diff --git a/av/video/frame.pyi b/av/video/frame.pyi index d650cf44c..bba60cc5d 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -15,6 +15,8 @@ _SupportedNDarray = Union[ np.ndarray[Any, np.dtype[np.float32]], ] +supported_np_pix_fmts: set[str] + class PictureType(IntEnum): NONE = 0 I = 1 diff --git a/av/video/frame.pyx b/av/video/frame.pyx index b74d58d1c..bb8f83486 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -12,6 +12,19 @@ from av.video.plane cimport VideoPlane cdef object _cinit_bypass_sentinel +# `pix_fmt`s supported by Frame.to_ndarray() and Frame.from_ndarray() +supported_np_pix_fmts = { + "abgr", "argb", "bayer_bggr16be", "bayer_bggr16le", "bayer_bggr8", "bayer_gbrg16be", + "bayer_gbrg16le", "bayer_gbrg8", "bayer_grbg16be", "bayer_grbg16le", "bayer_grbg8", + "bayer_rggb16be", "bayer_rggb16le", "bayer_rggb8", "bgr24", "bgr8", "bgra", + "gbrapf32be", "gbrapf32le", "gbrp", "gbrp10be", "gbrp10le", "gbrp12be", "gbrp12le", + "gbrp14be", "gbrp14le", "gbrp16be", "gbrp16le", "gbrpf32be", "gbrpf32le", "gray", + "gray16be", "gray16le", "gray8", "grayf32be", "grayf32le", "nv12", "pal8", "rgb24", + "rgb48be", "rgb48le", "rgb8", "rgba", "rgba64be", "rgba64le", "yuv420p", "yuv444p", + "yuv444p16be", "yuv444p16le", "yuva444p16be", "yuva444p16le", "yuvj420p", + "yuvj444p", "yuyv422", +} + cdef VideoFrame alloc_video_frame(): """Get a mostly uninitialized VideoFrame. diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 9581a7c41..5ae603f64 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -209,6 +209,7 @@ def test_ndarray_gray() -> None: array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) for format in ("gray", "gray8"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == "gray" assertNdarraysEqual(frame.to_ndarray(), array) @@ -218,6 +219,7 @@ def test_ndarray_gray_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318), dtype=numpy.uint8) for format in ("gray", "gray8"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == "gray" assertNdarraysEqual(frame.to_ndarray(), array) @@ -227,6 +229,7 @@ def test_ndarray_grayf32() -> None: array = numpy.random.random_sample(size=(480, 640)).astype(numpy.float32) for format in ("grayf32be", "grayf32le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -236,6 +239,7 @@ def test_ndarray_grayf32_align() -> None: array = numpy.random.random_sample(size=(238, 318)).astype(numpy.float32) for format in ("grayf32be", "grayf32le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -245,6 +249,7 @@ def test_ndarray_rgb() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 3), dtype=numpy.uint8) for format in ("rgb24", "bgr24"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -254,6 +259,7 @@ def test_ndarray_rgb_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318, 3), dtype=numpy.uint8) for format in ("rgb24", "bgr24"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -263,6 +269,7 @@ def test_ndarray_rgba() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 4), dtype=numpy.uint8) for format in ("argb", "rgba", "abgr", "bgra"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -272,6 +279,7 @@ def test_ndarray_rgba_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318, 4), dtype=numpy.uint8) for format in ("argb", "rgba", "abgr", "bgra"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -297,6 +305,7 @@ def test_ndarray_gbrp10() -> None: array = numpy.random.randint(0, 1024, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp10be", "gbrp10le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -306,6 +315,7 @@ def test_ndarray_gbrp10_align() -> None: array = numpy.random.randint(0, 1024, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp10be", "gbrp10le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -315,6 +325,7 @@ def test_ndarray_gbrp12() -> None: array = numpy.random.randint(0, 4096, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp12be", "gbrp12le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -324,6 +335,7 @@ def test_ndarray_gbrp12_align() -> None: array = numpy.random.randint(0, 4096, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp12be", "gbrp12le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -333,6 +345,7 @@ def test_ndarray_gbrp14() -> None: array = numpy.random.randint(0, 16384, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp14be", "gbrp14le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -342,6 +355,7 @@ def test_ndarray_gbrp14_align() -> None: array = numpy.random.randint(0, 16384, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp14be", "gbrp14le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -351,6 +365,7 @@ def test_ndarray_gbrp16() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp16be", "gbrp16le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -360,6 +375,7 @@ def test_ndarray_gbrp16_align() -> None: array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp16be", "gbrp16le"): frame = VideoFrame.from_ndarray(array, format=format) + assert format in av.video.frame.supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -603,6 +619,7 @@ def test_ndarray_pal8(): frame = VideoFrame.from_ndarray((array, palette), format="pal8") assert frame.width == 640 and frame.height == 480 assert frame.format.name == "pal8" + assert frame.format.name in av.video.frame.supported_np_pix_fmts returned = frame.to_ndarray() assert type(returned) is tuple and len(returned) == 2 assertNdarraysEqual(returned[0], array) @@ -614,6 +631,7 @@ def test_ndarray_nv12() -> None: frame = VideoFrame.from_ndarray(array, format="nv12") assert frame.width == 640 and frame.height == 480 assert frame.format.name == "nv12" + assert frame.format.name in av.video.frame.supported_np_pix_fmts assertNdarraysEqual(frame.to_ndarray(), array)