Add yuv420 support

PyAV-Org · Apr 9, 2023 · 6c6d398 · 6c6d398
1 parent 13d5c9b
commit 6c6d398
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/av/video/frame.pyx b/av/video/frame.pyx
@@ -347,7 +347,7 @@ cdef class VideoFrame(Frame):
         elif format in ('gray', 'gray8', 'rgb8', 'bgr8'):
             check_ndarray(array, 'uint8', 2)
             height, width = array.shape[:2]
-        elif format == "nv12":
+        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)

diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py
@@ -526,6 +526,26 @@ def test_shares_memory_rgb24(self):
         # Make sure the frame reflects that
         self.assertNdarraysEqual(frame.to_ndarray(), array)
 
+    def test_shares_memory_yuv420p(self):
+        array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
+        frame = VideoFrame.from_numpy_buffer(array, "yuv420p")
+        self.assertNdarraysEqual(frame.to_ndarray(), array)
+
+        # overwrite the array, the contents thereof
+        array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8)
+        # Make sure the frame reflects that
+        self.assertNdarraysEqual(frame.to_ndarray(), array)
+
+    def test_shares_memory_yuvj420p(self):
+        array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
+        frame = VideoFrame.from_numpy_buffer(array, "yuvj420p")
+        self.assertNdarraysEqual(frame.to_ndarray(), array)
+
+        # overwrite the array, the contents thereof
+        array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8)
+        # Make sure the frame reflects that
+        self.assertNdarraysEqual(frame.to_ndarray(), array)
+
     def test_shares_memory_nv12(self):
         array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
         frame = VideoFrame.from_numpy_buffer(array, "nv12")