from __future__ import annotations from dataclasses import fields, replace from pathlib import Path from typing import Any, TypeVar from .settings import ( AlphaPostSettings, ChromaSettings, 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", "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")), 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"] 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"] despill_updates: dict[str, Any] = {} if overrides.get("despill") is not None: despill_updates["enabled"] = overrides["despill"] seg_updates: dict[str, Any] = {} if overrides.get("device") is not None: seg_updates["device"] = overrides["device"] 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=chroma, trimap=_update_dataclass(trimap, trimap_updates), alpha_post=alpha_post, foreground=foreground, despill=_update_dataclass(despill, despill_updates), model=_update_dataclass(model, model_updates), segmentation=_update_dataclass(segmentation, seg_updates), screen_color=screen_color, )