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:
2026-07-01 19:27:09 +08:00
parent 22ef1d6336
commit 5f48f7a2cd
6 changed files with 122 additions and 2 deletions
+75
View File
@@ -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