Support YAML pipeline configuration

2026-06-30 15:03:15 +08:00
parent 923dbe9f09
commit 10705d49df
5 changed files with 150 additions and 81 deletions
@@ -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__":
@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import fields, replace
+from pathlib import Path
+from typing import Any, TypeVar
+
+from .settings import (
+    AlphaPostSettings,
+    ChromaSettings,
+    DespillSettings,
+    ModelSettings,
+    PipelineSettings,
+    TrimapSettings,
+)
+
+T = TypeVar("T")
+
+
+def _load_yaml(path: str | Path) -> dict[str, Any]:
+    try:
+        import yaml
+    except ModuleNotFoundError as exc:
+        raise RuntimeError("Missing dependency 'PyYAML'. Install it with: pip install PyYAML") from exc
+
+    with Path(path).open("r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle) or {}
+    if not isinstance(data, dict):
+        raise ValueError(f"Config file must contain a YAML mapping: {path}")
+    return data
+
+
+def _update_dataclass(instance: T, values: dict[str, Any] | None) -> T:
+    if not values:
+        return instance
+    allowed = {field.name for field in fields(instance)}
+    unknown = sorted(set(values) - allowed)
+    if unknown:
+        cls_name = type(instance).__name__
+        raise ValueError(f"Unknown {cls_name} field(s): {', '.join(unknown)}")
+    return replace(instance, **values)
+
+
+def settings_from_dict(data: dict[str, Any]) -> PipelineSettings:
+    allowed_sections = {"chroma", "trimap", "alpha_post", "despill", "model"}
+    unknown_sections = sorted(set(data) - allowed_sections)
+    if unknown_sections:
+        raise ValueError(f"Unknown config section(s): {', '.join(unknown_sections)}")
+
+    return PipelineSettings(
+        chroma=_update_dataclass(ChromaSettings(), data.get("chroma")),
+        trimap=_update_dataclass(TrimapSettings(), data.get("trimap")),
+        alpha_post=_update_dataclass(AlphaPostSettings(), data.get("alpha_post")),
+        despill=_update_dataclass(DespillSettings(), data.get("despill")),
+        model=_update_dataclass(ModelSettings(), data.get("model")),
+    )
+
+
+def load_settings(path: str | Path | None = None) -> PipelineSettings:
+    if path is None:
+        return PipelineSettings()
+    return settings_from_dict(_load_yaml(path))
+
+
+def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineSettings:
+    chroma = settings.chroma
+    trimap = settings.trimap
+    alpha_post = settings.alpha_post
+    despill = settings.despill
+    model = settings.model
+
+    model_updates = {
+        key: overrides[key]
+        for key in ["model_name", "device", "matting_method", "fallback_to_chroma_alpha"]
+        if overrides.get(key) is not None
+    }
+    trimap_updates = {
+        key: overrides[key]
+        for key in [
+            "sure_bg_threshold",
+            "sure_fg_threshold",
+            "unknown_radius_ratio",
+            "fg_safe_radius_ratio",
+        ]
+        if overrides.get(key) is not None
+    }
+    despill_updates: dict[str, Any] = {}
+    if overrides.get("despill") is not None:
+        despill_updates["enabled"] = overrides["despill"]
+
+    return PipelineSettings(
+        chroma=chroma,
+        trimap=_update_dataclass(trimap, trimap_updates),
+        alpha_post=alpha_post,
+        despill=_update_dataclass(despill, despill_updates),
+        model=_update_dataclass(model, model_updates),
+    )