Skip to content

Commit

Permalink
Make parameters same, test channel_last
Browse files Browse the repository at this point in the history
  • Loading branch information
WyattBlue committed Jan 22, 2025
1 parent 73baa3d commit 787917b
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 10 deletions.
2 changes: 1 addition & 1 deletion av/video/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...
Expand Down
12 changes: 5 additions & 7 deletions av/video/frame.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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

Expand Down
15 changes: 13 additions & 2 deletions tests/test_videoframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down

0 comments on commit 787917b

Please sign in to comment.