Implement green screen matting pipeline

2026-06-30 14:14:15 +08:00
parent 4d2a2015fe
commit 8fc53925e0
17 changed files with 896 additions and 0 deletions
@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict
+from pathlib import Path
+
+import numpy as np
+
+from .alpha_post import clean_alpha, enforce_trimap
+from .chroma import compute_bg_confidence
+from .despill import despill_green
+from .io import load_rgb, save_gray, save_rgba, write_text
+from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews
+from .settings import PipelineSettings
+from .trimap import generate_trimap, trimap_to_alpha_seed
+from .vitmatte_infer import ViTMatteRunner
+
+
+class MattingPipeline:
+    def __init__(self, settings: PipelineSettings):
+        self.settings = settings
+        self._runner: ViTMatteRunner | None = None
+
+    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_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 _run_image(
+    input_path: str | Path,
+    output_path: str | Path,
+    debug_dir: str | Path | None,
+    pipeline: MattingPipeline,
+) -> dict:
+    settings = pipeline.settings
+    rgb = load_rgb(input_path)
+    bg_confidence, model = compute_bg_confidence(rgb, settings=settings.chroma)
+    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 = clean_alpha(alpha, trimap, settings.alpha_post)
+    corrected_rgb, despill_mask = despill_green(
+        rgb, alpha, bg_confidence, model, settings.despill
+    )
+    save_rgba(output_path, corrected_rgb, alpha)
+
+    result = {
+        "input": str(input_path),
+        "output": str(output_path),
+        "alpha_source": alpha_source,
+        "background_model": model.to_dict(),
+        "trimap": trimap_stats,
+        "settings": asdict(settings),
+    }
+
+    if debug_dir is not None:
+        debug = Path(debug_dir)
+        debug.mkdir(parents=True, exist_ok=True)
+        save_gray(debug / "bg_confidence.png", bg_confidence)
+        save_gray(debug / "trimap.png", trimap)
+        save_gray(debug / "alpha.png", alpha)
+        save_gray(debug / "despill_mask.png", despill_mask)
+        save_previews(debug, corrected_rgb, alpha)
+        tiles = {
+            "input": rgb,
+            "bg confidence": np.repeat((bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2),
+            "trimap": np.repeat(trimap[..., None], 3, axis=2),
+            "alpha": np.repeat((alpha[..., None] * 255).astype(np.uint8), 3, axis=2),
+        }
+        for name, color in BACKGROUND_COLORS.items():
+            tiles[f"preview {name}"] = composite(corrected_rgb, alpha, color)
+        make_qa_grid(debug, tiles)
+        write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))
+
+    return result