Support yuv420p10le with ndarray conversions

CHANGELOG.rst

@@ -38,6 +38,7 @@ Breaking:
 Features:
 
 - Add ``options`` parameter to ``AudioResampler`` for passing ``libswresample`` options (e.g. ``resampler``, ``filter_size``, ``cutoff``) by :gh-user:`WyattBlue` (:issue:`2262`).
+- Support ``yuv420p10le`` in ``VideoFrame.to_ndarray`` and ``VideoFrame.from_ndarray`` by :gh-user:`WyattBlue` (:issue:`1981`).
 
 Fixes:
 

av/video/frame.py

@@ -254,6 +254,7 @@ def _numpy_avbuffer_free(
     "rgbf32be",
     "rgbf32le",
     "yuv420p",
+    "yuv420p10le",
     "yuv422p10le",
     "yuv444p",
     "yuv444p16be",
@@ -736,6 +737,7 @@ def to_image(self, **kwargs):
             "RGB", (plane.width, plane.height), bytes(o_buf), "raw", "RGB", 0, 1
         )
 
+    @cython.cdivision(True)
     def to_ndarray(self, channel_last=False, **kwargs):
         """Get a numpy array of this frame.
 
@@ -802,14 +804,19 @@ def to_ndarray(self, channel_last=False, **kwargs):
             return array
 
         # special cases
-        if format_name in {"yuv420p", "yuvj420p", "yuv422p"}:
+        if format_name in {"yuv420p", "yuvj420p", "yuv422p", "yuv420p10le"}:
             assert frame.ptr.width % 2 == 0, "width has to be even for this format"
             assert frame.ptr.height % 2 == 0, "height has to be even for this format"
+            is_10bit: cython.bint = format_name == "yuv420p10le"
+            itemsize = 2 if is_10bit else 1
+            dtype = "uint16" if is_10bit else "uint8"
+            if frame.ptr.width == 0 or frame.ptr.height == 0:
+                return np.empty((0, frame.ptr.width), dtype=dtype)
             return np.hstack(
                 [
-                    useful_array(planes[0]).reshape(-1),
-                    useful_array(planes[1]).reshape(-1),
-                    useful_array(planes[2]).reshape(-1),
+                    useful_array(planes[0], itemsize, dtype).reshape(-1),
+                    useful_array(planes[1], itemsize, dtype).reshape(-1),
+                    useful_array(planes[2], itemsize, dtype).reshape(-1),
                 ]
             ).reshape(-1, frame.ptr.width)
         if format_name == "yuv422p10le":
@@ -1143,6 +1150,7 @@ def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format):
         self._init_user_attributes()
 
     @staticmethod
+    @cython.cdivision(True)
     def from_ndarray(array, format="rgb24", channel_last=False):
         """
         Construct a frame from a numpy array.
@@ -1261,19 +1269,38 @@ def from_ndarray(array, format="rgb24", channel_last=False):
             check_ndarray_shape(array, array.shape[1] % 2 == 0)
 
             frame = VideoFrame(array.shape[1], (array.shape[0] * 2) // 3, format)
+            if frame.width == 0 or frame.height == 0:
+                return frame
             u_start = frame.width * frame.height
             v_start = 5 * u_start // 4
             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 == "yuv420p10le":
+            check_ndarray(array, "uint16", 2)
+            check_ndarray_shape(array, array.shape[0] % 3 == 0)
+            check_ndarray_shape(array, array.shape[1] % 2 == 0)
+
+            frame = VideoFrame(array.shape[1], (array.shape[0] * 2) // 3, format)
+            if frame.width == 0 or frame.height == 0:
+                return frame
+            u_start = frame.width * frame.height
+            v_start = 5 * u_start // 4
+            flat = array.reshape(-1)
+            copy_array_to_plane(flat[0:u_start], frame.planes[0], 2)
+            copy_array_to_plane(flat[u_start:v_start], frame.planes[1], 2)
+            copy_array_to_plane(flat[v_start:], frame.planes[2], 2)
+            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)
+            if frame.width == 0 or frame.height == 0:
+                return frame
             u_start = frame.width * frame.height
             v_start = u_start + u_start // 2
             flat = array.reshape(-1)

tests/test_videoframe.py

@@ -688,6 +688,39 @@ def test_ndarray_yuv420p() -> None:
     assertNdarraysEqual(frame.to_ndarray(), array)
 
 
+def test_ndarray_yuv420p10le() -> None:
+    array = numpy.random.randint(0, 1024, size=(720, 640), dtype=numpy.uint16)
+    frame = VideoFrame.from_ndarray(array, format="yuv420p10le")
+    assert frame.width == 640 and frame.height == 480
+    assert frame.format.name == "yuv420p10le"
+    assert "yuv420p10le" in supported_np_pix_fmts
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+
+def test_ndarray_yuv420p10le_align() -> None:
+    array = numpy.random.randint(0, 1024, size=(357, 318), dtype=numpy.uint16)
+    frame = VideoFrame.from_ndarray(array, format="yuv420p10le")
+    assert frame.width == 318 and frame.height == 238
+    assert frame.format.name == "yuv420p10le"
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+
+def test_ndarray_yuv420p10le_zero_size() -> None:
+    # A frame with zero width and/or height has no allocated planes; conversions
+    # should degrade gracefully instead of raising IndexError.
+    for w, h in ((0, 0), (0, 480), (640, 0)):
+        frame = VideoFrame(w, h, "yuv420p10le")
+        array = frame.to_ndarray()
+        assert array.dtype == numpy.uint16
+        assert array.size == 0
+
+    empty = numpy.empty((0, 0), dtype=numpy.uint16)
+    frame = VideoFrame.from_ndarray(empty, format="yuv420p10le")
+    assert frame.width == 0 and frame.height == 0
+    assert frame.format.name == "yuv420p10le"
+    assertNdarraysEqual(frame.to_ndarray(), empty)
+
+
 def test_ndarray_yuv422p() -> None:
     array = numpy.random.randint(0, 256, size=(960, 640), dtype=numpy.uint8)
     frame = VideoFrame.from_ndarray(array, format="yuv422p")