BGfilter/bgfilter/pipeline.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np

from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma
from .chroma import compute_bg_confidence, parse_hex_color
from .despill import despill
from .foreground import estimate_foreground_rgb
from .io import load_rgb, save_gray, save_rgb, save_rgba, write_text
from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews
from .settings import PipelineSettings
from .trimap import (
    fuse_trimap,
    fuse_trimap_directional,
    generate_trimap,
    trimap_to_alpha_seed,
)
from .vitmatte_infer import ViTMatteRunner


@dataclass
class MattingResult:
    """In-memory result of one matting run.

    ``rgb`` is the despilled foreground (uint8 HxWx3) and ``alpha`` is float32 HxW in
    0..1; together they compose the output RGBA. The remaining arrays are the
    intermediate maps used for debug previews.
    """

    rgb: np.ndarray
    alpha: np.ndarray
    metadata: dict
    bg_confidence: np.ndarray
    trimap: np.ndarray
    seg_mask: np.ndarray | None
    color_mask: np.ndarray
    input_rgb: np.ndarray


class MattingPipeline:
    def __init__(
        self,
        settings: PipelineSettings,
        *,
        runner: ViTMatteRunner | None = None,
        segmenter=None,
    ):
        self.settings = settings
        # Pre-loaded models can be injected so a long-lived service shares them
        # across requests instead of reloading per call; None = lazy-load on first use.
        self._runner: ViTMatteRunner | None = runner
        self._segmenter = segmenter

    def _segment(self, rgb: np.ndarray) -> np.ndarray:
        if self._segmenter is None:
            from .segmentation import make_segmenter

            self._segmenter = make_segmenter(self.settings.segmentation)
        return self._segmenter.mask(rgb)

    def _predict_alpha(self, rgb: np.ndarray, trimap: np.ndarray, bg_confidence: np.ndarray) -> tuple[np.ndarray, str]:
        if self.settings.model.matting_method == "chroma":
            return trimap_to_alpha_seed(trimap, bg_confidence), "chroma"
        if self.settings.model.matting_method != "vitmatte":
            raise RuntimeError(
                f"Unsupported matting method '{self.settings.model.matting_method}'. "
                "Use 'vitmatte' or 'chroma'."
            )

        try:
            if self._runner is None:
                self._runner = ViTMatteRunner(self.settings.model)
            return self._runner.predict(rgb, trimap), "vitmatte"
        except RuntimeError:
            if not self.settings.model.fallback_to_chroma_alpha:
                raise
            return trimap_to_alpha_seed(trimap, bg_confidence), "chroma_fallback"

    def run_rgb(self, rgb: np.ndarray) -> MattingResult:
        """Run the full pipeline on an in-memory RGB array (no disk I/O)."""
        return _process_rgb(rgb, self)

    def run_image(
        self,
        input_path: str | Path,
        output_path: str | Path,
        debug_dir: str | Path | None,
    ) -> dict:
        return _run_image(input_path, output_path, debug_dir, self)


def run_image(
    input_path: str | Path,
    output_path: str | Path,
    debug_dir: str | Path | None,
    settings: PipelineSettings,
) -> dict:
    return MattingPipeline(settings).run_image(input_path, output_path, debug_dir)


def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult:
    settings = pipeline.settings
    screen_color = parse_hex_color(settings.screen_color) if settings.screen_color else None
    bg_confidence, model = compute_bg_confidence(
        rgb, settings=settings.chroma, screen_color=screen_color
    )
    if settings.segmentation.enabled:
        seg_mask = pipeline._segment(rgb)
        mode = settings.trimap.mode
        if mode == "seg":
            trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
        elif mode in ("directional", "directional-hard-bg"):
            from .chroma import convert_color_spaces

            lab = convert_color_spaces(rgb)[1]
            trimap, trimap_stats = fuse_trimap_directional(
                seg_mask, bg_confidence, lab, model.lab_center, settings.trimap
            )
        else:
            raise RuntimeError(
                f"Unknown trimap.mode '{mode}'. "
                "Use 'directional', 'seg', or 'directional-hard-bg'."
            )
    else:
        seg_mask = None
        trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
    alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)

    alpha = enforce_trimap(alpha, trimap)
    alpha = suppress_alpha_by_chroma(alpha, bg_confidence, trimap, settings.alpha_post)
    alpha = clean_alpha(alpha, trimap, settings.alpha_post)
    foreground = estimate_foreground_rgb(
        rgb, alpha, bg_confidence, model, settings.foreground
    )
    corrected_rgb, color_mask = despill(foreground.rgb, alpha, model, settings.despill)

    metadata = {
        "alpha_source": alpha_source,
        "background_model": model.to_dict(),
        "trimap": trimap_stats,
        "settings": asdict(settings),
    }
    return MattingResult(
        rgb=corrected_rgb,
        alpha=alpha,
        metadata=metadata,
        bg_confidence=bg_confidence,
        trimap=trimap,
        seg_mask=seg_mask,
        color_mask=color_mask,
        input_rgb=rgb,
    )


def _run_image(
    input_path: str | Path,
    output_path: str | Path,
    debug_dir: str | Path | None,
    pipeline: MattingPipeline,
) -> dict:
    res = _process_rgb(load_rgb(input_path), pipeline)
    save_rgba(output_path, res.rgb, res.alpha)

    result = {"input": str(input_path), "output": str(output_path), **res.metadata}

    if debug_dir is not None:
        debug = Path(debug_dir)
        debug.mkdir(parents=True, exist_ok=True)
        save_gray(debug / "bg_confidence.png", res.bg_confidence)
        save_gray(debug / "trimap.png", res.trimap)
        if res.seg_mask is not None:
            save_gray(debug / "seg_mask.png", res.seg_mask)
        save_gray(debug / "alpha.png", res.alpha)
        save_rgb(debug / "foreground_rgb.png", res.rgb)
        save_gray(debug / "color_mask.png", res.color_mask)
        save_previews(debug, res.rgb, res.alpha)
        tiles = {
            "input": res.input_rgb,
            "bg confidence": np.repeat((res.bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2),
            "trimap": np.repeat(res.trimap[..., None], 3, axis=2),
            "alpha": np.repeat((res.alpha[..., None] * 255).astype(np.uint8), 3, axis=2),
        }
        for name, color in BACKGROUND_COLORS.items():
            tiles[f"preview {name}"] = composite(res.rgb, res.alpha, color)
        make_qa_grid(debug, tiles)
        write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))

    return result