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,108 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+from rich import print
+
+from .pipeline import MattingPipeline, run_image
+from .settings import (
+    AlphaPostSettings,
+    ChromaSettings,
+    DespillSettings,
+    ModelSettings,
+    PipelineSettings,
+    TrimapSettings,
+)
+
+app = typer.Typer(help="Offline green screen character matting.")
+
+
+def _settings(
+    model_name: str,
+    device: str,
+    matting_method: str,
+    fallback_to_chroma_alpha: bool,
+    sure_bg_threshold: float,
+    sure_fg_threshold: float,
+    unknown_radius_ratio: float,
+    fg_safe_radius_ratio: float,
+    despill: bool,
+) -> PipelineSettings:
+    return PipelineSettings(
+        chroma=ChromaSettings(),
+        trimap=TrimapSettings(
+            sure_bg_threshold=sure_bg_threshold,
+            sure_fg_threshold=sure_fg_threshold,
+            unknown_radius_ratio=unknown_radius_ratio,
+            fg_safe_radius_ratio=fg_safe_radius_ratio,
+        ),
+        alpha_post=AlphaPostSettings(),
+        despill=DespillSettings(enabled=despill),
+        model=ModelSettings(
+            model_name=model_name,
+            device=device,
+            matting_method=matting_method,
+            fallback_to_chroma_alpha=fallback_to_chroma_alpha,
+        ),
+    )
+
+
+@app.command()
+def main(
+    input: Path | None = typer.Option(None, "--input", "-i", exists=True, file_okay=True, dir_okay=False),
+    output: Path | None = typer.Option(None, "--output", "-o", file_okay=True, dir_okay=False),
+    input_dir: Path | None = typer.Option(None, "--input-dir", exists=True, file_okay=False, dir_okay=True),
+    output_dir: Path | None = typer.Option(None, "--output-dir", file_okay=False, dir_okay=True),
+    debug_dir: Path | None = typer.Option(None, "--debug-dir", file_okay=False, dir_okay=True),
+    model_name: str = typer.Option("hustvl/vitmatte-small-composition-1k", "--model-name"),
+    device: str = typer.Option("cuda", "--device"),
+    matting_method: str = typer.Option("vitmatte", "--matting-method"),
+    fallback_to_chroma_alpha: bool = typer.Option(False, "--fallback-to-chroma-alpha/--no-fallback-to-chroma-alpha"),
+    sure_bg_threshold: float = typer.Option(0.92, "--sure-bg-threshold", min=0.0, max=1.0),
+    sure_fg_threshold: float = typer.Option(0.12, "--sure-fg-threshold", min=0.0, max=1.0),
+    unknown_radius_ratio: float = typer.Option(0.012, "--unknown-radius-ratio", min=0.0),
+    fg_safe_radius_ratio: float = typer.Option(0.006, "--fg-safe-radius-ratio", min=0.0),
+    despill: bool = typer.Option(True, "--despill/--no-despill"),
+) -> None:
+    settings = _settings(
+        model_name=model_name,
+        device=device,
+        matting_method=matting_method,
+        fallback_to_chroma_alpha=fallback_to_chroma_alpha,
+        sure_bg_threshold=sure_bg_threshold,
+        sure_fg_threshold=sure_fg_threshold,
+        unknown_radius_ratio=unknown_radius_ratio,
+        fg_safe_radius_ratio=fg_safe_radius_ratio,
+        despill=despill,
+    )
+
+    try:
+        if input_dir is not None:
+            if output_dir is None:
+                raise typer.BadParameter("--output-dir is required when --input-dir is used")
+            images = sorted(
+                p for p in input_dir.iterdir() if p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}
+            )
+            if not images:
+                raise typer.BadParameter(f"No supported images found in {input_dir}")
+            pipeline = MattingPipeline(settings)
+            for image_path in images:
+                out_path = output_dir / f"{image_path.stem}_rgba.png"
+                dbg = None if debug_dir is None else debug_dir / image_path.stem
+                result = pipeline.run_image(image_path, out_path, dbg)
+                print(f"[green]wrote[/green] {result['output']}")
+            return
+
+        if input is None or output is None:
+            raise typer.BadParameter("Use either --input/--output or --input-dir/--output-dir")
+
+        result = run_image(input, output, debug_dir, settings)
+        print(f"[green]wrote[/green] {result['output']}")
+    except RuntimeError as exc:
+        print(f"[red]error:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+
+if __name__ == "__main__":
+    app()