From 32837a59d6d0118ab0e594d34ef99574ebb53fb7 Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Mon, 23 Jun 2025 22:24:04 -0400 Subject: [PATCH] video/frame: apply BasswoodAV changes This also applies "add GBRP, GBRAP, RGBA formats" by z-khan Co-authored-by: z-khan --- av/video/frame.pxd | 1 + av/video/{frame.pyx => frame.py} | 766 +++++++++++++++++++++++-------- av/video/frame.pyi | 9 +- tests/test_videoframe.py | 469 ++++++++++++------- 4 files changed, 886 insertions(+), 359 deletions(-) rename av/video/{frame.pyx => frame.py} (54%) diff --git a/av/video/frame.pxd b/av/video/frame.pxd index 779b23977..d352a7ab8 100644 --- a/av/video/frame.pxd +++ b/av/video/frame.pxd @@ -17,5 +17,6 @@ cdef class VideoFrame(Frame): cdef _init(self, lib.AVPixelFormat format, unsigned int width, unsigned int height) cdef _init_user_attributes(self) + cpdef save(self, object filepath) cdef VideoFrame alloc_video_frame() diff --git a/av/video/frame.pyx b/av/video/frame.py similarity index 54% rename from av/video/frame.pyx rename to av/video/frame.py index 04a14fd8b..bc4bf9154 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.py @@ -1,38 +1,115 @@ import sys from enum import IntEnum -from libc.stdint cimport uint8_t +import cython +from cython.cimports.av.error import err_check +from cython.cimports.av.sidedata.sidedata import get_display_rotation +from cython.cimports.av.utils import check_ndarray +from cython.cimports.av.video.format import get_pix_fmt, get_video_format +from cython.cimports.av.video.plane import VideoPlane +from cython.cimports.libc.stdint import uint8_t -from av.error cimport err_check -from av.sidedata.sidedata cimport get_display_rotation -from av.utils cimport check_ndarray -from av.video.format cimport get_pix_fmt, get_video_format -from av.video.plane cimport VideoPlane - - -cdef object _cinit_bypass_sentinel +_cinit_bypass_sentinel = object() # `pix_fmt`s supported by Frame.to_ndarray() and Frame.from_ndarray() supported_np_pix_fmts = { - "abgr", "argb", "bayer_bggr16be", "bayer_bggr16le", "bayer_bggr8", "bayer_gbrg16be", - "bayer_gbrg16le", "bayer_gbrg8", "bayer_grbg16be", "bayer_grbg16le", "bayer_grbg8", - "bayer_rggb16be", "bayer_rggb16le", "bayer_rggb8", "bgr24", "bgr48be", "bgr48le", "bgr8", "bgra", "bgra64be", "bgra64le", - "gbrapf32be", "gbrapf32le", "gbrp", "gbrp10be", "gbrp10le", "gbrp12be", "gbrp12le", - "gbrp14be", "gbrp14le", "gbrp16be", "gbrp16le", "gbrpf32be", "gbrpf32le", "gray", - "gray16be", "gray16le", "gray8", "grayf32be", "grayf32le", "nv12", "pal8", "rgb24", - "rgb48be", "rgb48le", "rgb8", "rgba", "rgba64be", "rgba64le", "yuv420p", - "yuv420p10le", "yuv422p10le", "yuv444p", "yuv444p16be", "yuv444p16le", "yuva444p16be", - "yuva444p16le", "yuvj420p", "yuvj444p", "yuyv422", + "abgr", + "argb", + "bayer_bggr16be", + "bayer_bggr16le", + "bayer_bggr8", + "bayer_gbrg16be", + "bayer_gbrg16le", + "bayer_gbrg8", + "bayer_grbg16be", + "bayer_grbg16le", + "bayer_grbg8", + "bayer_rggb16be", + "bayer_rggb16le", + "bayer_rggb8", + "bgr24", + "bgr48be", + "bgr48le", + "bgr8", + "bgra", + "bgra64be", + "bgra64le", + "gbrap", + "gbrap10be", + "gbrap10le", + "gbrap12be", + "gbrap12le", + "gbrap14be", + "gbrap14le", + "gbrap16be", + "gbrap16le", + "gbrapf32be", + "gbrapf32le", + "gbrp", + "gbrp10be", + "gbrp10le", + "gbrp12be", + "gbrp12le", + "gbrp14be", + "gbrp14le", + "gbrp16be", + "gbrp16le", + "gbrp9be", + "gbrp9le", + "gbrpf32be", + "gbrpf32le", + "gray", + "gray10be", + "gray10le", + "gray12be", + "gray12le", + "gray14be", + "gray14le", + "gray16be", + "gray16le", + "gray8", + "gray9be", + "gray9le", + "grayf32be", + "grayf32le", + "nv12", + "pal8", + "rgb24", + "rgb48be", + "rgb48le", + "rgb8", + "rgba", + "rgba64be", + "rgba64le", + "rgbaf16be", + "rgbaf16le", + "rgbaf32be", + "rgbaf32le", + "rgbf32be", + "rgbf32le", + "yuv420p", + "yuv422p10le", + "yuv444p", + "yuv444p16be", + "yuv444p16le", + "yuva444p16be", + "yuva444p16le", + "yuvj420p", + "yuvj444p", + "yuyv422", } -cdef VideoFrame alloc_video_frame(): + +@cython.cfunc +def alloc_video_frame() -> VideoFrame: """Get a mostly uninitialized VideoFrame. You MUST call VideoFrame._init(...) or VideoFrame._init_user_attributes() before exposing to the user. """ - return VideoFrame.__new__(VideoFrame, _cinit_bypass_sentinel) + return VideoFrame(_cinit_bypass_sentinel) + class PictureType(IntEnum): NONE = lib.AV_PICTURE_TYPE_NONE # Undefined @@ -44,23 +121,31 @@ class PictureType(IntEnum): SP = lib.AV_PICTURE_TYPE_SP # Switching predicted BI = lib.AV_PICTURE_TYPE_BI # BI type -cdef byteswap_array(array, bint big_endian): + +@cython.cfunc +def byteswap_array(array, big_endian: cython.bint): if (sys.byteorder == "big") != big_endian: return array.byteswap() - else: - return array - - -cdef copy_bytes_to_plane(img_bytes, VideoPlane plane, unsigned int bytes_per_pixel, bint flip_horizontal, bint flip_vertical): - cdef const uint8_t[:] i_buf = img_bytes - cdef size_t i_pos = 0 - cdef size_t i_stride = plane.width * bytes_per_pixel - - cdef uint8_t[:] o_buf = plane - cdef size_t o_pos = 0 - cdef size_t o_stride = abs(plane.line_size) - - cdef int start_row, end_row, step + return array + + +@cython.cfunc +def copy_bytes_to_plane( + img_bytes, + plane: VideoPlane, + bytes_per_pixel: cython.uint, + flip_horizontal: cython.bint, + flip_vertical: cython.bint, +): + i_buf: cython.const[uint8_t][:] = img_bytes + i_pos: cython.size_t = 0 + i_stride: cython.size_t = plane.width * bytes_per_pixel + + o_buf: uint8_t[:] = plane + o_pos: cython.size_t = 0 + o_stride: cython.size_t = abs(plane.line_size) + + start_row, end_row, step = cython.declare(cython.int) if flip_vertical: start_row = plane.height - 1 end_row = -1 @@ -70,56 +155,65 @@ class PictureType(IntEnum): end_row = plane.height step = 1 - cdef int i, j + i, j = cython.declare(cython.int) for row in range(start_row, end_row, step): i_pos = row * i_stride if flip_horizontal: for i in range(0, i_stride, bytes_per_pixel): for j in range(bytes_per_pixel): - o_buf[o_pos + i + j] = i_buf[i_pos + i_stride - i - bytes_per_pixel + j] + o_buf[o_pos + i + j] = i_buf[ + i_pos + i_stride - i - bytes_per_pixel + j + ] else: - o_buf[o_pos:o_pos + i_stride] = i_buf[i_pos:i_pos + i_stride] + o_buf[o_pos : o_pos + i_stride] = i_buf[i_pos : i_pos + i_stride] o_pos += o_stride -cdef copy_array_to_plane(array, VideoPlane plane, unsigned int bytes_per_pixel): - cdef bytes imgbytes = array.tobytes() +@cython.cfunc +def copy_array_to_plane(array, plane: VideoPlane, bytes_per_pixel: cython.uint): + imgbytes: bytes = array.tobytes() copy_bytes_to_plane(imgbytes, plane, bytes_per_pixel, False, False) -cdef useful_array(VideoPlane plane, unsigned int bytes_per_pixel=1, str dtype="uint8"): +@cython.cfunc +def useful_array( + plane: VideoPlane, bytes_per_pixel: cython.uint = 1, dtype: str = "uint8" +): """ Return the useful part of the VideoPlane as a single dimensional array. We are simply discarding any padding which was added for alignment. """ import numpy as np - cdef size_t total_line_size = abs(plane.line_size) - cdef size_t useful_line_size = plane.width * bytes_per_pixel + + total_line_size: cython.size_t = abs(plane.line_size) + useful_line_size: cython.size_t = plane.width * bytes_per_pixel arr = np.frombuffer(plane, np.uint8) if total_line_size != useful_line_size: arr = arr.reshape(-1, total_line_size)[:, 0:useful_line_size].reshape(-1) return arr.view(np.dtype(dtype)) -cdef check_ndarray_shape(object array, bint ok): +@cython.cfunc +def check_ndarray_shape(array: object, ok: cython.bint): if not ok: raise ValueError(f"Unexpected numpy array shape `{array.shape}`") -cdef class VideoFrame(Frame): +@cython.cclass +class VideoFrame(Frame): def __cinit__(self, width=0, height=0, format="yuv420p"): if width is _cinit_bypass_sentinel: return - cdef lib.AVPixelFormat c_format = get_pix_fmt(format) - + c_format: lib.AVPixelFormat = get_pix_fmt(format) self._init(c_format, width, height) - cdef _init(self, lib.AVPixelFormat format, unsigned int width, unsigned int height): - cdef int res = 0 + @cython.cfunc + def _init(self, format: lib.AVPixelFormat, width: cython.uint, height: cython.uint): + res: cython.int = 0 - with nogil: + with cython.nogil: self.ptr.width = width self.ptr.height = height self.ptr.format = format @@ -137,13 +231,18 @@ def __cinit__(self, width=0, height=0, format="yuv420p"): self._init_user_attributes() - cdef _init_user_attributes(self): - self.format = get_video_format(self.ptr.format, self.ptr.width, self.ptr.height) + @cython.cfunc + def _init_user_attributes(self): + self.format = get_video_format( + cython.cast(lib.AVPixelFormat, self.ptr.format), + self.ptr.width, + self.ptr.height, + ) def __dealloc__(self): # The `self._buffer` member is only set if *we* allocated the buffer in `_init`, # as opposed to a buffer allocated by a decoder. - lib.av_freep(&self._buffer) + lib.av_freep(cython.address(self._buffer)) # Let go of the reference from the numpy buffers if we made one self._np_buffer = None @@ -158,11 +257,10 @@ def planes(self): """ A tuple of :class:`.VideoPlane` objects. """ - # We need to detect which planes actually exist, but also contrain - # ourselves to the maximum plane count (as determined only by VideoFrames - # so far), in case the library implementation does not set the last - # plane to NULL. - cdef int max_plane_count = 0 + # We need to detect which planes actually exist, but also constrain ourselves to + # the maximum plane count (as determined only by VideoFrames so far), in case + # the library implementation does not set the last plane to NULL. + max_plane_count: cython.int = 0 for i in range(self.format.ptr.nb_components): count = self.format.ptr.comp[i].plane + 1 if max_plane_count < count: @@ -170,7 +268,7 @@ def planes(self): if self.format.name == "pal8": max_plane_count = 2 - cdef int plane_count = 0 + plane_count: cython.int = 0 while plane_count < max_plane_count and self.ptr.extended_data[plane_count]: plane_count += 1 return tuple([VideoPlane(self, i) for i in range(plane_count)]) @@ -267,6 +365,34 @@ def to_rgb(self, **kwargs): """ return self.reformat(format="rgb24", **kwargs) + @cython.ccall + def save(self, filepath: object): + """Save a VideoFrame as a JPG or PNG. + + :param filepath: str | Path + """ + is_jpg: cython.bint + + if filepath.endswith(".png"): + is_jpg = False + elif filepath.endswith(".jpg") or filepath.endswith(".jpeg"): + is_jpg = True + else: + raise ValueError("filepath must end with png or jpg.") + + encoder: str = "mjpeg" if is_jpg else "png" + pix_fmt: str = "yuvj420p" if is_jpg else "rgb24" + + from av.container.core import open + + with open(filepath, "w", options={"update": "1"}) as output: + output_stream = output.add_stream(encoder, pix_fmt=pix_fmt) + output_stream.width = self.width + output_stream.height = self.height + + output.mux(output_stream.encode(self.reformat(format=pix_fmt))) + output.mux(output_stream.encode(None)) + def to_image(self, **kwargs): """Get an RGB ``PIL.Image`` of this frame. @@ -276,23 +402,26 @@ def to_image(self, **kwargs): """ from PIL import Image - cdef VideoPlane plane = self.reformat(format="rgb24", **kwargs).planes[0] - cdef const uint8_t[:] i_buf = plane - cdef size_t i_pos = 0 - cdef size_t i_stride = plane.line_size + plane: VideoPlane = self.reformat(format="rgb24", **kwargs).planes[0] - cdef size_t o_pos = 0 - cdef size_t o_stride = plane.width * 3 - cdef size_t o_size = plane.height * o_stride - cdef bytearray o_buf = bytearray(o_size) + i_buf: cython.const[uint8_t][:] = plane + i_pos: cython.size_t = 0 + i_stride: cython.size_t = plane.line_size + + o_pos: cython.size_t = 0 + o_stride: cython.size_t = plane.width * 3 + o_size: cython.size_t = plane.height * o_stride + o_buf: bytearray = bytearray(o_size) while o_pos < o_size: - o_buf[o_pos:o_pos + o_stride] = i_buf[i_pos:i_pos + o_stride] + o_buf[o_pos : o_pos + o_stride] = i_buf[i_pos : i_pos + o_stride] i_pos += i_stride o_pos += o_stride - return Image.frombytes("RGB", (plane.width, plane.height), bytes(o_buf), "raw", "RGB", 0, 1) + return Image.frombytes( + "RGB", (plane.width, plane.height), bytes(o_buf), "raw", "RGB", 0, 1 + ) def to_ndarray(self, channel_last=False, **kwargs): """Get a numpy array of this frame. @@ -307,7 +436,7 @@ def to_ndarray(self, channel_last=False, **kwargs): .. note:: Numpy must be installed. - .. note:: For formats which return an array of ``uint16`` or ``float32``, + .. note:: For formats which return an array of ``uint16``, ``float16`` or ``float32``, the samples will be in the system's native byte order. .. note:: For ``pal8``, an ``(image, palette)`` tuple will be returned, @@ -316,14 +445,18 @@ def to_ndarray(self, channel_last=False, **kwargs): .. note:: For ``gbrp`` formats, channels are flipped to RGB order. """ - cdef VideoFrame frame = self.reformat(**kwargs) + frame: VideoFrame = self.reformat(**kwargs) import numpy as np # check size - if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422", "yuv420p10le", "yuv422p10le"}: - assert frame.width % 2 == 0, "the width has to be even for this pixel format" - assert frame.height % 2 == 0, "the height has to be even for this pixel format" + if frame.format.name in {"yuv420p", "yuvj420p", "yuyv422", "yuv422p10le"}: + assert frame.width % 2 == 0, ( + "the width has to be even for this pixel format" + ) + assert frame.height % 2 == 0, ( + "the height has to be even for this pixel format" + ) # cases planes are simply concatenated in shape (height, width, channels) itemsize, dtype = { @@ -346,6 +479,17 @@ def to_ndarray(self, channel_last=False, **kwargs): "bgr48le": (6, "uint16"), "bgr8": (1, "uint8"), "bgra": (4, "uint8"), + "bgra64be": (8, "uint16"), + "bgra64le": (8, "uint16"), + "gbrap": (1, "uint8"), + "gbrap10be": (2, "uint16"), + "gbrap10le": (2, "uint16"), + "gbrap12be": (2, "uint16"), + "gbrap12le": (2, "uint16"), + "gbrap14be": (2, "uint16"), + "gbrap14le": (2, "uint16"), + "gbrap16be": (2, "uint16"), + "gbrap16le": (2, "uint16"), "gbrapf32be": (4, "float32"), "gbrapf32le": (4, "float32"), "gbrp": (1, "uint8"), @@ -357,12 +501,22 @@ def to_ndarray(self, channel_last=False, **kwargs): "gbrp14le": (2, "uint16"), "gbrp16be": (2, "uint16"), "gbrp16le": (2, "uint16"), + "gbrp9be": (2, "uint16"), + "gbrp9le": (2, "uint16"), "gbrpf32be": (4, "float32"), "gbrpf32le": (4, "float32"), "gray": (1, "uint8"), + "gray10be": (2, "uint16"), + "gray10le": (2, "uint16"), + "gray12be": (2, "uint16"), + "gray12le": (2, "uint16"), + "gray14be": (2, "uint16"), + "gray14le": (2, "uint16"), "gray16be": (2, "uint16"), "gray16le": (2, "uint16"), "gray8": (1, "uint8"), + "gray9be": (2, "uint16"), + "gray9le": (2, "uint16"), "grayf32be": (4, "float32"), "grayf32le": (4, "float32"), "rgb24": (3, "uint8"), @@ -372,8 +526,12 @@ def to_ndarray(self, channel_last=False, **kwargs): "rgba": (4, "uint8"), "rgba64be": (8, "uint16"), "rgba64le": (8, "uint16"), - "bgra64be": (8, "uint16"), - "bgra64le": (8, "uint16"), + "rgbaf16be": (8, "float16"), + "rgbaf16le": (8, "float16"), + "rgbaf32be": (16, "float32"), + "rgbaf32le": (16, "float32"), + "rgbf32be": (12, "float32"), + "rgbf32le": (12, "float32"), "yuv444p": (1, "uint8"), "yuv444p16be": (2, "uint16"), "yuv444p16le": (2, "uint16"), @@ -384,8 +542,9 @@ def to_ndarray(self, channel_last=False, **kwargs): }.get(frame.format.name, (None, None)) if itemsize is not None: layers = [ - useful_array(plan, itemsize, dtype) - .reshape(frame.height, frame.width, -1) + useful_array(plan, itemsize, dtype).reshape( + frame.height, frame.width, -1 + ) for plan in frame.planes ] if len(layers) == 1: # shortcut, avoid memory copy @@ -406,26 +565,24 @@ def to_ndarray(self, channel_last=False, **kwargs): # special cases if frame.format.name in {"yuv420p", "yuvj420p"}: - return np.hstack([ - useful_array(frame.planes[0]), - useful_array(frame.planes[1]), - useful_array(frame.planes[2]), - ]).reshape(-1, frame.width) - if frame.format.name == "yuv420p10le": - # Read planes as uint16: - y = useful_array(frame.planes[0], 2, "uint16").reshape(frame.height, frame.width) - u = useful_array(frame.planes[1], 2, "uint16").reshape(frame.height // 2, frame.width // 2) - v = useful_array(frame.planes[2], 2, "uint16").reshape(frame.height // 2, frame.width // 2) - u_full = np.repeat(np.repeat(u, 2, axis=1), 2, axis=0) - v_full = np.repeat(np.repeat(u, 2, axis=1), 2, axis=0) - if channel_last: - return np.stack([y, u_full, v_full], axis=2) - return np.stack([y, u_full, v_full], axis=0) + return np.hstack( + [ + useful_array(frame.planes[0]), + useful_array(frame.planes[1]), + useful_array(frame.planes[2]), + ] + ).reshape(-1, frame.width) if frame.format.name == "yuv422p10le": # Read planes as uint16 at their original width - y = useful_array(frame.planes[0], 2, "uint16").reshape(frame.height, frame.width) - u = useful_array(frame.planes[1], 2, "uint16").reshape(frame.height, frame.width // 2) - v = useful_array(frame.planes[2], 2, "uint16").reshape(frame.height, frame.width // 2) + y = useful_array(frame.planes[0], 2, "uint16").reshape( + frame.height, frame.width + ) + u = useful_array(frame.planes[1], 2, "uint16").reshape( + frame.height, frame.width // 2 + ) + v = useful_array(frame.planes[2], 2, "uint16").reshape( + frame.height, frame.width // 2 + ) # Double the width of U and V by repeating each value u_full = np.repeat(u, 2, axis=1) @@ -435,13 +592,20 @@ def to_ndarray(self, channel_last=False, **kwargs): return np.stack([y, u_full, v_full], axis=0) if frame.format.name == "pal8": image = useful_array(frame.planes[0]).reshape(frame.height, frame.width) - palette = np.frombuffer(frame.planes[1], "i4").astype(">i4").reshape(-1, 1).view(np.uint8) + palette = ( + np.frombuffer(frame.planes[1], "i4") + .astype(">i4") + .reshape(-1, 1) + .view(np.uint8) + ) return image, palette if frame.format.name == "nv12": - return np.hstack([ - useful_array(frame.planes[0]), - useful_array(frame.planes[1], 2), - ]).reshape(-1, frame.width) + return np.hstack( + [ + useful_array(frame.planes[0]), + useful_array(frame.planes[1], 2), + ] + ).reshape(-1, frame.width) raise ValueError( f"Conversion to numpy array with format `{frame.format.name}` is not yet supported" @@ -455,18 +619,29 @@ def from_image(img): if img.mode != "RGB": img = img.convert("RGB") - cdef VideoFrame frame = VideoFrame(img.size[0], img.size[1], "rgb24") + frame: VideoFrame = VideoFrame(img.size[0], img.size[1], "rgb24") copy_array_to_plane(img, frame.planes[0], 3) return frame @staticmethod 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. + """ + Construct a frame from a numpy buffer. + + :param int width: optional width of actual image, if different from the array width. + + .. note:: For formats which expect an array of ``uint16``, ``float16`` or ``float32``, + the samples must be in the system's native byte order. + + .. note:: for ``gbrp`` formats, channels are assumed to be given in RGB order. + + .. note:: For formats where width of the array is not the same as the width of the image, + for example with yuv420p images 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. + """ + import numpy as np + height = array.shape[0] if not width: width = array.shape[1] @@ -476,40 +651,169 @@ def from_numpy_buffer(array, format="rgb24", width=0): check_ndarray_shape(array, array.shape[2] == 3) if array.strides[1:] != (3, 1): raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) + linesizes = (array.strides[0],) elif format in {"rgb48le", "rgb48be", "bgr48le", "bgr48be"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 3) if array.strides[1:] != (6, 2): raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) + linesizes = (array.strides[0],) + elif format in {"rgbf32le", "rgbf32be"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 3) + if array.strides[1:] != (12, 4): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0],) elif format in {"rgba", "bgra", "argb", "abgr"}: 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], ) + linesizes = (array.strides[0],) elif format in {"rgba64le", "rgba64be", "bgra64le", "bgra64be"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 4) if array.strides[1:] != (8, 2): raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) - elif format in {"gray", "gray8", "rgb8", "bgr8","bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8"}: + linesizes = (array.strides[0],) + elif format in {"rgbaf16le", "rgbaf16be"}: + check_ndarray(array, "float16", 3) + check_ndarray_shape(array, array.shape[2] == 4) + if array.strides[1:] != (8, 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0],) + elif format in {"rgbaf32le", "rgbaf32be"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 4) + if array.strides[1:] != (16, 4): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = (array.strides[0],) + elif format in { + "gray", + "gray8", + "rgb8", + "bgr8", + "bayer_bggr8", + "bayer_gbrg8", + "bayer_grbg8", + "bayer_rggb8", + }: check_ndarray(array, "uint8", 2) if array.strides[1] != 1: raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) - elif format in {"gray16le", "gray16be", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"}: + linesizes = (array.strides[0],) + elif format in { + "gray9be", + "gray9le", + "gray10be", + "gray10le", + "gray12be", + "gray12le", + "gray14be", + "gray14le", + "gray16be", + "gray16le", + "bayer_bggr16be", + "bayer_bggr16le", + "bayer_gbrg16be", + "bayer_gbrg16le", + "bayer_grbg16be", + "bayer_grbg16le", + "bayer_rggb16be", + "bayer_rggb16le", + }: check_ndarray(array, "uint16", 2) if array.strides[1] != 2: raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) + linesizes = (array.strides[0],) elif format in {"grayf32le", "grayf32be"}: check_ndarray(array, "float32", 2) if array.strides[1] != 4: raise ValueError("provided array does not have C_CONTIGUOUS rows") - linesizes = (array.strides[0], ) + linesizes = (array.strides[0],) + elif format in {"gbrp"}: + check_ndarray(array, "uint8", 3) + check_ndarray_shape(array, array.shape[2] == 3) + if array.strides[1:] != (3, 1): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = ( + array.strides[0] // 3, + array.strides[0] // 3, + array.strides[0] // 3, + ) + elif format in { + "gbrp9be", + "gbrp9le", + "gbrp10be", + "gbrp10le", + "gbrp12be", + "gbrp12le", + "gbrp14be", + "gbrp14le", + "gbrp16be", + "gbrp16le", + }: + check_ndarray(array, "uint16", 3) + check_ndarray_shape(array, array.shape[2] == 3) + if array.strides[1:] != (6, 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = ( + array.strides[0] // 3, + array.strides[0] // 3, + array.strides[0] // 3, + ) + elif format in {"gbrpf32be", "gbrpf32le"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 3) + if array.strides[1:] != (12, 4): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = ( + array.strides[0] // 3, + array.strides[0] // 3, + array.strides[0] // 3, + ) + elif format in {"gbrap"}: + 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] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + ) + elif format in { + "gbrap10be", + "gbrap10le", + "gbrap12be", + "gbrap12le", + "gbrap14be", + "gbrap14le", + "gbrap16be", + "gbrap16le", + }: + check_ndarray(array, "uint16", 3) + check_ndarray_shape(array, array.shape[2] == 4) + if array.strides[1:] != (8, 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = ( + array.strides[0] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + ) + elif format in {"gbrapf32be", "gbrapf32le"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 4) + if array.strides[1:] != (16, 4): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + linesizes = ( + array.strides[0] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + array.strides[0] // 4, + ) elif format in {"yuv420p", "yuvj420p", "nv12"}: check_ndarray(array, "uint8", 2) check_ndarray_shape(array, array.shape[0] % 3 == 0) @@ -519,26 +823,37 @@ def from_numpy_buffer(array, format="rgb24", width=0): 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) + 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") + raise ValueError( + f"Conversion from numpy array with format `{format}` is not yet supported" + ) + + if format.startswith("gbrap"): # rgba -> gbra + array = np.ascontiguousarray(np.moveaxis(array[..., [1, 2, 0, 3]], -1, 0)) + elif format.startswith("gbrp"): # rgb -> gbr + array = np.ascontiguousarray(np.moveaxis(array[..., [1, 2, 0]], -1, 0)) - frame = alloc_video_frame() + frame = VideoFrame(_cinit_bypass_sentinel) frame._image_fill_pointers_numpy(array, width, height, linesizes, format) return frame 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 + c_format: lib.AVPixelFormat + c_ptr: cython.pointer[uint8_t] + c_data: cython.size_t - # If you want to use the numpy notation - # then you need to include the following two lines at the top of the file + # If you want to use the numpy notation, then you need to include the following lines at the top of the file: # cimport numpy as cnp # cnp.import_array() + # And add the numpy include directories to the setup.py files # hint np.get_include() # cdef cnp.ndarray[ @@ -548,19 +863,15 @@ def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format): # c_ptr = &c_buffer[0] # c_ptr = ((buffer.ctypes.data)) - # Using buffer.ctypes.data helps avoid any kind of - # usage of the c-api from numpy, which avoid the need to add numpy - # as a compile time dependency - # Without this double cast, you get an error that looks like - # c_ptr = (buffer.ctypes.data) - # TypeError: expected bytes, int found + # Using buffer.ctypes.data helps avoid any kind of usage of the c-api from + # numpy, which avoid the need to add numpy as a compile time dependency. + c_data = buffer.ctypes.data - c_ptr = (c_data) + c_ptr = cython.cast(cython.pointer[uint8_t], c_data) c_format = get_pix_fmt(format) - lib.av_freep(&self._buffer) + lib.av_freep(cython.address(self._buffer)) - # Hold on to a reference for the numpy buffer - # so that it doesn't get accidentally garbage collected + # Hold on to a reference for the numpy buffer so that it doesn't get accidentally garbage collected self._np_buffer = buffer self.ptr.format = c_format self.ptr.width = width @@ -570,7 +881,7 @@ def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format): res = lib.av_image_fill_pointers( self.ptr.data, - self.ptr.format, + cython.cast(lib.AVPixelFormat, self.ptr.format), self.ptr.height, c_ptr, self.ptr.linesize, @@ -588,10 +899,11 @@ def from_ndarray(array, format="rgb24", channel_last=False): :param bool channel_last: If False (default), the shape for the yuv444p and yuvj444p is given by (channels, height, width) rather than (height, width, channels). - .. note:: For formats which expect an array of ``uint16``, + .. note:: For formats which expect an array of ``uint16``, ``float16`` or ``float32``, the samples must be in the system's native byte order. - .. note:: for ``pal8``, an ``(image, palette)`` pair must be passed. `palette` must have shape (256, 4) and is given in ARGB format (PyAV will swap bytes if needed). + .. note:: for ``pal8``, an ``(image, palette)`` pair must be passed. `palette` must + have shape (256, 4) and is given in ARGB format (PyAV will swap bytes if needed). .. note:: for ``gbrp`` formats, channels are assumed to be given in RGB order. @@ -600,45 +912,64 @@ def from_ndarray(array, format="rgb24", channel_last=False): # case layers are concatenated channels, itemsize, dtype = { - "yuv444p": (3, 1, "uint8"), - "yuvj444p": (3, 1, "uint8"), + "bayer_bggr16be": (1, 2, "uint16"), + "bayer_bggr16le": (1, 2, "uint16"), + "bayer_bggr8": (1, 1, "uint8"), + "bayer_gbrg16be": (1, 2, "uint16"), + "bayer_gbrg16le": (1, 2, "uint16"), + "bayer_gbrg8": (1, 1, "uint8"), + "bayer_grbg16be": (1, 2, "uint16"), + "bayer_grbg16le": (1, 2, "uint16"), + "bayer_grbg8": (1, 1, "uint8"), + "bayer_rggb16be": (1, 2, "uint16"), + "bayer_rggb16le": (1, 2, "uint16"), + "bayer_rggb8": (1, 1, "uint8"), + "bgr8": (1, 1, "uint8"), + "gbrap": (4, 1, "uint8"), + "gbrap10be": (4, 2, "uint16"), + "gbrap10le": (4, 2, "uint16"), + "gbrap12be": (4, 2, "uint16"), + "gbrap12le": (4, 2, "uint16"), + "gbrap14be": (4, 2, "uint16"), + "gbrap14le": (4, 2, "uint16"), + "gbrap16be": (4, 2, "uint16"), + "gbrap16le": (4, 2, "uint16"), + "gbrapf32be": (4, 4, "float32"), + "gbrapf32le": (4, 4, "float32"), "gbrp": (3, 1, "uint8"), "gbrp10be": (3, 2, "uint16"), - "gbrp12be": (3, 2, "uint16"), - "gbrp14be": (3, 2, "uint16"), - "gbrp16be": (3, 2, "uint16"), "gbrp10le": (3, 2, "uint16"), + "gbrp12be": (3, 2, "uint16"), "gbrp12le": (3, 2, "uint16"), + "gbrp14be": (3, 2, "uint16"), "gbrp14le": (3, 2, "uint16"), + "gbrp16be": (3, 2, "uint16"), "gbrp16le": (3, 2, "uint16"), + "gbrp9be": (3, 2, "uint16"), + "gbrp9le": (3, 2, "uint16"), "gbrpf32be": (3, 4, "float32"), "gbrpf32le": (3, 4, "float32"), "gray": (1, 1, "uint8"), - "gray8": (1, 1, "uint8"), - "rgb8": (1, 1, "uint8"), - "bgr8": (1, 1, "uint8"), + "gray10be": (1, 2, "uint16"), + "gray10le": (1, 2, "uint16"), + "gray12be": (1, 2, "uint16"), + "gray12le": (1, 2, "uint16"), + "gray14be": (1, 2, "uint16"), + "gray14le": (1, 2, "uint16"), "gray16be": (1, 2, "uint16"), "gray16le": (1, 2, "uint16"), + "gray8": (1, 1, "uint8"), + "gray9be": (1, 2, "uint16"), + "gray9le": (1, 2, "uint16"), "grayf32be": (1, 4, "float32"), "grayf32le": (1, 4, "float32"), - "gbrapf32be": (4, 4, "float32"), - "gbrapf32le": (4, 4, "float32"), + "rgb8": (1, 1, "uint8"), + "yuv444p": (3, 1, "uint8"), "yuv444p16be": (3, 2, "uint16"), "yuv444p16le": (3, 2, "uint16"), "yuva444p16be": (4, 2, "uint16"), - "yuva444p16le": (4, 2, "uint16"), - "bayer_bggr8": (1, 1, "uint8"), - "bayer_rggb8": (1, 1, "uint8"), - "bayer_grbg8": (1, 1, "uint8"), - "bayer_gbrg8": (1, 1, "uint8"), - "bayer_bggr16be": (1, 2, "uint16"), - "bayer_bggr16le": (1, 2, "uint16"), - "bayer_rggb16be": (1, 2, "uint16"), - "bayer_rggb16le": (1, 2, "uint16"), - "bayer_grbg16be": (1, 2, "uint16"), - "bayer_grbg16le": (1, 2, "uint16"), - "bayer_gbrg16be": (1, 2, "uint16"), - "bayer_gbrg16le": (1, 2, "uint16"), + "yuva444p16le": (4, 2, "uint16"), + "yuvj444p": (3, 1, "uint8"), }.get(format, (None, None, None)) if channels is not None: if array.ndim == 2: # (height, width) -> (height, width, 1) @@ -650,9 +981,14 @@ def from_ndarray(array, format="rgb24", channel_last=False): array = byteswap_array(array, format.endswith("be")) frame = VideoFrame(array.shape[1], array.shape[0], format) if frame.format.name.startswith("gbr"): # rgb -> gbr - array = np.concatenate([ # not inplace to avoid bad surprises - array[:, :, 1:3], array[:, :, 0:1], array[:, :, 3:], - ], axis=2) + array = np.concatenate( + [ # not inplace to avoid bad surprises + array[:, :, 1:3], + array[:, :, 0:1], + array[:, :, 3:], + ], + axis=2, + ) for i in range(channels): copy_array_to_plane(array[:, :, i], frame.planes[i], itemsize) return frame @@ -681,28 +1017,6 @@ 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": - if not isinstance(array, np.ndarray) or array.dtype != np.uint16: - raise ValueError("Array must be uint16 type") - - # Convert to channel-first if needed: - if channel_last and array.shape[2] == 3: - array = np.moveaxis(array, 2, 0) - elif not (array.shape[0] == 3): - raise ValueError("Array must have shape (3, height, width) or (height, width, 3)") - - height, width = array.shape[1:] - if width % 2 != 0 or height % 2 != 0: - raise ValueError("Width and height must be even") - - frame = VideoFrame(width, height, format) - copy_array_to_plane(array[0], frame.planes[0], 2) - # Subsample U and V by taking every other row and column: - u = array[1, ::2, ::2].copy() # Need copy to ensure C-contiguous - v = array[2, ::2, ::2].copy() # Need copy to ensure C-contiguous - copy_array_to_plane(u, frame.planes[1], 2) - copy_array_to_plane(v, frame.planes[2], 2) - return frame elif format == "yuv422p10le": if not isinstance(array, np.ndarray) or array.dtype != np.uint16: raise ValueError("Array must be uint16 type") @@ -711,7 +1025,9 @@ def from_ndarray(array, format="rgb24", channel_last=False): if channel_last and array.shape[2] == 3: array = np.moveaxis(array, 2, 0) elif not (array.shape[0] == 3): - raise ValueError("Array must have shape (3, height, width) or (height, width, 3)") + raise ValueError( + "Array must have shape (3, height, width) or (height, width, 3)" + ) height, width = array.shape[1:] if width % 2 != 0 or height % 2 != 0: @@ -736,17 +1052,45 @@ def from_ndarray(array, format="rgb24", channel_last=False): elif format in {"argb", "rgba", "abgr", "bgra"}: check_ndarray(array, "uint8", 3) check_ndarray_shape(array, array.shape[2] == 4) - elif format in {"rgb48be", "rgb48le","bgr48be", "bgr48le"}: + elif format in {"rgb48be", "rgb48le", "bgr48be", "bgr48le"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 3) frame = VideoFrame(array.shape[1], array.shape[0], format) - copy_array_to_plane(byteswap_array(array, format.endswith("be")), frame.planes[0], 6) + copy_array_to_plane( + byteswap_array(array, format.endswith("be")), frame.planes[0], 6 + ) + return frame + elif format in {"rgbf32be", "rgbf32le"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 3) + frame = VideoFrame(array.shape[1], array.shape[0], format) + copy_array_to_plane( + byteswap_array(array, format.endswith("be")), frame.planes[0], 12 + ) return frame elif format in {"rgba64be", "rgba64le", "bgra64be", "bgra64le"}: check_ndarray(array, "uint16", 3) check_ndarray_shape(array, array.shape[2] == 4) frame = VideoFrame(array.shape[1], array.shape[0], format) - copy_array_to_plane(byteswap_array(array, format.endswith("be")), frame.planes[0], 8) + copy_array_to_plane( + byteswap_array(array, format.endswith("be")), frame.planes[0], 8 + ) + return frame + elif format in {"rgbaf16be", "rgbaf16le"}: + check_ndarray(array, "float16", 3) + check_ndarray_shape(array, array.shape[2] == 4) + frame = VideoFrame(array.shape[1], array.shape[0], format) + copy_array_to_plane( + byteswap_array(array, format.endswith("be")), frame.planes[0], 8 + ) + return frame + elif format in {"rgbaf32be", "rgbaf32le"}: + check_ndarray(array, "float32", 3) + check_ndarray_shape(array, array.shape[2] == 4) + frame = VideoFrame(array.shape[1], array.shape[0], format) + copy_array_to_plane( + byteswap_array(array, format.endswith("be")), frame.planes[0], 16 + ) return frame elif format == "nv12": check_ndarray(array, "uint8", 2) @@ -760,20 +1104,52 @@ def from_ndarray(array, format="rgb24", channel_last=False): copy_array_to_plane(flat[uv_start:], frame.planes[1], 2) return frame else: - raise ValueError(f"Conversion from numpy array with format `{format}` is not yet supported") + raise ValueError( + f"Conversion from numpy array with format `{format}` is not yet supported" + ) frame = VideoFrame(array.shape[1], array.shape[0], format) - copy_array_to_plane(array, frame.planes[0], 1 if array.ndim == 2 else array.shape[2]) + copy_array_to_plane( + array, frame.planes[0], 1 if array.ndim == 2 else array.shape[2] + ) return frame @staticmethod - def from_bytes(img_bytes: bytes, width: int, height: int, format="rgba", flip_horizontal=False, flip_vertical=False): + def from_bytes( + img_bytes: bytes, + width: int, + height: int, + format="rgba", + flip_horizontal=False, + flip_vertical=False, + ): frame = VideoFrame(width, height, format) if format == "rgba": - copy_bytes_to_plane(img_bytes, frame.planes[0], 4, flip_horizontal, flip_vertical) - elif format in ("bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8","bayer_bggr16le", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"): - copy_bytes_to_plane(img_bytes, frame.planes[0], 1 if format.endswith("8") else 2, flip_horizontal, flip_vertical) + copy_bytes_to_plane( + img_bytes, frame.planes[0], 4, flip_horizontal, flip_vertical + ) + elif format in { + "bayer_bggr8", + "bayer_rggb8", + "bayer_gbrg8", + "bayer_grbg8", + "bayer_bggr16le", + "bayer_rggb16le", + "bayer_gbrg16le", + "bayer_grbg16le", + "bayer_bggr16be", + "bayer_rggb16be", + "bayer_gbrg16be", + "bayer_grbg16be", + }: + copy_bytes_to_plane( + img_bytes, + frame.planes[0], + 1 if format.endswith("8") else 2, + flip_horizontal, + flip_vertical, + ) else: raise NotImplementedError(f"Format '{format}' is not supported.") return frame diff --git a/av/video/frame.pyi b/av/video/frame.pyi index 313e184f9..c6868aabc 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -1,8 +1,8 @@ from enum import IntEnum +from pathlib import Path from typing import Any, ClassVar, Union import numpy as np -from PIL import Image from av.frame import Frame @@ -12,6 +12,7 @@ from .plane import VideoPlane _SupportedNDarray = Union[ np.ndarray[Any, np.dtype[np.uint8]], np.ndarray[Any, np.dtype[np.uint16]], + np.ndarray[Any, np.dtype[np.float16]], np.ndarray[Any, np.dtype[np.float32]], ] @@ -30,7 +31,6 @@ class PictureType(IntEnum): class VideoFrame(Frame): format: VideoFormat pts: int - duration: int planes: tuple[VideoPlane, ...] pict_type: int colorspace: int @@ -61,12 +61,13 @@ class VideoFrame(Frame): dst_color_range: int | str | None = None, ) -> VideoFrame: ... def to_rgb(self, **kwargs: Any) -> VideoFrame: ... - def to_image(self, **kwargs: Any) -> Image.Image: ... + def save(self, filepath: str | Path) -> None: ... + def to_image(self, **kwargs): ... def to_ndarray( self, channel_last: bool = False, **kwargs: Any ) -> _SupportedNDarray: ... @staticmethod - def from_image(img: Image.Image) -> VideoFrame: ... + def from_image(img): ... @staticmethod def from_numpy_buffer( array: _SupportedNDarray, format: str = "rgb24", width: int = 0 diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 677bc1fc1..9c9773ae8 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -7,16 +7,11 @@ import av from av import VideoFrame +from av.frame import Frame +from av.video.frame import supported_np_pix_fmts from av.video.reformatter import ColorRange, Colorspace, Interpolation -from .common import ( - TestCase, - assertImagesAlmostEqual, - assertNdarraysEqual, - fate_png, - fate_suite, - has_pillow, -) +from .common import TestCase, assertNdarraysEqual, fate_png, fate_suite def assertPixelValue16(plane, expected, byteorder: str) -> None: @@ -29,26 +24,6 @@ def assertPixelValue16(plane, expected, byteorder: str) -> None: assert view[1] == (expected >> 8 & 0xFF) -def test_opaque() -> None: - with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: - video_stream = container.streams.video[0] - - ctx = video_stream.codec_context - ctx.flags |= av.codec.context.Flags.copy_opaque - - assert video_stream.codec_context.copy_opaque - - for packet_idx, packet in enumerate(container.demux()): - packet.opaque = (time.time(), packet_idx) - for frame in packet.decode(): - assert isinstance(frame, av.frame.Frame) - - if frame.opaque is None: - continue - - assert type(frame.opaque) is tuple and len(frame.opaque) == 2 - - def test_frame_duration_matches_packet() -> None: with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: packet_durations = [ @@ -147,64 +122,24 @@ def test_memoryview_read() -> None: assert mem[:7] == b"0.234xx" -class TestVideoFrameImage(TestCase): - def setUp(self) -> None: - if not has_pillow: - pytest.skip() +def test_interpolation() -> None: + container = av.open(fate_png()) + for _ in container.decode(video=0): + frame = _ + break - def test_roundtrip(self) -> None: - import PIL.Image as Image + assert frame.width == 330 and frame.height == 330 - image = Image.open(fate_png()) - frame = VideoFrame.from_image(image) - img = frame.to_image() - img.save(self.sandboxed("roundtrip-high.jpg")) - assertImagesAlmostEqual(image, img) + img = frame.reformat(width=200, height=100, interpolation=Interpolation.BICUBIC) + assert img.width == 200 and img.height == 100 - def test_interpolation(self) -> None: - import PIL.Image as Image + img = frame.reformat(width=200, height=100, interpolation="BICUBIC") + assert img.width == 200 and img.height == 100 - image = Image.open(fate_png()) - frame = VideoFrame.from_image(image) - assert frame.width == 330 and frame.height == 330 - - img = frame.to_image(width=200, height=100, interpolation=Interpolation.BICUBIC) - assert img.width == 200 and img.height == 100 - - img = frame.to_image(width=200, height=100, interpolation="BICUBIC") - assert img.width == 200 and img.height == 100 - - img = frame.to_image( - width=200, height=100, interpolation=int(Interpolation.BICUBIC) - ) - assert img.width == 200 and img.height == 100 - - def test_to_image_rgb24(self) -> None: - sizes = [(318, 238), (320, 240), (500, 500)] - for width, height in sizes: - frame = VideoFrame(width, height, format="rgb24") - - # fill video frame data - for plane in frame.planes: - ba = bytearray(plane.buffer_size) - pos = 0 - for row in range(height): - for i in range(plane.line_size): - ba[pos] = i % 256 - pos += 1 - plane.update(ba) - - # construct expected image data - expected = bytearray(height * width * 3) - pos = 0 - for row in range(height): - for i in range(width * 3): - expected[pos] = i % 256 - pos += 1 - - img = frame.to_image() - assert img.size == (width, height) - assert img.tobytes() == expected + img = frame.reformat( + width=200, height=100, interpolation=int(Interpolation.BICUBIC) + ) + assert img.width == 200 and img.height == 100 def test_basic_to_ndarray() -> None: @@ -212,19 +147,11 @@ def test_basic_to_ndarray() -> None: assert array.shape == (480, 640, 3) -def test_to_image_with_dimensions() -> None: - if not has_pillow: - pytest.skip() - - img = VideoFrame(640, 480, format="rgb24").to_image(width=320, height=240) - assert img.size == (320, 240) - - def test_ndarray_gray() -> None: array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) for format in ("gray", "gray8"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == "gray" assertNdarraysEqual(frame.to_ndarray(), array) @@ -234,17 +161,127 @@ def test_ndarray_gray_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318), dtype=numpy.uint8) for format in ("gray", "gray8"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == "gray" assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_gray9be() -> None: + array = numpy.random.randint(0, 512, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray9be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray9be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "big") + + +def test_ndarray_gray9le() -> None: + array = numpy.random.randint(0, 512, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray9le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray9le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "little") + + +def test_ndarray_gray10be() -> None: + array = numpy.random.randint(0, 1024, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray10be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray10be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "big") + + +def test_ndarray_gray10le() -> None: + array = numpy.random.randint(0, 1024, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray10le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray10le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "little") + + +def test_ndarray_gray12be() -> None: + array = numpy.random.randint(0, 4096, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray12be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray12be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "big") + + +def test_ndarray_gray12le() -> None: + array = numpy.random.randint(0, 4096, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray12le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray12le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "little") + + +def test_ndarray_gray14be() -> None: + array = numpy.random.randint(0, 16384, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray14be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray14be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "big") + + +def test_ndarray_gray14le() -> None: + array = numpy.random.randint(0, 16384, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray14le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray14le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "little") + + +def test_ndarray_gray16be() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray16be") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray16be" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "big") + + +def test_ndarray_gray16le() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16) + frame = VideoFrame.from_ndarray(array, format="gray16le") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gray16le" + assertNdarraysEqual(frame.to_ndarray(), array) + + # check endianness by examining value of first pixel + assertPixelValue16(frame.planes[0], array[0][0], "little") + + def test_ndarray_grayf32() -> None: array = numpy.random.random_sample(size=(480, 640)).astype(numpy.float32) for format in ("grayf32be", "grayf32le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -254,7 +291,7 @@ def test_ndarray_grayf32_align() -> None: array = numpy.random.random_sample(size=(238, 318)).astype(numpy.float32) for format in ("grayf32be", "grayf32le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -264,7 +301,7 @@ def test_ndarray_rgb() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 3), dtype=numpy.uint8) for format in ("rgb24", "bgr24"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -274,17 +311,26 @@ def test_ndarray_rgb_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318, 3), dtype=numpy.uint8) for format in ("rgb24", "bgr24"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_rgbf32() -> None: + array = numpy.random.random_sample(size=(480, 640, 3)).astype(numpy.float32) + for format in ("rgbf32be", "rgbf32le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_rgba() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 4), dtype=numpy.uint8) for format in ("argb", "rgba", "abgr", "bgra"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -294,7 +340,7 @@ def test_ndarray_rgba_align() -> None: array = numpy.random.randint(0, 256, size=(238, 318, 4), dtype=numpy.uint8) for format in ("argb", "rgba", "abgr", "bgra"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -304,7 +350,7 @@ def test_ndarray_bayer8() -> None: array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) for format in ("bayer_bggr8", "bayer_gbrg8", "bayer_grbg8", "bayer_rggb8"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -323,12 +369,126 @@ def test_ndarray_bayer16() -> None: "bayer_rggb16le", ): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_gbrap() -> None: + array = numpy.random.randint(0, 256, size=(480, 640, 4), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="gbrap") + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == "gbrap" + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap_align() -> None: + array = numpy.random.randint(0, 256, size=(238, 318, 4), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="gbrap") + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == "gbrap" + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap10() -> None: + array = numpy.random.randint(0, 1024, size=(480, 640, 4), dtype=numpy.uint16) + for format in ("gbrap10be", "gbrap10le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap10_align() -> None: + array = numpy.random.randint(0, 1024, size=(238, 318, 4), dtype=numpy.uint16) + for format in ("gbrap10be", "gbrap10le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap12() -> None: + array = numpy.random.randint(0, 4096, size=(480, 640, 4), dtype=numpy.uint16) + for format in ("gbrap12be", "gbrap12le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap12_align() -> None: + array = numpy.random.randint(0, 4096, size=(238, 318, 4), dtype=numpy.uint16) + for format in ("gbrap12be", "gbrap12le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap14() -> None: + array = numpy.random.randint(0, 16384, size=(480, 640, 4), dtype=numpy.uint16) + for format in ("gbrap14be", "gbrap14le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap14_align() -> None: + array = numpy.random.randint(0, 16384, size=(238, 318, 4), dtype=numpy.uint16) + for format in ("gbrap14be", "gbrap14le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap16() -> None: + array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16) + for format in ("gbrap16be", "gbrap16le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrap16_align() -> None: + array = numpy.random.randint(0, 65536, size=(238, 318, 4), dtype=numpy.uint16) + for format in ("gbrap16be", "gbrap16le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrapf32() -> None: + array = numpy.random.random_sample(size=(480, 640, 4)).astype(numpy.float32) + for format in ("gbrapf32be", "gbrapf32le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrapf32_align() -> None: + array = numpy.random.random_sample(size=(238, 318, 4)).astype(numpy.float32) + for format in ("gbrapf32be", "gbrapf32le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_gbrp() -> None: array = numpy.random.randint(0, 256, size=(480, 640, 3), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="gbrp") @@ -345,11 +505,31 @@ def test_ndarray_gbrp_align() -> None: assertNdarraysEqual(frame.to_ndarray(), array) +def test_ndarray_gbrp9() -> None: + array = numpy.random.randint(0, 512, size=(480, 640, 3), dtype=numpy.uint16) + for format in ("gbrp9be", "gbrp9le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_gbrp9_align() -> None: + array = numpy.random.randint(0, 512, size=(238, 318, 3), dtype=numpy.uint16) + for format in ("gbrp9be", "gbrp9le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert format in supported_np_pix_fmts + assert frame.width == 318 and frame.height == 238 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_gbrp10() -> None: array = numpy.random.randint(0, 1024, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp10be", "gbrp10le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -359,7 +539,7 @@ def test_ndarray_gbrp10_align() -> None: array = numpy.random.randint(0, 1024, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp10be", "gbrp10le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -369,7 +549,7 @@ def test_ndarray_gbrp12() -> None: array = numpy.random.randint(0, 4096, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp12be", "gbrp12le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -379,7 +559,7 @@ def test_ndarray_gbrp12_align() -> None: array = numpy.random.randint(0, 4096, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp12be", "gbrp12le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -389,7 +569,7 @@ def test_ndarray_gbrp14() -> None: array = numpy.random.randint(0, 16384, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp14be", "gbrp14le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -399,7 +579,7 @@ def test_ndarray_gbrp14_align() -> None: array = numpy.random.randint(0, 16384, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp14be", "gbrp14le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -409,7 +589,7 @@ def test_ndarray_gbrp16() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) for format in ("gbrp16be", "gbrp16le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 640 and frame.height == 480 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -419,7 +599,7 @@ def test_ndarray_gbrp16_align() -> None: array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16) for format in ("gbrp16be", "gbrp16le"): frame = VideoFrame.from_ndarray(array, format=format) - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts assert frame.width == 318 and frame.height == 238 assert frame.format.name == format assertNdarraysEqual(frame.to_ndarray(), array) @@ -443,24 +623,6 @@ def test_ndarray_gbrpf32_align() -> None: assertNdarraysEqual(frame.to_ndarray(), array) -def test_ndarray_gbrapf32() -> None: - array = numpy.random.random_sample(size=(480, 640, 4)).astype(numpy.float32) - for format in ("gbrapf32be", "gbrapf32le"): - frame = VideoFrame.from_ndarray(array, format=format) - assert frame.width == 640 and frame.height == 480 - assert frame.format.name == format - assertNdarraysEqual(frame.to_ndarray(), array) - - -def test_ndarray_gbrapf32_allign() -> None: - array = numpy.random.random_sample(size=(238, 318, 4)).astype(numpy.float32) - for format in ("gbrapf32be", "gbrapf32le"): - frame = VideoFrame.from_ndarray(array, format=format) - assert frame.width == 318 and frame.height == 238 - assert frame.format.name == format - assertNdarraysEqual(frame.to_ndarray(), array) - - def test_ndarray_yuv420p() -> None: array = numpy.random.randint(0, 256, size=(720, 640), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="yuv420p") @@ -529,22 +691,13 @@ def test_ndarray_yuv444p16() -> None: assertNdarraysEqual(frame.to_ndarray(), array) -def test_ndarray_yuv420p10le() -> None: - array = numpy.random.randint(0, 65536, size=(3, 480, 640), dtype=numpy.uint16) - for format in ("yuv420p10le",): - frame = VideoFrame.from_ndarray(array, format=format) - assert frame.width == 640 and frame.height == 480 - assert frame.format.name == format - assert format in av.video.frame.supported_np_pix_fmts - - def test_ndarray_yuv422p10le() -> None: array = numpy.random.randint(0, 65536, size=(3, 480, 640), dtype=numpy.uint16) for format in ("yuv422p10le",): frame = VideoFrame.from_ndarray(array, format=format) assert frame.width == 640 and frame.height == 480 assert frame.format.name == format - assert format in av.video.frame.supported_np_pix_fmts + assert format in supported_np_pix_fmts def test_ndarray_yuv444p16_align() -> None: @@ -582,28 +735,6 @@ def test_ndarray_yuyv422_align() -> None: assertNdarraysEqual(frame.to_ndarray(), array) -def test_ndarray_gray16be() -> None: - array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16) - frame = VideoFrame.from_ndarray(array, format="gray16be") - assert frame.width == 640 and frame.height == 480 - assert frame.format.name == "gray16be" - assertNdarraysEqual(frame.to_ndarray(), array) - - # check endianness by examining value of first pixel - assertPixelValue16(frame.planes[0], array[0][0], "big") - - -def test_ndarray_gray16le() -> None: - array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16) - frame = VideoFrame.from_ndarray(array, format="gray16le") - assert frame.width == 640 and frame.height == 480 - assert frame.format.name == "gray16le" - assertNdarraysEqual(frame.to_ndarray(), array) - - # check endianness by examining value of first pixel - assertPixelValue16(frame.planes[0], array[0][0], "little") - - def test_ndarray_rgb48be() -> None: array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16) frame = VideoFrame.from_ndarray(array, format="rgb48be") @@ -714,6 +845,24 @@ def test_ndarray_bgra64le() -> None: assertPixelValue16(frame.planes[0], array[0][0][0], "little") +def test_ndarray_rgbaf16() -> None: + array = numpy.random.random_sample(size=(480, 640, 4)).astype(numpy.float16) + for format in ("rgbaf16be", "rgbaf16le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + +def test_ndarray_rgbaf32() -> None: + array = numpy.random.random_sample(size=(480, 640, 4)).astype(numpy.float32) + for format in ("rgbaf32be", "rgbaf32le"): + frame = VideoFrame.from_ndarray(array, format=format) + assert frame.width == 640 and frame.height == 480 + assert frame.format.name == format + assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_rgb8() -> None: array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="rgb8") @@ -736,7 +885,7 @@ def test_ndarray_pal8(): frame = VideoFrame.from_ndarray((array, palette), format="pal8") assert frame.width == 640 and frame.height == 480 assert frame.format.name == "pal8" - assert frame.format.name in av.video.frame.supported_np_pix_fmts + assert frame.format.name in supported_np_pix_fmts returned = frame.to_ndarray() assert type(returned) is tuple and len(returned) == 2 assertNdarraysEqual(returned[0], array) @@ -748,7 +897,7 @@ def test_ndarray_nv12() -> None: frame = VideoFrame.from_ndarray(array, format="nv12") assert frame.width == 640 and frame.height == 480 assert frame.format.name == "nv12" - assert frame.format.name in av.video.frame.supported_np_pix_fmts + assert frame.format.name in supported_np_pix_fmts assertNdarraysEqual(frame.to_ndarray(), array)