Make VideoFrame.from_numpy_buffer support buffers with padding

Some devices have hardware that creates image buffers with padding, so
adding support here means less frame buffer copying is required.

Specifically, we extend the support to buffers where the pixel rows
are contiguous, though the image doesn't comprise all the pixels on
the row (and is therefore not strictly contiguous).

We also support yuv420p images with padding. These have padding in the
middle of the UV rows as well as at the end, so can't be trimmed by
the application before being passed in. Instead, the true image width
must be passed.

Tests are also added to ensure all these cases now avoid copying.

av/video/frame.pyi

@@ -62,7 +62,7 @@ class VideoFrame(Frame):
     def from_image(img: Image.Image) -> VideoFrame: ...
     @staticmethod
     def from_numpy_buffer(
-        array: _SupportedNDarray, format: str = "rgb24"
+        array: _SupportedNDarray, format: str = "rgb24", width: int = 0
     ) -> VideoFrame: ...
     @staticmethod
     def from_ndarray(array: _SupportedNDarray, format: str = "rgb24") -> VideoFrame: ...

av/video/frame.pyx

@@ -374,31 +374,54 @@ cdef class VideoFrame(Frame):
         return frame
 
     @staticmethod
-    def from_numpy_buffer(array, format="rgb24"):
+    def from_numpy_buffer(array, format="rgb24", width=0):
+        # Usually the width of the array is the same as the width of the image. But sometimes
+        # this is not possible, for example with yuv420p images that have padding. These are
+        # awkward because the UV rows at the bottom have padding bytes in the middle of the
+        # row as well as at the end. To cope with these, callers need to be able to pass the
+        # actual width to us.
+        height = array.shape[0]
+        if not width:
+            width = array.shape[1]
+
         if format in ("rgb24", "bgr24"):
             check_ndarray(array, "uint8", 3)
             check_ndarray_shape(array, array.shape[2] == 3)
-            height, width = array.shape[:2]
+            if array.strides[1:] != (3, 1):
+                raise ValueError("provided array does not have C_CONTIGUOUS rows")
+            linesizes = (array.strides[0], )
+        elif format in ("rgba", "bgra"):
+            check_ndarray(array, "uint8", 3)
+            check_ndarray_shape(array, array.shape[2] == 4)
+            if array.strides[1:] != (4, 1):
+                raise ValueError("provided array does not have C_CONTIGUOUS rows")
+            linesizes = (array.strides[0], )
         elif format in ("gray", "gray8", "rgb8", "bgr8"):
             check_ndarray(array, "uint8", 2)
-            height, width = array.shape[:2]
+            if array.strides[1] != 1:
+                raise ValueError("provided array does not have C_CONTIGUOUS rows")
+            linesizes = (array.strides[0], )
         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)
-            height, width = array.shape[:2]
             height = height // 6 * 4
+            if array.strides[1] != 1:
+                raise ValueError("provided array does not have C_CONTIGUOUS rows")
+            if format in ("yuv420p", "yuvj420p"):
+                # For YUV420 planar formats, the UV plane stride is always half the Y stride.
+                linesizes = (array.strides[0], array.strides[0] // 2, array.strides[0] // 2)
+            else:
+                # Planes where U and V are interleaved have the same stride as Y.
+                linesizes = (array.strides[0], array.strides[0])
         else:
             raise ValueError(f"Conversion from numpy array with format `{format}` is not yet supported")
 
-        if not array.flags["C_CONTIGUOUS"]:
-            raise ValueError("provided array must be C_CONTIGUOUS")
-
         frame = alloc_video_frame()
-        frame._image_fill_pointers_numpy(array, width, height, format)
+        frame._image_fill_pointers_numpy(array, width, height, linesizes, format)
         return frame
 
-    def _image_fill_pointers_numpy(self, buffer, width, height, format):
+    def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format):
         cdef lib.AVPixelFormat c_format
         cdef uint8_t * c_ptr
         cdef size_t c_data
@@ -433,13 +456,8 @@ cdef class VideoFrame(Frame):
         self.ptr.format = c_format
         self.ptr.width = width
         self.ptr.height = height
-        res = lib.av_image_fill_linesizes(
-            self.ptr.linesize,
-            <lib.AVPixelFormat>self.ptr.format,
-            width,
-        )
-        if res:
-          err_check(res)
+        for i, linesize in enumerate(linesizes):
+            self.ptr.linesize[i] = linesize
 
         res = lib.av_image_fill_pointers(
             self.ptr.data,

tests/test_videoframe.py

@@ -528,6 +528,19 @@ def test_shares_memory_gray() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
+    array = array[:, :300]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "gray")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_gray8() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
@@ -539,6 +552,19 @@ def test_shares_memory_gray8() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
+    array = array[:, :300]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "gray8")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_rgb8() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
@@ -550,6 +576,19 @@ def test_shares_memory_rgb8() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
+    array = array[:, :300]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "rgb8")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_bgr8() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
@@ -561,6 +600,19 @@ def test_shares_memory_bgr8() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
+    array = array[:, :300]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "bgr8")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_rgb24() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318, 3), dtype=numpy.uint8)
@@ -572,6 +624,43 @@ def test_shares_memory_rgb24() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318, 3), dtype=numpy.uint8)
+    array = array[:, :300, :]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "rgb24")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+
+def test_shares_memory_rgba() -> None:
+    array = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    frame = VideoFrame.from_numpy_buffer(array, "rgba")
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+    # overwrite the array, the contents thereof
+    array[...] = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    # Make sure the frame reflects that
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    array = array[:, :300, :]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "rgba")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_yuv420p() -> None:
     array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
