cv_arrays.py

 1"""Image contracts: uint8 HWC BGR input; float32 RGB/CHW model input."""
 2from __future__ import annotations
 3
 4from typing import cast
 5
 6import cv2
 7import numpy as np
 8from numpy.typing import NDArray
 9
10ByteImage = NDArray[np.uint8]
11FloatImage = NDArray[np.float32]
12
13# region pixels
14def pixel_demo() -> tuple[ByteImage, ByteImage, ByteImage]:
15    """Return a (2,2,3) BGR image, its RGB copy, and a copied one-pixel crop."""
16    bgr: ByteImage = np.array(
17        [[[0, 0, 255], [0, 255, 0]], [[255, 0, 0], [30, 60, 90]]],
18        dtype=np.uint8,
19    )
20    rgb: ByteImage = cast(ByteImage, cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
21    crop: ByteImage = rgb[1:2, 1:2, :].copy()
22    crop[0, 0, 0] = 0
23    return bgr, rgb, crop
24# endregion pixels
25
26# region normalize
27def to_unit_rgb(bgr: ByteImage) -> FloatImage:
28    """Nonempty uint8 (H,W,3) BGR -> float32 (H,W,3) RGB in [0,1]."""
29    if bgr.dtype != np.uint8 or bgr.ndim != 3 or bgr.shape[2] != 3:
30        raise ValueError("Expected uint8 HWC with three BGR channels")
31    if min(bgr.shape[:2]) == 0:
32        raise ValueError("Image must have pixels")
33    rgb: ByteImage = cast(ByteImage, cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
34    unit_rgb: FloatImage = rgb.astype(np.float32) / np.float32(255.0)
35    return unit_rgb
36# endregion normalize
37
38# region layout
39def model_input(unit_rgb: FloatImage) -> FloatImage:
40    """Finite float32 HWC RGB [0,1] -> contiguous NCHW batch of size one."""
41    if unit_rgb.dtype != np.float32 or unit_rgb.ndim != 3 or unit_rgb.shape[2] != 3:
42        raise ValueError("Expected float32 HWC with three RGB channels")
43    if unit_rgb.size == 0 or not np.isfinite(unit_rgb).all():
44        raise ValueError("Expected nonempty finite image")
45    if np.any(unit_rgb < 0) or np.any(unit_rgb > 1):
46        raise ValueError("Expected values in [0,1]")
47    chw: FloatImage = unit_rgb.transpose(2, 0, 1)
48    batch: FloatImage = np.ascontiguousarray(chw[None, ...])
49    return batch
50# endregion layout