From 82370473c32554942b74a1da3130e2e38fab0c67 Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Wed, 22 Jan 2025 09:03:58 -0500 Subject: [PATCH] Make parameters same, test channel_last --- av/video/frame.pyi | 2 +- av/video/frame.pyx | 12 +++++------- tests/test_videoframe.py | 15 +++++++++++++-- 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/av/video/frame.pyi b/av/video/frame.pyi index f58554b8f..3533c8e3f 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -60,7 +60,7 @@ class VideoFrame(Frame): def to_rgb(self, **kwargs: Any) -> VideoFrame: ... def to_image(self, **kwargs: Any) -> Image.Image: ... def to_ndarray( - self, force_channel_last: bool = False, **kwargs: Any + self, channel_last: bool = False, **kwargs: Any ) -> _SupportedNDarray: ... @staticmethod def from_image(img: Image.Image) -> VideoFrame: ... diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 5cdfb233b..36f4deae1 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -282,18 +282,16 @@ cdef class VideoFrame(Frame): return Image.frombytes("RGB", (plane.width, plane.height), bytes(o_buf), "raw", "RGB", 0, 1) - def to_ndarray(self, force_channel_last=False, **kwargs): + def to_ndarray(self, channel_last=False, **kwargs): """Get a numpy array of this frame. Any ``**kwargs`` are passed to :meth:`.VideoReformatter.reformat`. The array returned is generally of dimension (height, width, channels). - :param bool force_channel_last: If False (default), the shape for the yuv444p and yuvj444p - will be (channels, height, width) rather than (height, width, channels) as usual. - This is for backward compatibility and also for keeping that - `bytes(to_ndarray(frame))` should be the same as the ffmpeg cli - when returning the pix_fmt with `-c:v rawvideo`. + :param bool channel_last: If True, the shape of array will be + (height, width, channels) rather than (channels, height, width) for + the "yuv444p" and "yuvj444p" formats. .. note:: Numpy must be installed. @@ -374,7 +372,7 @@ cdef class VideoFrame(Frame): array[:, :, 0] = array[:, :, 2] array[:, :, 2] = array[:, :, 1] array[:, :, 1] = buffer - if not force_channel_last and frame.format.name in {"yuv444p", "yuvj444p"}: + if not channel_last and frame.format.name in {"yuv444p", "yuvj444p"}: array = np.moveaxis(array, 2, 0) return array diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 8cb520b82..9581a7c41 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -440,6 +440,17 @@ def test_ndarray_yuv444p() -> None: assert frame.format.name == "yuv444p" assertNdarraysEqual(frame.to_ndarray(), array) + array = numpy.random.randint(0, 256, size=(3, 480, 640), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, channel_last=False, format="yuv444p") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "yuv444p" + assertNdarraysEqual(frame.to_ndarray(channel_last=False), array) + assert array.shape != frame.to_ndarray(channel_last=True).shape + assert ( + frame.to_ndarray(channel_last=False).shape + != frame.to_ndarray(channel_last=True).shape + ) + def test_ndarray_yuvj444p() -> None: array = numpy.random.randint(0, 256, size=(3, 480, 640), dtype=numpy.uint8) @@ -458,7 +469,7 @@ def test_ndarray_yuv444p16() -> None: assertNdarraysEqual(frame.to_ndarray(), array) -def test_ndarray_yuv444p16_allign() -> None: +def test_ndarray_yuv444p16_align() -> None: array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16) for format in ("yuv444p16be", "yuv444p16le"): frame = VideoFrame.from_ndarray(array, format=format) @@ -476,7 +487,7 @@ def test_ndarray_yuva444p16() -> None: assertNdarraysEqual(frame.to_ndarray(), array) -def test_ndarray_yuva444p16_allign() -> None: +def test_ndarray_yuva444p16_align() -> None: array = numpy.random.randint(0, 65536, size=(238, 318, 4), dtype=numpy.uint16) for format in ("yuva444p16be", "yuva444p16le"): frame = VideoFrame.from_ndarray(array, format=format)