lhk229
7555321cfc
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 <noreply@anthropic.com>
191 lines
7.0 KiB
Python
191 lines
7.0 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass
|
|
|
|
import numpy as np
|
|
|
|
from .deps import require_cv2
|
|
from .settings import ChromaSettings
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
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
|
|
|
|
def to_dict(self) -> dict:
|
|
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)
|
|
|
|
|
|
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)
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
mask[:border, :] = True
|
|
mask[-border:, :] = True
|
|
mask[:, :border] = True
|
|
mask[:, -border:] = True
|
|
return mask
|
|
|
|
|
|
def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
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
|
|
|
|
|
|
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:
|
|
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 estimate_background_model(
|
|
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)
|
|
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)
|
|
|
|
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)
|
|
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)
|
|
rgb_sigma = np.maximum(
|
|
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(
|
|
float(np.percentile(lab_dists, 75) * 1.4826), settings.lab_sigma_min
|
|
)
|
|
|
|
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)
|
|
|
|
|
|
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, screen_color)
|
|
else:
|
|
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)
|
|
|
|
rgb_center = np.asarray(model.rgb_center, dtype=np.float32)
|
|
rgb_sigma = np.asarray(model.rgb_sigma, dtype=np.float32)
|
|
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)
|
|
return np.clip(conf, 0.0, 1.0).astype(np.float32), model
|