Add directional trimap mode (default); keep seg/hard-bg via trimap.mode
Replace the two experimental bools (directional, bg_hued_to_bg) with a single
TrimapSettings.mode selector for the segmentation pipeline:
- "directional" (new default): chroma magnitude + seg + a Lab hue-direction
split. In the chroma-unknown zone a confidently-segmented pixel stays
foreground unless it is displaced toward the background hue, so a neutral
background-coloured garment (e.g. a white shirt) is kept while a background-
hued residual (blue between hair strands) is left unknown for ViTMatte /
chroma-suppress to clear.
- "seg": original fuse_trimap, unchanged.
- "directional-hard-bg": aggressive variant that hard-removes background-hued
pixels (can eat cool/shadowed white cloth).
Selectable via configs/default.yaml (trimap.mode) or CLI --trimap-mode; unknown
modes raise. Default CLI output verified byte-identical to the reviewed UNK
result on the pastel sample.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -100,6 +100,81 @@ def fuse_trimap(
|
||||
return trimap, stats
|
||||
|
||||
|
||||
def fuse_trimap_directional(
|
||||
seg_mask: np.ndarray,
|
||||
bg_confidence: np.ndarray,
|
||||
lab: np.ndarray,
|
||||
lab_center: tuple[float, float, float],
|
||||
settings: TrimapSettings,
|
||||
) -> tuple[np.ndarray, dict[str, int]]:
|
||||
"""Three-signal trimap: chroma magnitude + seg + a directional chroma test.
|
||||
|
||||
Per pixel (``bgc`` = chroma bg-confidence, ``seg`` = subject confidence)::
|
||||
|
||||
BG if bgc >= sure_bg_threshold (clean screen colour)
|
||||
OR (seg < seg_low AND bgc > sure_fg_threshold) (seg says not-subject,
|
||||
but never override a pixel chroma is sure is foreground -> keeps
|
||||
fine wisps the segmenter underestimates)
|
||||
FG if bgc <= sure_fg_threshold AND seg >= seg_low (clearly non-screen colour)
|
||||
else (chroma-unknown): if seg >= seg_core_threshold, split by hue --
|
||||
a pixel *not* displaced toward the background hue (projection <
|
||||
bg_hue_proj_min: e.g. a neutral white shirt) becomes FG, while one
|
||||
strongly displaced toward it (e.g. blue between hair strands) is left
|
||||
unknown for ViTMatte / chroma-suppress. Lower-seg unknowns stay unknown.
|
||||
|
||||
A thin unknown band at the segmentation silhouette is preserved so ViTMatte can
|
||||
anti-alias the boundary (it is deliberately small so interior holes survive).
|
||||
"""
|
||||
cv2 = require_cv2()
|
||||
shape = seg_mask.shape
|
||||
fg_safe_radius = radius_from_ratio(
|
||||
shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius
|
||||
)
|
||||
band_radius = max(settings.min_fg_safe_radius, fg_safe_radius // 2)
|
||||
|
||||
chroma_bg = bg_confidence >= settings.sure_bg_threshold
|
||||
chroma_fg = bg_confidence <= settings.sure_fg_threshold
|
||||
chroma_unknown = ~chroma_bg & ~chroma_fg
|
||||
core = seg_mask >= settings.seg_core_threshold
|
||||
|
||||
# Directional chroma: projection of (a*, b*) onto the background chroma
|
||||
# direction. Positive = colour displaced toward the background hue.
|
||||
lab_c = np.asarray(lab_center, dtype=np.float32)
|
||||
mag = float(np.hypot(lab_c[1], lab_c[2]))
|
||||
if mag < 1e-3:
|
||||
proj = np.zeros(shape, dtype=np.float32)
|
||||
else:
|
||||
proj = (lab[..., 1] * lab_c[1] + lab[..., 2] * lab_c[2]) / mag
|
||||
bg_hued = proj >= settings.bg_hue_proj_min
|
||||
|
||||
bg = chroma_bg | ((seg_mask < settings.seg_low) & ~chroma_fg)
|
||||
fg = chroma_fg & (seg_mask >= settings.seg_low)
|
||||
rule3_fg = chroma_unknown & core & ~bg_hued
|
||||
if settings.mode == "directional-hard-bg":
|
||||
# Aggressive variant: hard-remove background-hued pixels instead of leaving
|
||||
# them unknown. Can eat cool/shadowed white cloth, so it is not the default.
|
||||
bg = bg | (chroma_unknown & core & bg_hued)
|
||||
|
||||
trimap = np.full(shape, 128, dtype=np.uint8)
|
||||
trimap[bg] = 0
|
||||
trimap[fg | rule3_fg] = 255
|
||||
|
||||
loose = (seg_mask >= settings.seg_loose_threshold).astype(np.uint8)
|
||||
kernel = elliptical_kernel(band_radius)
|
||||
band = cv2.dilate(loose, kernel).astype(bool) & ~cv2.erode(loose, kernel).astype(bool)
|
||||
trimap[band & ~chroma_bg] = 128
|
||||
|
||||
stats = {
|
||||
"sure_bg_pixels": int((trimap == 0).sum()),
|
||||
"unknown_pixels": int((trimap == 128).sum()),
|
||||
"sure_fg_pixels": int((trimap == 255).sum()),
|
||||
"band_radius": int(band_radius),
|
||||
"rule3_fg_pixels": int(rule3_fg.sum()),
|
||||
"bg_hued_pixels": int((bg_hued & chroma_unknown & core).sum()),
|
||||
}
|
||||
return trimap, stats
|
||||
|
||||
|
||||
def trimap_to_alpha_seed(trimap: np.ndarray, bg_confidence: np.ndarray) -> np.ndarray:
|
||||
alpha = np.clip(1.0 - bg_confidence, 0.0, 1.0).astype(np.float32)
|
||||
alpha[trimap == 0] = 0.0
|
||||
|
||||
Reference in New Issue
Block a user