diff --git a/README.md b/README.md index f6e067d..c793047 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,19 @@ background-hued residual such as blue trapped between hair strands). Switch with `--trimap-mode seg` (topology only, no hue split) or `directional-hard-bg` (aggressive — hard-removes background-hued pixels; can eat cool/shadowed white cloth). +## Non-flat backgrounds + +Pass `--no-chroma` (or `chroma.enabled: false`) to matte images whose background is +**not** one flat colour — gradients, textures, scenes. There is no colour key in +this mode: the segmentation mask alone drives the trimap (mode forced to `seg`), +ViTMatte still refines the unknown band at full resolution, and pymatting still +estimates edge foreground colour. Every colour-keyed stage is bypassed: background +auto-detection, the directional hue split, chroma alpha suppression, despill, and +the cross-check veto (with `--cross-check-as-seg` the second model still provides +the segmentation mask). Requires the segmentation pipeline; quality then rests +entirely on the segmenter's mask, so expect flat-background results to stay +stronger on hair-level detail. + ## Background colour By default (`screen_color: null`) the background colour is **auto-detected** from the diff --git a/bgfilter/cli.py b/bgfilter/cli.py index ad4d299..b9b58d1 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -23,6 +23,7 @@ def main( device: str | None = typer.Option(None, "--device"), precision: str | None = typer.Option(None, "--precision", help="Compute precision for all models (ViTMatte cast + BiRefNet autocast): fp32 (default) | bf16 (faster + halves matting activation memory; needs bf16-capable hardware, else falls back to fp32)"), screen_color: str | None = typer.Option(None, "--screen-color", help="Background colour prior as #RRGGBB (default: auto-detect the flat background colour)"), + chroma: bool | None = typer.Option(None, "--chroma/--no-chroma", help="Colour-key the flat background (default: on). --no-chroma = non-flat background mode: segmentation alone drives the trimap; auto-detect, hue split, chroma suppression, despill and the cross-check veto are bypassed"), 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), @@ -42,6 +43,7 @@ def main( device=device, precision=precision, screen_color=screen_color, + chroma=chroma, matting_method=matting_method, fallback_to_chroma_alpha=fallback_to_chroma_alpha, sure_bg_threshold=sure_bg_threshold, diff --git a/bgfilter/config.py b/bgfilter/config.py index 688aa1d..d58574f 100644 --- a/bgfilter/config.py +++ b/bgfilter/config.py @@ -99,6 +99,9 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS } if overrides.get("trimap_mode") is not None: trimap_updates["mode"] = overrides["trimap_mode"] + chroma_updates: dict[str, Any] = {} + if overrides.get("chroma") is not None: + chroma_updates["enabled"] = overrides["chroma"] despill_updates: dict[str, Any] = {} if overrides.get("despill") is not None: despill_updates["enabled"] = overrides["despill"] @@ -130,7 +133,7 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS else settings.screen_color ) return PipelineSettings( - chroma=chroma, + chroma=_update_dataclass(chroma, chroma_updates), trimap=_update_dataclass(trimap, trimap_updates), alpha_post=alpha_post, cross_check=_update_dataclass(settings.cross_check, cross_check_updates), diff --git a/bgfilter/despill.py b/bgfilter/despill.py index cb86144..6c9c815 100644 --- a/bgfilter/despill.py +++ b/bgfilter/despill.py @@ -10,7 +10,7 @@ from .settings import DespillSettings def despill( rgb: np.ndarray, alpha: np.ndarray, - model: BackgroundModel, + model: BackgroundModel | None, settings: DespillSettings, ) -> tuple[np.ndarray, np.ndarray]: """Remove background-colour spill from the foreground. @@ -24,7 +24,8 @@ def despill( Returns the corrected RGB (uint8) and the per-pixel despill weight (float32). """ - if not settings.enabled: + # model is None with chroma disabled: no key colour means no spill direction. + if not settings.enabled or model is None: return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32) bg_lab = np.asarray(model.lab_center, dtype=np.float32) diff --git a/bgfilter/foreground.py b/bgfilter/foreground.py index b898761..41f0de7 100644 --- a/bgfilter/foreground.py +++ b/bgfilter/foreground.py @@ -11,7 +11,7 @@ def estimate_foreground_rgb( rgb: np.ndarray, alpha: np.ndarray, bg_confidence: np.ndarray, - model: BackgroundModel, + model: BackgroundModel | None, settings: ForegroundSettings, ) -> np.ndarray: """Estimated foreground colour F (uint8 HxWx3) for compositing over alpha.""" @@ -23,10 +23,15 @@ def estimate_foreground_rgb( try: return _estimate_ml(rgb, alpha, settings) except RuntimeError: - if not settings.fallback_to_unmix: + if not settings.fallback_to_unmix or model is None: raise method = "unmix" if method == "unmix": + if model is None: + raise RuntimeError( + "Foreground method 'unmix' needs the chroma background model, which " + "does not exist with chroma disabled; use method 'ml'." + ) return _estimate_unmix(rgb, alpha, bg_confidence, model, settings) raise RuntimeError( f"Unsupported foreground method '{settings.method}'. Use 'ml' or 'unmix'." diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index f1906aa..f6d2880 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -132,10 +132,26 @@ def run_image( def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: settings = pipeline.settings - screen_color = parse_hex_color(settings.screen_color) if settings.screen_color else None - bg_confidence, model = compute_bg_confidence( - rgb, settings=settings.chroma, screen_color=screen_color - ) + if settings.chroma.enabled: + screen_color = parse_hex_color(settings.screen_color) if settings.screen_color else None + bg_confidence, model = compute_bg_confidence( + rgb, settings=settings.chroma, screen_color=screen_color + ) + else: + if not settings.segmentation.enabled: + raise RuntimeError( + "chroma.enabled: false (non-flat background mode) needs the segmentation " + "pipeline; enable segmentation or re-enable chroma." + ) + if settings.model.matting_method == "chroma": + raise RuntimeError( + "matting_method 'chroma' needs the chroma key; use 'vitmatte' or re-enable chroma." + ) + # Non-flat background: no colour key exists. A zero background-confidence + # map makes every chroma-fused formula degrade to its seg-only form; the + # colour-keyed stages (hue split, suppression, veto, despill) are skipped. + bg_confidence = np.zeros(rgb.shape[:2], dtype=np.float32) + model = None second_alpha = None if settings.segmentation.enabled: reuse = ( @@ -154,7 +170,9 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: seg_mask = second_alpha else: seg_mask = pipeline._segment(rgb) - mode = settings.trimap.mode + # The directional modes are colour tests against the key colour; without + # one, the seg-topology trimap is the only meaningful choice. + mode = settings.trimap.mode if settings.chroma.enabled else "seg" if mode == "seg": trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap) elif mode in ("directional", "directional-hard-bg"): @@ -173,14 +191,17 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: raw_alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence) alpha = enforce_trimap(raw_alpha, trimap) - alpha = suppress_alpha_by_chroma( - alpha, bg_confidence, trimap, settings.alpha_post, - # The matte-confidence gate only makes sense for a real matting prediction; - # a chroma-seeded alpha is itself colour evidence, so no gate there. - raw_alpha=raw_alpha if alpha_source == "vitmatte" else None, - ) + if settings.chroma.enabled: + alpha = suppress_alpha_by_chroma( + alpha, bg_confidence, trimap, settings.alpha_post, + # The matte-confidence gate only makes sense for a real matting prediction; + # a chroma-seeded alpha is itself colour evidence, so no gate there. + raw_alpha=raw_alpha if alpha_source == "vitmatte" else None, + ) alpha = clean_alpha(alpha, trimap, settings.alpha_post) - if settings.cross_check.enabled: + # The veto's suspect zone is defined by background hue, so it needs the key + # colour; without chroma the second opinion still serves as seg (reuse above). + if settings.cross_check.enabled and settings.chroma.enabled: lab = convert_color_spaces(rgb)[1] if second_alpha is None: second_alpha = pipeline._second_opinion(rgb) @@ -199,7 +220,7 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: metadata = { "alpha_source": alpha_source, - "background_model": model.to_dict(), + "background_model": model.to_dict() if model is not None else None, "trimap": trimap_stats, "settings": asdict(settings), } diff --git a/bgfilter/settings.py b/bgfilter/settings.py index fe95a47..eb516b9 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -5,6 +5,14 @@ from dataclasses import dataclass @dataclass(frozen=True) class ChromaSettings: + # False = non-flat-background mode (--no-chroma): no colour key at all. + # Segmentation alone drives the trimap (mode forced to "seg"), and every + # colour-keyed stage is bypassed: background auto-detect, the directional + # hue split, chroma alpha suppression, despill, the cross-check veto (its + # suspect zone is background-hued by definition; with reuse_as_seg the + # second opinion still serves as the segmenter), and the unmix foreground + # fallback. Requires segmentation.enabled and matting_method 'vitmatte'. + enabled: bool = True border_ratio: float = 0.04 min_samples: int = 2048 lab_sigma_min: float = 10.0 diff --git a/configs/default.yaml b/configs/default.yaml index b9b5192..38daf31 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -3,6 +3,13 @@ # like "#CFEFFF" sets it explicitly. Also overridable via --screen-color. screen_color: null +chroma: + # false = non-flat background mode (--no-chroma): no colour key; segmentation + # alone drives the trimap (mode forced to "seg") and every colour-keyed stage + # is bypassed (auto-detect, hue split, chroma suppression, despill, the + # cross-check veto). Needs segmentation.enabled and matting_method vitmatte. + enabled: true + model: model_name: hustvl/vitmatte-base-composition-1k device: cpu