Add yuv422p10le pix_fmt

Author: WyattBlue

av/video/frame.pyx

@@ -20,9 +20,9 @@ supported_np_pix_fmts = {
     "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",
+    "rgb48be", "rgb48le", "rgb8", "rgba", "rgba64be", "rgba64le", "yuv420p",
+    "yuv422p10le", "yuv444p", "yuv444p16be", "yuv444p16le", "yuva444p16be",
+    "yuva444p16le", "yuvj420p", "yuvj444p", "yuyv422",
 }
 
 cdef VideoFrame alloc_video_frame():
@@ -322,7 +322,7 @@ cdef class VideoFrame(Frame):
         import numpy as np
 
         # check size
-        if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422"}:
+        if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422", "yuv422p10le"}:
             assert frame.width % 2 == 0, "the width has to be even for this pixel format"
             assert frame.height % 2 == 0, "the height has to be even for this pixel format"
 
@@ -408,6 +408,18 @@ cdef class VideoFrame(Frame):
                 useful_array(frame.planes[1]),
                 useful_array(frame.planes[2]),
             ]).reshape(-1, frame.width)
+        if frame.format.name == "yuv422p10le":
+            # Read planes as uint16 at their original width
+            y = useful_array(frame.planes[0], 2, "uint16").reshape(frame.height, frame.width)
+            u = useful_array(frame.planes[1], 2, "uint16").reshape(frame.height, frame.width // 2)
+            v = useful_array(frame.planes[2], 2, "uint16").reshape(frame.height, frame.width // 2)
+
+            # Double the width of U and V by repeating each value
+            u_full = np.repeat(u, 2, axis=1)
+            v_full = np.repeat(v, 2, axis=1)
+            if channel_last:
+                return np.stack([y, u_full, v_full], axis=2)
+            return np.stack([y, u_full, v_full], axis=0)
         if frame.format.name == "pal8":
             image = useful_array(frame.planes[0]).reshape(frame.height, frame.width)
             palette = np.frombuffer(frame.planes[1], "i4").astype(">i4").reshape(-1, 1).view(np.uint8)
@@ -641,6 +653,28 @@ cdef class VideoFrame(Frame):
             copy_array_to_plane(flat[u_start:v_start], frame.planes[1], 1)
             copy_array_to_plane(flat[v_start:], frame.planes[2], 1)
             return frame
+        elif format == "yuv422p10le":
+            if not isinstance(array, np.ndarray) or array.dtype != np.uint16:
+                raise ValueError("Array must be uint16 type")
+
+            # Convert to channel-first if needed
+            if channel_last and array.shape[2] == 3:
+                array = np.moveaxis(array, 2, 0)
+            elif not (array.shape[0] == 3):
+                raise ValueError("Array must have shape (3, height, width) or (height, width, 3)")
+
+            height, width = array.shape[1:]
+            if width % 2 != 0 or height % 2 != 0:
+                raise ValueError("Width and height must be even")
+
+            frame = VideoFrame(width, height, format)
+            copy_array_to_plane(array[0], frame.planes[0], 2)
+            # Subsample U and V by taking every other column
+            u = array[1, :, ::2].copy()  # Need copy to ensure C-contiguous
+            v = array[2, :, ::2].copy()  # Need copy to ensure C-contiguous
+            copy_array_to_plane(u, frame.planes[1], 2)
+            copy_array_to_plane(v, frame.planes[2], 2)
+            return frame
         elif format == "yuyv422":
             check_ndarray(array, "uint8", 3)
             check_ndarray_shape(array, array.shape[0] % 2 == 0)

tests/test_videoframe.py

@@ -485,6 +485,15 @@ def test_ndarray_yuv444p16() -> None:
         assertNdarraysEqual(frame.to_ndarray(), array)
 
 
+def test_ndarray_yuv422p10le() -> None:
+    array = numpy.random.randint(0, 65536, size=(3, 480, 640), dtype=numpy.uint16)
+    for format in ("yuv422p10le",):
+        frame = VideoFrame.from_ndarray(array, format=format)
+        assert frame.width == 640 and frame.height == 480
+        assert frame.format.name == format
+        assert format in av.video.frame.supported_np_pix_fmts
+
+
 def test_ndarray_yuv444p16_align() -> None:
     array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16)
     for format in ("yuv444p16be", "yuv444p16le"):