xiaoxia-saas/packages/domain/noise_reduction_config.py

"""音频降噪配置领域模型 — 纯逻辑，无FFmpeg依赖.

抽离自 noise_reduction_engine.py，包含：
- NoiseReductionLevel 枚举（low/medium/high/custom）
- NoiseReductionConfig 数据类（解析/钳制/效果判断）
- 等级预设参数
- afftdn / arnndn 滤镜构建
- 便捷函数（apply_noise_reduction_if_needed）
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


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


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

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


# 各等级对应的降噪参数（afftdn 的 noise floor，单位 dB）
# 值越大（越接近 0），降噪越强；值越小（越负），降噪越弱
_LEVEL_PARAMS: dict[NoiseReductionLevel, dict[str, float]] = {
    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,
    },
}

# 参数范围
MIN_NOISE_FLOOR = -60.0
MAX_NOISE_FLOOR = -5.0

# 默认值
DEFAULT_LEVEL = NoiseReductionLevel.MEDIUM
DEFAULT_NOISE_FLOOR = -25.0


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


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

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

    enabled: bool = False
    level: NoiseReductionLevel = DEFAULT_LEVEL
    noise_floor: float = DEFAULT_NOISE_FLOOR  # dB
    voice_enhance: bool = False

    @classmethod
    def from_dict(cls, data: dict[str, Any] | 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 = DEFAULT_LEVEL

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

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

        # 钳制到合法范围
        noise_floor = max(MIN_NOISE_FLOOR, min(MAX_NOISE_FLOOR, 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[DEFAULT_LEVEL])
        return float(params["nf"])

    def get_level_params(self) -> dict[str, float]:
        """获取当前等级的完整参数字典."""
        if self.level == NoiseReductionLevel.CUSTOM:
            return {
                "nf": self.noise_floor,
                "tn": -10.0,
                "tr": 50.0,
            }
        params = _LEVEL_PARAMS.get(self.level, _LEVEL_PARAMS[DEFAULT_LEVEL])
        return {k: float(v) for k, v in params.items()}

    def validate(self) -> tuple[bool, str]:
        """校验配置是否有效."""
        if not self.enabled:
            return True, ""

        if not (MIN_NOISE_FLOOR <= self.noise_floor <= MAX_NOISE_FLOOR):
            return False, f"noise_floor 必须在 {MIN_NOISE_FLOOR}~{MAX_NOISE_FLOOR} dB 之间"

        return True, ""


# ── 滤镜构建 ────────────────────────────────────────────────────────────────


def build_afftdn_filter(
    config: NoiseReductionConfig,
    input_label: str,
    output_label: str,
) -> str:
    """构建 afftdn 音频降噪滤镜字符串.

    Args:
        config: 降噪配置
        input_label: 输入标签，如 "[0:a]" 或 "[a0]"
        output_label: 输出标签，如 "[nr0]"

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

    params = config.get_level_params()
    nf = params["nf"]
    tn = params["tn"]
    tr = params["tr"]

    # 构建 afftdn 滤镜
    filter_parts = [f"afftdn=nf={nf}:tn={tn}:tr={tr}"]

    # 人声增强：通过 highpass + 压缩 + 响度归一化实现
    if config.voice_enhance:
        filter_parts.append("highpass=f=80")
        filter_parts.append("acompressor=threshold=-20:ratio=2:attack=5:release=50")
        filter_parts.append("loudnorm=I=-16:TP=-1.5:LRA=11")

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


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

    注意：需要额外下载 RNNNoise 模型文件，默认使用 afftdn（无需额外依赖）。
    """
    if not 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[str, Any] | None,
    input_label: str,
    output_label: str,
) -> str | None:
    """便捷函数：根据配置判断是否需要应用音频降噪.

    Args:
        config_data: 降噪配置字典
        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

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


def get_level_names() -> list[str]:
    """获取所有降噪等级名称列表."""
    return [level.value for level in NoiseReductionLevel]