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:
@@ -29,6 +29,7 @@ def main(
|
||||
unknown_radius_ratio: float | None = typer.Option(None, "--unknown-radius-ratio", min=0.0),
|
||||
fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0),
|
||||
despill: bool | None = typer.Option(None, "--despill/--no-despill"),
|
||||
trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"),
|
||||
) -> None:
|
||||
try:
|
||||
settings = override_settings(
|
||||
@@ -43,6 +44,7 @@ def main(
|
||||
unknown_radius_ratio=unknown_radius_ratio,
|
||||
fg_safe_radius_ratio=fg_safe_radius_ratio,
|
||||
despill=despill,
|
||||
trimap_mode=trimap_mode,
|
||||
)
|
||||
|
||||
if input_dir is not None:
|
||||
|
||||
@@ -92,6 +92,8 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
|
||||
]
|
||||
if overrides.get(key) is not None
|
||||
}
|
||||
if overrides.get("trimap_mode") is not None:
|
||||
trimap_updates["mode"] = overrides["trimap_mode"]
|
||||
despill_updates: dict[str, Any] = {}
|
||||
if overrides.get("despill") is not None:
|
||||
despill_updates["enabled"] = overrides["despill"]
|
||||
|
||||
+21
-2
@@ -13,7 +13,12 @@ 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
|
||||
from .settings import PipelineSettings
|
||||
from .trimap import fuse_trimap, generate_trimap, trimap_to_alpha_seed
|
||||
from .trimap import (
|
||||
fuse_trimap,
|
||||
fuse_trimap_directional,
|
||||
generate_trimap,
|
||||
trimap_to_alpha_seed,
|
||||
)
|
||||
from .vitmatte_infer import ViTMatteRunner
|
||||
|
||||
|
||||
@@ -80,7 +85,21 @@ def _run_image(
|
||||
)
|
||||
if settings.segmentation.enabled:
|
||||
seg_mask = pipeline._segment(rgb)
|
||||
trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
|
||||
mode = settings.trimap.mode
|
||||
if mode == "seg":
|
||||
trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
|
||||
elif mode in ("directional", "directional-hard-bg"):
|
||||
from .chroma import convert_color_spaces
|
||||
|
||||
lab = convert_color_spaces(rgb)[2]
|
||||
trimap, trimap_stats = fuse_trimap_directional(
|
||||
seg_mask, bg_confidence, lab, model.lab_center, settings.trimap
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Unknown trimap.mode '{mode}'. "
|
||||
"Use 'directional', 'seg', or 'directional-hard-bg'."
|
||||
)
|
||||
else:
|
||||
seg_mask = None
|
||||
trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
|
||||
|
||||
@@ -28,6 +28,19 @@ class TrimapSettings:
|
||||
# Semantic fusion thresholds (used by fuse_trimap when segmentation is on).
|
||||
seg_core_threshold: float = 0.60
|
||||
seg_loose_threshold: float = 0.08
|
||||
# Trimap mode for the segmentation pipeline (the chroma-only pipeline ignores it):
|
||||
# "directional" (default) -- chroma magnitude + seg + a hue-direction split. In
|
||||
# the chroma-unknown zone a confidently-segmented pixel stays foreground
|
||||
# unless it is displaced toward the background hue: a neutral background-
|
||||
# coloured garment (e.g. a white shirt) is kept, while a background-hued
|
||||
# residual (e.g. blue trapped between hair strands) is left UNKNOWN for
|
||||
# ViTMatte / chroma-suppress to clear.
|
||||
# "seg" -- original fuse_trimap: seg topology + chroma veto, no hue split.
|
||||
# "directional-hard-bg" -- like "directional" but hard-marks background-hued
|
||||
# pixels as background; more aggressive, can eat cool/shadowed white cloth.
|
||||
mode: str = "directional"
|
||||
seg_low: float = 0.15 # seg below this -> background-eligible; at/above -> foreground-eligible
|
||||
bg_hue_proj_min: float = 4.0 # Lab a*/b* projection onto bg direction above which a pixel is background-hued
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -25,6 +25,15 @@ trimap:
|
||||
fg_safe_radius_ratio: 0.006
|
||||
seg_core_threshold: 0.60
|
||||
seg_loose_threshold: 0.08
|
||||
# Trimap mode (segmentation pipeline only; also settable via --trimap-mode):
|
||||
# directional = chroma + seg + hue split; background-hued residual
|
||||
# (blue between hair) -> unknown, white shirt kept. (default)
|
||||
# seg = original fuse_trimap (no hue split).
|
||||
# directional-hard-bg = hard-remove background-hued pixels (aggressive; can eat
|
||||
# cool/shadowed white cloth).
|
||||
mode: directional
|
||||
seg_low: 0.15
|
||||
bg_hue_proj_min: 4.0
|
||||
|
||||
alpha_post:
|
||||
chroma_suppress: true
|
||||
|
||||
Reference in New Issue
Block a user