288 lines
9.0 KiB
Python
Executable File
288 lines
9.0 KiB
Python
Executable File
"""绿幕抠像配置领域模型 — 纯逻辑，无FFmpeg依赖.
|
||
|
||
抽离自 chroma_key_engine.py，包含：
|
||
- ChromaKeyConfig 数据类（解析/钳制/效果判断）
|
||
- 预设配置（绿幕/蓝幕/红幕等）
|
||
- 颜色归一化
|
||
- colorkey / chromakey 滤镜构建
|
||
- 便捷函数（apply_chroma_key_if_needed）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ── 预设配置 ──────────────────────────────────────────────────────────────────
|
||
|
||
# 常见绿幕/蓝幕预设
|
||
CHROMA_KEY_PRESETS: dict[str, dict[str, Any]] = {
|
||
"green_screen": {
|
||
"key_color": "#00FF00",
|
||
"similarity": 0.3,
|
||
"blend": 0.1,
|
||
"spill_suppress": 0.5,
|
||
},
|
||
"blue_screen": {
|
||
"key_color": "#0000FF",
|
||
"similarity": 0.3,
|
||
"blend": 0.1,
|
||
"spill_suppress": 0.5,
|
||
},
|
||
"red_screen": {
|
||
"key_color": "#FF0000",
|
||
"similarity": 0.3,
|
||
"blend": 0.1,
|
||
"spill_suppress": 0.0,
|
||
},
|
||
"precise_green": {
|
||
"key_color": "#00FF00",
|
||
"similarity": 0.2,
|
||
"blend": 0.05,
|
||
"spill_suppress": 0.3,
|
||
},
|
||
"soft_green": {
|
||
"key_color": "#00FF00",
|
||
"similarity": 0.45,
|
||
"blend": 0.2,
|
||
"spill_suppress": 0.5,
|
||
},
|
||
}
|
||
|
||
VALID_PRESETS = set(CHROMA_KEY_PRESETS.keys())
|
||
|
||
# 参数范围
|
||
MIN_SIMILARITY = 0.01
|
||
MAX_SIMILARITY = 1.0
|
||
MIN_BLEND = 0.0
|
||
MAX_BLEND = 1.0
|
||
MIN_SPILL_SUPPRESS = 0.0
|
||
MAX_SPILL_SUPPRESS = 1.0
|
||
|
||
# 默认值
|
||
DEFAULT_KEY_COLOR = "#00FF00"
|
||
DEFAULT_SIMILARITY = 0.3
|
||
DEFAULT_BLEND = 0.1
|
||
DEFAULT_SPILL_SUPPRESS = 0.0
|
||
|
||
|
||
# ── 配置模型 ──────────────────────────────────────────────────────────────────
|
||
|
||
|
||
@dataclass
|
||
class ChromaKeyConfig:
|
||
"""绿幕抠像配置.
|
||
|
||
Attributes:
|
||
enabled: 是否启用抠像
|
||
key_color: 要抠除的颜色，支持 hex 格式（如 "#00FF00"）或颜色名
|
||
similarity: 颜色相似度阈值 0.01~1.0，值越大抠除范围越大
|
||
blend: 边缘平滑/混合度 0.0~1.0，值越大边缘越柔和
|
||
spill_suppress: 溢色抑制 0.0~1.0，减少边缘的绿幕反光
|
||
"""
|
||
|
||
enabled: bool = False
|
||
key_color: str = DEFAULT_KEY_COLOR
|
||
similarity: float = DEFAULT_SIMILARITY
|
||
blend: float = DEFAULT_BLEND
|
||
spill_suppress: float = DEFAULT_SPILL_SUPPRESS
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict[str, Any] | None) -> ChromaKeyConfig:
|
||
"""从字典解析配置，参数越界自动钳制."""
|
||
if not data or not data.get("enabled", False):
|
||
return cls(enabled=False)
|
||
|
||
key_color = str(data.get("key_color", DEFAULT_KEY_COLOR)).strip()
|
||
|
||
def _safe_float(val: Any, default: float) -> float:
|
||
try:
|
||
return float(val)
|
||
except (TypeError, ValueError):
|
||
return default
|
||
|
||
similarity = _safe_float(data.get("similarity", DEFAULT_SIMILARITY), DEFAULT_SIMILARITY)
|
||
blend = _safe_float(data.get("blend", DEFAULT_BLEND), DEFAULT_BLEND)
|
||
spill_suppress = _safe_float(data.get("spill_suppress", DEFAULT_SPILL_SUPPRESS), DEFAULT_SPILL_SUPPRESS)
|
||
|
||
# 钳制到合法范围
|
||
similarity = max(MIN_SIMILARITY, min(MAX_SIMILARITY, similarity))
|
||
blend = max(MIN_BLEND, min(MAX_BLEND, blend))
|
||
spill_suppress = max(MIN_SPILL_SUPPRESS, min(MAX_SPILL_SUPPRESS, spill_suppress))
|
||
|
||
return cls(
|
||
enabled=True,
|
||
key_color=key_color,
|
||
similarity=similarity,
|
||
blend=blend,
|
||
spill_suppress=spill_suppress,
|
||
)
|
||
|
||
@classmethod
|
||
def from_preset(cls, preset_name: str) -> ChromaKeyConfig | None:
|
||
"""从预设名称创建配置."""
|
||
preset = CHROMA_KEY_PRESETS.get(preset_name)
|
||
if not preset:
|
||
return None
|
||
return cls(
|
||
enabled=True,
|
||
key_color=preset["key_color"],
|
||
similarity=preset["similarity"],
|
||
blend=preset["blend"],
|
||
spill_suppress=preset["spill_suppress"],
|
||
)
|
||
|
||
def has_effect(self) -> bool:
|
||
"""判断是否有实际抠像效果."""
|
||
return self.enabled and self.similarity > 0
|
||
|
||
def validate(self) -> tuple[bool, str]:
|
||
"""校验配置是否有效."""
|
||
if not self.enabled:
|
||
return True, ""
|
||
|
||
if not self.key_color:
|
||
return False, "key_color 不能为空"
|
||
|
||
if not (MIN_SIMILARITY <= self.similarity <= MAX_SIMILARITY):
|
||
return False, f"similarity 必须在 {MIN_SIMILARITY}~{MAX_SIMILARITY} 之间"
|
||
|
||
if not (MIN_BLEND <= self.blend <= MAX_BLEND):
|
||
return False, f"blend 必须在 {MIN_BLEND}~{MAX_BLEND} 之间"
|
||
|
||
if not (MIN_SPILL_SUPPRESS <= self.spill_suppress <= MAX_SPILL_SUPPRESS):
|
||
return False, f"spill_suppress 必须在 {MIN_SPILL_SUPPRESS}~{MAX_SPILL_SUPPRESS} 之间"
|
||
|
||
return True, ""
|
||
|
||
|
||
# ── 颜色归一化 ────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def normalize_color(color_str: str) -> str:
|
||
"""将颜色字符串转为 FFmpeg colorkey 接受的格式.
|
||
|
||
支持:
|
||
- "#RRGGBB" / "#RRGGBBAA" → 0xRRGGBB
|
||
- "0xRRGGBB" → 直接使用
|
||
- 颜色名（green/blue/red/black/white 等）→ 直接透传
|
||
"""
|
||
color = color_str.strip()
|
||
|
||
# hex 格式
|
||
hex_match = re.match(r"^#?([0-9a-fA-F]{6})([0-9a-fA-F]{2})?$", color)
|
||
if hex_match:
|
||
return f"0x{hex_match.group(1).upper()}"
|
||
|
||
# 已经是 0x 格式
|
||
if color.lower().startswith("0x"):
|
||
return color.upper()
|
||
|
||
# 颜色名直接透传（FFmpeg 支持常见颜色名）
|
||
return color
|
||
|
||
|
||
# ── 滤镜构建 ────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def build_colorkey_filter(
|
||
config: ChromaKeyConfig,
|
||
input_label: str,
|
||
output_label: str,
|
||
) -> str:
|
||
"""构建 colorkey 滤镜字符串.
|
||
|
||
Args:
|
||
config: 抠像配置
|
||
input_label: 输入标签，如 "[0:v]" 或 "[v0]"
|
||
output_label: 输出标签，如 "[ck0]"
|
||
|
||
Returns:
|
||
FFmpeg 滤镜字符串
|
||
"""
|
||
if not config.has_effect():
|
||
return f"{input_label}copy{output_label}"
|
||
|
||
color = normalize_color(config.key_color)
|
||
similarity = config.similarity
|
||
blend = config.blend
|
||
|
||
# 基础 colorkey 滤镜
|
||
parts = [f"colorkey=color={color}:similarity={similarity}:blend={blend}"]
|
||
|
||
# 溢色抑制（通过 colorchannelmixer 降低绿色通道增益）
|
||
if config.spill_suppress > 0:
|
||
spill = config.spill_suppress
|
||
g_gain = max(0.3, 1.0 - spill * 0.7)
|
||
r_gain = 1.0 + spill * 0.15
|
||
b_gain = 1.0 + spill * 0.15
|
||
parts.append(f"colorchannelmixer=rr={r_gain}:gg={g_gain}:bb={b_gain}:aa=1")
|
||
|
||
return f"{input_label}{','.join(parts)}{output_label}"
|
||
|
||
|
||
def build_chromakey_filter(
|
||
config: ChromaKeyConfig,
|
||
input_label: str,
|
||
output_label: str,
|
||
) -> str:
|
||
"""使用 chromakey 滤镜（更高级的版本，支持更多参数）.
|
||
|
||
注意：并非所有 FFmpeg 版本都支持 chromakey 滤镜，
|
||
优先使用 colorkey（兼容性更好）。
|
||
"""
|
||
if not config.has_effect():
|
||
return f"{input_label}copy{output_label}"
|
||
|
||
color = normalize_color(config.key_color)
|
||
similarity = config.similarity
|
||
blend = config.blend
|
||
|
||
return f"{input_label}chromakey=color={color}:similarity={similarity}:blend={blend}{output_label}"
|
||
|
||
|
||
# ── 工具函数 ────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def apply_chroma_key_if_needed(
|
||
clip_config: dict[str, Any] | None,
|
||
input_label: str,
|
||
output_label: str,
|
||
) -> str | None:
|
||
"""便捷函数：根据 clip 配置判断是否需要应用绿幕抠像.
|
||
|
||
Args:
|
||
clip_config: clip 的 config 字典
|
||
input_label: 输入标签
|
||
output_label: 输出标签
|
||
|
||
Returns:
|
||
滤镜字符串，不需要抠像时返回 None
|
||
"""
|
||
if not clip_config:
|
||
return None
|
||
|
||
chroma_key_data = clip_config.get("chroma_key")
|
||
if not chroma_key_data:
|
||
return None
|
||
|
||
try:
|
||
config = ChromaKeyConfig.from_dict(chroma_key_data)
|
||
if not config.has_effect():
|
||
return None
|
||
|
||
return build_colorkey_filter(config, input_label, output_label)
|
||
except Exception as e:
|
||
logger.warning("[chroma-key] 应用抠像失败，跳过: %s", e)
|
||
return None
|
||
|
||
|
||
def get_preset_names() -> list[str]:
|
||
"""获取所有预设名称列表."""
|
||
return sorted(list(CHROMA_KEY_PRESETS.keys()))
|