Add yuv422p support for video frame to_ndarray and from_ndarray

Author: wader

av/video/frame.py

@@ -449,13 +449,19 @@ def to_ndarray(self, channel_last=False, **kwargs):
         import numpy as np
 
         # check size
-        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"
-            )
+        if frame.format.name in {
+            "yuv420p",
+            "yuvj420p",
+            "yuyv422",
+            "yuv422p10le",
+            "yuv422p",
+        }:
+            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"
 
         # cases planes are simply concatenated in shape (height, width, channels)
         itemsize, dtype = {
@@ -563,7 +569,7 @@ def to_ndarray(self, channel_last=False, **kwargs):
             return array
 
         # special cases
-        if frame.format.name in {"yuv420p", "yuvj420p"}:
+        if frame.format.name in {"yuv420p", "yuvj420p", "yuv422p"}:
             return np.hstack(
                 [
                     useful_array(frame.planes[0]),
@@ -1022,6 +1028,19 @@ def from_ndarray(array, format="rgb24", channel_last=False):
             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 == "yuv422p":
+            check_ndarray(array, "uint8", 2)
+            check_ndarray_shape(array, array.shape[0] % 4 == 0)
+            check_ndarray_shape(array, array.shape[1] % 2 == 0)
+
+            frame = VideoFrame(array.shape[1], array.shape[0] // 2, format)
+            u_start = frame.width * frame.height
+            v_start = u_start + u_start // 2
+            flat = array.reshape(-1)
+            copy_array_to_plane(flat[0:u_start], frame.planes[0], 1)
+            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")

tests/test_videoframe.py

@@ -648,6 +648,14 @@ def test_ndarray_yuv420p() -> None:
     assertNdarraysEqual(frame.to_ndarray(), array)
 
 
+def test_ndarray_yuv422p() -> None:
+    array = numpy.random.randint(0, 256, size=(960, 640), dtype=numpy.uint8)
+    frame = VideoFrame.from_ndarray(array, format="yuv422p")
+    assert frame.width == 640 and frame.height == 480
+    assert frame.format.name == "yuv422p"
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+
 def test_ndarray_yuv420p_align() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
     frame = VideoFrame.from_ndarray(array, format="yuv420p")