lhk229
5f48f7a2cd
Replace the two experimental bools (directional, bg_hued_to_bg) with a single
TrimapSettings.mode selector for the segmentation pipeline:
- "directional" (new default): chroma magnitude + seg + a Lab hue-direction
split. In the chroma-unknown zone a confidently-segmented pixel stays
foreground unless it is displaced toward the background hue, so a neutral
background-coloured garment (e.g. a white shirt) is kept while a background-
hued residual (blue between hair strands) is left unknown for ViTMatte /
chroma-suppress to clear.
- "seg": original fuse_trimap, unchanged.
- "directional-hard-bg": aggressive variant that hard-removes background-hued
pixels (can eat cool/shadowed white cloth).
Selectable via configs/default.yaml (trimap.mode) or CLI --trimap-mode; unknown
modes raise. Default CLI output verified byte-identical to the reviewed UNK
result on the pastel sample.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
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")
|
|
|
|
|
|
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"]
|
|
|
|
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,
|
|
)
|