From 10705d49dfc7576ffb2a35e391a5623c36e4daa6 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 30 Jun 2026 15:03:15 +0800 Subject: [PATCH] Support YAML pipeline configuration --- README.md | 9 ++++ bgfilter/cli.py | 86 +++++++++++------------------------ bgfilter/config.py | 96 ++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + scripts/smoke_samples.py | 39 +++++++--------- 5 files changed, 150 insertions(+), 81 deletions(-) create mode 100644 bgfilter/config.py diff --git a/README.md b/README.md index c28f765..871a03a 100644 --- a/README.md +++ b/README.md @@ -31,9 +31,14 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --input Samples\TestImage.png ` --output Outputs\TestImage_rgba.png ` --debug-dir Outputs\TestImage_debug ` + --config configs\default.yaml ` --device cuda ``` +The CLI reads `configs/default.yaml` when `--config` is provided. Command-line +options override config values, so tuning can usually happen in YAML while +runtime choices such as `--device cpu` stay on the command line. + Use CPU for validation when CUDA is unavailable: ```powershell @@ -41,6 +46,7 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --input Samples\TestImage.png ` --output Outputs\TestImage_rgba.png ` --debug-dir Outputs\TestImage_debug ` + --config configs\default.yaml ` --device cpu ``` @@ -51,6 +57,7 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --input-dir Samples ` --output-dir Outputs ` --debug-dir Outputs\debug ` + --config configs\default.yaml ` --device cuda ``` @@ -64,6 +71,7 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --input-dir Samples ` --output-dir Outputs\chroma ` --debug-dir Outputs\chroma_debug ` + --config configs\default.yaml ` --matting-method chroma ` --device cpu ``` @@ -103,6 +111,7 @@ Run the bundled sample smoke check: D:\MiniConda\envs\lightML\python.exe scripts\smoke_samples.py ` --samples-dir Samples ` --output-dir Outputs\smoke_samples ` + --config configs\default.yaml ` --device cpu ` --fallback-to-chroma-alpha ` --max-edge-green-excess-p95 0.50 diff --git a/bgfilter/cli.py b/bgfilter/cli.py index dfbb1a6..c509ae8 100644 --- a/bgfilter/cli.py +++ b/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__": diff --git a/bgfilter/config.py b/bgfilter/config.py new file mode 100644 index 0000000..1b99194 --- /dev/null +++ b/bgfilter/config.py @@ -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), + ) diff --git a/requirements.txt b/requirements.txt index a410475..35d5373 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ huggingface_hub tqdm typer rich +PyYAML diff --git a/scripts/smoke_samples.py b/scripts/smoke_samples.py index 0061e3b..3d77342 100644 --- a/scripts/smoke_samples.py +++ b/scripts/smoke_samples.py @@ -8,16 +8,10 @@ REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) +from bgfilter.config import load_settings, override_settings from bgfilter.pipeline import MattingPipeline from bgfilter.quality import assert_quality, measure_rgba -from bgfilter.settings import ( - AlphaPostSettings, - ChromaSettings, - DespillSettings, - ModelSettings, - PipelineSettings, - TrimapSettings, -) +from bgfilter.settings import PipelineSettings REQUIRED_DEBUG_FILES = { @@ -36,17 +30,12 @@ REQUIRED_DEBUG_FILES = { def build_settings(args: argparse.Namespace) -> PipelineSettings: - return PipelineSettings( - chroma=ChromaSettings(), - trimap=TrimapSettings(), - alpha_post=AlphaPostSettings(), - despill=DespillSettings(), - model=ModelSettings( - model_name=args.model_name, - device=args.device, - matting_method=args.matting_method, - fallback_to_chroma_alpha=args.fallback_to_chroma_alpha, - ), + return override_settings( + load_settings(args.config), + model_name=args.model_name, + device=args.device, + matting_method=args.matting_method, + fallback_to_chroma_alpha=args.fallback_to_chroma_alpha, ) @@ -63,10 +52,16 @@ def main() -> int: parser = argparse.ArgumentParser(description="Run sample-image smoke checks.") parser.add_argument("--samples-dir", type=Path, default=Path("Samples")) parser.add_argument("--output-dir", type=Path, default=Path("Outputs/smoke_samples")) + parser.add_argument("--config", type=Path, default=Path("configs/default.yaml")) parser.add_argument("--device", default="cpu") - parser.add_argument("--matting-method", default="vitmatte", choices=["vitmatte", "chroma"]) - parser.add_argument("--model-name", default="hustvl/vitmatte-small-composition-1k") - parser.add_argument("--fallback-to-chroma-alpha", action="store_true") + parser.add_argument("--matting-method", default=None, choices=["vitmatte", "chroma"]) + parser.add_argument("--model-name", default=None) + parser.add_argument( + "--fallback-to-chroma-alpha", + action="store_true", + default=None, + help="Override config and allow chroma alpha when ViTMatte fails.", + ) parser.add_argument("--max-edge-green-excess-p95", type=float, default=0.50) args = parser.parse_args()