BGfilter/bgfilter/chroma.py

from __future__ import annotations

from dataclasses import asdict, dataclass

import numpy as np

from .deps import require_cv2
from .settings import ChromaSettings


@dataclass(frozen=True)
class BackgroundModel:
    rgb_center: tuple[float, float, float]
    rgb_sigma: tuple[float, float, float]
    hsv_center: tuple[float, float, float]
    hue_sigma: float
    lab_center: tuple[float, float, float]
    lab_sigma: float
    sample_count: int

    def to_dict(self) -> dict:
        return asdict(self)


def _smoothstep(x: np.ndarray, edge0: float, edge1: float) -> np.ndarray:
    t = np.clip((x - edge0) / max(edge1 - edge0, 1e-6), 0.0, 1.0)
    return t * t * (3.0 - 2.0 * t)


def _hue_distance(hue: np.ndarray, center: float) -> np.ndarray:
    diff = np.abs(hue.astype(np.float32) - float(center))
    return np.minimum(diff, 180.0 - diff)


def _border_mask(height: int, width: int, ratio: float) -> np.ndarray:
    border = max(8, int(round(max(height, width) * ratio)))
    border = min(border, height // 2, width // 2)
    mask = np.zeros((height, width), dtype=bool)
    mask[:border, :] = True
    mask[-border:, :] = True
    mask[:, :border] = True
    mask[:, -border:] = True
    return mask


def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    cv2 = require_cv2()
    rgb_f = rgb.astype(np.float32) / 255.0
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV).astype(np.float32)
    hsv[..., 1:] /= 255.0
    lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB).astype(np.float32)
    return rgb_f, hsv, lab


def initial_green_candidates(
    rgb_f: np.ndarray, hsv: np.ndarray, settings: ChromaSettings
) -> np.ndarray:
    r = rgb_f[..., 0]
    g = rgb_f[..., 1]
    b = rgb_f[..., 2]
    dominance = g - np.maximum(r, b)
    hue = hsv[..., 0]
    sat = hsv[..., 1]
    return (
        (dominance >= settings.green_margin)
        & (g >= settings.min_green_value)
        & (sat >= settings.min_saturation)
        & (hue >= settings.hue_low)
        & (hue <= settings.hue_high)
    )


def estimate_background_model(
    rgb: np.ndarray, settings: ChromaSettings
) -> tuple[BackgroundModel, tuple[np.ndarray, np.ndarray, np.ndarray]]:
    rgb_f, hsv, lab = convert_color_spaces(rgb)
    candidates = initial_green_candidates(rgb_f, hsv, settings)
    h, w = candidates.shape
    border_candidates = candidates & _border_mask(h, w, settings.border_ratio)

    sample_mask = border_candidates
    if int(sample_mask.sum()) < settings.min_samples:
        sample_mask = candidates
    if int(sample_mask.sum()) < max(64, settings.min_samples // 16):
        # Last-resort fallback: choose pixels with strongest green dominance.
        dominance = rgb_f[..., 1] - np.maximum(rgb_f[..., 0], rgb_f[..., 2])
        cutoff = np.percentile(dominance, 90.0)
        sample_mask = dominance >= cutoff

    rgb_samples = rgb_f[sample_mask]
    hsv_samples = hsv[sample_mask]
    lab_samples = lab[sample_mask]

    rgb_center = np.median(rgb_samples, axis=0)
    rgb_sigma = np.maximum(
        np.percentile(np.abs(rgb_samples - rgb_center), 75, axis=0) * 1.4826,
        settings.rgb_sigma_min,
    )
    hsv_center = np.median(hsv_samples, axis=0)
    hue_center = float(hsv_center[0])
    hue_dists = _hue_distance(hsv_samples[:, 0], hue_center)
    hue_sigma = max(
        float(np.percentile(hue_dists, 75) * 1.4826), settings.hue_sigma_min
    )
    lab_center = np.median(lab_samples, axis=0)
    lab_dists = np.linalg.norm(lab_samples - lab_center, axis=1)
    lab_sigma = max(
        float(np.percentile(lab_dists, 75) * 1.4826), settings.lab_sigma_min
    )

    model = BackgroundModel(
        rgb_center=tuple(float(x) for x in rgb_center),
        rgb_sigma=tuple(float(x) for x in rgb_sigma),
        hsv_center=tuple(float(x) for x in hsv_center),
        hue_sigma=float(hue_sigma),
        lab_center=tuple(float(x) for x in lab_center),
        lab_sigma=float(lab_sigma),
        sample_count=int(sample_mask.sum()),
    )
    return model, (rgb_f, hsv, lab)


def compute_bg_confidence(
    rgb: np.ndarray,
    model: BackgroundModel | None = None,
    settings: ChromaSettings | None = None,
) -> tuple[np.ndarray, BackgroundModel]:
    settings = settings or ChromaSettings()
    if model is None:
        model, spaces = estimate_background_model(rgb, settings)
    else:
        spaces = convert_color_spaces(rgb)
    rgb_f, hsv, lab = spaces

    r = rgb_f[..., 0]
    g = rgb_f[..., 1]
    b = rgb_f[..., 2]
    dominance = g - np.maximum(r, b)
    hue = hsv[..., 0]
    sat = hsv[..., 1]

    hue_conf = np.exp(-0.5 * (_hue_distance(hue, model.hsv_center[0]) / model.hue_sigma) ** 2)
    sat_conf = _smoothstep(sat, max(0.05, model.hsv_center[1] * 0.45), max(0.2, model.hsv_center[1] * 0.85))
    dominance_conf = _smoothstep(dominance, settings.green_margin * 0.4, settings.green_margin * 1.6)

    lab_center = np.asarray(model.lab_center, dtype=np.float32)
    lab_dist = np.linalg.norm(lab - lab_center, axis=2)
    lab_conf = np.exp(-0.5 * (lab_dist / model.lab_sigma) ** 2)

    rgb_center = np.asarray(model.rgb_center, dtype=np.float32)
    rgb_sigma = np.asarray(model.rgb_sigma, dtype=np.float32)
    rgb_dist = np.linalg.norm((rgb_f - rgb_center) / rgb_sigma, axis=2)
    rgb_conf = np.exp(-0.5 * (rgb_dist / 2.5) ** 2)

    conf = hue_conf * sat_conf * dominance_conf * np.maximum(lab_conf, rgb_conf * 0.85)
    return np.clip(conf, 0.0, 1.0).astype(np.float32), model