Add scripts/analyze_audio.py

scripts/analyze_audio.py

@@ -0,0 +1,279 @@
+"""
+analyze_audio.py
+
+Comprehensive audio quality comparison between two WAV files.
+Designed for comparing PyTorch TTS output vs ONNX/browser output.
+
+Metrics:
+  1. Objective (librosa): mel cosine similarity, MFCC similarity, duration, pitch contour
+  2. Groq Whisper: transcription + WER
+  3. Gemini Flash: MOS score (1-5) with reasoning
+
+Usage:
+    conda run -n chatterbox-onnx python analyze_audio.py <file_a.wav> <file_b.wav> [--reference-text "..."]
+    conda run -n chatterbox-onnx python analyze_audio.py _cmp/pytorch_output.wav _cmp/onnx_output.wav
+
+    # Compare against the perfect baseline:
+    conda run -n chatterbox-onnx python analyze_audio.py \
+        Chatterbox-Finnish/output_finnish.wav \
+        _cmp/browser_sim_output.wav
+"""
+
+import sys, os, base64, json, argparse
+import numpy as np
+import librosa
+import soundfile as sf
+import requests
+from pathlib import Path
+
+# Load from .env
+def load_env():
+    env = {}
+    env_path = Path(__file__).parent / ".env"
+    if env_path.exists():
+        for line in env_path.read_text().splitlines():
+            if "=" in line and not line.startswith("#"):
+                k, v = line.split("=", 1)
+                env[k.strip()] = v.strip()
+    return env
+
+ENV = load_env()
+GROQ_KEY   = os.environ.get("GROQ_API_KEY",   ENV.get("QROQ_API_KEY", ""))
+GEMINI_KEY = os.environ.get("GEMINI_API_KEY",  ENV.get("GEMINI_API_KEY", ""))
+
+
+# ── Objective metrics ─────────────────────────────────────────────────────────
+
+def load_mono(path: str, target_sr: int = 22050) -> tuple[np.ndarray, int]:
+    y, sr = librosa.load(path, sr=target_sr, mono=True)
+    return y, sr
+
+
+def cosine(a: np.ndarray, b: np.ndarray) -> float:
+    a, b = a.flatten(), b.flatten()
+    denom = np.linalg.norm(a) * np.linalg.norm(b)
+    return float(np.dot(a, b) / denom) if denom > 0 else 0.0
+
+
+def mel_similarity(y_a, y_b, sr) -> float:
+    """Cosine similarity of mean mel spectrograms (overall timbre match)."""
+    mel_a = librosa.feature.melspectrogram(y=y_a, sr=sr, n_mels=128)
+    mel_b = librosa.feature.melspectrogram(y=y_b, sr=sr, n_mels=128)
+    # Mean over time
+    return cosine(mel_a.mean(axis=1), mel_b.mean(axis=1))
+
+
+def mfcc_similarity(y_a, y_b, sr, n_mfcc=20) -> float:
+    """Cosine similarity of mean MFCCs (phonetic content match)."""
+    mfcc_a = librosa.feature.mfcc(y=y_a, sr=sr, n_mfcc=n_mfcc).mean(axis=1)
+    mfcc_b = librosa.feature.mfcc(y=y_b, sr=sr, n_mfcc=n_mfcc).mean(axis=1)
+    return cosine(mfcc_a, mfcc_b)
+
+
+def pitch_correlation(y_a, y_b, sr) -> float:
+    """Correlation of F0 contours (prosody match). NaN frames excluded."""
+    f0_a = librosa.yin(y_a, fmin=60, fmax=400)
+    f0_b = librosa.yin(y_b, fmin=60, fmax=400)
+    # Resample to same length
+    length = min(len(f0_a), len(f0_b))
+    f0_a, f0_b = f0_a[:length], f0_b[:length]
+    voiced = (f0_a > 0) & (f0_b > 0)
+    if voiced.sum() < 10:
+        return float("nan")
+    a, b = f0_a[voiced], f0_b[voiced]
+    corr = np.corrcoef(a, b)[0, 1]
+    return float(corr)
+
+
+def spectral_flux_similarity(y_a, y_b, sr) -> float:
+    """How similar the rhythm/energy flow is (pacing match)."""
+    flux_a = np.diff(librosa.feature.rms(y=y_a)[0])
+    flux_b = np.diff(librosa.feature.rms(y=y_b)[0])
+    length = min(len(flux_a), len(flux_b))
+    return cosine(flux_a[:length], flux_b[:length])
+
+
+def objective_metrics(path_a: str, path_b: str) -> dict:
+    SR = 22050
+    y_a, _ = load_mono(path_a, SR)
+    y_b, _ = load_mono(path_b, SR)
+
+    dur_a = librosa.get_duration(y=y_a, sr=SR)
+    dur_b = librosa.get_duration(y=y_b, sr=SR)
+
+    return {
+        "duration_a_s":         round(dur_a, 2),
+        "duration_b_s":         round(dur_b, 2),
+        "duration_ratio":       round(dur_b / dur_a if dur_a > 0 else 0, 3),
+        "mel_cosine":           round(mel_similarity(y_a, y_b, SR), 4),
+        "mfcc_cosine":          round(mfcc_similarity(y_a, y_b, SR), 4),
+        "pitch_correlation":    round(pitch_correlation(y_a, y_b, SR), 4),
+        "energy_flux_cosine":   round(spectral_flux_similarity(y_a, y_b, SR), 4),
+    }
+
+
+# ── WER helper ────────────────────────────────────────────────────────────────
+
+def simple_wer(ref: str, hyp: str) -> float:
+    """Token-level WER."""
+    ref_words = ref.lower().split()
+    hyp_words = hyp.lower().split()
+    n, m = len(ref_words), len(hyp_words)
+    dp = list(range(m + 1))
+    for i in range(1, n + 1):
+        prev = dp.copy()
+        dp[0] = i
+        for j in range(1, m + 1):
+            dp[j] = min(prev[j] + 1, dp[j - 1] + 1,
+                        prev[j - 1] + (0 if ref_words[i-1] == hyp_words[j-1] else 1))
+    return dp[m] / max(n, 1)
+
+
+# ── Groq transcription ────────────────────────���───────────────────────────────
+
+def transcribe_groq(wav_path: str, lang: str = "fi") -> str:
+    if not GROQ_KEY:
+        return "(no GROQ_API_KEY)"
+    with open(wav_path, "rb") as f:
+        r = requests.post(
+            "https://api.groq.com/openai/v1/audio/transcriptions",
+            headers={"Authorization": f"Bearer {GROQ_KEY}"},
+            files={"file": (os.path.basename(wav_path), f, "audio/wav")},
+            data={"model": "whisper-large-v3", "language": lang, "response_format": "text"},
+        )
+    if r.ok:
+        return r.text.strip()
+    return f"(error {r.status_code})"
+
+
+# ── Gemini MOS ────────────────────────────────────────────────────────────────
+
+def gemini_mos(wav_path: str, label: str = "") -> dict:
+    """
+    Uses Gemini 2.0 Flash to give a MOS score + reasoning for a TTS audio file.
+    Matches methodology used in the Chatterbox Finnish fine-tuning evaluation.
+    """
+    if not GEMINI_KEY:
+        return {"score": None, "comment": "(no GEMINI_API_KEY)"}
+
+    audio_bytes = open(wav_path, "rb").read()
+    audio_b64   = base64.b64encode(audio_bytes).decode()
+
+    prompt = (
+        "You are an expert speech quality evaluator. "
+        "Listen to this Finnish text-to-speech audio sample and evaluate its naturalness.\n\n"
+        "Rate on MOS (Mean Opinion Score) scale 1-5:\n"
+        "  1.0 = Completely unintelligible or robotic\n"
+        "  2.0 = Very poor quality, hard to understand\n"
+        "  3.0 = Acceptable but clearly synthetic\n"
+        "  4.0 = Good quality, natural-sounding\n"
+        "  5.0 = Excellent, indistinguishable from human speech\n\n"
+        "Return ONLY valid JSON: {\"mos\": <float 1.0-5.0>, \"reason\": \"<one sentence>\"}"
+    )
+
+    body = {
+        "contents": [{
+            "parts": [
+                {"inline_data": {"mime_type": "audio/wav", "data": audio_b64}},
+                {"text": prompt},
+            ]
+        }],
+        "generationConfig": {"temperature": 0.1, "maxOutputTokens": 1024},
+    }
+
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={GEMINI_KEY}"
+    r = requests.post(url, json=body, timeout=30)
+    if not r.ok:
+        return {"score": None, "comment": f"(Gemini error {r.status_code}: {r.text[:200]})"}
+
+    try:
+        text = r.json()["candidates"][0]["content"]["parts"][0]["text"]
+        # Strip markdown fences if present
+        text = text.strip().lstrip("```json").lstrip("```").rstrip("```").strip()
+        data = json.loads(text)
+        return {"score": data.get("mos"), "comment": data.get("reason", "")}
+    except Exception as e:
+        return {"score": None, "comment": f"(parse error: {e} | raw: {r.text[:200]})"}
+
+
+# ── Main report ───────────────────────────────────────────────────────────────
+
+def report(path_a: str, path_b: str, label_a: str = "A", label_b: str = "B",
+           ref_text: str = "", lang: str = "fi"):
+
+    BAR = "=" * 65
+    print(f"\n{BAR}")
+    print(f"  AUDIO COMPARISON REPORT")
+    print(f"  A: {path_a}")
+    print(f"  B: {path_b}")
+    print(BAR)
+
+    # ── Objective metrics ──
+    print("\n── Objective metrics ──────────────────────────────────────────")
+    obj = objective_metrics(path_a, path_b)
+    print(f"  Duration        A={obj['duration_a_s']}s  B={obj['duration_b_s']}s  "
+          f"ratio(B/A)={obj['duration_ratio']}")
+    print(f"  Mel cosine      {obj['mel_cosine']:.4f}   (timbre match, 1.0=identical)")
+    print(f"  MFCC cosine     {obj['mfcc_cosine']:.4f}   (phonetic match, 1.0=identical)")
+    print(f"  Pitch corr      {obj['pitch_correlation']:.4f}   (prosody match, 1.0=identical)")
+    print(f"  Energy flux     {obj['energy_flux_cosine']:.4f}   (pacing match, 1.0=identical)")
+
+    mel = obj["mel_cosine"]
+    mfcc = obj["mfcc_cosine"]
+    quality = "excellent (near-identical)" if mel > 0.98 and mfcc > 0.98 \
+         else "good" if mel > 0.95 and mfcc > 0.95 \
+         else "fair" if mel > 0.90 \
+         else "poor — significant differences"
+    print(f"\n  → Waveform match: {quality}")
+
+    # ── Transcription ──
+    print("\n── Groq Whisper transcription ─────────────────────────────────")
+    tx_a = transcribe_groq(path_a, lang)
+    tx_b = transcribe_groq(path_b, lang)
+    print(f"  {label_a}: '{tx_a}'")
+    print(f"  {label_b}: '{tx_b}'")
+    if ref_text:
+        wer_a = simple_wer(ref_text, tx_a)
+        wer_b = simple_wer(ref_text, tx_b)
+        print(f"  Ref:  '{ref_text}'")
+        print(f"  WER   {label_a}: {wer_a:.1%}   {label_b}: {wer_b:.1%}")
+
+    # ── Gemini MOS ──
+    print("\n── Gemini 2.0 Flash MOS ───────────────────────────────────────")
+    mos_a = gemini_mos(path_a, label_a)
+    mos_b = gemini_mos(path_b, label_b)
+    print(f"  {label_a}: MOS={mos_a['score']}  — {mos_a['comment']}")
+    print(f"  {label_b}: MOS={mos_b['score']}  — {mos_b['comment']}")
+
+    # ── Summary ──
+    print(f"\n{BAR}")
+    print("  SUMMARY")
+    print(BAR)
+    print(f"  Mel cosine:   {obj['mel_cosine']:.4f}  (target: >0.95 for 'good match')")
+    print(f"  MFCC cosine:  {obj['mfcc_cosine']:.4f}  (target: >0.95)")
+    print(f"  MOS {label_a}: {mos_a['score']}   MOS {label_b}: {mos_b['score']}")
+    if ref_text:
+        wer_a = simple_wer(ref_text, tx_a)
+        wer_b = simple_wer(ref_text, tx_b)
+        print(f"  WER {label_a}: {wer_a:.1%}   WER {label_b}: {wer_b:.1%}")
+
+    return {
+        "objective": obj,
+        "transcription": {"a": tx_a, "b": tx_b},
+        "mos": {"a": mos_a, "b": mos_b},
+    }
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description="Compare two TTS audio files")
+    p.add_argument("file_a", help="Reference/baseline WAV (e.g. pytorch output)")
+    p.add_argument("file_b", help="Target WAV to compare against (e.g. ONNX/browser output)")
+    p.add_argument("--label-a", default="PyTorch",  help="Label for file A")
+    p.add_argument("--label-b", default="ONNX",     help="Label for file B")
+    p.add_argument("--ref-text", default="", help="Reference transcript for WER")
+    p.add_argument("--lang", default="fi", help="Language code for transcription")
+    args = p.parse_args()
+
+    report(args.file_a, args.file_b,
+           label_a=args.label_a, label_b=args.label_b,
+           ref_text=args.ref_text, lang=args.lang)