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 <noreply@anthropic.com>
This commit is contained in:
+73
-56
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user