diff --git a/README.md b/README.md index f1ffaff..badb898 100644 --- a/README.md +++ b/README.md @@ -20,9 +20,9 @@ Two pipelines, selected by `segmentation.enabled` in the config: Lightest / fastest; leans entirely on the colour key for topology. - **Single segmenter** (`enabled: true`, default) — one segmentation model drives the trimap topology (holes, hair), ViTMatte then refines the soft edges. Backend is - switchable: `anime-seg` (default, ONNX, for anime characters) or `birefnet` (general - salient objects — text, logos, photos; needs `trust_remote_code`). Use - `--config configs\birefnet.yaml` for the BiRefNet backend. + switchable: `birefnet` (default, general salient objects — text, logos, photos; needs + `trust_remote_code`) or `anime-seg` (ONNX, tuned for anime characters). Switch at + runtime with `--seg-backend anime-seg` (it also selects the matching weights). Both share pymatting foreground estimation and a colour de-spill. There is **no** green-contamination rescue / recolour layer — with clean source images it is @@ -74,15 +74,15 @@ 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 + --config configs\default.yaml ``` No `--screen-color` is needed — the background colour is auto-detected. Add `--screen-color "#CFEFFF"` only to override it. The CLI reads `configs/default.yaml` when `--config` is provided; command-line options override config values, so tuning usually happens in YAML while runtime choices (`--device`, `--screen-color`, -`--trimap-mode`) stay on the command line. +`--trimap-mode`, `--seg-backend`) stay on the command line. The default device is +**CPU**; add `--device cuda` for GPU. ## Batch @@ -91,8 +91,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 + --config configs\default.yaml ``` ## Chroma-alpha debug mode diff --git a/bgfilter/cli.py b/bgfilter/cli.py index e551d23..08c0928 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -30,6 +30,7 @@ def main( 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"), trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"), + seg_backend: str | None = typer.Option(None, "--seg-backend", help="Segmentation backend: birefnet (default) | anime-seg"), ) -> None: try: settings = override_settings( @@ -45,6 +46,7 @@ def main( fg_safe_radius_ratio=fg_safe_radius_ratio, despill=despill, trimap_mode=trimap_mode, + seg_backend=seg_backend, ) if input_dir is not None: diff --git a/bgfilter/config.py b/bgfilter/config.py index 32af54a..bd38528 100644 --- a/bgfilter/config.py +++ b/bgfilter/config.py @@ -17,6 +17,9 @@ from .settings import ( T = TypeVar("T") +# Canonical weights per segmentation backend (selected at runtime via --seg-backend). +_SEG_MODELS = {"birefnet": "ZhengPeng7/BiRefNet", "anime-seg": "skytnt/anime-seg"} + def _load_yaml(path: str | Path) -> dict[str, Any]: try: @@ -100,6 +103,14 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS seg_updates: dict[str, Any] = {} if overrides.get("device") is not None: seg_updates["device"] = overrides["device"] + if overrides.get("seg_backend") is not None: + backend = overrides["seg_backend"] + if backend not in _SEG_MODELS: + raise ValueError( + f"Unknown --seg-backend '{backend}'. Use one of: {', '.join(_SEG_MODELS)}." + ) + seg_updates["backend"] = backend + seg_updates["model_name"] = _SEG_MODELS[backend] screen_color = ( overrides["screen_color"] diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 6ddd337..dcdfaf1 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -93,7 +93,7 @@ class DespillSettings: @dataclass(frozen=True) class ModelSettings: model_name: str = "hustvl/vitmatte-small-composition-1k" - device: str = "cuda" + device: str = "cpu" matting_method: str = "vitmatte" fallback_to_chroma_alpha: bool = False @@ -101,9 +101,9 @@ class ModelSettings: @dataclass(frozen=True) class SegmentationSettings: enabled: bool = True # False = chroma-only pipeline (no seg model) - backend: str = "anime-seg" # "anime-seg" or "birefnet" - model_name: str = "skytnt/anime-seg" - device: str = "cuda" + backend: str = "birefnet" # "birefnet" (general) or "anime-seg" (anime characters) + model_name: str = "ZhengPeng7/BiRefNet" + device: str = "cpu" input_size: int = 1024 diff --git a/configs/birefnet.yaml b/configs/birefnet.yaml index abcaac1..923cb7c 100644 --- a/configs/birefnet.yaml +++ b/configs/birefnet.yaml @@ -7,7 +7,7 @@ screen_color: null model: model_name: hustvl/vitmatte-base-composition-1k - device: cuda + device: cpu matting_method: vitmatte fallback_to_chroma_alpha: false @@ -15,7 +15,7 @@ segmentation: enabled: true backend: birefnet model_name: ZhengPeng7/BiRefNet - device: cuda + device: cpu input_size: 1024 trimap: diff --git a/configs/default.yaml b/configs/default.yaml index 0ed326a..d5fa03e 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -5,7 +5,7 @@ screen_color: null model: model_name: hustvl/vitmatte-base-composition-1k - device: cuda + device: cpu matting_method: vitmatte fallback_to_chroma_alpha: false @@ -13,9 +13,9 @@ segmentation: # enabled: true = single-segmenter pipeline (one seg model drives the trimap) # enabled: false = no seg model, chroma-only pipeline (Chroma + ViTMatte) enabled: true - backend: anime-seg # single trimap segmenter; switchable: anime-seg | birefnet - model_name: skytnt/anime-seg # birefnet -> ZhengPeng7/BiRefNet - device: cuda + backend: birefnet # single trimap segmenter; switchable at runtime: --seg-backend birefnet|anime-seg + model_name: ZhengPeng7/BiRefNet # matches backend; --seg-backend anime-seg auto-selects skytnt/anime-seg + device: cpu input_size: 1024 trimap: diff --git a/docs/green_screen_matting_workflow.md b/docs/green_screen_matting_workflow.md index 66cd524..5278cda 100644 --- a/docs/green_screen_matting_workflow.md +++ b/docs/green_screen_matting_workflow.md @@ -8,8 +8,8 @@ 本文档是**第一阶段的绿幕方案原始规格**,仍准确描述 “Chroma + ViTMatte” 这条管线。实现此后已泛化,**当前架构以 [README](../README.md) 为准**,主要变化: -- **两条管线**(`segmentation.enabled` 切换):无分割模型的 Chroma + ViTMatte;或单个语义分割模型(默认 anime-seg,可切 birefnet)定 trimap 拓扑,ViTMatte 细化。 -- **背景色泛化**:`screen_color` 先验(如 `#CFEFFF`),chroma / despill 不再写死绿色;`null` 时才走本文档的绿幕自动探测。 +- **两条管线**(`segmentation.enabled` 切换):无分割模型的 Chroma + ViTMatte;或单个语义分割模型(默认 birefnet,可切 anime-seg,用 `--seg-backend`)定 trimap 拓扑,ViTMatte 细化。 +- **背景色泛化**:`screen_color` 先验(如 `#CFEFFF`);`null` 时自动从图像边框探测平背景色(探测不到干净平背景会报错),不再写死绿幕。 - **已移除**:染绿修复(recolor)、发丝补救(dropout-fill)、BiRefNet+anime-seg 结合方案——它们只是给“生图即染色/断裂”的坏素材擦屁股;源头改用干净背景生图后不再需要。 下文其余章节描述绿幕基线流程,仍然有效。