BGfilter/bgfilter/settings.py

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ChromaSettings:
    green_margin: float = 0.08
    min_green_value: float = 0.22
    min_saturation: float = 0.20
    hue_low: float = 35.0
    hue_high: float = 95.0
    border_ratio: float = 0.04
    min_samples: int = 2048
    hue_sigma_min: float = 7.0
    lab_sigma_min: float = 10.0
    rgb_sigma_min: float = 0.08


@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 (green) 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


@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
    green_excess_margin: float = 0.015
    edge_expand_radius: int = 2
    alpha_weight_floor: float = 0.35
    # Custom-screen despill (when screen_color is set): margin, in Lab a*/b*
    # units, of chroma allowed along the background colour direction before it is
    # treated 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-small-composition-1k"
    device: str = "cuda"
    matting_method: str = "vitmatte"
    fallback_to_chroma_alpha: bool = False


@dataclass(frozen=True)
class SegmentationSettings:
    enabled: bool = True       # False = chroma-only pipeline (no seg model)
    backend: str = "anime-seg"  # "anime-seg" or "birefnet"
    model_name: str = "skytnt/anime-seg"
    device: str = "cuda"
    input_size: int = 1024


@dataclass(frozen=True)
class PipelineSettings:
    chroma: ChromaSettings = ChromaSettings()
    trimap: TrimapSettings = TrimapSettings()
    alpha_post: AlphaPostSettings = AlphaPostSettings()
    foreground: ForegroundSettings = ForegroundSettings()
    despill: DespillSettings = DespillSettings()
    model: ModelSettings = ModelSettings()
    segmentation: SegmentationSettings = SegmentationSettings()
    # Background-colour prior. None = auto-detect a green screen (current default,
    # unchanged). A hex string like "#CFEFFF" keys chroma + despill off that colour
    # instead, for arbitrary (e.g. pastel) flat backgrounds.
    screen_color: str | None = None