From 8bf5d03ca7b479ba1ed45d9e0d0600d99532b85d Mon Sep 17 00:00:00 2001 From: David Plowman Date: Wed, 13 Nov 2024 16:16:05 +0000 Subject: [PATCH] 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 | 2 +- av/video/frame.pyx | 50 ++++++---- tests/test_videoframe.py | 201 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 236 insertions(+), 17 deletions(-) diff --git a/av/video/frame.pyi b/av/video/frame.pyi index d837ed606..0739010c1 100644 --- a/av/video/frame.pyi +++ b/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: ... diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 80cf266f8..862db8513 100644 --- a/av/video/frame.pyx +++ b/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, - 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, diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 6396f1f45..32b6e5482 100644 --- a/tests/test_videoframe.py +++ b/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")