conv/opencv: Add stride support

Signed-off-by: Tomi Valkeinen <tomi.valkeinen@ideasonboard.com>

Author: tomba

pixutils/conv/conv.py

@@ -70,7 +70,7 @@ def to_bgr888(
         if backend == 'opencv':
             from .opencv import opencv_to_bgr888
 
-            result = opencv_to_bgr888(fmt, width, height, arr, options)
+            result = opencv_to_bgr888(fmt, width, height, bytesperline, arr, options)
             if result is not None:
                 return result
             # opencv couldn't handle this format/options, try next backend

pixutils/conv/opencv.py

@@ -61,7 +61,12 @@ def _can_use_opencv_rgb(fmt: PixelFormat) -> bool:
 
 
 def opencv_to_bgr888(
-    fmt: PixelFormat, width: int, height: int, arr: npt.NDArray[np.uint8], options: dict | None
+    fmt: PixelFormat,
+    width: int,
+    height: int,
+    bytesperline: int,
+    arr: npt.NDArray[np.uint8],
+    options: dict | None,
 ) -> npt.NDArray[np.uint8] | None:
     if fmt.color == PixelColorEncoding.YUV:
         if not _can_use_opencv_yuv(fmt, options):
@@ -78,4 +83,4 @@ def opencv_to_bgr888(
     # Import and call implementation only if format is supported
     from .opencv_impl import opencv_convert
 
-    return opencv_convert(fmt, width, height, arr)
+    return opencv_convert(fmt, width, height, bytesperline, arr)

pixutils/conv/opencv_impl.py

@@ -3,11 +3,12 @@
 
 from __future__ import annotations
 
-from typing import Callable, cast
+from typing import cast
 
 import cv2  # type: ignore[import-not-found]
 import numpy as np
 import numpy.typing as npt
+from numpy.lib.stride_tricks import as_strided
 
 from pixutils.formats import PixelFormat, PixelColorEncoding
 
@@ -20,95 +21,126 @@
     'GBRG': cv2.COLOR_BAYER_GB2BGR,
 }
 
-# Tuple is (cv2 color code or None, reshape function)
-RGB_FORMAT_MAP: dict[str, tuple[int | None, Callable]] = {
+RGB_FORMAT_MAP: dict[str, int | None] = {
     # 32-bit BGRA formats
-    'XRGB8888': (cv2.COLOR_BGRA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'BGRX8888': (cv2.COLOR_BGRA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'ARGB8888': (cv2.COLOR_BGRA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'BGRA8888': (cv2.COLOR_BGRA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
+    'XRGB8888': cv2.COLOR_BGRA2BGR,
+    'BGRX8888': cv2.COLOR_BGRA2BGR,
+    'ARGB8888': cv2.COLOR_BGRA2BGR,
+    'BGRA8888': cv2.COLOR_BGRA2BGR,
     # 32-bit RGBA formats
-    'XBGR8888': (cv2.COLOR_RGBA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'RGBX8888': (cv2.COLOR_RGBA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'ABGR8888': (cv2.COLOR_RGBA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
-    'RGBA8888': (cv2.COLOR_RGBA2BGR, lambda b, w, h: b.reshape(h, w, 4)),
+    'XBGR8888': cv2.COLOR_RGBA2BGR,
+    'RGBX8888': cv2.COLOR_RGBA2BGR,
+    'ABGR8888': cv2.COLOR_RGBA2BGR,
+    'RGBA8888': cv2.COLOR_RGBA2BGR,
     # 24-bit formats
-    'RGB888': (cv2.COLOR_RGB2BGR, lambda b, w, h: b.reshape(h, w, 3)),
-    'BGR888': (None, lambda b, w, h: b.reshape(h, w, 3)),
+    'RGB888': cv2.COLOR_RGB2BGR,
+    'BGR888': None,  # Already BGR
 }
 
-# Tuple is (cv2 color code or None, reshape function)
-YUV_FORMAT_MAP: dict[str, tuple[int, Callable]] = {
-    'YUYV': (cv2.COLOR_YUV2BGR_YUY2, lambda b, w, h: b.reshape(h, w, 2)),
-    'UYVY': (cv2.COLOR_YUV2BGR_UYVY, lambda b, w, h: b.reshape(h, w, 2)),
-    'YVYU': (cv2.COLOR_YUV2BGR_YVYU, lambda b, w, h: b.reshape(h, w, 2)),
-    'NV12': (cv2.COLOR_YUV2BGR_NV12, lambda b, w, h: b.reshape(h * 3 // 2, w)),
-    'NV21': (cv2.COLOR_YUV2BGR_NV21, lambda b, w, h: b.reshape(h * 3 // 2, w)),
+YUV_FORMAT_MAP: dict[str, int] = {
+    'YUYV': cv2.COLOR_YUV2BGR_YUY2,
+    'UYVY': cv2.COLOR_YUV2BGR_UYVY,
+    'YVYU': cv2.COLOR_YUV2BGR_YVYU,
+    'NV12': cv2.COLOR_YUV2BGR_NV12,
+    'NV21': cv2.COLOR_YUV2BGR_NV21,
 }
 
 
 def _convert_yuv(
-    fmt: PixelFormat, width: int, height: int, arr: npt.NDArray[np.uint8]
+    fmt: PixelFormat, width: int, height: int, stride: int, arr: npt.NDArray[np.uint8]
 ) -> npt.NDArray[np.uint8]:
-    cv_code, reshape_func = YUV_FORMAT_MAP[fmt.name]
-    reshaped = reshape_func(arr, width, height)
+    cv_code = YUV_FORMAT_MAP[fmt.name]
+
+    if len(fmt.planes) == 1:
+        # Packed formats (YUYV, UYVY, YVYU)
+        plane = fmt.planes[0]
+        bytes_per_pixel = plane.bytes_per_block // plane.pixels_per_block
+
+        # OpenCV requires 3D array with channel dimension
+        reshaped = as_strided(
+            arr,
+            shape=(height, width, bytes_per_pixel),
+            strides=(stride, bytes_per_pixel, 1),
+            writeable=False,
+        )
+    else:
+        # Multi-plane formats (NV12, NV21)
+        # OpenCV expects concatenated layout: (h * 3/2, w)
+        reshaped = arr.reshape(height * 3 // 2, width)
+
     return cv2.cvtColor(reshaped, cv_code)
 
 
 def _convert_rgb(
-    fmt: PixelFormat, width: int, height: int, arr: npt.NDArray[np.uint8]
+    fmt: PixelFormat, width: int, height: int, stride: int, arr: npt.NDArray[np.uint8]
 ) -> npt.NDArray[np.uint8]:
-    cv_code, reshape_func = RGB_FORMAT_MAP[fmt.name]
-    reshaped = reshape_func(arr, width, height)
+    cv_code = RGB_FORMAT_MAP[fmt.name]
+
+    # Generic bytes_per_pixel from plane info
+    plane = fmt.planes[0]
+    bytes_per_pixel = plane.bytes_per_block // plane.pixels_per_block
+
+    # OpenCV requires 3D array with channel dimension
+    reshaped = as_strided(
+        arr,
+        shape=(height, width, bytes_per_pixel),
+        strides=(stride, bytes_per_pixel, 1),
+        writeable=False,
+    )
 
     if cv_code is None:
-        # Already BGR, just return a copy
         return reshaped.copy()
 
     return cv2.cvtColor(reshaped, cv_code)
 
 
 def _convert_raw(
-    fmt: PixelFormat, width: int, height: int, arr: npt.NDArray[np.uint8]
+    fmt: PixelFormat, width: int, height: int, stride: int, arr: npt.NDArray[np.uint8]
 ) -> npt.NDArray[np.uint8] | None:
     pattern = fmt.bayer_pattern
     assert pattern is not None
     cv_code = BAYER_PATTERN_MAP[pattern]
 
     name = fmt.name
+    plane = fmt.planes[0]
+
+    # Determine element size from plane info
+    bytes_per_pixel = plane.bytes_per_block // plane.pixels_per_block
 
-    # Determine bit depth from format name
-    if '8' in name:
-        # 8-bit: reshape to (h, w) uint8
-        bayer = arr.reshape(height, width)
+    if bytes_per_pixel == 1:
+        # 8-bit formats
+        bayer = as_strided(arr, shape=(height, width), strides=(stride, 1), writeable=False)
         return cast(npt.NDArray[np.uint8], cv2.cvtColor(bayer, cv_code))
-    elif '16' in name:
-        # 16-bit: reshape to (h, w) uint16, convert, then scale to 8-bit
-        bayer = arr.view(np.uint16).reshape(height, width)
-        bgr16 = cast(npt.NDArray[np.uint16], cv2.cvtColor(bayer, cv_code))
-        return (bgr16 >> 8).astype(np.uint8)
-    elif '10' in name or '12' in name:
-        # 10/12-bit unpacked (stored in 16-bit): shift up to use full 16-bit range
-        bits = 10 if '10' in name else 12
-        bayer = arr.view(np.uint16).reshape(height, width)
-        bayer = bayer << (16 - bits)
-        bgr16 = cast(npt.NDArray[np.uint16], cv2.cvtColor(bayer, cv_code))
-        return (bgr16 >> 8).astype(np.uint8)
-    else:
-        # Unknown bit depth
-        return None
+    elif bytes_per_pixel == 2:
+        # 16-bit formats (could be 10, 12, or 16 bit stored in 16)
+        arr16 = arr.view(np.uint16)
+        bayer = as_strided(arr16, shape=(height, width), strides=(stride, 2), writeable=False)
+
+        # Detect actual bit depth from format name for scaling
+        if '16' in name:
+            bgr16 = cast(npt.NDArray[np.uint16], cv2.cvtColor(bayer, cv_code))
+            return (bgr16 >> 8).astype(np.uint8)
+        elif '10' in name or '12' in name:
+            bits = 10 if '10' in name else 12
+            bayer = bayer << (16 - bits)
+            bgr16 = cast(npt.NDArray[np.uint16], cv2.cvtColor(bayer, cv_code))
+            return (bgr16 >> 8).astype(np.uint8)
+
+    return None
 
 
 def opencv_convert(
-    fmt: PixelFormat, width: int, height: int, arr: npt.NDArray[np.uint8]
+    fmt: PixelFormat, width: int, height: int, bytesperline: int, arr: npt.NDArray[np.uint8]
 ) -> npt.NDArray[np.uint8] | None:
+    stride = bytesperline if bytesperline > 0 else fmt.stride(width, 0)
+
     if fmt.color == PixelColorEncoding.YUV:
-        return _convert_yuv(fmt, width, height, arr)
+        return _convert_yuv(fmt, width, height, stride, arr)
 
     if fmt.color == PixelColorEncoding.RAW:
-        return _convert_raw(fmt, width, height, arr)
+        return _convert_raw(fmt, width, height, stride, arr)
 
     if fmt.color == PixelColorEncoding.RGB:
-        return _convert_rgb(fmt, width, height, arr)
+        return _convert_rgb(fmt, width, height, stride, arr)
 
     return None