Support YAML pipeline configuration

bgfilter/cli.py

@@ -5,49 +5,12 @@ from pathlib import Path
 import typer
 from rich import print
 
+from .config import load_settings, override_settings
 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),
@@ -55,29 +18,31 @@ def main(
     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"),
+    config: Path | None = typer.Option(None, "--config", exists=True, file_okay=True, dir_okay=False),
+    model_name: str | None = typer.Option(None, "--model-name"),
+    device: str | None = typer.Option(None, "--device"),
+    matting_method: str | None = typer.Option(None, "--matting-method"),
+    fallback_to_chroma_alpha: bool | None = typer.Option(None, "--fallback-to-chroma-alpha/--no-fallback-to-chroma-alpha"),
+    sure_bg_threshold: float | None = typer.Option(None, "--sure-bg-threshold", min=0.0, max=1.0),
+    sure_fg_threshold: float | None = typer.Option(None, "--sure-fg-threshold", min=0.0, max=1.0),
+    unknown_radius_ratio: float | None = typer.Option(None, "--unknown-radius-ratio", min=0.0),
+    fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0),
+    despill: bool | None = typer.Option(None, "--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:
+        settings = override_settings(
+            load_settings(config),
+            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,
+        )
+
         if input_dir is not None:
             if output_dir is None:
                 raise typer.BadParameter("--output-dir is required when --input-dir is used")
@@ -102,6 +67,9 @@ def main(
     except RuntimeError as exc:
         print(f"[red]error:[/red] {exc}")
         raise typer.Exit(code=1) from exc
+    except ValueError as exc:
+        print(f"[red]config error:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
 
 
 if __name__ == "__main__":