diff --git a/av/video/format.pyi b/av/video/format.pyi index ee16b85b8..e102ef4c0 100644 --- a/av/video/format.pyi +++ b/av/video/format.pyi @@ -6,7 +6,10 @@ class VideoFormat: has_palette: bool is_bit_stream: bool is_planar: bool - is_rgb: bool + @property + def is_rgb(self) -> bool: ... + @property + def is_bayer(self) -> bool: ... width: int height: int components: tuple[VideoFormatComponent, ...] diff --git a/av/video/format.pyx b/av/video/format.pyx index 6ae66c3a2..8779f7151 100644 --- a/av/video/format.pyx +++ b/av/video/format.pyx @@ -104,7 +104,12 @@ cdef class VideoFormat: def is_rgb(self): """The pixel format contains RGB-like data (as opposed to YUV/grayscale).""" return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_RGB) + + @property + def is_bayer(self): + """The pixel format contains Bayer data.""" + return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_BAYER) cpdef chroma_width(self, int luma_width=0): """chroma_width(luma_width=0) diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 36f4deae1..b74d58d1c 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -317,6 +317,18 @@ cdef class VideoFrame(Frame): itemsize, dtype = { "abgr": (4, "uint8"), "argb": (4, "uint8"), + "bayer_bggr8": (1, "uint8"), + "bayer_gbrg8": (1, "uint8"), + "bayer_grbg8": (1, "uint8"), + "bayer_rggb8": (1, "uint8"), + "bayer_bggr16le": (2, "uint16"), + "bayer_bggr16be": (2, "uint16"), + "bayer_gbrg16le": (2, "uint16"), + "bayer_gbrg16be": (2, "uint16"), + "bayer_grbg16le": (2, "uint16"), + "bayer_grbg16be": (2, "uint16"), + "bayer_rggb16le": (2, "uint16"), + "bayer_rggb16be": (2, "uint16"), "bgr24": (3, "uint8"), "bgr8": (1, "uint8"), "bgra": (4, "uint8"), @@ -451,6 +463,13 @@ cdef class VideoFrame(Frame): else: # Planes where U and V are interleaved have the same stride as Y. linesizes = (array.strides[0], array.strides[0]) + 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"}: + check_ndarray(array, "uint8" if format.endswith("8") else "uint16", 2) + + if array.strides[1] != (1 if format.endswith("8") else 2): + raise ValueError("provided array does not have C_CONTIGUOUS rows") + + linesizes = (array.strides[0],) else: raise ValueError(f"Conversion from numpy array with format `{format}` is not yet supported") @@ -554,7 +573,19 @@ cdef class VideoFrame(Frame): "yuv444p16be": (3, 2, "uint16"), "yuv444p16le": (3, 2, "uint16"), "yuva444p16be": (4, 2, "uint16"), - "yuva444p16le": (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"), }.get(format, (None, None, None)) if channels is not None: if array.ndim == 2: # (height, width) -> (height, width, 1) @@ -644,6 +675,8 @@ cdef class VideoFrame(Frame): 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) else: raise NotImplementedError(f"Format '{format}' is not supported.") return frame