From 62f175948ee7b08bbba727dd35e66be3b7768fc7 Mon Sep 17 00:00:00 2001 From: lingh Date: Wed, 1 Jul 2026 20:34:31 +0800 Subject: [PATCH] Auto-detect the background colour (default); drop the green heuristic Replace the hardcoded green-screen auto-detection with a general flat-colour detector. When screen_color is null the background colour is now detected from the image border (dominant colour of the border strip) and used directly as the chroma model, with a failure gate that raises when the border is not one clean flat colour (gradient / texture / subject filling the frame). - chroma: add detect_background_color + _background_border_cluster; the null case of estimate_background_model samples the detected cluster directly (no seed search); compute_bg_confidence is now always perceptual (Lab/RGB distance). Removed the green heuristic (initial_green_candidates), the green multiplicative gating, the now-unused _smoothstep, and five green-only ChromaSettings fields. - pipeline: reuse the detected colour for de-spill in the auto case. - settings / default.yaml: add detect_* tuning fields; refresh screen_color docs. - README: document auto-detection + the failure gate, trimap modes, birefnet config. A supplied --screen-color hex still uses the seed-search refinement. Validated: detects green / pastel / text samples correctly, raises on gradient / two-colour, CLI exits 1 on failure, green and pastel mattes unchanged in quality. Co-Authored-By: Claude Opus 4.8 --- README.md | 40 +++++++++----- bgfilter/chroma.py | 129 ++++++++++++++++++++++++------------------- bgfilter/cli.py | 4 +- bgfilter/pipeline.py | 5 +- bgfilter/settings.py | 16 +++--- configs/default.yaml | 6 +- 6 files changed, 117 insertions(+), 83 deletions(-) diff --git a/README.md b/README.md index 4e1de3b..f1ffaff 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # BgFilter Offline character matting for AI-generated images on a flat-colour background. -Green screen is the default; any flat colour (e.g. a pastel `#CFEFFF`) works via a -`screen_color` prior. +The background colour is **auto-detected** from the image border (green, pastel, any +flat colour); pass `--screen-color` to set it explicitly. ```text RGB input - -> chroma bg-confidence (keyed to the screen colour) + -> chroma bg-confidence (keyed to the auto-detected or given colour) -> [optional] semantic segmentation mask -> trimap -> ViTMatte -> alpha cleanup -> pymatting foreground -> despill -> RGBA PNG -> QA previews @@ -20,24 +20,37 @@ 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) or `birefnet` (needs `trust_remote_code`). + 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. Both share pymatting foreground estimation and a colour de-spill. There is **no** green-contamination rescue / recolour layer — with clean source images it is unnecessary, so it was removed. +The segmentation trimap defaults to `directional` mode (chroma + seg + a hue-direction +split: it keeps a background-coloured garment such as a white shirt while dropping a +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). + ## Background colour -Default is green auto-detection. For any other flat background, pass the colour as a -hex prior so both the chroma key and the de-spill work off it: +By default (`screen_color: null`) the background colour is **auto-detected** from the +image border: the dominant flat colour of the border strip becomes the key colour. If +the border is not one clean flat colour — a gradient, texture, or a subject filling the +frame — detection **fails with an error**; pass `--screen-color` explicitly in that case. + +To set it yourself, give a hex prior: ```powershell ... --screen-color "#CFEFFF" ``` -or set `screen_color: "#CFEFFF"` in the config (`null` = green auto-detect). The -chroma key then scores pixels by perceptual (Lab/RGB) distance to that colour, and -de-spill removes chroma along that colour's direction instead of the green channel. +or `screen_color: "#CFEFFF"` in the config. Either way the chroma key scores pixels by +perceptual (Lab/RGB) distance to the colour, and de-spill removes chroma along that +colour's direction. (A supplied hex is refined against nearby border pixels; an +auto-detected colour is used directly.) ## Environment @@ -62,13 +75,14 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --output Outputs\TestImage_rgba.png ` --debug-dir Outputs\TestImage_debug ` --config configs\default.yaml ` - --screen-color "#CFEFFF" ` --device cuda ``` -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 such as `--device cpu` or `--screen-color` stay on the command line. +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. ## Batch diff --git a/bgfilter/chroma.py b/bgfilter/chroma.py index ca2522d..b890123 100644 --- a/bgfilter/chroma.py +++ b/bgfilter/chroma.py @@ -22,11 +22,6 @@ class BackgroundModel: return asdict(self) -def _smoothstep(x: np.ndarray, edge0: float, edge1: float) -> np.ndarray: - t = np.clip((x - edge0) / max(edge1 - edge0, 1e-6), 0.0, 1.0) - return t * t * (3.0 - 2.0 * t) - - def _hue_distance(hue: np.ndarray, center: float) -> np.ndarray: diff = np.abs(hue.astype(np.float32) - float(center)) return np.minimum(diff, 180.0 - diff) @@ -67,22 +62,63 @@ def rgb01_to_lab(color: tuple[float, float, float]) -> np.ndarray: return cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)[0, 0] -def initial_green_candidates( - rgb_f: np.ndarray, hsv: np.ndarray, settings: ChromaSettings -) -> np.ndarray: - r = rgb_f[..., 0] - g = rgb_f[..., 1] - b = rgb_f[..., 2] - dominance = g - np.maximum(r, b) - hue = hsv[..., 0] - sat = hsv[..., 1] - return ( - (dominance >= settings.green_margin) - & (g >= settings.min_green_value) - & (sat >= settings.min_saturation) - & (hue >= settings.hue_low) - & (hue <= settings.hue_high) - ) +def _background_border_cluster( + lab: np.ndarray, settings: ChromaSettings +) -> tuple[np.ndarray, float, int]: + """Border pixels belonging to the dominant flat colour, plus quality metrics. + + Raises RuntimeError if the border is not dominated by a single flat colour, or + the corners disagree (gradient / vignette / a subject occupying a corner) -- + i.e. there is no clean flat background to key against. + """ + h, w = lab.shape[:2] + border = _border_mask(h, w, settings.border_ratio) + bl = lab[border] + + # Dominant colour = mode of a coarse Lab histogram over the border strip. + bins = np.floor(bl / 6.0).astype(np.int64) + bins -= bins.min(axis=0) + key = (bins[:, 0] * 4096 + bins[:, 1]) * 4096 + bins[:, 2] + vals, counts = np.unique(key, return_counts=True) + peak_center = bl[key == vals[counts.argmax()]].mean(axis=0) + cluster_local = np.linalg.norm(bl - peak_center, axis=1) <= settings.detect_cluster_radius + share = float(cluster_local.mean()) + center = np.median(bl[cluster_local], axis=0) + + # Corner agreement: catches gradients / vignettes / a subject in one corner. + cs = max(8, int(round(min(h, w) * 0.03))) + corner_meds = [ + np.median(c.reshape(-1, 3), axis=0) + for c in (lab[:cs, :cs], lab[:cs, -cs:], lab[-cs:, :cs], lab[-cs:, -cs:]) + ] + agree = int(sum(np.linalg.norm(m - center) <= settings.detect_corner_tol for m in corner_meds)) + + if share < settings.detect_min_border_share or agree < 3: + raise RuntimeError( + "Could not auto-detect a flat background colour: dominant-colour border " + f"share={share:.2f} (need >= {settings.detect_min_border_share:.2f}), corners " + f"agreeing={agree}/4 (need >= 3). Pass --screen-color to set it explicitly." + ) + + full = np.zeros((h, w), dtype=bool) + full[border] = cluster_local + return full, share, agree + + +def detect_background_color( + rgb: np.ndarray, settings: ChromaSettings | None = None +) -> tuple[float, float, float]: + """Detect a single flat background colour from the image border. + + Returns the colour as a 0..1 RGB triple. Raises RuntimeError (via + :func:`_background_border_cluster`) if the border is not dominated by one flat + colour, so callers fail loudly rather than matting against a wrong colour. + Pure: reads pixels only, with no matting side effects. + """ + settings = settings or ChromaSettings() + rgb_f, _, lab = convert_color_spaces(rgb) + mask, _, _ = _background_border_cluster(lab, settings) + return tuple(float(x) for x in np.median(rgb_f[mask], axis=0)) def estimate_background_model( @@ -92,25 +128,21 @@ def estimate_background_model( ) -> tuple[BackgroundModel, tuple[np.ndarray, np.ndarray, np.ndarray]]: rgb_f, hsv, lab = convert_color_spaces(rgb) if screen_color is None: - candidates = initial_green_candidates(rgb_f, hsv, settings) + # AUTO: sample the flat colour that dominates the border (raises if there is + # no clean flat background). Its own pixels are the model -- no seed search. + sample_mask, _, _ = _background_border_cluster(lab, settings) else: - # Candidates = pixels near the given colour; the border sample then refines - # the exact observed shade (compression / lighting drift) around it. + # Seed search: pixels near the given colour, border-preferred; the border + # sample then refines the exact observed shade (compression / drift). target_lab = rgb01_to_lab(screen_color) candidates = np.linalg.norm(lab - target_lab, axis=2) <= 25.0 - h, w = candidates.shape - border_candidates = candidates & _border_mask(h, w, settings.border_ratio) - - sample_mask = border_candidates - if int(sample_mask.sum()) < settings.min_samples: - sample_mask = candidates - if int(sample_mask.sum()) < max(64, settings.min_samples // 16): - # Last-resort fallback: the pixels closest to the target (or greenest). - if screen_color is None: - score = rgb_f[..., 1] - np.maximum(rgb_f[..., 0], rgb_f[..., 2]) - sample_mask = score >= np.percentile(score, 90.0) - else: - dist = np.linalg.norm(lab - rgb01_to_lab(screen_color), axis=2) + h, w = candidates.shape + border_candidates = candidates & _border_mask(h, w, settings.border_ratio) + sample_mask = border_candidates + if int(sample_mask.sum()) < settings.min_samples: + sample_mask = candidates + if int(sample_mask.sum()) < max(64, settings.min_samples // 16): + dist = np.linalg.norm(lab - target_lab, axis=2) sample_mask = dist <= np.percentile(dist, 10.0) rgb_samples = rgb_f[sample_mask] @@ -159,17 +191,6 @@ def compute_bg_confidence( spaces = convert_color_spaces(rgb) rgb_f, hsv, lab = spaces - r = rgb_f[..., 0] - g = rgb_f[..., 1] - b = rgb_f[..., 2] - dominance = g - np.maximum(r, b) - hue = hsv[..., 0] - sat = hsv[..., 1] - - hue_conf = np.exp(-0.5 * (_hue_distance(hue, model.hsv_center[0]) / model.hue_sigma) ** 2) - sat_conf = _smoothstep(sat, max(0.05, model.hsv_center[1] * 0.45), max(0.2, model.hsv_center[1] * 0.85)) - dominance_conf = _smoothstep(dominance, settings.green_margin * 0.4, settings.green_margin * 1.6) - lab_center = np.asarray(model.lab_center, dtype=np.float32) lab_dist = np.linalg.norm(lab - lab_center, axis=2) lab_conf = np.exp(-0.5 * (lab_dist / model.lab_sigma) ** 2) @@ -179,12 +200,8 @@ def compute_bg_confidence( rgb_dist = np.linalg.norm((rgb_f - rgb_center) / rgb_sigma, axis=2) rgb_conf = np.exp(-0.5 * (rgb_dist / 2.5) ** 2) - if screen_color is None: - # Green screen: gate closeness by green hue / saturation / channel dominance. - conf = hue_conf * sat_conf * dominance_conf * np.maximum(lab_conf, rgb_conf * 0.85) - else: - # Arbitrary colour: rely on perceptual closeness to the estimated colour. - # Downstream thresholds (trimap 0.92, suppress 0.35-0.80) keep the weaker - # pastel signal from touching the seg-driven subject. - conf = np.maximum(lab_conf, rgb_conf * 0.85) + # Perceptual closeness to the background colour (auto-detected when none was + # given, otherwise the supplied colour). Downstream thresholds (trimap 0.92, + # suppress 0.35-0.80) keep the weaker pastel signal off the seg-driven subject. + conf = np.maximum(lab_conf, rgb_conf * 0.85) return np.clip(conf, 0.0, 1.0).astype(np.float32), model diff --git a/bgfilter/cli.py b/bgfilter/cli.py index c939863..e551d23 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -8,7 +8,7 @@ from rich import print from .config import load_settings, override_settings from .pipeline import MattingPipeline, run_image -app = typer.Typer(help="Offline green screen character matting.") +app = typer.Typer(help="Offline flat-background character matting.") @app.command() @@ -21,7 +21,7 @@ def main( 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"), - screen_color: str | None = typer.Option(None, "--screen-color", help="Background colour prior as #RRGGBB (default: auto-detect green screen)"), + screen_color: str | None = typer.Option(None, "--screen-color", help="Background colour prior as #RRGGBB (default: auto-detect the flat background colour)"), 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), diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index b5b66d8..2b1f9dc 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -83,6 +83,9 @@ def _run_image( bg_confidence, model = compute_bg_confidence( rgb, settings=settings.chroma, screen_color=screen_color ) + # With no colour given, chroma auto-detected one (model.rgb_center); reuse it for + # despill so both key off the same background colour instead of the green path. + despill_color = screen_color if screen_color is not None else tuple(model.rgb_center) if settings.segmentation.enabled: seg_mask = pipeline._segment(rgb) mode = settings.trimap.mode @@ -112,7 +115,7 @@ def _run_image( rgb, alpha, bg_confidence, model, settings.foreground ) corrected_rgb, color_mask = despill_green( - foreground.rgb, alpha, model, settings.despill, screen_color=screen_color + foreground.rgb, alpha, model, settings.despill, screen_color=despill_color ) save_rgba(output_path, corrected_rgb, alpha) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 778d4e7..6ddd337 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -5,16 +5,16 @@ from dataclasses import dataclass @dataclass(frozen=True) class ChromaSettings: - green_margin: float = 0.08 - min_green_value: float = 0.22 - min_saturation: float = 0.20 - hue_low: float = 35.0 - hue_high: float = 95.0 border_ratio: float = 0.04 min_samples: int = 2048 hue_sigma_min: float = 7.0 lab_sigma_min: float = 10.0 rgb_sigma_min: float = 0.08 + # Background-colour auto-detection (used when screen_color is None): find the + # flat colour that dominates the image border; raise if none is clean enough. + detect_cluster_radius: float = 12.0 # Lab radius grouping border pixels into the background cluster + detect_min_border_share: float = 0.55 # min fraction of the border that must be this one colour + detect_corner_tol: float = 20.0 # max Lab distance a corner may sit from the cluster centre @dataclass(frozen=True) @@ -116,7 +116,7 @@ class PipelineSettings: despill: DespillSettings = DespillSettings() model: ModelSettings = ModelSettings() segmentation: SegmentationSettings = SegmentationSettings() - # Background-colour prior. None = auto-detect a green screen (current default, - # unchanged). A hex string like "#CFEFFF" keys chroma + despill off that colour - # instead, for arbitrary (e.g. pastel) flat backgrounds. + # Background-colour prior. None = auto-detect the flat background colour from the + # image border (raises if there is no clean flat background). A hex string like + # "#CFEFFF" sets it explicitly. screen_color: str | None = None diff --git a/configs/default.yaml b/configs/default.yaml index 29d216d..0ed326a 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -1,6 +1,6 @@ -# Background-colour prior. null = auto-detect a green screen (default behaviour). -# A hex string like "#CFEFFF" keys chroma + despill off that colour instead, for -# arbitrary (e.g. pastel) flat backgrounds. Also overridable via --screen-color. +# Background-colour prior. null = auto-detect the flat background colour from the +# image border (raises if the border is not one clean flat colour). A hex string +# like "#CFEFFF" sets it explicitly. Also overridable via --screen-color. screen_color: null model: