lhk229
711669ea51
Wire BiRefNet into the pipeline (lazy-loaded) and build the trimap from the semantic subject mask refined by the green key, used when segmentation is enabled. generate_trimap (chroma-only) stays as the no-segmentation fallback. Authority split: the mask decides subject topology (keeps green-tinted hair as foreground, drops see-through holes), the bright-green key sharpens the screen boundary, ViTMatte refines a small protective band. Background follows the mask directly so interior holes survive; the band is small so it does not swallow finger gaps. End-to-end on the samples: green hair preserved, finger-gap green 99% transparent (BiRefNet marks it background; was 100% foreground before), leg edges smooth -- topology/colour/edge problems local cues could not solve, with no hard alpha surgery. - SegmentationSettings wired into PipelineSettings/config; --device propagates. - TrimapSettings: seg_core_threshold, seg_loose_threshold. - New debug output seg_mask.png. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from .deps import require_cv2
|
|
from .settings import TrimapSettings
|
|
|
|
|
|
def radius_from_ratio(shape: tuple[int, int], ratio: float, minimum: int) -> int:
|
|
return max(minimum, int(round(max(shape) * ratio)))
|
|
|
|
|
|
def elliptical_kernel(radius: int) -> np.ndarray:
|
|
cv2 = require_cv2()
|
|
size = radius * 2 + 1
|
|
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))
|
|
|
|
|
|
def generate_trimap(
|
|
bg_confidence: np.ndarray, settings: TrimapSettings
|
|
) -> tuple[np.ndarray, dict[str, int]]:
|
|
cv2 = require_cv2()
|
|
shape = bg_confidence.shape
|
|
unknown_radius = radius_from_ratio(
|
|
shape, settings.unknown_radius_ratio, settings.min_unknown_radius
|
|
)
|
|
fg_safe_radius = radius_from_ratio(
|
|
shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius
|
|
)
|
|
|
|
sure_bg = bg_confidence >= settings.sure_bg_threshold
|
|
low_bg = bg_confidence <= settings.sure_fg_threshold
|
|
|
|
bg_u8 = sure_bg.astype(np.uint8)
|
|
unknown_band = cv2.dilate(bg_u8, elliptical_kernel(unknown_radius)).astype(bool)
|
|
fg_safe = ~cv2.dilate(bg_u8, elliptical_kernel(fg_safe_radius)).astype(bool)
|
|
sure_fg = low_bg & fg_safe
|
|
|
|
trimap = np.full(shape, 128, dtype=np.uint8)
|
|
trimap[sure_bg] = 0
|
|
trimap[sure_fg] = 255
|
|
|
|
# Keep a protective unknown band around all sure background, including holes.
|
|
trimap[unknown_band & ~sure_bg & ~sure_fg] = 128
|
|
|
|
stats = {
|
|
"sure_bg_pixels": int((trimap == 0).sum()),
|
|
"unknown_pixels": int((trimap == 128).sum()),
|
|
"sure_fg_pixels": int((trimap == 255).sum()),
|
|
"unknown_radius": int(unknown_radius),
|
|
"fg_safe_radius": int(fg_safe_radius),
|
|
}
|
|
return trimap, stats
|
|
|
|
|
|
def fuse_trimap(
|
|
seg_mask: np.ndarray, bg_confidence: np.ndarray, settings: TrimapSettings
|
|
) -> tuple[np.ndarray, dict[str, int]]:
|
|
"""Build a trimap from a semantic subject mask, refined by the green key.
|
|
|
|
Authority split: the segmentation mask decides subject *topology* (it keeps
|
|
colour-contaminated hair as foreground and drops see-through holes), while the
|
|
bright-green chroma key sharpens the screen boundary. ViTMatte then refines
|
|
the unknown band.
|
|
"""
|
|
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)
|
|
|
|
green_bg = bg_confidence >= settings.sure_bg_threshold # bright clean screen
|
|
core = (seg_mask >= settings.seg_core_threshold).astype(np.uint8)
|
|
loose = (seg_mask >= settings.seg_loose_threshold).astype(np.uint8)
|
|
|
|
# A small protective band around the mask boundary lets ViTMatte anti-alias
|
|
# crisp edges; hair gets extra width from BiRefNet's own soft region. The band
|
|
# is deliberately small so it does not swallow small holes (finger gaps).
|
|
kernel = elliptical_kernel(band_radius)
|
|
band = cv2.dilate(loose, kernel).astype(bool) & ~cv2.erode(loose, kernel).astype(bool)
|
|
|
|
# Background follows the mask directly so interior holes stay background, minus
|
|
# the boundary band; foreground is the confident subject, never bright screen.
|
|
sure_fg = core.astype(bool) & ~band & ~green_bg
|
|
sure_bg = ((loose == 0) & ~band) | (green_bg & (loose == 0))
|
|
|
|
trimap = np.full(shape, 128, dtype=np.uint8)
|
|
trimap[sure_bg] = 0
|
|
trimap[sure_fg] = 255
|
|
|
|
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),
|
|
"fg_safe_radius": int(fg_safe_radius),
|
|
"seg_subject_pixels": int((core > 0).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
|
|
alpha[trimap == 255] = 1.0
|
|
return alpha
|