xiaoxia-saas/tests/render_compare/audio_diff.py

"""音频对比工具 — 基于 FFmpeg 的音频质量对比.

使用以下指标评估两段音频的相似度：
1. 波形差异（RMS 差值）
2. 频谱相似度（FFT 分帧比较）
3. 时长差异

对比方式：
- 直接对两个音频做 `ametadata=select='gt(scene\\,0.3)'` 过于复杂
- 简化方案：用 `amerge` + `astats` 计算差值音频的 RMS

更精确的方案（已实现）：
- 将两轨音频做差（amix=0:weights='1 -1' → 实际上用 pan 更简单）
- 对差值音频做 astats，获取差值的 RMS、峰值等指标
"""

from __future__ import annotations

import json
import re
import shutil
import subprocess  # nosec B404
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

FFMPEG_BIN: str = shutil.which("ffmpeg") or "ffmpeg"
FFPROBE_BIN: str = shutil.which("ffprobe") or "ffprobe"


@dataclass
class AudioDiffResult:
    """音频对比结果."""

    audio_a: str
    audio_b: str
    duration_a: float
    duration_b: float
    duration_diff: float
    sample_rate_match: bool
    channels_match: bool
    diff_rms_db: float  # 差值音频的 RMS（dB，越低越相似）
    diff_peak_db: float  # 差值音频的峰值（dB，越低越相似）
    similarity_score: float  # 综合相似度评分 [0, 1]，1 = 完全一致
    passed: bool

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def probe_duration(file_path: str) -> float:
    """探测文件时长（秒），失败返回 0."""
    try:
        result = subprocess.run(  # nosec B603
            [
                FFPROBE_BIN,
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(file_path),
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=10,
        )
        return round(float(result.stdout.strip()), 3)
    except Exception:
        return 0.0


def probe_has_audio(file_path: str | Path) -> bool:
    """探测文件是否包含音频流."""
    try:
        result = subprocess.run(  # nosec B603
            [
                FFPROBE_BIN,
                "-v",
                "error",
                "-select_streams",
                "a:0",
                "-show_entries",
                "stream=codec_type",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(file_path),
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=10,
        )
        return result.stdout.strip() == "audio"
    except Exception:
        return False  # 探测失败保守返回 False，避免误判有音频


def compute_audio_diff(
    audio_a: str | Path,
    audio_b: str | Path,
    *,
    similarity_threshold: float = 0.90,
    duration_tolerance: float = 0.1,
) -> AudioDiffResult:
    """计算两段音频的差异.

    方案：用 pan 滤镜将两轨相减，对差值音频做 astats 分析。

    Args:
        audio_a: 音频A（基线）
        audio_b: 音频B（对比）
        similarity_threshold: 相似度合格阈值
        duration_tolerance: 时长容忍度（秒）

    Returns:
        AudioDiffResult 对比结果
    """
    dur_a = probe_duration(str(audio_a))
    dur_b = probe_duration(str(audio_b))
    duration_diff = abs(dur_a - dur_b)

    # 获取音频元信息
    info_a = _probe_audio_info(str(audio_a))
    info_b = _probe_audio_info(str(audio_b))

    sample_rate_match = info_a["sample_rate"] == info_b["sample_rate"]
    channels_match = info_a["channels"] == info_b["channels"]

    # 相减后分析差值
    # 取较短时长做对比
    min_dur = min(dur_a, dur_b)
    if min_dur <= 0:
        return AudioDiffResult(
            audio_a=str(audio_a),
            audio_b=str(audio_b),
            duration_a=dur_a,
            duration_b=dur_b,
            duration_diff=duration_diff,
            sample_rate_match=sample_rate_match,
            channels_match=channels_match,
            diff_rms_db=-999.0,
            diff_peak_db=-999.0,
            similarity_score=0.0,
            passed=False,
        )

    # 做差值音频：a - b
    # 注意：amix 会自动按输入数归一化音量（除以N），
    # 所以 a + (-1)*b 经过 amix=inputs=2 后整体音量会减半（-6dB）。
    # 加 volume=2 补偿回来，确保差值 RMS 反映真实差异幅度。
    command = [
        FFMPEG_BIN,
        "-i",
        str(audio_a),
        "-i",
        str(audio_b),
        "-filter_complex",
        # 第2轨反相 → amix混合 → volume=2补偿amix的自动缩放
        "[1:a]volume=-1[inv];[0:a][inv]amix=inputs=2:duration=shortest:dropout_transition=0,volume=2[diff]",
        "-map",
        "[diff]",
        "-f",
        "null",
        "-af",
        "astats=metadata=1:reset=0",
        "-",
    ]

    try:
        result = subprocess.run(  # nosec B603
            command,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=120,
        )
        stderr = result.stderr or ""
    except subprocess.CalledProcessError:
        # 如果音频格式不兼容，返回失败
        return AudioDiffResult(
            audio_a=str(audio_a),
            audio_b=str(audio_b),
            duration_a=dur_a,
            duration_b=dur_b,
            duration_diff=duration_diff,
            sample_rate_match=sample_rate_match,
            channels_match=channels_match,
            diff_rms_db=999.0,
            diff_peak_db=999.0,
            similarity_score=0.0,
            passed=False,
        )

    diff_rms_db, diff_peak_db = _parse_astats(stderr)

    # 相似度评分：基于差值 RMS
    # 差值 RMS -60dB → 相似度 ~1.0（几乎无声差）
    # 差值 RMS -20dB → 相似度 ~0.5（有明显差异）
    # 差值 RMS 0dB → 相似度 ~0.0（完全相反）
    if diff_rms_db <= -60:
        similarity_score = 1.0
    elif diff_rms_db >= 0:
        similarity_score = 0.0
    else:
        # 线性映射：-60dB → 1.0, 0dB → 0.0
        similarity_score = max(0.0, min(1.0, 1.0 + diff_rms_db / 60.0))

    passed = (
        duration_diff <= duration_tolerance
        and sample_rate_match
        and channels_match
        and similarity_score >= similarity_threshold
    )

    return AudioDiffResult(
        audio_a=str(audio_a),
        audio_b=str(audio_b),
        duration_a=round(dur_a, 3),
        duration_b=round(dur_b, 3),
        duration_diff=round(duration_diff, 3),
        sample_rate_match=sample_rate_match,
        channels_match=channels_match,
        diff_rms_db=round(diff_rms_db, 2),
        diff_peak_db=round(diff_peak_db, 2),
        similarity_score=round(similarity_score, 4),
        passed=passed,
    )


def _probe_audio_info(file_path: str) -> dict[str, int]:
    """探测音频元信息."""
    try:
        result = subprocess.run(  # nosec B603
            [
                FFPROBE_BIN,
                "-v",
                "error",
                "-select_streams",
                "a:0",
                "-show_entries",
                "stream=sample_rate,channels",
                "-of",
                "json",
                file_path,
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=10,
        )

        info = json.loads(result.stdout)
        stream = info.get("streams", [{}])[0]
        return {
            "sample_rate": int(stream.get("sample_rate", 44100)),
            "channels": int(stream.get("channels", 2)),
        }
    except Exception:
        return {"sample_rate": 0, "channels": 0}


def _parse_astats(stderr: str) -> tuple[float, float]:
    """从 astats 输出中解析 RMS 和峰值.

    astats 输出格式（在 stderr 中）:
    [Parsed_astats_1 @ 0x...] Channel: 1
    [Parsed_astats_1 @ 0x...] ...
    [Parsed_astats_1 @ 0x...] Overall
    [Parsed_astats_1 @ 0x...] DC offset: 0.000000
    [Parsed_astats_1 @ 0x...] Min level: -0.123456
    [Parsed_astats_1 @ 0x...] Max level: 0.789012
    [Parsed_astats_1 @ 0x...] Peak level dB: -2.01
    [Parsed_astats_1 @ 0x...] RMS level dB: -10.56
    ...
    """
    lines = stderr.split("\n")
    rms_db = -999.0
    peak_db = -999.0

    for line in lines:
        # 找 Overall 部分的统计（双声道时取整体值）
        rms_match = re.search(r"RMS level dB:\s*(-?\d+\.?\d*)", line)
        peak_match = re.search(r"Peak level dB:\s*(-?\d+\.?\d*)", line)
        if rms_match:
            rms_db = float(rms_match.group(1))
        if peak_match:
            peak_db = float(peak_match.group(1))

    return rms_db, peak_db


def extract_audio(video_path: str | Path, output_path: str | Path) -> Path:
    """从视频中提取音频（AAC 格式）.

    Args:
        video_path: 视频文件路径
        output_path: 输出音频路径

    Returns:
        输出音频文件路径
    """
    command = [
        FFMPEG_BIN,
        "-y",
        "-i",
        str(video_path),
        "-vn",
        "-acodec",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]
    subprocess.run(command, check=True, capture_output=True, timeout=120)  # nosec B603
    return Path(output_path)