From 7555321cfc9c8a1c9ab4050f57bb07409b3b5884 Mon Sep 17 00:00:00 2001 From: lingh Date: Wed, 1 Jul 2026 15:33:58 +0800 Subject: [PATCH] Simplify to single-segmenter pipelines; add screen_color prior; drop rescue + recolor Strategic pivot: the dropout-fill rescue, the anime-seg fill backend, the BiRefNet+anime-seg intersection, and the recolor pass all existed to repair source images whose hair was already green-contaminated or broken at generation time. Fixing the source instead (clean pastel-background generation) makes the matte high-contrast, so that whole compensation layer is unnecessary. Remove it. - Two pipelines via segmentation.enabled: false -> chroma-only (Chroma + ViTMatte) true -> single segmenter (default anime-seg, switchable birefnet) -> trimap then ViTMatte refines, pymatting estimates foreground, despill cleans spill. - Generalise the green-hardcoded colour logic to a screen_color prior (hex, e.g. "#CFEFFF"; default null = green auto-detect, byte-identical). chroma keys off Lab/RGB distance to the colour; despill removes chroma along its Lab direction. Added --screen-color CLI flag. - Removed dead code: fill_seg_dropouts + seg_fill*/intersection params, SegmentationSettings.fill_backend/fill_model_name/sharpen, unsharp_mask, recolor.py + RecolorSettings. Co-Authored-By: Claude Opus 4.8 --- bgfilter/alpha_post.py | 39 ------------------- bgfilter/chroma.py | 50 +++++++++++++++++++++---- bgfilter/cli.py | 2 + bgfilter/config.py | 15 +++++--- bgfilter/despill.py | 59 +++++++++++++++++++++++++++++ bgfilter/pipeline.py | 52 ++++---------------------- bgfilter/recolor.py | 81 ---------------------------------------- bgfilter/segmentation.py | 13 ------- bgfilter/settings.py | 46 ++++++----------------- configs/default.yaml | 30 +++++---------- 10 files changed, 142 insertions(+), 245 deletions(-) delete mode 100644 bgfilter/recolor.py diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py index 206457d..5151c32 100644 --- a/bgfilter/alpha_post.py +++ b/bgfilter/alpha_post.py @@ -85,42 +85,3 @@ def clean_alpha( out[hole & (trimap == 255)] = 1.0 return enforce_trimap(out, trimap) - - -def fill_seg_dropouts( - alpha: np.ndarray, - seg_mask: np.ndarray, - settings: AlphaPostSettings, - gate_mask: np.ndarray | None = None, -) -> np.ndarray: - """Rescue alpha where ViTMatte dropped a region the semantic mask asserts. - - Green-contaminated wisps are nearly the background colour, so ViTMatte (a - contrast-based matter) cuts them even though BiRefNet still sees them. Where - ViTMatte's alpha is near zero but the segmentation mask is confidently - foreground, restore a faint alpha from the mask. Limited to those dropouts so - ViTMatte's sharper edges elsewhere are untouched, and gated by seg_min so the - mask's faint halo and the finger gap (low mask) are not filled. - - ``gate_mask`` (when ``seg_fill_intersect`` is on) is a second, reliable mask - -- BiRefNet -- required to also endorse the rescue (>= ``seg_fill_gate_min``). - ``seg_mask`` (anime-seg, high recall) then only proposes wisps; the - intersection with BiRefNet drops its over-inclusive halo, while the rescued - *value* still comes from anime-seg's smooth mask (no grainy BiRefNet value). - """ - if not settings.seg_fill: - return alpha - out = alpha.astype(np.float32).copy() - rescued = settings.seg_fill_strength * np.clip(seg_mask, 0.0, 1.0) - dropout = (out < settings.seg_fill_alpha_max) & (seg_mask >= settings.seg_fill_seg_min) - # Intersection mode: keep only rescues the reliable mask also endorses. - if settings.seg_fill_intersect and gate_mask is not None: - dropout &= gate_mask >= settings.seg_fill_gate_min - # The segmentation mask carries scattered low-confidence noise; rescuing it - # verbatim sprinkles isolated specks (burrs) over the matte. Keep only - # connected structures (genuine wisps), drop the isolated specks. - if settings.seg_fill_min_area_ratio > 0: - min_area = max(4, int(round(alpha.size * settings.seg_fill_min_area_ratio))) - dropout = _remove_small_components(dropout, min_area) - out[dropout] = np.maximum(out[dropout], rescued[dropout]) - return np.clip(out, 0.0, 1.0) diff --git a/bgfilter/chroma.py b/bgfilter/chroma.py index 3ddc77b..ca2522d 100644 --- a/bgfilter/chroma.py +++ b/bgfilter/chroma.py @@ -52,6 +52,21 @@ def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.nd return rgb_f, hsv, lab +def parse_hex_color(text: str) -> tuple[float, float, float]: + """Parse ``#RRGGBB`` (or ``RRGGBB``) into an RGB triple normalised to 0..1.""" + s = text.strip().lstrip("#") + if len(s) != 6: + raise ValueError(f"screen_color must be a #RRGGBB hex string, got '{text}'") + return tuple(int(s[i : i + 2], 16) / 255.0 for i in (0, 2, 4)) # type: ignore[return-value] + + +def rgb01_to_lab(color: tuple[float, float, float]) -> np.ndarray: + """Lab (L 0..100, a/b ~ -127..127) of a single normalised-RGB colour.""" + cv2 = require_cv2() + arr = np.array([[list(color)]], dtype=np.float32) + return cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)[0, 0] + + def initial_green_candidates( rgb_f: np.ndarray, hsv: np.ndarray, settings: ChromaSettings ) -> np.ndarray: @@ -71,10 +86,18 @@ def initial_green_candidates( def estimate_background_model( - rgb: np.ndarray, settings: ChromaSettings + rgb: np.ndarray, + settings: ChromaSettings, + screen_color: tuple[float, float, float] | None = None, ) -> tuple[BackgroundModel, tuple[np.ndarray, np.ndarray, np.ndarray]]: rgb_f, hsv, lab = convert_color_spaces(rgb) - candidates = initial_green_candidates(rgb_f, hsv, settings) + if screen_color is None: + candidates = initial_green_candidates(rgb_f, hsv, settings) + else: + # Candidates = pixels near the given colour; the border sample then refines + # the exact observed shade (compression / lighting drift) around it. + 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) @@ -82,10 +105,13 @@ def estimate_background_model( 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: choose pixels with strongest green dominance. - dominance = rgb_f[..., 1] - np.maximum(rgb_f[..., 0], rgb_f[..., 2]) - cutoff = np.percentile(dominance, 90.0) - sample_mask = dominance >= cutoff + # 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) + sample_mask = dist <= np.percentile(dist, 10.0) rgb_samples = rgb_f[sample_mask] hsv_samples = hsv[sample_mask] @@ -124,10 +150,11 @@ def compute_bg_confidence( rgb: np.ndarray, model: BackgroundModel | None = None, settings: ChromaSettings | None = None, + screen_color: tuple[float, float, float] | None = None, ) -> tuple[np.ndarray, BackgroundModel]: settings = settings or ChromaSettings() if model is None: - model, spaces = estimate_background_model(rgb, settings) + model, spaces = estimate_background_model(rgb, settings, screen_color) else: spaces = convert_color_spaces(rgb) rgb_f, hsv, lab = spaces @@ -152,5 +179,12 @@ 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) - conf = hue_conf * sat_conf * dominance_conf * np.maximum(lab_conf, rgb_conf * 0.85) + 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) return np.clip(conf, 0.0, 1.0).astype(np.float32), model diff --git a/bgfilter/cli.py b/bgfilter/cli.py index c509ae8..ec77410 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -21,6 +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)"), 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), @@ -34,6 +35,7 @@ def main( load_settings(config), model_name=model_name, device=device, + screen_color=screen_color, 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 73c3f63..2a3e73c 100644 --- a/bgfilter/config.py +++ b/bgfilter/config.py @@ -11,7 +11,6 @@ from .settings import ( ForegroundSettings, ModelSettings, PipelineSettings, - RecolorSettings, SegmentationSettings, TrimapSettings, ) @@ -44,7 +43,9 @@ def _update_dataclass(instance: T, values: dict[str, Any] | None) -> T: def settings_from_dict(data: dict[str, Any]) -> PipelineSettings: - allowed_sections = {"chroma", "trimap", "alpha_post", "foreground", "despill", "recolor", "model", "segmentation"} + data = dict(data) + screen_color = data.pop("screen_color", None) + allowed_sections = {"chroma", "trimap", "alpha_post", "foreground", "despill", "model", "segmentation"} unknown_sections = sorted(set(data) - allowed_sections) if unknown_sections: raise ValueError(f"Unknown config section(s): {', '.join(unknown_sections)}") @@ -55,9 +56,9 @@ def settings_from_dict(data: dict[str, Any]) -> PipelineSettings: alpha_post=_update_dataclass(AlphaPostSettings(), data.get("alpha_post")), foreground=_update_dataclass(ForegroundSettings(), data.get("foreground")), despill=_update_dataclass(DespillSettings(), data.get("despill")), - recolor=_update_dataclass(RecolorSettings(), data.get("recolor")), model=_update_dataclass(ModelSettings(), data.get("model")), segmentation=_update_dataclass(SegmentationSettings(), data.get("segmentation")), + screen_color=screen_color, ) @@ -73,7 +74,6 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS alpha_post = settings.alpha_post foreground = settings.foreground despill = settings.despill - recolor = settings.recolor model = settings.model segmentation = settings.segmentation @@ -99,13 +99,18 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS if overrides.get("device") is not None: seg_updates["device"] = overrides["device"] + screen_color = ( + overrides["screen_color"] + if overrides.get("screen_color") is not None + else settings.screen_color + ) return PipelineSettings( chroma=chroma, trimap=_update_dataclass(trimap, trimap_updates), alpha_post=alpha_post, foreground=foreground, despill=_update_dataclass(despill, despill_updates), - recolor=recolor, model=_update_dataclass(model, model_updates), segmentation=_update_dataclass(segmentation, seg_updates), + screen_color=screen_color, ) diff --git a/bgfilter/despill.py b/bgfilter/despill.py index 35544f1..3517e2d 100644 --- a/bgfilter/despill.py +++ b/bgfilter/despill.py @@ -11,10 +11,14 @@ def despill_green( alpha: np.ndarray, model: BackgroundModel, settings: DespillSettings, + screen_color: tuple[float, float, float] | None = None, ) -> tuple[np.ndarray, np.ndarray]: rgb_f = rgb.astype(np.float32) / 255.0 if not settings.enabled: return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32) + if screen_color is not None: + # Arbitrary background colour: despill along its Lab chroma direction. + return _despill_color(rgb, alpha, model, settings) # (1) Despill across the whole foreground, not only the soft edge band. # Green frequently bleeds into pixels ViTMatte marks as solid foreground @@ -65,3 +69,58 @@ def despill_green( out[..., 2] = np.clip(out[..., 2] + compensation * 0.5, 0.0, 1.0) return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32) + + +def _despill_color( + rgb: np.ndarray, + alpha: np.ndarray, + model: BackgroundModel, + settings: DespillSettings, +) -> tuple[np.ndarray, np.ndarray]: + """Despill toward an arbitrary background colour. + + Generalises the green channel algebra: the background's Lab a*/b* vector gives + a chroma direction; foreground chroma projecting along it beyond a margin is + treated as spill and pulled back, keeping luminance so texture survives. Reduces + to a green-style suppression when the background is a saturated primary, but also + handles low-chroma pastels where no single RGB channel dominates. + """ + from .deps import require_cv2 + + cv2 = require_cv2() + rgb_f = rgb.astype(np.float32) / 255.0 + + bg_lab = np.asarray(model.lab_center, dtype=np.float32) + va, vb = float(bg_lab[1]), float(bg_lab[2]) + mag = float(np.hypot(va, vb)) + if mag < 1e-3: + return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32) + da, db = va / mag, vb / mag + + foreground = alpha > settings.edge_low + if settings.edge_expand_radius > 0: + radius = settings.edge_expand_radius + kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1) + ) + foreground = cv2.dilate(foreground.astype(np.uint8), kernel).astype(bool) + + lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB) + A, B = lab[..., 1], lab[..., 2] + proj = A * da + B * db # chroma toward the background hue + excess = np.maximum(proj - settings.color_margin_lab, 0.0) + + alpha_weight = np.maximum( + np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor + ) + color_weight = np.clip( + excess / max(settings.color_margin_lab * 8.0, 1e-6), 0.0, 1.0 + ) + mask = foreground.astype(np.float32) * alpha_weight * color_weight + reduce = excess * mask * settings.strength # pull chroma back along the bg axis + + lab_out = lab.copy() + lab_out[..., 1] = A - reduce * da + lab_out[..., 2] = B - reduce * db + out = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB) + return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 198a92b..ad03ea5 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -6,16 +6,10 @@ from pathlib import Path import numpy as np -from .alpha_post import ( - clean_alpha, - enforce_trimap, - fill_seg_dropouts, - suppress_alpha_by_chroma, -) -from .chroma import compute_bg_confidence +from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma +from .chroma import compute_bg_confidence, parse_hex_color from .despill import despill_green from .foreground import estimate_foreground_rgb -from .recolor import recolor_foreground from .io import load_rgb, save_gray, save_rgb, save_rgba, write_text from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews from .settings import PipelineSettings @@ -28,36 +22,13 @@ class MattingPipeline: self.settings = settings self._runner: ViTMatteRunner | None = None self._segmenter = None - self._fill_segmenter = None def _segment(self, rgb: np.ndarray) -> np.ndarray: if self._segmenter is None: from .segmentation import make_segmenter self._segmenter = make_segmenter(self.settings.segmentation) - mask = self._segmenter.mask(rgb) - seg = self.settings.segmentation - if seg.sharpen > 0: - from .segmentation import unsharp_mask - - mask = unsharp_mask(mask, seg.sharpen, seg.sharpen_radius) - return mask - - def _segment_fill(self, rgb: np.ndarray, default_mask: np.ndarray) -> np.ndarray: - """Mask for the dropout-fill; may use a different (higher-recall) backend - than the trimap, since the fill only adds and cannot break anything.""" - seg = self.settings.segmentation - if not seg.fill_backend or seg.fill_backend == seg.backend: - return default_mask - if self._fill_segmenter is None: - from dataclasses import replace - - from .segmentation import make_segmenter - - self._fill_segmenter = make_segmenter( - replace(seg, backend=seg.fill_backend, model_name=seg.fill_model_name) - ) - return self._fill_segmenter.mask(rgb) + return self._segmenter.mask(rgb) def _predict_alpha(self, rgb: np.ndarray, trimap: np.ndarray, bg_confidence: np.ndarray) -> tuple[np.ndarray, str]: if self.settings.model.matting_method == "chroma": @@ -103,7 +74,10 @@ def _run_image( ) -> dict: settings = pipeline.settings rgb = load_rgb(input_path) - bg_confidence, model = compute_bg_confidence(rgb, settings=settings.chroma) + 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.segmentation.enabled: seg_mask = pipeline._segment(rgb) trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap) @@ -115,22 +89,12 @@ def _run_image( alpha = enforce_trimap(alpha, trimap) alpha = suppress_alpha_by_chroma(alpha, bg_confidence, trimap, settings.alpha_post) alpha = clean_alpha(alpha, trimap, settings.alpha_post) - if seg_mask is not None: - fill_mask = pipeline._segment_fill(rgb, seg_mask) - # seg_mask is the reliable trimap backend (BiRefNet); it gates the fill - # when seg_fill_intersect is on (fill_mask ∧ seg_mask). - alpha = fill_seg_dropouts( - alpha, fill_mask, settings.alpha_post, gate_mask=seg_mask - ) foreground = estimate_foreground_rgb( rgb, alpha, bg_confidence, model, settings.foreground ) corrected_rgb, color_mask = despill_green( - foreground.rgb, alpha, model, settings.despill + foreground.rgb, alpha, model, settings.despill, screen_color=screen_color ) - if settings.recolor.enabled: - recolored = recolor_foreground(corrected_rgb, alpha, seg_mask, settings.recolor) - corrected_rgb, color_mask = recolored.rgb, recolored.mask save_rgba(output_path, corrected_rgb, alpha) result = { diff --git a/bgfilter/recolor.py b/bgfilter/recolor.py deleted file mode 100644 index 5abc5a8..0000000 --- a/bgfilter/recolor.py +++ /dev/null @@ -1,81 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np - -from .deps import require_cv2 -from .settings import RecolorSettings - - -@dataclass(frozen=True) -class RecolorResult: - rgb: np.ndarray # recoloured foreground colour (uint8) - mask: np.ndarray # per-pixel recolour weight applied (float32, 0..1) - - -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 recolor_foreground( - rgb: np.ndarray, - alpha: np.ndarray, - seg_mask: np.ndarray | None, - settings: RecolorSettings, -) -> RecolorResult: - """Repair AI green-contamination by propagating clean neighbour colour. - - Green that the generator bled onto the subject is recoloured by pulling the - chroma from nearby *clean* (non-green) subject pixels, while keeping each - pixel's luminance so hair texture and shading survive. Because it installs a - real observed colour rather than subtracting green, it does not overshoot into - magenta the way a despill/unmix does. No assumption that the contaminated - region is hair -- the target is inferred locally per pixel. - """ - rgb_f = rgb.astype(np.float32) / 255.0 - if not settings.enabled: - return RecolorResult(rgb=rgb.copy(), mask=np.zeros(alpha.shape, dtype=np.float32)) - - cv2 = require_cv2() - r, g, b = rgb_f[..., 0], rgb_f[..., 1], rgb_f[..., 2] - # Residual off-colour after despill: leftover green or the magenta that an - # unmix can overshoot into. Despill already neutralises the matte edge, so the - # clean edge has off ~ 0 and is left alone; this targets contaminated interior. - off = np.maximum(g - np.maximum(r, b), np.minimum(r, b) - g) - - subject = np.clip(seg_mask if seg_mask is not None else alpha, 0.0, 1.0) - coverage = _smoothstep(np.clip(alpha, 0.0, 1.0), 0.0, settings.coverage_soft) - contam = _smoothstep(off, settings.contam_low, settings.contam_high) * coverage - clean = (subject >= settings.person_threshold).astype(np.float32) * ( - off < settings.clean_dom_max - ).astype(np.float32) - - lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB) - L, A, B = lab[..., 0], lab[..., 1], lab[..., 2] - radius = max(3, int(round(max(alpha.shape) * settings.propagate_radius_ratio))) - k = (radius * 2 + 1, radius * 2 + 1) - den = cv2.GaussianBlur(clean, k, 0) - inv_den = 1.0 / np.maximum(den, 1e-6) - target_a = cv2.GaussianBlur(A * clean, k, 0) * inv_den - target_b = cv2.GaussianBlur(B * clean, k, 0) * inv_den - target_l = cv2.GaussianBlur(L * clean, k, 0) * inv_den - - # Fade the recolour out where there is too little clean reference nearby - # (a fully-contaminated pocket has no trustworthy colour to borrow). - avail = _smoothstep(den, settings.min_reference_weight * 0.5, settings.min_reference_weight) - w = contam * avail - - out_a = (1.0 - w) * A + w * target_a - out_b = (1.0 - w) * B + w * target_b - # Green is bright, so pull luminance partway to the clean neighbour to avoid a - # bright fringe, but only partway so hair texture is not flattened. - wl = w * settings.luminance_strength - out_l = (1.0 - wl) * L + wl * target_l - lab_out = np.stack([out_l, out_a, out_b], axis=-1) - out = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB) - return RecolorResult( - rgb=np.clip(out * 255.0, 0, 255).astype(np.uint8), - mask=w.astype(np.float32), - ) diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py index fbfcae5..e6269bf 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -111,19 +111,6 @@ class AnimeSegSegmenter: return np.asarray(mask_img, dtype=np.float32) / 255.0 -def unsharp_mask(mask: np.ndarray, amount: float, radius: int) -> np.ndarray: - """Sharpen a soft mask: raise thin-strand peaks and crisp the blur so strands - reach the sure-foreground threshold instead of dissolving into the unknown band.""" - if amount <= 0: - return mask - from .deps import require_cv2 - - cv2 = require_cv2() - k = max(1, radius) * 2 + 1 - blur = cv2.GaussianBlur(mask, (k, k), 0) - return np.clip(mask + amount * (mask - blur), 0.0, 1.0).astype(np.float32) - - def make_segmenter(settings: SegmentationSettings): if settings.backend == "birefnet": return BiRefNetSegmenter(settings) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index a9927a9..9f9bcdf 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -42,20 +42,6 @@ class AlphaPostSettings: chroma_suppress_bg_low: float = 0.35 chroma_suppress_bg_high: float = 0.80 chroma_suppress_strength: float = 1.0 - # Segmentation dropout-fill: rescue alpha where ViTMatte cut a region the - # semantic mask asserts is foreground (green-contaminated wisps it cannot see). - seg_fill: bool = True - seg_fill_alpha_max: float = 0.10 - seg_fill_seg_min: float = 0.10 - seg_fill_strength: float = 0.70 - seg_fill_min_area_ratio: float = 0.00002 # drop rescued specks below this area - # Intersection fill: gate the (anime-seg) dropout-fill by a second, reliable - # mask (BiRefNet). anime-seg proposes freely at a low seg_fill_seg_min; only - # rescues BiRefNet also endorses (>= seg_fill_gate_min, its threshold left - # unchanged) survive. Trims anime-seg's over-inclusive halo to the structure - # both backends agree on, without reintroducing BiRefNet's grainy value. - seg_fill_intersect: bool = False - seg_fill_gate_min: float = 0.10 @dataclass(frozen=True) @@ -85,19 +71,10 @@ class DespillSettings: green_excess_margin: float = 0.015 edge_expand_radius: int = 2 alpha_weight_floor: float = 0.35 - - -@dataclass(frozen=True) -class RecolorSettings: - enabled: bool = True - person_threshold: float = 0.5 # clean reference = subject above this seg/alpha - coverage_soft: float = 0.08 # recolour ramps in over alpha 0..this - contam_low: float = 0.02 # residual off-colour where recolour starts - contam_high: float = 0.12 # residual off-colour for full recolour - clean_dom_max: float = 0.03 # reference = subject pixels below this off-colour - propagate_radius_ratio: float = 0.02 - min_reference_weight: float = 0.05 - luminance_strength: float = 0.5 # how far to pull luminance to the clean neighbour + # Custom-screen despill (when screen_color is set): margin, in Lab a*/b* + # units, of chroma allowed along the background colour direction before it is + # treated as spill. Larger keeps more of the subject's own colour. + color_margin_lab: float = 4.0 @dataclass(frozen=True) @@ -110,15 +87,11 @@ class ModelSettings: @dataclass(frozen=True) class SegmentationSettings: - enabled: bool = True - backend: str = "birefnet" # "birefnet" or "anime-seg" - model_name: str = "ZhengPeng7/BiRefNet" - fill_backend: str = "" # "" = reuse `backend` for dropout-fill; else e.g. "anime-seg" - fill_model_name: str = "" + 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" input_size: int = 1024 - sharpen: float = 0.0 # unsharp-mask amount on the mask (0 = off) - sharpen_radius: int = 4 @dataclass(frozen=True) @@ -128,6 +101,9 @@ class PipelineSettings: alpha_post: AlphaPostSettings = AlphaPostSettings() foreground: ForegroundSettings = ForegroundSettings() despill: DespillSettings = DespillSettings() - recolor: RecolorSettings = RecolorSettings() 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. + screen_color: str | None = None diff --git a/configs/default.yaml b/configs/default.yaml index 3fc21e9..6387393 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -1,3 +1,8 @@ +# 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. +screen_color: null + model: model_name: hustvl/vitmatte-base-composition-1k device: cuda @@ -5,11 +10,11 @@ model: fallback_to_chroma_alpha: false 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: birefnet # birefnet | anime-seg (anime-seg needs model_name: skytnt/anime-seg) - model_name: ZhengPeng7/BiRefNet - fill_backend: anime-seg # dropout-fill uses anime-seg (higher wisp recall; add-only, harmless) - fill_model_name: skytnt/anime-seg + backend: anime-seg # single trimap segmenter; switchable: anime-seg | birefnet + model_name: skytnt/anime-seg # birefnet -> ZhengPeng7/BiRefNet device: cuda input_size: 1024 @@ -26,11 +31,6 @@ alpha_post: chroma_suppress_bg_low: 0.35 chroma_suppress_bg_high: 0.80 chroma_suppress_strength: 1.0 - seg_fill: true - seg_fill_alpha_max: 0.10 - seg_fill_seg_min: 0.10 - seg_fill_strength: 0.70 - seg_fill_min_area_ratio: 0.00002 foreground: enabled: true @@ -56,14 +56,4 @@ despill: green_excess_margin: 0.015 edge_expand_radius: 2 alpha_weight_floor: 0.35 - -recolor: - enabled: true - person_threshold: 0.5 - coverage_soft: 0.08 - contam_low: 0.02 - contam_high: 0.12 - clean_dom_max: 0.03 - propagate_radius_ratio: 0.02 - min_reference_weight: 0.05 - luminance_strength: 0.5 + color_margin_lab: 4.0 # custom-screen despill: Lab chroma margin (screen_color set)