"""绿幕抠像引擎 — 基于 FFmpeg colorkey / chromakey 滤镜. 支持将指定颜色(默认绿色)变为透明,可用于虚拟背景、画中画背景替换等场景。 使用方式: config = ChromaKeyConfig(key_color="#00FF00", similarity=0.3, blend=0.1) engine = ChromaKeyEngine(config) filter_str = engine.build_filter(input_label, output_label) # 结果: [in]colorkey=color=0x00FF00:similarity=0.3:blend=0.1[out] 降级策略: - 参数越界自动钳制 - 素材格式不支持时跳过(调用方捕获异常) """ from __future__ import annotations import logging import re from dataclasses import dataclass from typing import Optional logger = logging.getLogger(__name__) # ── 配置模型 ────────────────────────────────────────────────────────────────── @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 = "#00FF00" similarity: float = 0.3 blend: float = 0.1 spill_suppress: float = 0.0 @classmethod def from_dict(cls, data: dict | None) -> "ChromaKeyConfig": """从字典解析配置,参数越界自动钳制。""" if not data or not data.get("enabled", False): return cls(enabled=False) key_color = str(data.get("key_color", "#00FF00")).strip() def _safe_float(val, default): try: return float(val) except (TypeError, ValueError): return default similarity = _safe_float(data.get("similarity", 0.3), 0.3) blend = _safe_float(data.get("blend", 0.1), 0.1) spill_suppress = _safe_float(data.get("spill_suppress", 0.0), 0.0) # 钳制到合法范围 similarity = max(0.01, min(1.0, similarity)) blend = max(0.0, min(1.0, blend)) spill_suppress = max(0.0, min(1.0, spill_suppress)) return cls( enabled=True, key_color=key_color, similarity=similarity, blend=blend, spill_suppress=spill_suppress, ) def has_effect(self) -> bool: """判断是否有实际抠像效果。""" return self.enabled and self.similarity > 0 # ── 预设配置 ────────────────────────────────────────────────────────────────── # 常见绿幕/蓝幕预设 CHROMA_KEY_PRESETS = { "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, }, } # ── 引擎实现 ────────────────────────────────────────────────────────────────── class ChromaKeyEngine: """绿幕抠像引擎。 基于 FFmpeg colorkey 滤镜实现,将指定颜色变为透明。 适用于绿幕/蓝幕视频的背景去除,配合画中画或 overlay 实现虚拟背景。 """ def __init__(self, config: ChromaKeyConfig): self.config = config @staticmethod 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_filter(self, input_label: str, output_label: str) -> str: """构建 colorkey 滤镜字符串。 Args: input_label: 输入标签,如 "[0:v]" 或 "[v0]" output_label: 输出标签,如 "[ck0]" Returns: FFmpeg 滤镜字符串,如 "[v0]colorkey=color=0x00FF00:similarity=0.3:blend=0.1[ck0]" Raises: ValueError: 配置无效时抛出(调用方应捕获并降级) """ if not self.config.has_effect(): # 无效果,直接直通 return f"{input_label}copy{output_label}" color = self._normalize_color(self.config.key_color) similarity = self.config.similarity blend = self.config.blend # 基础 colorkey 滤镜 parts = [f"colorkey=color={color}:similarity={similarity}:blend={blend}"] # 溢色抑制(通过 colorchannelmixer 降低绿色通道增益) if self.config.spill_suppress > 0: # 降低绿通道增益,减少绿幕反光溢出 spill = self.config.spill_suppress # 绿通道增益 = 1 - spill_factor 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=" f"rr={r_gain}:" f"gg={g_gain}:" f"bb={b_gain}:" f"aa=1") filter_str = f"{input_label}{','.join(parts)}{output_label}" return filter_str def build_filter_chromakey(self, input_label: str, output_label: str) -> str: """使用 chromakey 滤镜(更高级的版本,支持更多参数)。 注意:并非所有 FFmpeg 版本都支持 chromakey 滤镜, 优先使用 colorkey(兼容性更好)。 Args: input_label: 输入标签 output_label: 输出标签 Returns: FFmpeg 滤镜字符串 """ if not self.config.has_effect(): return f"{input_label}copy{output_label}" color = self._normalize_color(self.config.key_color) similarity = self.config.similarity blend = self.config.blend return f"{input_label}" f"chromakey=color={color}:similarity={similarity}:blend={blend}" f"{output_label}" def apply_chroma_key_if_needed( clip_config: dict | None, input_label: str, output_label: str, ) -> Optional[str]: """便捷函数:根据 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 engine = ChromaKeyEngine(config) return engine.build_filter(input_label, output_label) except Exception as e: logger.warning("[chroma-key] 应用抠像失败,跳过: %s", e) return None