xiaoxia-saas/apps/worker/video_processing/chroma_key_engine.py

"""绿幕抠像引擎 — 基于 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