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

"""音频降噪引擎 — 基于 FFmpeg afftdn 滤镜.

支持对音频进行背景噪音消除、人声增强，适用于语音录制、采访等场景。

使用方式：
    config = NoiseReductionConfig(level="medium")
    engine = NoiseReductionEngine(config)
    filter_str = engine.build_filter(input_label, output_label)
    # 结果: [0:a]afftdn=nf=-25[out]

降级策略：
    - 参数越界自动钳制
    - FFmpeg 不支持 afftdn 时，调用方可捕获异常并跳过
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import Enum
from typing import Optional

logger = logging.getLogger(__name__)


# ── 降噪等级 ──────────────────────────────────────────────────────────────────


class NoiseReductionLevel(str, Enum):
    """降噪等级预设。"""

    LOW = "low"  # 轻度降噪，保留细节，适合轻微背景噪音
    MEDIUM = "medium"  # 中度降噪，平衡效果和音质
    HIGH = "high"  # 高度降噪，适合嘈杂环境，可能轻微影响音质
    CUSTOM = "custom"  # 自定义参数


# 各等级对应的降噪参数（afftdn 的 noise floor，单位 dB）
# 值越大（越接近 0），降噪越强；值越小（越负），降噪越弱
_LEVEL_PARAMS = {
    NoiseReductionLevel.LOW: {
        "nf": -35,  # 噪音阈值（dB），越负越保守
        "tn": -10,  # 噪音频谱平滑度
        "tr": 50,  # 时间分辨率（ms）
    },
    NoiseReductionLevel.MEDIUM: {
        "nf": -25,
        "tn": -10,
        "tr": 50,
    },
    NoiseReductionLevel.HIGH: {
        "nf": -15,
        "tn": -5,
        "tr": 30,
    },
}


# ── 配置模型 ──────────────────────────────────────────────────────────────────


@dataclass
class NoiseReductionConfig:
    """音频降噪配置。

    Attributes:
        enabled: 是否启用降噪
        level: 降噪等级 low/medium/high/custom
        noise_floor: 自定义噪音阈值（dB），仅 level=custom 时有效，范围 -60 ~ -5
        voice_enhance: 是否启用人声增强
        output_format: 输出格式描述（内部使用）
    """

    enabled: bool = False
    level: NoiseReductionLevel = NoiseReductionLevel.MEDIUM
    noise_floor: float = -25.0  # dB
    voice_enhance: bool = False

    @classmethod
    def from_dict(cls, data: dict | None) -> "NoiseReductionConfig":
        """从字典解析配置，参数越界自动钳制。"""
        if not data or not data.get("enabled", False):
            return cls(enabled=False)

        level_str = str(data.get("level", "medium")).lower()
        try:
            level = NoiseReductionLevel(level_str)
        except ValueError:
            level = NoiseReductionLevel.MEDIUM

        try:
            noise_floor = float(data.get("noise_floor", -25.0))
        except (TypeError, ValueError):
            noise_floor = -25.0

        voice_enhance = bool(data.get("voice_enhance", False))

        # 钳制到合法范围
        noise_floor = max(-60.0, min(-5.0, noise_floor))

        return cls(
            enabled=True,
            level=level,
            noise_floor=noise_floor,
            voice_enhance=voice_enhance,
        )

    def has_effect(self) -> bool:
        """判断是否有实际降噪效果。"""
        return self.enabled

    def get_effective_noise_floor(self) -> float:
        """获取实际生效的噪音阈值（dB）。"""
        if self.level == NoiseReductionLevel.CUSTOM:
            return self.noise_floor
        params = _LEVEL_PARAMS.get(self.level, _LEVEL_PARAMS[NoiseReductionLevel.MEDIUM])
        return float(params["nf"])


# ── 引擎实现 ──────────────────────────────────────────────────────────────────


class NoiseReductionEngine:
    """音频降噪引擎。

    基于 FFmpeg afftdn（Audio FFt Denoiser）滤镜实现：
    - 使用短时傅里叶变换分析音频频谱
    - 识别并消除稳态背景噪音
    - 保留人声等非稳态信号
    """

    def __init__(self, config: NoiseReductionConfig):
        self.config = config

    def build_filter(self, input_label: str, output_label: str) -> str:
        """构建音频降噪滤镜字符串。

        Args:
            input_label: 输入标签，如 "[0:a]" 或 "[a0]"
            output_label: 输出标签，如 "[nr0]"

        Returns:
            FFmpeg 滤镜字符串，如 "[a0]afftdn=nf=-25:tn=-10:tr=50[nr0]"

        Raises:
            ValueError: 配置无效时抛出（调用方应捕获并降级）
        """
        if not self.config.has_effect():
            return f"{input_label}anull{output_label}"

        # 获取参数
        if self.config.level == NoiseReductionLevel.CUSTOM:
            nf = self.config.noise_floor
            tn = -10  # 默认频谱平滑度
            tr = 50  # 默认时间分辨率
        else:
            params = _LEVEL_PARAMS.get(
                self.config.level,
                _LEVEL_PARAMS[NoiseReductionLevel.MEDIUM],
            )
            nf = float(params["nf"])
            tn = float(params["tn"])
            tr = float(params["tr"])

        # 构建 afftdn 滤镜
        # nf: noise floor (dB)
        # tn: temporal noise floor smoothing (dB)
        # tr: time resolution (ms)
        filter_parts = [f"afftdn=nf={nf}:tn={tn}:tr={tr}"]

        # 人声增强：通过 highpass + 轻微压缩实现
        if self.config.voice_enhance:
            # 1. 高通滤波，去除低频噪音
            filter_parts.append("highpass=f=80")
            # 2. 轻微压缩，提升人声清晰度
            filter_parts.append("acompressor=threshold=-20:ratio=2:attack=5:release=50")
            # 3. 响度归一化
            filter_parts.append("loudnorm=I=-16:TP=-1.5:LRA=11")

        filter_str = f"{input_label}{','.join(filter_parts)}{output_label}"
        return filter_str

    def build_filter_arnndn(self, input_label: str, output_label: str, model_file: str) -> str:
        """使用 RNN 降噪滤镜（arnndn，效果更好但需要模型文件）。

        注意：需要额外下载 RNNNoise 模型文件，默认使用 afftdn（无需额外依赖）。

        Args:
            input_label: 输入标签
            output_label: 输出标签
            model_file: RNNNoise 模型文件路径（.rnnn 格式）

        Returns:
            FFmpeg 滤镜字符串
        """
        if not self.config.has_effect():
            return f"{input_label}anull{output_label}"

        return f"{input_label}arnndn=m={model_file}{output_label}"


def apply_noise_reduction_if_needed(
    config_data: dict | None,
    input_label: str,
    output_label: str,
) -> Optional[str]:
    """便捷函数：根据配置判断是否需要应用音频降噪。

    Args:
        config_data: 降噪配置字典（从 plan.config.audio_noise_reduction 或 clip.config.noise_reduction 读取）
        input_label: 输入标签
        output_label: 输出标签

    Returns:
        滤镜字符串，不需要降噪时返回 None
    """
    if not config_data:
        return None

    try:
        config = NoiseReductionConfig.from_dict(config_data)
        if not config.has_effect():
            return None

        engine = NoiseReductionEngine(config)
        return engine.build_filter(input_label, output_label)
    except Exception as e:
        logger.warning("[noise-reduction] 应用降噪失败，跳过: %s", e)
        return None