h3-blackwell-runtime/tools/compare_audio_pcm.py
2026-08-22 14:09:45 +07:00

64 lines
2.3 KiB
Python

"""Compare decoded audio streams using aligned float PCM arrays."""
import argparse
import json
import subprocess
from pathlib import Path
import numpy as np
def decode(path: Path) -> np.ndarray:
raw = subprocess.check_output([
"ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(path),
"-map", "0:a:0", "-f", "f32le", "-acodec", "pcm_f32le", "-",
])
return np.frombuffer(raw, dtype="<f4").reshape(-1, 2)
def metrics(reference: np.ndarray, candidate: np.ndarray) -> dict:
count = min(len(reference), len(candidate))
reference = reference[:count]
candidate = candidate[:count]
error = candidate - reference
signal_power = np.maximum(np.mean(reference**2, axis=0), 1e-30)
noise_power = np.maximum(np.mean(error**2, axis=0), 1e-30)
first_250ms = error[:8000]
first_signal_power = np.maximum(np.mean(reference[:8000] ** 2, axis=0), 1e-30)
first_noise_power = np.maximum(np.mean(first_250ms**2, axis=0), 1e-30)
return {
"reference_samples": len(reference),
"candidate_samples": len(candidate),
"compared_samples": count,
"snr_db_by_channel": (10.0 * np.log10(signal_power / noise_power)).tolist(),
"rmse": float(np.sqrt(np.mean(error**2))),
"first_250ms_snr_db_by_channel": (
10.0 * np.log10(first_signal_power / first_noise_power)
).tolist(),
"first_250ms_rmse": float(np.sqrt(np.mean(first_250ms**2))),
"max_abs_error": float(np.max(np.abs(error))),
}
parser = argparse.ArgumentParser()
parser.add_argument("--reference", type=Path, required=True)
parser.add_argument("--candidate", action="append", default=[], metavar="NAME=PATH")
parser.add_argument("--output", type=Path)
args = parser.parse_args()
reference = decode(args.reference)
report = {"reference": str(args.reference), "candidates": {}}
for value in args.candidate:
if "=" not in value:
raise ValueError(f"candidate must be NAME=PATH, got {value!r}")
name, raw_path = value.split("=", 1)
report["candidates"][name] = {
"path": raw_path,
**metrics(reference, decode(Path(raw_path))),
}
serialized = json.dumps(report, indent=2)
if args.output is not None:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(serialized + "\n", encoding="utf-8")
print(serialized)