diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a18a60f07..67e265cf6 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,6 +38,7 @@ Breaking: Features: - Add ``options`` parameter to ``AudioResampler`` for passing ``libswresample`` options (e.g. ``resampler``, ``filter_size``, ``cutoff``) by :gh-user:`WyattBlue` (:issue:`2262`). +- Support ``yuv420p10le`` in ``VideoFrame.to_ndarray`` and ``VideoFrame.from_ndarray`` by :gh-user:`WyattBlue` (:issue:`1981`). Fixes: diff --git a/av/video/frame.py b/av/video/frame.py index 1a7b5364d..e530f9026 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -254,6 +254,7 @@ def _numpy_avbuffer_free( "rgbf32be", "rgbf32le", "yuv420p", + "yuv420p10le", "yuv422p10le", "yuv444p", "yuv444p16be", @@ -736,6 +737,7 @@ def to_image(self, **kwargs): "RGB", (plane.width, plane.height), bytes(o_buf), "raw", "RGB", 0, 1 ) + @cython.cdivision(True) def to_ndarray(self, channel_last=False, **kwargs): """Get a numpy array of this frame. @@ -802,14 +804,19 @@ def to_ndarray(self, channel_last=False, **kwargs): return array # special cases - if format_name in {"yuv420p", "yuvj420p", "yuv422p"}: + if format_name in {"yuv420p", "yuvj420p", "yuv422p", "yuv420p10le"}: assert frame.ptr.width % 2 == 0, "width has to be even for this format" assert frame.ptr.height % 2 == 0, "height has to be even for this format" + is_10bit: cython.bint = format_name == "yuv420p10le" + itemsize = 2 if is_10bit else 1 + dtype = "uint16" if is_10bit else "uint8" + if frame.ptr.width == 0 or frame.ptr.height == 0: + return np.empty((0, frame.ptr.width), dtype=dtype) return np.hstack( [ - useful_array(planes[0]).reshape(-1), - useful_array(planes[1]).reshape(-1), - useful_array(planes[2]).reshape(-1), + useful_array(planes[0], itemsize, dtype).reshape(-1), + useful_array(planes[1], itemsize, dtype).reshape(-1), + useful_array(planes[2], itemsize, dtype).reshape(-1), ] ).reshape(-1, frame.ptr.width) if format_name == "yuv422p10le": @@ -1143,6 +1150,7 @@ def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format): self._init_user_attributes() @staticmethod + @cython.cdivision(True) def from_ndarray(array, format="rgb24", channel_last=False): """ Construct a frame from a numpy array. @@ -1261,6 +1269,8 @@ def from_ndarray(array, format="rgb24", channel_last=False): check_ndarray_shape(array, array.shape[1] % 2 == 0) frame = VideoFrame(array.shape[1], (array.shape[0] * 2) // 3, format) + if frame.width == 0 or frame.height == 0: + return frame u_start = frame.width * frame.height v_start = 5 * u_start // 4 flat = array.reshape(-1) @@ -1268,12 +1278,29 @@ def from_ndarray(array, format="rgb24", channel_last=False): 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": + check_ndarray(array, "uint16", 2) + check_ndarray_shape(array, array.shape[0] % 3 == 0) + check_ndarray_shape(array, array.shape[1] % 2 == 0) + + frame = VideoFrame(array.shape[1], (array.shape[0] * 2) // 3, format) + if frame.width == 0 or frame.height == 0: + return frame + u_start = frame.width * frame.height + v_start = 5 * u_start // 4 + flat = array.reshape(-1) + copy_array_to_plane(flat[0:u_start], frame.planes[0], 2) + copy_array_to_plane(flat[u_start:v_start], frame.planes[1], 2) + copy_array_to_plane(flat[v_start:], frame.planes[2], 2) + return frame elif format == "yuv422p": check_ndarray(array, "uint8", 2) check_ndarray_shape(array, array.shape[0] % 4 == 0) check_ndarray_shape(array, array.shape[1] % 2 == 0) frame = VideoFrame(array.shape[1], array.shape[0] // 2, format) + if frame.width == 0 or frame.height == 0: + return frame u_start = frame.width * frame.height v_start = u_start + u_start // 2 flat = array.reshape(-1) diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 26386adb0..826a13c5d 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -688,6 +688,39 @@ def test_ndarray_yuv420p() -> None: assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_yuv420p10le() -> None: + array = numpy.random.randint(0, 1024, size=(720, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="yuv420p10le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "yuv420p10le" + assert "yuv420p10le" in supported_np_pix_fmts + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_yuv420p10le_align() -> None: + array = numpy.random.randint(0, 1024, size=(357, 318), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="yuv420p10le") + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == "yuv420p10le" + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_yuv420p10le_zero_size() -> None: + # A frame with zero width and/or height has no allocated planes; conversions + # should degrade gracefully instead of raising IndexError. + for w, h in ((0, 0), (0, 480), (640, 0)): + frame = VideoFrame(w, h, "yuv420p10le") + array = frame.to_ndarray() + assert array.dtype == numpy.uint16 + assert array.size == 0 + + empty = numpy.empty((0, 0), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(empty, format="yuv420p10le") + assert frame.width == 0 and frame.height == 0 + assert frame.format.name == "yuv420p10le" + assertNdarraysEqual(frame.to_ndarray(), empty) + + def test_ndarray_yuv422p() -> None: array = numpy.random.randint(0, 256, size=(960, 640), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="yuv422p")