xiaoxia-saas/tests/render_compare/video_diff.py

"""视频对比工具 — 基于 FFmpeg 的像素级质量对比.

使用 SSIM + PSNR 双指标评估两个视频的相似度：
- SSIM (Structural Similarity): 结构相似性，范围 [0, 1]，越接近 1 越相似
- PSNR (Peak Signal-to-Noise Ratio): 峰值信噪比，单位 dB，越高越好

灰度验收标准：
- 平均 SSIM >= 0.95 → 视觉上几乎无差异（P0 场景必达）
- 最低 SSIM >= 0.90 → 最严重帧差异可接受
- 平均 PSNR >= 28dB → 质量达标
"""

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 VideoDiffResult:
    """视频对比结果."""

    video_a: str
    video_b: str
    width: int
    height: int
    duration_a: float
    duration_b: float
    avg_ssim: float
    min_ssim: float
    avg_psnr: float  # dB
    min_psnr: float
    frame_count: int
    duration_diff: float  # 时长差（秒）
    resolution_match: bool
    passed: bool  # 是否通过阈值

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


def probe_video_info(video_path: str) -> dict[str, Any]:
    """获取视频信息（宽、高、时长、fps）."""
    try:
        result = subprocess.run(  # nosec B603
            [
                FFPROBE_BIN,
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=width,height,r_frame_rate,duration",
                "-show_entries",
                "format=duration",
                "-of",
                "json",
                video_path,
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=10,
        )
        info = json.loads(result.stdout)
        stream = info.get("streams", [{}])[0]
        fmt = info.get("format", {})

        width = int(stream.get("width", 1280))
        height = int(stream.get("height", 720))
        fps_str = stream.get("r_frame_rate", "25/1")
        if "/" in fps_str:
            num, den = fps_str.split("/")
            fps = float(num) / float(den) if float(den) > 0 else 25.0
        else:
            fps = float(fps_str) if fps_str else 25.0
        duration = float(fmt.get("duration", 0)) or float(stream.get("duration", 0))

        return {"width": width, "height": height, "duration": duration, "fps": round(fps, 2)}
    except Exception:
        return {"width": 1280, "height": 720, "duration": 0.0, "fps": 25.0}


def compute_video_diff(
    video_a: str | Path,
    video_b: str | Path,
    *,
    ssim_threshold: float = 0.95,
    psnr_threshold: float = 28.0,
    duration_tolerance: float = 0.1,
) -> VideoDiffResult:
    """计算两个视频的像素差异.

    使用 FFmpeg ssim + psnr 滤镜一次性计算两个指标。

    Args:
        video_a: 视频A路径（基线）
        video_b: 视频B路径（对比）
        ssim_threshold: SSIM 合格阈值（默认 0.90）
        psnr_threshold: PSNR 合格阈值（默认 25dB）
        duration_tolerance: 时长容忍度（秒，默认 0.1s）

    Returns:
        VideoDiffResult 对比结果

    Raises:
        subprocess.CalledProcessError: FFmpeg 执行失败
    """
    info_a = probe_video_info(str(video_a))
    info_b = probe_video_info(str(video_b))

    duration_diff = abs(info_a["duration"] - info_b["duration"])
    resolution_match = info_a["width"] == info_b["width"] and info_a["height"] == info_b["height"]

    # ssim 和 psnr 的 stats_file 都输出到 stdout
    # 用行格式区分：SSIM 行含 "All:"，PSNR 行含 "psnr_avg:"
    command = [
        FFMPEG_BIN,
        "-i",
        str(video_a),
        "-i",
        str(video_b),
        "-lavfi",
        "[0:v][1:v]ssim=stats_file=-[out1];[0:v][1:v]psnr=stats_file=-[out2]",
        "-f",
        "null",
        "-",
    ]

    result = subprocess.run(  # nosec B603
        command,
        check=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        timeout=300,
    )

    # 逐帧统计在 stdout（stats_file=-），汇总日志在 stderr
    stats_stdout = result.stdout or ""

    avg_ssim, min_ssim = _parse_ssim_stats(stats_stdout)
    avg_psnr, min_psnr = _parse_psnr_stats(stats_stdout)
    frame_count = _count_frames(result.stderr or "")

    passed = (
        resolution_match
        and duration_diff <= duration_tolerance
        and avg_ssim >= ssim_threshold
        and avg_psnr >= psnr_threshold
    )

    return VideoDiffResult(
        video_a=str(video_a),
        video_b=str(video_b),
        width=info_a["width"],
        height=info_a["height"],
        duration_a=round(info_a["duration"], 3),
        duration_b=round(info_b["duration"], 3),
        avg_ssim=round(avg_ssim, 6),
        min_ssim=round(min_ssim, 6),
        avg_psnr=round(avg_psnr, 3),
        min_psnr=round(min_psnr, 3),
        frame_count=frame_count,
        duration_diff=round(duration_diff, 3),
        resolution_match=resolution_match,
        passed=passed,
    )


def _parse_ssim_stats(stats_output: str) -> tuple[float, float]:
    """从 SSIM stats_file 输出中解析逐帧 SSIM.

    FFmpeg ssim 滤镜 stats_file 输出格式（每行一帧）:
        n:1 Y:0.987654 U:0.991234 V:0.990000 All:0.989000 (19.585642)
        n:2 Y:0.986543 U:0.990123 V:0.988888 All:0.987654 (19.123456)
        ...

    Returns:
        (avg_ssim, min_ssim)
    """
    ssim_values: list[float] = []

    for line in stats_output.split("\n"):
        # 匹配 stats_file 格式：n:数字 ... All:数字
        if not line.startswith("n:"):
            continue
        match = re.search(r"All:(\d+\.\d+)", line)
        if match:
            ssim_values.append(float(match.group(1)))

    if not ssim_values:
        return 0.0, 0.0

    avg_ssim = sum(ssim_values) / len(ssim_values)
    min_ssim = min(ssim_values)
    return avg_ssim, min_ssim


def _parse_psnr_stats(stats_output: str) -> tuple[float, float]:
    """从 PSNR stats_file 输出中解析逐帧 PSNR.

    FFmpeg psnr 滤镜 stats_file 输出格式（每行一帧）:
        n:1 mse_avg:100.23 mse_y:150.12 mse_u:50.34 mse_v:80.56 psnr_avg:28.12 psnr_y:26.34 psnr_u:31.12 psnr_v:29.08
        n:2 ...

    Returns:
        (avg_psnr, min_psnr)  —  avg_psnr 是逐帧 psnr_avg 的均值，min_psnr 是逐帧最小值
    """
    psnr_values: list[float] = []

    for line in stats_output.split("\n"):
        if not line.startswith("n:"):
            continue
        match = re.search(r"psnr_avg:(\d+\.\d+)", line)
        if match:
            psnr_values.append(float(match.group(1)))

    if not psnr_values:
        return 0.0, 0.0

    avg_psnr = sum(psnr_values) / len(psnr_values)
    min_psnr = min(psnr_values)
    return avg_psnr, min_psnr


def _count_frames(stderr: str) -> int:
    """从 FFmpeg 输出中统计帧数."""
    match = re.search(r"frame=\s*(\d+)", stderr)
    return int(match.group(1)) if match else 0


def save_diff_frame(
    video_a: str | Path,
    video_b: str | Path,
    output_path: str | Path,
    *,
    timestamp: float = 1.0,
) -> Path:
    """生成差异帧可视化图（红绿色差）.

    使用 blend 滤镜生成差异可视化图，差异越大越亮。

    Args:
        video_a: 视频A
        video_b: 视频B
        output_path: 输出图片路径
        timestamp: 截取的时间点（秒）

    Returns:
        输出图片路径
    """
    command = [
        FFMPEG_BIN,
        "-y",
        "-ss",
        str(timestamp),
        "-i",
        str(video_a),
        "-ss",
        str(timestamp),
        "-i",
        str(video_b),
        "-lavfi",
        "[0:v][1:v]blend=all_mode=difference,eq=contrast=5:brightness=0.5[diff]",
        "-map",
        "[diff]",
        "-vframes",
        "1",
        str(output_path),
    ]

    subprocess.run(command, check=True, capture_output=True, timeout=60)  # nosec B603
    return Path(output_path)