@@ -583,6 +672,38 @@ def test_shares_memory_yuv420p() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array where there are some padding bytes
+    # note that the uv rows have half the padding in the middle of a row, and the
+    # other half at the end
+    height = 512
+    stride = 256
+    width = 200
+    array = numpy.random.randint(
+        0, 256, size=(height * 6 // 4, stride), dtype=numpy.uint8
+    )
+    uv_width = width // 2
+    uv_stride = stride // 2
+
+    # compare carefully, avoiding all the padding bytes which to_ndarray strips out
+    frame = VideoFrame.from_numpy_buffer(array, "yuv420p", width=width)
+    frame_array = frame.to_ndarray()
+    assertNdarraysEqual(frame_array[:height, :width], array[:height, :width])
+    assertNdarraysEqual(frame_array[height:, :uv_width], array[height:, :uv_width])
+    assertNdarraysEqual(
+        frame_array[height:, uv_width:],
+        array[height:, uv_stride : uv_stride + uv_width],
+    )
+
+    # overwrite the array, and check the shared frame buffer changed too!
+    array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8)
+    frame_array = frame.to_ndarray()
+    assertNdarraysEqual(frame_array[:height, :width], array[:height, :width])
+    assertNdarraysEqual(frame_array[height:, :uv_width], array[height:, :uv_width])
+    assertNdarraysEqual(
+        frame_array[height:, uv_width:],
+        array[height:, uv_stride : uv_stride + uv_width],
+    )
+
 
 def test_shares_memory_yuvj420p() -> None:
     array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
@@ -594,6 +715,36 @@ def test_shares_memory_yuvj420p() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test with padding, just as we did in the yuv420p case
+    height = 512
+    stride = 256
+    width = 200
+    array = numpy.random.randint(
+        0, 256, size=(height * 6 // 4, stride), dtype=numpy.uint8
+    )
+    uv_width = width // 2
+    uv_stride = stride // 2
+
+    # compare carefully, avoiding all the padding bytes which to_ndarray strips out
+    frame = VideoFrame.from_numpy_buffer(array, "yuvj420p", width=width)
+    frame_array = frame.to_ndarray()
+    assertNdarraysEqual(frame_array[:height, :width], array[:height, :width])
+    assertNdarraysEqual(frame_array[height:, :uv_width], array[height:, :uv_width])
+    assertNdarraysEqual(
+        frame_array[height:, uv_width:],
+        array[height:, uv_stride : uv_stride + uv_width],
+    )
+
+    # overwrite the array, and check the shared frame buffer changed too!
+    array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8)
+    frame_array = frame.to_ndarray()
+    assertNdarraysEqual(frame_array[:height, :width], array[:height, :width])
+    assertNdarraysEqual(frame_array[height:, :uv_width], array[height:, :uv_width])
+    assertNdarraysEqual(
+        frame_array[height:, uv_width:],
+        array[height:, uv_stride : uv_stride + uv_width],
+    )
+
 
 def test_shares_memory_nv12() -> None:
     array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
@@ -605,6 +756,19 @@ def test_shares_memory_nv12() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
+    array = array[:, :200]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "nv12")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_shares_memory_bgr24() -> None:
     array = numpy.random.randint(0, 256, size=(357, 318, 3), dtype=numpy.uint8)
@@ -616,6 +780,43 @@ def test_shares_memory_bgr24() -> None:
     # Make sure the frame reflects that
     assertNdarraysEqual(frame.to_ndarray(), array)
 
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318, 3), dtype=numpy.uint8)
+    array = array[:, :300, :]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "bgr24")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+
+def test_shares_memory_bgra() -> None:
+    array = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    frame = VideoFrame.from_numpy_buffer(array, "bgra")
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+    # overwrite the array, the contents thereof
+    array[...] = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    # Make sure the frame reflects that
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
+    # repeat the test, but with an array that is not fully contiguous, though the
+    # pixels in a row are
+    array = numpy.random.randint(0, 256, size=(357, 318, 4), dtype=numpy.uint8)
+    array = array[:, :300, :]
+    assert not array.data.c_contiguous
+    frame = VideoFrame.from_numpy_buffer(array, "bgra")
+    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
+    assertNdarraysEqual(frame.to_ndarray(), array)
+
 
 def test_reformat_pts() -> None:
     frame = VideoFrame(640, 480, "rgb24")