diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 785603a1a..04a14fd8b 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -16,7 +16,7 @@ cdef object _cinit_bypass_sentinel 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", + "bayer_rggb16be", "bayer_rggb16le", "bayer_rggb8", "bgr24", "bgr48be", "bgr48le", "bgr8", "bgra", "bgra64be", "bgra64le", "gbrapf32be", "gbrapf32le", "gbrp", "gbrp10be", "gbrp10le", "gbrp12be", "gbrp12le", "gbrp14be", "gbrp14le", "gbrp16be", "gbrp16le", "gbrpf32be", "gbrpf32le", "gray", "gray16be", "gray16le", "gray8", "grayf32be", "grayf32le", "nv12", "pal8", "rgb24", @@ -342,6 +342,8 @@ cdef class VideoFrame(Frame): "bayer_rggb16le": (2, "uint16"), "bayer_rggb16be": (2, "uint16"), "bgr24": (3, "uint8"), + "bgr48be": (6, "uint16"), + "bgr48le": (6, "uint16"), "bgr8": (1, "uint8"), "bgra": (4, "uint8"), "gbrapf32be": (4, "float32"), @@ -370,6 +372,8 @@ cdef class VideoFrame(Frame): "rgba": (4, "uint8"), "rgba64be": (8, "uint16"), "rgba64le": (8, "uint16"), + "bgra64be": (8, "uint16"), + "bgra64le": (8, "uint16"), "yuv444p": (1, "uint8"), "yuv444p16be": (2, "uint16"), "yuv444p16le": (2, "uint16"), @@ -467,43 +471,58 @@ cdef class VideoFrame(Frame): if not width: width = array.shape[1] - if format in ("rgb24", "bgr24"): + if format in {"rgb24", "bgr24"}: check_ndarray(array, "uint8", 3) check_ndarray_shape(array, array.shape[2] == 3) if array.strides[1:] != (3, 1): raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0], ) + elif format in {"rgb48le", "rgb48be", "bgr48le", "bgr48be"}: + check_ndarray(array, "uint16", 3) + check_ndarray_shape(array, array.shape[2] == 3) + if array.strides[1:] != (6, 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") linesizes = (array.strides[0], ) - elif format in ("rgba", "bgra"): + elif format in {"rgba", "bgra", "argb", "abgr"}: check_ndarray(array, "uint8", 3) check_ndarray_shape(array, array.shape[2] == 4) if array.strides[1:] != (4, 1): raise ValueError("provided array does not have C_CONTIGUOUS rows") linesizes = (array.strides[0], ) - elif format in ("gray", "gray8", "rgb8", "bgr8"): + elif format in {"rgba64le", "rgba64be", "bgra64le", "bgra64be"}: + check_ndarray(array, "uint16", 3) + check_ndarray_shape(array, array.shape[2] == 4) + if array.strides[1:] != (8, 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0], ) + elif format in {"gray", "gray8", "rgb8", "bgr8","bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8"}: check_ndarray(array, "uint8", 2) if array.strides[1] != 1: raise ValueError("provided array does not have C_CONTIGUOUS rows") linesizes = (array.strides[0], ) - elif format in ("yuv420p", "yuvj420p", "nv12"): + elif format in {"gray16le", "gray16be", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"}: + check_ndarray(array, "uint16", 2) + if array.strides[1] != 2: + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0], ) + elif format in {"grayf32le", "grayf32be"}: + check_ndarray(array, "float32", 2) + if array.strides[1] != 4: + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0], ) + elif format in {"yuv420p", "yuvj420p", "nv12"}: check_ndarray(array, "uint8", 2) check_ndarray_shape(array, array.shape[0] % 3 == 0) check_ndarray_shape(array, array.shape[1] % 2 == 0) height = height // 6 * 4 if array.strides[1] != 1: raise ValueError("provided array does not have C_CONTIGUOUS rows") - if format in ("yuv420p", "yuvj420p"): + if format in {"yuv420p", "yuvj420p"}: # For YUV420 planar formats, the UV plane stride is always half the Y stride. linesizes = (array.strides[0], array.strides[0] // 2, array.strides[0] // 2) else: # Planes where U and V are interleaved have the same stride as Y. linesizes = (array.strides[0], array.strides[0]) - elif format in {"bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8","bayer_bggr16le", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"}: - check_ndarray(array, "uint8" if format.endswith("8") else "uint16", 2) - - if array.strides[1] != (1 if format.endswith("8") else 2): - raise ValueError("provided array does not have C_CONTIGUOUS rows") - - linesizes = (array.strides[0],) else: raise ValueError(f"Conversion from numpy array with format `{format}` is not yet supported") @@ -717,13 +736,13 @@ cdef class VideoFrame(Frame): elif format in {"argb", "rgba", "abgr", "bgra"}: check_ndarray(array, "uint8", 3) check_ndarray_shape(array, array.shape[2] == 4) - elif format in {"rgb48be", "rgb48le"}: + elif format in {"rgb48be", "rgb48le","bgr48be", "bgr48le"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 3) frame = VideoFrame(array.shape[1], array.shape[0], format) copy_array_to_plane(byteswap_array(array, format.endswith("be")), frame.planes[0], 6) return frame - elif format in {"rgba64be", "rgba64le"}: + elif format in {"rgba64be", "rgba64le", "bgra64be", "bgra64le"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 4) frame = VideoFrame(array.shape[1], array.shape[0], format) diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 9327a4ccd..26549b31b 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -285,6 +285,35 @@ def test_ndarray_rgba_align() -> None: assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_bayer8() -> None: + array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) + for format in ("bayer_bggr8", "bayer_gbrg8", "bayer_grbg8", "bayer_rggb8"): + 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) + + +def test_ndarray_bayer16() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16) + for format in ( + "bayer_bggr16be", + "bayer_bggr16le", + "bayer_gbrg16be", + "bayer_gbrg16le", + "bayer_grbg16be", + "bayer_grbg16le", + "bayer_rggb16be", + "bayer_rggb16le", + ): + 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) + + def test_ndarray_gbrp() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 3), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="gbrp") @@ -571,6 +600,17 @@ def test_ndarray_rgb48be() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "big") +def test_ndarray_bgr48be() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="bgr48be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "bgr48be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining blue value of first pixel + assertPixelValue16(frame.planes[0], array[0][0][0], "big") + + def test_ndarray_rgb48le() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) frame = VideoFrame.from_ndarray(array, format="rgb48le") @@ -582,6 +622,17 @@ def test_ndarray_rgb48le() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "little") +def test_ndarray_bgr48le() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="bgr48le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "bgr48le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining blue value of first pixel + assertPixelValue16(frame.planes[0], array[0][0][0], "little") + + def test_ndarray_rgb48le_align() -> None: array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16) frame = VideoFrame.from_ndarray(array, format="rgb48le") @@ -593,6 +644,17 @@ def test_ndarray_rgb48le_align() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "little") +def test_ndarray_bgr48le_align() -> None: + array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="bgr48le") + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == "bgr48le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining blue value of first pixel + assertPixelValue16(frame.planes[0], array[0][0][0], "little") + + def test_ndarray_rgba64be() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16) frame = VideoFrame.from_ndarray(array, format="rgba64be") @@ -604,6 +666,17 @@ def test_ndarray_rgba64be() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "big") +def test_ndarray_bgra64be() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="bgra64be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "bgra64be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining blue value of first pixel + assertPixelValue16(frame.planes[0], array[0][0][0], "big") + + def test_ndarray_rgba64le() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16) frame = VideoFrame.from_ndarray(array, format="rgba64le") @@ -615,6 +688,17 @@ def test_ndarray_rgba64le() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "little") +def test_ndarray_bgra64le() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="bgra64le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "bgra64le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining blue value of first pixel + assertPixelValue16(frame.planes[0], array[0][0][0], "little") + + def test_ndarray_rgb8() -> None: array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="rgb8") @@ -805,6 +889,25 @@ def test_shares_memory_rgba() -> None: assertNdarraysEqual(frame.to_ndarray(), array) +def test_shares_memory_bayer8() -> None: + for format in ("bayer_rggb8", "bayer_bggr8", "bayer_grbg8", "bayer_gbrg8"): + array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8) + frame = VideoFrame.from_numpy_buffer(array, format) + assertNdarraysEqual(frame.to_ndarray(), array) + + array[...] = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8) + assertNdarraysEqual(frame.to_ndarray(), array) + + array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8) + array = array[:, :300] + assert not array.data.c_contiguous + frame = VideoFrame.from_numpy_buffer(array, format) + assertNdarraysEqual(frame.to_ndarray(), array) + + array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8) + assertNdarraysEqual(frame.to_ndarray(), array) + + def test_shares_memory_yuv420p() -> None: array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8) frame = VideoFrame.from_numpy_buffer(array, "yuv420p")