helpers/background.py

"""
Step 1 — background removal.

Two backends, both producing a single 8-bit alpha matte at the ORIGINAL image
resolution:

  "hf"    briaai/RMBG-2.0 via the HuggingFace Inference API. Best quality.
  "local" Offline fallback: Lab chroma-distance seeding + GrabCut. No network,
          no token, no GPU. Good enough for a specimen on a plain backdrop, and
          it makes the pipeline testable in CI.

Why the fallback keys on *chroma* rather than brightness
--------------------------------------------------------
A fossil photographed against a wall casts a soft shadow. The shadow has almost
the same hue as the wall — it is just darker. Thresholding on brightness pulls
the shadow into the foreground; thresholding on Lab (a, b) chroma distance does
not, because the shadow's chroma barely shifts. We therefore weight chroma
heavily and luminance lightly. Validated on IMG_3105.

NOTE ON A BUG INHERITED FROM `organoid`
---------------------------------------
`organoid/services/processing.py` does:

    result = client.image_segmentation(...)
    bg_removed = result[0]["mask"]
    if bg_removed.mode == "RGBA": ...   # never true — a mask is mode "L"

...so it saves the *silhouette mask* as the "background-removed image" and every
downstream step (including the vision model's coordinate guess) operates on a
flat grey blob rather than the specimen. Fixed here: the mask is used as an
alpha channel and composited against the original RGB.
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
from PIL import Image

from config.settings import BG_BACKEND, BG_STRICT, HF_TOKEN
from helpers.images import composite_on_white, load_rgb, specimen_dir, to_rgba


# ---------------------------------------------------------------------------
# Backend: HuggingFace RMBG-2.0
# ---------------------------------------------------------------------------

def _alpha_from_hf(image_path: Path, size: tuple[int, int]) -> np.ndarray:
    from huggingface_hub import InferenceClient

    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN is not set — export it or use FOSSIL_BG_BACKEND=local")

    client = InferenceClient(token=HF_TOKEN)
    segments = client.image_segmentation(image=str(image_path), model="briaai/RMBG-2.0")
    if not segments:
        raise RuntimeError("RMBG-2.0 returned no segments")

    # RMBG returns a single foreground segment; take the largest if several.
    masks = [np.array(s["mask"].convert("L")) for s in segments]
    alpha = max(masks, key=lambda m: int((m > 127).sum()))

    if (alpha.shape[1], alpha.shape[0]) != size:
        alpha = cv2.resize(alpha, size, interpolation=cv2.INTER_LINEAR)
    return alpha


# ---------------------------------------------------------------------------
# Backend: offline GrabCut fallback
# ---------------------------------------------------------------------------

def _alpha_from_local(bgr: np.ndarray) -> np.ndarray:
    h, w = bgr.shape[:2]
    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
    L, A, B = lab[..., 0], lab[..., 1], lab[..., 2]

    # Background colour model sampled from a border band.
    band = max(4, int(0.02 * min(h, w)))
    border = np.zeros((h, w), bool)
    border[:band, :] = border[-band:, :] = border[:, :band] = border[:, -band:] = True
    L0, A0, B0 = L[border].mean(), A[border].mean(), B[border].mean()

    d_chroma = np.sqrt((A - A0) ** 2 + (B - B0) ** 2)
    d_lum = np.abs(L - L0)
    score = cv2.normalize(d_chroma + 0.20 * d_lum, None, 0, 255, cv2.NORM_MINMAX)
    _, rough = cv2.threshold(score.astype(np.uint8), 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    gc = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
    gc[rough > 0] = cv2.GC_PR_FGD
    gc[cv2.erode(rough, np.ones((25, 25), np.uint8)) > 0] = cv2.GC_FGD
    gc[cv2.dilate(rough, np.ones((45, 45), np.uint8)) == 0] = cv2.GC_BGD

    bgd, fgd = np.zeros((1, 65), np.float64), np.zeros((1, 65), np.float64)
    cv2.grabCut(bgr, gc, None, bgd, fgd, 5, cv2.GC_INIT_WITH_MASK)
    return np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 255, 0).astype(np.uint8)


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

def clean_alpha(alpha: np.ndarray) -> np.ndarray:
    """Binarise, keep the largest connected component, fill interior holes.

    Fossils are single rigid objects: anything disconnected is dust, a label
    card, or a shadow fragment, and any hole is a pore or a dark recess that
    still belongs to the specimen.
    """
    binary = (alpha > 127).astype(np.uint8) * 255

    n, labels, stats, _ = cv2.connectedComponentsWithStats(binary, 8)
    if n > 1:
        largest = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
        binary = np.where(labels == largest, 255, 0).astype(np.uint8)

    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    filled = np.zeros_like(binary)
    cv2.drawContours(filled, contours, -1, 255, -1)
    return filled


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def execute_background_removal(image_path: str | Path, backend: str | None = None) -> dict:
    """
    Remove the background and write four artefacts to `output/<stem>/`.

    Returns a dict matching `schemas.models.BackgroundRemovalResult`.
    """
    image_path = Path(image_path)
    backend = (backend or BG_BACKEND).lower()
    out = specimen_dir(image_path)

    rgb = load_rgb(image_path)
    bgr = cv2.cvtColor(np.array(rgb), cv2.COLOR_RGB2BGR)
    print(f"  [BG] backend={backend} size={rgb.size}")

    if backend == "hf":
        try:
            alpha = _alpha_from_hf(image_path, rgb.size)
        except Exception as exc:  # noqa: BLE001
            # STRICT: refuse to quietly downgrade. GrabCut produces a *worse but
            # plausible* matte — the run would look successful and the masks would
            # just be a bit off, which is the kind of failure nobody notices until
            # a mesh is already built from it. In a deployment that picked RMBG on
            # purpose, a loud stop beats a quiet degradation.
            if BG_STRICT:
                raise RuntimeError(
                    f"RMBG-2.0 background removal failed and FOSSIL_BG_STRICT is on, "
                    f"so no fallback was attempted: {exc}\n"
                    f"Usually this is a missing or unauthorised HF_TOKEN. Fix the token, "
                    f"or set FOSSIL_BG_STRICT=0 to allow the offline GrabCut fallback "
                    f"(which will give visibly worse mattes)."
                ) from exc
            print(f"  [BG] HF backend failed ({exc}); falling back to local GrabCut")
            alpha, backend = _alpha_from_local(bgr), "local (hf-fallback)"
    else:
        alpha = _alpha_from_local(bgr)

    alpha = clean_alpha(alpha)
    coverage = float((alpha > 0).mean())
    if coverage < 0.005:
        raise RuntimeError(
            f"Foreground is only {coverage:.3%} of the frame — background removal "
            "almost certainly failed. Check lighting/contrast against the backdrop."
        )

    alpha_img = Image.fromarray(alpha)
    stem = image_path.stem

    alpha_path = out / f"{stem}_alpha.png"
    cutout_path = out / f"{stem}_cutout_white.png"
    rgba_path = out / f"{stem}_cutout_rgba.png"

    alpha_img.save(alpha_path)
    composite_on_white(rgb, alpha_img).save(cutout_path)
    to_rgba(rgb, alpha_img).save(rgba_path)

    print(f"  [BG] coverage={coverage:.1%} -> {out}")

    return {
        "status": "success",
        "original_path": str(image_path),
        "cutout_path": str(cutout_path),
        "rgba_path": str(rgba_path),
        "alpha_mask_path": str(alpha_path),
        "dimensions": {"width": rgb.width, "height": rgb.height},
        "coverage": round(coverage, 4),
        "backend": backend,
        "message": f"Background removed with '{backend}'. Alpha matte saved to {alpha_path}.",
    }