From 3238ab43643484452d55b9f6979d0141ffb7b55c Mon Sep 17 00:00:00 2001
From: Nikhil Idiculla <nikhilidiculla@gmail.com>
Date: Wed, 26 Mar 2025 12:17:22 -0700
Subject: [PATCH] Add yuv420p10le pix_fmt

---
 av/video/frame.pyx       | 36 ++++++++++++++++++++++++++++++++++--
 tests/test_videoframe.py |  9 +++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/av/video/frame.pyx b/av/video/frame.pyx
index 84a35cc49..785603a1a 100644
--- a/av/video/frame.pyx
+++ b/av/video/frame.pyx
@@ -21,7 +21,7 @@ supported_np_pix_fmts = {
     "gbrp14be", "gbrp14le", "gbrp16be", "gbrp16le", "gbrpf32be", "gbrpf32le", "gray",
     "gray16be", "gray16le", "gray8", "grayf32be", "grayf32le", "nv12", "pal8", "rgb24",
     "rgb48be", "rgb48le", "rgb8", "rgba", "rgba64be", "rgba64le", "yuv420p",
-    "yuv422p10le", "yuv444p", "yuv444p16be", "yuv444p16le", "yuva444p16be",
+    "yuv420p10le", "yuv422p10le", "yuv444p", "yuv444p16be", "yuv444p16le", "yuva444p16be",
     "yuva444p16le", "yuvj420p", "yuvj444p", "yuyv422",
 }
 
@@ -321,7 +321,7 @@ cdef class VideoFrame(Frame):
         import numpy as np
 
         # check size
-        if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422", "yuv422p10le"}:
+        if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422", "yuv420p10le", "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"
 
@@ -407,6 +407,16 @@ cdef class VideoFrame(Frame):
                 useful_array(frame.planes[1]),
                 useful_array(frame.planes[2]),
             ]).reshape(-1, frame.width)
+        if frame.format.name == "yuv420p10le":
+            # Read planes as uint16:
+            y = useful_array(frame.planes[0], 2, "uint16").reshape(frame.height, frame.width)
+            u = useful_array(frame.planes[1], 2, "uint16").reshape(frame.height // 2, frame.width // 2)
+            v = useful_array(frame.planes[2], 2, "uint16").reshape(frame.height // 2, frame.width // 2)
+            u_full = np.repeat(np.repeat(u, 2, axis=1), 2, axis=0)
+            v_full = np.repeat(np.repeat(u, 2, axis=1), 2, axis=0)
+            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 == "yuv422p10le":
             # Read planes as uint16 at their original width
             y = useful_array(frame.planes[0], 2, "uint16").reshape(frame.height, frame.width)
@@ -652,6 +662,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 == "yuv420p10le":
+            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 row and column:
+            u = array[1, ::2, ::2].copy()  # Need copy to ensure C-contiguous
+            v = array[2, ::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 == "yuv422p10le":
             if not isinstance(array, np.ndarray) or array.dtype != np.uint16:
                 raise ValueError("Array must be uint16 type")
diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py
index bef69550b..9327a4ccd 100644
--- a/tests/test_videoframe.py
+++ b/tests/test_videoframe.py
@@ -485,6 +485,15 @@ def test_ndarray_yuv444p16() -> None:
         assertNdarraysEqual(frame.to_ndarray(), array)
 
 
+def test_ndarray_yuv420p10le() -> None:
+    array = numpy.random.randint(0, 65536, size=(3, 480, 640), dtype=numpy.uint16)
+    for format in ("yuv420p10le",):
+        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_yuv422p10le() -> None:
     array = numpy.random.randint(0, 65536, size=(3, 480, 640), dtype=numpy.uint16)
     for format in ("yuv422p10le",):