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")