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]
    lab_center: tuple[float, float, float]
    lab_sigma: float
    sample_count: int

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


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]:
    """Return ``(rgb_f, lab)``: RGB normalised to 0..1 and Lab (L 0..100, a/b signed)."""
    cv2 = require_cv2()
    rgb_f = rgb.astype(np.float32) / 255.0
    lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB).astype(np.float32)
    return rgb_f, lab


def bg_hue_projection(
    lab: np.ndarray, lab_center: tuple[float, float, float]
) -> np.ndarray:
    """Per-pixel projection of Lab chroma (a*, b*) onto the background hue direction.

    Positive = displaced toward the background hue, ~0 for neutral pixels. Returns
    zeros when the background is near-neutral (no meaningful hue direction).
    """
    center = np.asarray(lab_center, dtype=np.float32)
    mag = float(np.hypot(center[1], center[2]))
    if mag < 1e-3:
        return np.zeros(lab.shape[:2], dtype=np.float32)
    return (lab[..., 1] * center[1] + lab[..., 2] * center[2]) / mag


def parse_hex_color(text: str) -> tuple[float, float, float]:
    """Parse ``#RRGGBB`` (or ``RRGGBB``) into an RGB triple normalised to 0..1."""
    s = text.strip().lstrip("#")
    if len(s) != 6:
        raise ValueError(f"screen_color must be a #RRGGBB hex string, got '{text}'")
    return tuple(int(s[i : i + 2], 16) / 255.0 for i in (0, 2, 4))  # type: ignore[return-value]


def rgb01_to_lab(color: tuple[float, float, float]) -> np.ndarray:
    """Lab (L 0..100, a/b ~ -127..127) of a single normalised-RGB colour."""
    cv2 = require_cv2()
    arr = np.array([[list(color)]], dtype=np.float32)
    return cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)[0, 0]


def _background_border_cluster(
    lab: np.ndarray, settings: ChromaSettings
) -> tuple[np.ndarray, float, int]:
    """Border pixels belonging to the dominant flat colour, plus quality metrics.

    Raises RuntimeError if the border is not dominated by a single flat colour, or
    the corners disagree (gradient / vignette / a subject occupying a corner) --
    i.e. there is no clean flat background to key against.
    """
    h, w = lab.shape[:2]
    border = _border_mask(h, w, settings.border_ratio)
    bl = lab[border]

    # Dominant colour = mode of a coarse Lab histogram over the border strip.
    bins = np.floor(bl / 6.0).astype(np.int64)
    bins -= bins.min(axis=0)
    key = (bins[:, 0] * 4096 + bins[:, 1]) * 4096 + bins[:, 2]
    vals, counts = np.unique(key, return_counts=True)
    peak_center = bl[key == vals[counts.argmax()]].mean(axis=0)
    cluster_local = np.linalg.norm(bl - peak_center, axis=1) <= settings.detect_cluster_radius
    share = float(cluster_local.mean())
    center = np.median(bl[cluster_local], axis=0)

    # Corner agreement: catches gradients / vignettes / a subject in one corner.
    cs = max(8, int(round(min(h, w) * 0.03)))
    corner_meds = [
        np.median(c.reshape(-1, 3), axis=0)
        for c in (lab[:cs, :cs], lab[:cs, -cs:], lab[-cs:, :cs], lab[-cs:, -cs:])
    ]
    agree = int(sum(np.linalg.norm(m - center) <= settings.detect_corner_tol for m in corner_meds))

    if share < settings.detect_min_border_share or agree < 3:
        raise RuntimeError(
            "Could not auto-detect a flat background colour: dominant-colour border "
            f"share={share:.2f} (need >= {settings.detect_min_border_share:.2f}), corners "
            f"agreeing={agree}/4 (need >= 3). Pass --screen-color to set it explicitly."
        )

    full = np.zeros((h, w), dtype=bool)
    full[border] = cluster_local
    return full, share, agree


def estimate_background_model(
    rgb: np.ndarray,
    settings: ChromaSettings,
    screen_color: tuple[float, float, float] | None = None,
) -> tuple[BackgroundModel, tuple[np.ndarray, np.ndarray, np.ndarray]]:
    rgb_f, lab = convert_color_spaces(rgb)
    if screen_color is None:
        # AUTO: sample the flat colour that dominates the border (raises if there is
        # no clean flat background). Its own pixels are the model -- no seed search.
        sample_mask, _, _ = _background_border_cluster(lab, settings)
    else:
        # Seed search: pixels near the given colour, border-preferred; the border
        # sample then refines the exact observed shade (compression / drift).
        target_lab = rgb01_to_lab(screen_color)
        candidates = np.linalg.norm(lab - target_lab, axis=2) <= 25.0
        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):
            dist = np.linalg.norm(lab - target_lab, axis=2)
            sample_mask = dist <= np.percentile(dist, 10.0)

    rgb_samples = rgb_f[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,
    )
    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),
        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, lab)


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

    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)

    # Perceptual closeness to the background colour (auto-detected when none was
    # given, otherwise the supplied colour). Downstream thresholds (trimap 0.92,
    # suppress 0.35-0.80) keep the weaker pastel signal off the seg-driven subject.
    conf = np.maximum(lab_conf, rgb_conf * 0.85)
    return np.clip(conf, 0.0, 1.0).astype(np.float32), model