diff --git a/bgfilter/__init__.py b/bgfilter/__init__.py index 2e425e6..8d43e99 100644 --- a/bgfilter/__init__.py +++ b/bgfilter/__init__.py @@ -1,3 +1,3 @@ -"""Green screen character matting pipeline.""" +"""Flat-background character matting pipeline.""" __version__ = "0.1.0" diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py index 5151c32..c0f3634 100644 --- a/bgfilter/alpha_post.py +++ b/bgfilter/alpha_post.py @@ -24,12 +24,12 @@ def suppress_alpha_by_chroma( trimap: np.ndarray, settings: AlphaPostSettings, ) -> np.ndarray: - """Pull alpha toward 0 where background (green) confidence is high. + """Pull alpha toward 0 where background-colour confidence is high. Only the unknown band (trimap == 128) is touched, so confidently-foreground - pixels are never thinned. Green that survived in hair gaps / hole pockets - becomes transparent; this also clears the magenta fringe that over-unmixing - leaves on those low-alpha pixels, since alpha ~ 0 hides the foreground colour. + pixels are never thinned. Background colour that survived in hair gaps / hole + pockets becomes transparent; this also clears the colour fringe that + over-unmixing leaves on those low-alpha pixels, since alpha ~ 0 hides it. """ if not settings.chroma_suppress: return alpha diff --git a/bgfilter/chroma.py b/bgfilter/chroma.py index b890123..817af75 100644 --- a/bgfilter/chroma.py +++ b/bgfilter/chroma.py @@ -12,8 +12,6 @@ from .settings import ChromaSettings class BackgroundModel: rgb_center: tuple[float, float, float] rgb_sigma: tuple[float, float, float] - hsv_center: tuple[float, float, float] - hue_sigma: float lab_center: tuple[float, float, float] lab_sigma: float sample_count: int @@ -22,11 +20,6 @@ class BackgroundModel: return asdict(self) -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) - - def _border_mask(height: int, width: int, ratio: float) -> np.ndarray: border = max(8, int(round(max(height, width) * ratio))) border = min(border, height // 2, width // 2) @@ -38,13 +31,12 @@ def _border_mask(height: int, width: int, ratio: float) -> np.ndarray: return mask -def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return ``(rgb_f, lab)``: RGB normalised to 0..1 and Lab (L 0..100, a/b signed).""" cv2 = require_cv2() rgb_f = rgb.astype(np.float32) / 255.0 - hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV).astype(np.float32) - hsv[..., 1:] /= 255.0 lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB).astype(np.float32) - return rgb_f, hsv, lab + return rgb_f, lab def parse_hex_color(text: str) -> tuple[float, float, float]: @@ -116,7 +108,7 @@ def detect_background_color( Pure: reads pixels only, with no matting side effects. """ settings = settings or ChromaSettings() - rgb_f, _, lab = convert_color_spaces(rgb) + 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)) @@ -126,7 +118,7 @@ def estimate_background_model( 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) + rgb_f, lab = convert_color_spaces(rgb) if screen_color is None: # 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. @@ -146,7 +138,6 @@ def estimate_background_model( sample_mask = dist <= np.percentile(dist, 10.0) rgb_samples = rgb_f[sample_mask] - hsv_samples = hsv[sample_mask] lab_samples = lab[sample_mask] rgb_center = np.median(rgb_samples, axis=0) @@ -154,12 +145,6 @@ def estimate_background_model( np.percentile(np.abs(rgb_samples - rgb_center), 75, axis=0) * 1.4826, settings.rgb_sigma_min, ) - hsv_center = np.median(hsv_samples, axis=0) - hue_center = float(hsv_center[0]) - hue_dists = _hue_distance(hsv_samples[:, 0], hue_center) - hue_sigma = max( - float(np.percentile(hue_dists, 75) * 1.4826), settings.hue_sigma_min - ) lab_center = np.median(lab_samples, axis=0) lab_dists = np.linalg.norm(lab_samples - lab_center, axis=1) lab_sigma = max( @@ -169,13 +154,11 @@ def estimate_background_model( model = BackgroundModel( rgb_center=tuple(float(x) for x in rgb_center), rgb_sigma=tuple(float(x) for x in rgb_sigma), - hsv_center=tuple(float(x) for x in hsv_center), - hue_sigma=float(hue_sigma), lab_center=tuple(float(x) for x in lab_center), lab_sigma=float(lab_sigma), sample_count=int(sample_mask.sum()), ) - return model, (rgb_f, hsv, lab) + return model, (rgb_f, lab) def compute_bg_confidence( @@ -189,7 +172,7 @@ def compute_bg_confidence( model, spaces = estimate_background_model(rgb, settings, screen_color) else: spaces = convert_color_spaces(rgb) - rgb_f, hsv, lab = spaces + rgb_f, lab = spaces lab_center = np.asarray(model.lab_center, dtype=np.float32) lab_dist = np.linalg.norm(lab - lab_center, axis=2) diff --git a/bgfilter/despill.py b/bgfilter/despill.py index 3517e2d..cb86144 100644 --- a/bgfilter/despill.py +++ b/bgfilter/despill.py @@ -3,106 +3,43 @@ from __future__ import annotations import numpy as np from .chroma import BackgroundModel +from .deps import require_cv2 from .settings import DespillSettings -def despill_green( +def despill( rgb: np.ndarray, 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 + """Remove background-colour spill from the foreground. + + 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 toward + neutral, keeping luminance so hair texture and shading survive. This works for + any flat background: a saturated primary (green/blue screen) reduces to classic + channel suppression, while low-chroma pastels -- where no single RGB channel + dominates -- are handled by the same projection. + + Returns the corrected RGB (uint8) and the per-pixel despill weight (float32). + """ 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 - # (alpha ~ 1); the old upper bound left those uncorrected. Neutral pixels - # stay untouched because the green-excess term below is zero for them. - foreground = alpha > settings.edge_low - if settings.edge_expand_radius > 0: - from .deps import require_cv2 - - cv2 = require_cv2() - 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) - - r = rgb_f[..., 0] - g = rgb_f[..., 1] - b = rgb_f[..., 2] - neutral_green = np.maximum(r, b) + settings.green_excess_margin - excess = np.maximum(g - neutral_green, 0.0) - - # (2) Keep full strength across the whole opaque range (alpha >= 0.5) and - # ease off only where the pixel is mostly transparent. The previous tent - # peaked at alpha 0.5 and starved exactly the solid hair that needs despill. - alpha_weight = np.maximum( - np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor - ) - - # (3) Localise purely by green excess. bg_confidence is intentionally not a - # gate: it is low on the character edge where the spill lives, so using it - # suppressed the very correction we want. - green_weight = np.clip(excess / max(settings.green_excess_margin * 8.0, 1e-6), 0.0, 1.0) - mask = foreground.astype(np.float32) * alpha_weight * green_weight - - out = rgb_f.copy() - target_green = np.minimum(g, neutral_green) - out[..., 1] = g * (1.0 - mask * settings.strength) + target_green * ( - mask * settings.strength - ) - - # Light foreground color approximation: green spill usually removed too much - # luminance from thin hair edges, so redistribute part of the removed green - # into red/blue without changing fully opaque pixels. - bg_green = float(model.rgb_center[1]) - compensation = excess * mask * settings.strength * min(0.25, bg_green * 0.15) - out[..., 0] = np.clip(out[..., 0] + compensation * 0.5, 0.0, 1.0) - 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: + if mag < 1e-3: # near-neutral background: no chroma direction, nothing to despill return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32) da, db = va / mag, vb / mag + cv2 = require_cv2() + rgb_f = rgb.astype(np.float32) / 255.0 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) - ) + 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) @@ -110,12 +47,9 @@ def _despill_color( 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 - ) + # Full strength across the opaque range; ease off only where mostly transparent. + 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 diff --git a/bgfilter/foreground.py b/bgfilter/foreground.py index ff79dd3..a26f2b3 100644 --- a/bgfilter/foreground.py +++ b/bgfilter/foreground.py @@ -54,8 +54,8 @@ def _estimate_ml( Solves image = alpha * F + (1 - alpha) * B with spatial smoothness, so the foreground colour is *propagated* into semi-transparent edges from reliable high-alpha neighbours instead of being divided out pixel-wise. That removes - green spill on thin hair where a direct unmix (divide by small alpha) is - too noisy to trust. + background-colour spill on thin hair where a direct unmix (divide by small + alpha) is too noisy to trust. """ try: from pymatting import estimate_foreground_ml diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 2b1f9dc..4611dad 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -8,7 +8,7 @@ import numpy as np 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 .despill import despill from .foreground import estimate_foreground_rgb from .io import load_rgb, save_gray, save_rgb, save_rgba, write_text from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews @@ -83,9 +83,6 @@ 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 @@ -94,7 +91,7 @@ def _run_image( elif mode in ("directional", "directional-hard-bg"): from .chroma import convert_color_spaces - lab = convert_color_spaces(rgb)[2] + lab = convert_color_spaces(rgb)[1] trimap, trimap_stats = fuse_trimap_directional( seg_mask, bg_confidence, lab, model.lab_center, settings.trimap ) @@ -114,8 +111,8 @@ def _run_image( foreground = estimate_foreground_rgb( rgb, alpha, bg_confidence, model, settings.foreground ) - corrected_rgb, color_mask = despill_green( - foreground.rgb, alpha, model, settings.despill, screen_color=despill_color + corrected_rgb, color_mask = despill( + foreground.rgb, alpha, model, settings.despill ) save_rgba(output_path, corrected_rgb, alpha) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index dcdfaf1..5c881f1 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -7,7 +7,6 @@ from dataclasses import dataclass class ChromaSettings: 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 @@ -50,7 +49,7 @@ class AlphaPostSettings: alpha_floor: float = 0.002 alpha_ceil: float = 0.998 # Chroma-guided alpha suppression: in the unknown band, pull alpha toward 0 - # where background (green) confidence is high, clearing spill in gaps/holes. + # where background-colour confidence is high, clearing spill in gaps/holes. chroma_suppress: bool = True chroma_suppress_bg_low: float = 0.35 chroma_suppress_bg_high: float = 0.80 @@ -81,12 +80,11 @@ class DespillSettings: enabled: bool = True edge_low: float = 0.005 strength: float = 0.92 - green_excess_margin: float = 0.015 edge_expand_radius: int = 2 alpha_weight_floor: float = 0.35 - # 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. + # Margin, in Lab a*/b* units, of chroma allowed along the background-colour + # direction before it counts as spill. Larger keeps more of the subject's own + # colour. color_margin_lab: float = 4.0 diff --git a/bgfilter/trimap.py b/bgfilter/trimap.py index d647820..6c8b168 100644 --- a/bgfilter/trimap.py +++ b/bgfilter/trimap.py @@ -56,12 +56,12 @@ def generate_trimap( def fuse_trimap( seg_mask: np.ndarray, bg_confidence: np.ndarray, settings: TrimapSettings ) -> tuple[np.ndarray, dict[str, int]]: - """Build a trimap from a semantic subject mask, refined by the green key. + """Build a trimap from a semantic subject mask, refined by the chroma key. Authority split: the segmentation mask decides subject *topology* (it keeps colour-contaminated hair as foreground and drops see-through holes), while the - bright-green chroma key sharpens the screen boundary. ViTMatte then refines - the unknown band. + chroma key sharpens the flat-background boundary. ViTMatte then refines the + unknown band. """ cv2 = require_cv2() shape = seg_mask.shape @@ -70,7 +70,7 @@ def fuse_trimap( ) band_radius = max(settings.min_fg_safe_radius, fg_safe_radius // 2) - green_bg = bg_confidence >= settings.sure_bg_threshold # bright clean screen + screen_bg = bg_confidence >= settings.sure_bg_threshold # confident flat background core = (seg_mask >= settings.seg_core_threshold).astype(np.uint8) loose = (seg_mask >= settings.seg_loose_threshold).astype(np.uint8) @@ -81,9 +81,9 @@ def fuse_trimap( band = cv2.dilate(loose, kernel).astype(bool) & ~cv2.erode(loose, kernel).astype(bool) # Background follows the mask directly so interior holes stay background, minus - # the boundary band; foreground is the confident subject, never bright screen. - sure_fg = core.astype(bool) & ~band & ~green_bg - sure_bg = ((loose == 0) & ~band) | (green_bg & (loose == 0)) + # the boundary band; foreground is the confident subject, never the background. + sure_fg = core.astype(bool) & ~band & ~screen_bg + sure_bg = ((loose == 0) & ~band) | (screen_bg & (loose == 0)) trimap = np.full(shape, 128, dtype=np.uint8) trimap[sure_bg] = 0 diff --git a/configs/birefnet.yaml b/configs/birefnet.yaml index 923cb7c..c80932a 100644 --- a/configs/birefnet.yaml +++ b/configs/birefnet.yaml @@ -55,7 +55,6 @@ despill: enabled: true edge_low: 0.005 strength: 0.92 - green_excess_margin: 0.015 edge_expand_radius: 2 alpha_weight_floor: 0.35 color_margin_lab: 4.0 diff --git a/configs/default.yaml b/configs/default.yaml index d5fa03e..bda0325 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -62,7 +62,6 @@ despill: enabled: true edge_low: 0.005 strength: 0.92 - green_excess_margin: 0.015 edge_expand_radius: 2 alpha_weight_floor: 0.35 - color_margin_lab: 4.0 # custom-screen despill: Lab chroma margin (screen_color set) + color_margin_lab: 4.0 # Lab chroma margin along the background-colour direction