BGfilter/bgfilter/trimap.py

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 chroma key.

    Authority split: the segmentation mask decides subject *topology* (it keeps
    colour-contaminated hair as foreground and drops see-through holes), while the
    chroma key sharpens the flat-background 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)

    screen_bg = bg_confidence >= settings.sure_bg_threshold      # confident flat background
    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 the background.
    sure_fg = core.astype(bool) & ~band & ~screen_bg
    sure_bg = ((loose == 0) & ~band) | (screen_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 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
    alpha[trimap == 255] = 1.0
    return alpha