mosaic_generator/utils.py

"""
utils.py

Low-level utility functions used across the mosaic generator.
Includes:
- Numba-accelerated RGB → LAB conversion
- A safe wrapper for ensuring correct image dtype and shape
"""

import numpy as np
import cv2
from numba import njit


# ----------------------------------------------------------------------
# NUMBA RGB → LAB (HIGH SPEED)
# ----------------------------------------------------------------------
@njit
def fast_rgb2lab_numba(rgb):
    """
    Fast approximate RGB → LAB conversion using Numba JIT.

    Parameters
    ----------
    rgb : np.ndarray
        Float32 array of shape (H, W, 3) in [0, 255].

    Returns
    -------
    np.ndarray
        LAB array of shape (H, W, 3) (float32).
    """
    R = rgb[..., 0] / 255.0
    G = rgb[..., 1] / 255.0
    B = rgb[..., 2] / 255.0

    # sRGB → linear RGB
    def f(c):
        return np.where(c > 0.04045, ((c + 0.055) / 1.055) ** 2.4, c / 12.92)

    R = f(R); G = f(G); B = f(B)

    # Linear RGB → XYZ
    X = 0.4124 * R + 0.3576 * G + 0.1805 * B
    Y = 0.2126 * R + 0.7152 * G + 0.0722 * B
    Z = 0.0193 * R + 0.1192 * G + 0.9505 * B

    # Normalize by D65
    X /= 0.95047
    Z /= 1.08883

    # XYZ → LAB helper
    def g(t):
        return np.where(t > 0.008856, t ** (1/3), 7.787 * t + 16/116)

    fx = g(X); fy = g(Y); fz = g(Z)

    L = 116 * fy - 16
    a = 500 * (fx - fy)
    b = 200 * (fy - fz)

    out = np.empty(rgb.shape, dtype=np.float32)
    out[..., 0] = L
    out[..., 1] = a
    out[..., 2] = b
    return out


# ----------------------------------------------------------------------
# SAFE WRAPPER
# ----------------------------------------------------------------------
def fast_rgb2lab(img_rgb):
    """
    Safe wrapper for Numba LAB conversion.

    Parameters
    ----------
    img_rgb : np.ndarray
        RGB image, shape (H, W, 3), dtype uint8 or float32.

    Returns
    -------
    np.ndarray
        LAB image of shape (H, W, 3), dtype float32.

    Raises
    ------
    ValueError
        If the input image is not a valid RGB array.

    Notes
    -----
    - Numba does NOT allow Python exceptions inside the JIT function.
      Therefore, validation happens here before calling Numba.
    """
    if img_rgb is None or not isinstance(img_rgb, np.ndarray):
        raise ValueError("fast_rgb2lab(): expected a NumPy array.")

    if img_rgb.ndim != 3 or img_rgb.shape[2] != 3:
        raise ValueError(
            f"fast_rgb2lab(): expected image shape (H, W, 3), got {img_rgb.shape}"
        )

    # Ensure float32 for Numba kernel
    return fast_rgb2lab_numba(img_rgb.astype(np.float32))