diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 1c30a13..4f02fee 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -28,13 +28,36 @@ class MattingPipeline: self.settings = settings self._runner: ViTMatteRunner | None = None self._segmenter = None + self._fill_segmenter = None def _segment(self, rgb: np.ndarray) -> np.ndarray: if self._segmenter is None: from .segmentation import make_segmenter self._segmenter = make_segmenter(self.settings.segmentation) - return self._segmenter.mask(rgb) + mask = self._segmenter.mask(rgb) + seg = self.settings.segmentation + if seg.sharpen > 0: + from .segmentation import unsharp_mask + + mask = unsharp_mask(mask, seg.sharpen, seg.sharpen_radius) + return mask + + def _segment_fill(self, rgb: np.ndarray, default_mask: np.ndarray) -> np.ndarray: + """Mask for the dropout-fill; may use a different (higher-recall) backend + than the trimap, since the fill only adds and cannot break anything.""" + seg = self.settings.segmentation + if not seg.fill_backend or seg.fill_backend == seg.backend: + return default_mask + if self._fill_segmenter is None: + from dataclasses import replace + + from .segmentation import make_segmenter + + self._fill_segmenter = make_segmenter( + replace(seg, backend=seg.fill_backend, model_name=seg.fill_model_name) + ) + return self._fill_segmenter.mask(rgb) def _predict_alpha(self, rgb: np.ndarray, trimap: np.ndarray, bg_confidence: np.ndarray) -> tuple[np.ndarray, str]: if self.settings.model.matting_method == "chroma": @@ -93,7 +116,8 @@ def _run_image( alpha = suppress_alpha_by_chroma(alpha, bg_confidence, trimap, settings.alpha_post) alpha = clean_alpha(alpha, trimap, settings.alpha_post) if seg_mask is not None: - alpha = fill_seg_dropouts(alpha, seg_mask, settings.alpha_post) + fill_mask = pipeline._segment_fill(rgb, seg_mask) + alpha = fill_seg_dropouts(alpha, fill_mask, settings.alpha_post) foreground = estimate_foreground_rgb( rgb, alpha, bg_confidence, model, settings.foreground ) diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py index e6269bf..fbfcae5 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -111,6 +111,19 @@ class AnimeSegSegmenter: return np.asarray(mask_img, dtype=np.float32) / 255.0 +def unsharp_mask(mask: np.ndarray, amount: float, radius: int) -> np.ndarray: + """Sharpen a soft mask: raise thin-strand peaks and crisp the blur so strands + reach the sure-foreground threshold instead of dissolving into the unknown band.""" + if amount <= 0: + return mask + from .deps import require_cv2 + + cv2 = require_cv2() + k = max(1, radius) * 2 + 1 + blur = cv2.GaussianBlur(mask, (k, k), 0) + return np.clip(mask + amount * (mask - blur), 0.0, 1.0).astype(np.float32) + + def make_segmenter(settings: SegmentationSettings): if settings.backend == "birefnet": return BiRefNetSegmenter(settings) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 6c27e6e..03988d5 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -106,8 +106,12 @@ class SegmentationSettings: enabled: bool = True backend: str = "birefnet" # "birefnet" or "anime-seg" model_name: str = "ZhengPeng7/BiRefNet" + fill_backend: str = "" # "" = reuse `backend` for dropout-fill; else e.g. "anime-seg" + fill_model_name: str = "" device: str = "cuda" input_size: int = 1024 + sharpen: float = 0.0 # unsharp-mask amount on the mask (0 = off) + sharpen_radius: int = 4 @dataclass(frozen=True) diff --git a/configs/default.yaml b/configs/default.yaml index abae22e..3fc21e9 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -8,6 +8,8 @@ segmentation: enabled: true backend: birefnet # birefnet | anime-seg (anime-seg needs model_name: skytnt/anime-seg) model_name: ZhengPeng7/BiRefNet + fill_backend: anime-seg # dropout-fill uses anime-seg (higher wisp recall; add-only, harmless) + fill_model_name: skytnt/anime-seg device: cuda input_size: 1024