Add BGR48, BGRA64 pixel formats

Extend `supported_np_pix_fmts` pixel formats with `bgr48be` `bgr48le` `bgra64be` `bgra64le`.

Add handling for `argb` `abgr` `grayf32le` `grayf32be` pixel formats.
---------

Co-authored-by: WyattBlue <[email protected]>

Author: z-khan

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)

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")