Implement green screen matting pipeline
This commit is contained in:
@@ -0,0 +1,61 @@
|
||||
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
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DespillSettings:
|
||||
enabled: bool = True
|
||||
edge_low: float = 0.02
|
||||
edge_high: float = 0.98
|
||||
strength: float = 0.75
|
||||
green_excess_margin: float = 0.03
|
||||
|
||||
|
||||
@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 PipelineSettings:
|
||||
chroma: ChromaSettings = ChromaSettings()
|
||||
trimap: TrimapSettings = TrimapSettings()
|
||||
alpha_post: AlphaPostSettings = AlphaPostSettings()
|
||||
despill: DespillSettings = DespillSettings()
|
||||
model: ModelSettings = ModelSettings()
|
||||
Reference in New Issue
Block a user