BGfilter/bgfilter/config.py

from __future__ import annotations

from dataclasses import fields, replace
from pathlib import Path
from typing import Any, TypeVar

from .settings import (
    AlphaPostSettings,
    ChromaSettings,
    CrossCheckSettings,
    DespillSettings,
    ForegroundSettings,
    ModelSettings,
    PipelineSettings,
    SegmentationSettings,
    TrimapSettings,
)

T = TypeVar("T")

# Canonical weights per segmentation backend (selected at runtime via --seg-backend).
_SEG_MODELS = {"birefnet": "ZhengPeng7/BiRefNet", "anime-seg": "skytnt/anime-seg"}


def _load_yaml(path: str | Path) -> dict[str, Any]:
    try:
        import yaml
    except ModuleNotFoundError as exc:
        raise RuntimeError("Missing dependency 'PyYAML'. Install it with: pip install PyYAML") from exc

    with Path(path).open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle) or {}
    if not isinstance(data, dict):
        raise ValueError(f"Config file must contain a YAML mapping: {path}")
    return data


def _update_dataclass(instance: T, values: dict[str, Any] | None) -> T:
    if not values:
        return instance
    allowed = {field.name for field in fields(instance)}
    unknown = sorted(set(values) - allowed)
    if unknown:
        cls_name = type(instance).__name__
        raise ValueError(f"Unknown {cls_name} field(s): {', '.join(unknown)}")
    return replace(instance, **values)


def settings_from_dict(data: dict[str, Any]) -> PipelineSettings:
    data = dict(data)
    screen_color = data.pop("screen_color", None)
    allowed_sections = {"chroma", "trimap", "alpha_post", "cross_check", "foreground", "despill", "model", "segmentation"}
    unknown_sections = sorted(set(data) - allowed_sections)
    if unknown_sections:
        raise ValueError(f"Unknown config section(s): {', '.join(unknown_sections)}")

    return PipelineSettings(
        chroma=_update_dataclass(ChromaSettings(), data.get("chroma")),
        trimap=_update_dataclass(TrimapSettings(), data.get("trimap")),
        alpha_post=_update_dataclass(AlphaPostSettings(), data.get("alpha_post")),
        cross_check=_update_dataclass(CrossCheckSettings(), data.get("cross_check")),
        foreground=_update_dataclass(ForegroundSettings(), data.get("foreground")),
        despill=_update_dataclass(DespillSettings(), data.get("despill")),
        model=_update_dataclass(ModelSettings(), data.get("model")),
        segmentation=_update_dataclass(SegmentationSettings(), data.get("segmentation")),
        screen_color=screen_color,
    )


def load_settings(path: str | Path | None = None) -> PipelineSettings:
    if path is None:
        return PipelineSettings()
    return settings_from_dict(_load_yaml(path))


def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineSettings:
    chroma = settings.chroma
    trimap = settings.trimap
    alpha_post = settings.alpha_post
    foreground = settings.foreground
    despill = settings.despill
    model = settings.model
    segmentation = settings.segmentation

    model_updates = {
        key: overrides[key]
        for key in ["model_name", "device", "matting_method", "fallback_to_chroma_alpha", "precision"]
        if overrides.get(key) is not None
    }
    trimap_updates = {
        key: overrides[key]
        for key in [
            "sure_bg_threshold",
            "sure_fg_threshold",
            "unknown_radius_ratio",
            "fg_safe_radius_ratio",
        ]
        if overrides.get(key) is not None
    }
    if overrides.get("trimap_mode") is not None:
        trimap_updates["mode"] = overrides["trimap_mode"]
    chroma_updates: dict[str, Any] = {}
    if overrides.get("chroma") is not None:
        chroma_updates["enabled"] = overrides["chroma"]
    despill_updates: dict[str, Any] = {}
    if overrides.get("despill") is not None:
        despill_updates["enabled"] = overrides["despill"]
    cross_check_updates: dict[str, Any] = {}
    if overrides.get("cross_check") is not None:
        cross_check_updates["enabled"] = overrides["cross_check"]
    if overrides.get("cross_check_as_seg") is not None:
        cross_check_updates["reuse_as_seg"] = overrides["cross_check_as_seg"]
    seg_updates: dict[str, Any] = {}
    if overrides.get("device") is not None:
        seg_updates["device"] = overrides["device"]
    # --precision is a single knob for all three models (ViTMatte handled via
    # model_updates above; per-model values remain settable in the config file).
    if overrides.get("precision") is not None:
        seg_updates["precision"] = overrides["precision"]
        cross_check_updates["precision"] = overrides["precision"]
    if overrides.get("seg_backend") is not None:
        backend = overrides["seg_backend"]
        if backend not in _SEG_MODELS:
            raise ValueError(
                f"Unknown --seg-backend '{backend}'. Use one of: {', '.join(_SEG_MODELS)}."
            )
        seg_updates["backend"] = backend
        seg_updates["model_name"] = _SEG_MODELS[backend]

    screen_color = (
        overrides["screen_color"]
        if overrides.get("screen_color") is not None
        else settings.screen_color
    )
    return PipelineSettings(
        chroma=_update_dataclass(chroma, chroma_updates),
        trimap=_update_dataclass(trimap, trimap_updates),
        alpha_post=alpha_post,
        cross_check=_update_dataclass(settings.cross_check, cross_check_updates),
        foreground=foreground,
        despill=_update_dataclass(despill, despill_updates),
        model=_update_dataclass(model, model_updates),
        segmentation=_update_dataclass(segmentation, seg_updates),
        screen_color=screen_color,
    )