BGfilter/bgfilter/settings.py

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ChromaSettings:
    border_ratio: float = 0.04
    min_samples: int = 2048
    lab_sigma_min: float = 10.0
    rgb_sigma_min: float = 0.08
    # Background-colour auto-detection (used when screen_color is None): find the
    # flat colour that dominates the image border; raise if none is clean enough.
    detect_cluster_radius: float = 12.0    # Lab radius grouping border pixels into the background cluster
    detect_min_border_share: float = 0.55  # min fraction of the border that must be this one colour
    detect_corner_tol: float = 20.0        # max Lab distance a corner may sit from the cluster centre


@dataclass(frozen=True)
class TrimapSettings:
    sure_bg_threshold: float = 0.92
    sure_fg_threshold: float = 0.12
    unknown_radius_ratio: float = 0.012
    fg_safe_radius_ratio: float = 0.006
    min_unknown_radius: int = 4
    min_fg_safe_radius: int = 2
    # 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)
class AlphaPostSettings:
    min_component_area_ratio: float = 0.00001
    fill_hole_area_ratio: float = 0.00002
    alpha_floor: float = 0.002
    alpha_ceil: float = 0.998
    # Chroma-guided alpha suppression: in the unknown band, pull alpha toward 0
    # where background-colour confidence is high, clearing spill in gaps/holes.
    chroma_suppress: bool = True
    chroma_suppress_bg_low: float = 0.35
    chroma_suppress_bg_high: float = 0.80
    chroma_suppress_strength: float = 1.0
    # Matte-confidence gate: fade the suppression where ViTMatte itself is confident
    # the pixel is opaque (raw alpha ramps lo -> hi), so colour evidence only vetoes
    # the matte where the matte is unsure. Protects background-coloured subjects
    # (white shirt on pastel blue, blush on pink) that chroma alone cannot tell from
    # background residue. Set suppress_raw_lo >= 1.0 to disable the gate.
    suppress_raw_lo: float = 0.85
    suppress_raw_hi: float = 0.98


@dataclass(frozen=True)
class CrossCheckSettings:
    # Cross-model veto: a trimap-free matting model gives a second opinion that
    # may only LOWER alpha (min-fusion), only on background-hued, non-dark pixels
    # (the suspect zone) where the primary result is confident (gate_lo->gate_hi
    # alpha ramp). Clears colour-drifted background residue between hair strands
    # that chroma, segmentation and the primary matte all read as foreground;
    # already-soft wisps and dark hair are exempt by construction.
    enabled: bool = True
    model_name: str = "ZhengPeng7/BiRefNet_HR-matting"
    input_size: int = 2048
    # Reuse the cross-check forward as the segmentation mask: one HR-matting
    # inference serves both trimap topology and the second opinion, skipping the
    # primary seg model entirely (saves its load + forward). Pilot-validated on
    # TestImage3/FixImage1: trimap 99.8% identical, no structural change. Only
    # takes effect when both cross_check and segmentation are enabled.
    reuse_as_seg: bool = True
    # Compute precision of the cross-check forward: "fp32" or "bf16" (autocast;
    # needs bf16-capable hardware, else falls back to fp32 with a warning).
    precision: str = "fp32"
    proj_min: float = 3.0       # bg-hue projection above which a pixel is suspect
    l_min: float = 45.0         # Lab lightness below which a pixel is exempt (dark hair)
    feather_sigma: float = 2.0  # Gaussian feather of the zone boundary, in px
    gate_lo: float = 0.70       # primary alpha below this -> fully exempt
    gate_hi: float = 0.95       # primary alpha above this -> fully vetoable


@dataclass(frozen=True)
class ForegroundSettings:
    enabled: bool = True
    method: str = "ml"  # "ml" (pymatting) or "unmix" (legacy heuristic)
    fallback_to_unmix: bool = True
    ml_regularization: float = 1e-5
    # Legacy "unmix" method parameters (used only when method == "unmix").
    edge_low: float = 0.005
    edge_high: float = 0.995
    min_unmix_alpha: float = 0.08
    unmix_strength: float = 0.70
    local_strength: float = 0.45
    local_blur_radius: int = 11
    local_alpha_threshold: float = 0.92
    local_bg_confidence_max: float = 0.20
    green_excess_margin: float = 0.015
    bg_confidence_weight: float = 0.70


@dataclass(frozen=True)
class DespillSettings:
    enabled: bool = True
    edge_low: float = 0.005
    strength: float = 0.92
    edge_expand_radius: int = 2
    alpha_weight_floor: float = 0.35
    # Margin, in Lab a*/b* units, of chroma allowed along the background-colour
    # direction before it counts as spill. Larger keeps more of the subject's own
    # colour.
    color_margin_lab: float = 4.0


@dataclass(frozen=True)
class ModelSettings:
    model_name: str = "hustvl/vitmatte-base-composition-1k"
    device: str = "cpu"
    matting_method: str = "vitmatte"
    fallback_to_chroma_alpha: bool = False
    # ViTMatte compute precision: "fp32" or "bf16". bf16 halves the matting
    # model's activation memory (its full-resolution attention is the pipeline's
    # memory peak) with visually identical alpha, but needs bf16-capable
    # hardware (any modern GPU, or a CPU with AVX512-BF16/AMX) — otherwise it
    # falls back to fp32 with a warning. Segmentation always stays fp32.
    precision: str = "fp32"


@dataclass(frozen=True)
class SegmentationSettings:
    enabled: bool = True       # False = chroma-only pipeline (no seg model)
    backend: str = "birefnet"  # "birefnet" (general) or "anime-seg" (anime characters)
    model_name: str = "ZhengPeng7/BiRefNet"
    device: str = "cpu"
    input_size: int = 1024
    # Compute precision of the BiRefNet forward: "fp32" or "bf16" (autocast; needs
    # bf16-capable hardware, else falls back to fp32 with a warning). The ONNX
    # anime-seg backend ignores this.
    precision: str = "fp32"


@dataclass(frozen=True)
class PipelineSettings:
    chroma: ChromaSettings = ChromaSettings()
    trimap: TrimapSettings = TrimapSettings()
    alpha_post: AlphaPostSettings = AlphaPostSettings()
    cross_check: CrossCheckSettings = CrossCheckSettings()
    foreground: ForegroundSettings = ForegroundSettings()
    despill: DespillSettings = DespillSettings()
    model: ModelSettings = ModelSettings()
    segmentation: SegmentationSettings = SegmentationSettings()
    # Background-colour prior. None = auto-detect the flat background colour from the
    # image border (raises if there is no clean flat background). A hex string like
    # "#CFEFFF" sets it explicitly.
    screen_color: str | None = None