"""Validation suite for the surgical emotion module. Checks: 1. Frame-rate contract: T* == ceil(duration_s * 12.5) for varied durations. 2. Interface contract: shapes/dims/ranges of all outputs. 3. Emotion fidelity: top-1 agreement + JS divergence of student vs teacher utterance probabilities over the val set; per-RAVDESS-emotion breakdown. 4. Dynamics: correlation between student intensity curve and teacher max-prob proxy; frame-prob trajectory plot for one emotional utterance. Writes /tmp/work/surgical_out/eval_report.json + dynamics_plot.png """ from __future__ import annotations import glob import json import os import random import sys import time import numpy as np import torch import torchaudio sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from surgical_emotion import SurgicalEmotionEncoder FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt" OUT_DIR = "/tmp/work/surgical_out" RAVDESS_DIR = "/tmp/work/ravdess_kag" TEST_ACTORS = set(range(19, 25)) # held-out actors def _actor_of(path): try: return int(os.path.basename(path).split("-")[-1].split(".")[0]) except Exception: return 0 def _is_speech(path): try: return os.path.basename(path).split("-")[1] == "01" except Exception: return False def rms_envelope_50hz_t(wav: torch.Tensor, T_frames: int) -> torch.Tensor: frames = torch.zeros(T_frames) for i in range(T_frames): seg = wav[i * 320 : (i + 1) * 320] frames[i] = float(seg.pow(2).mean().sqrt()) if seg.numel() else (frames[i - 1] if i else 0.0) if T_frames >= 9: frames = torch.nn.functional.conv1d( frames[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0] lo, hi = frames.min(), frames.max() return torch.full_like(frames, 0.5) if hi - lo < 1e-6 else (frames - lo) / (hi - lo) def js_div(p, q, eps=1e-8): p = p.clamp(min=eps); q = q.clamp(min=eps) p = p.clamp(min=eps); q = q.clamp(min=eps) m = 0.5 * (p + q) kl = lambda a, b: (a * (a.log() - b.log())).sum(-1) return 0.5 * kl(p, m) + 0.5 * kl(q, m) def load_wav(path): wav, sr = torchaudio.load(path) wav = wav.mean(dim=0) if sr != 16000: wav = torchaudio.functional.resample(wav, sr, 16000) return wav def main(): device = "cuda:0" if torch.cuda.is_available() else "cpu" enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, CKPT, device=device) enc.eval() cfg = enc.cfg report = {"frame_rate_hz": cfg.frame_rate_hz, "checks": {}, "files": {}} # ---------- 1. rate & interface contract ---------- ok = True for sec in [1.0, 2.5, 5.0, 10.0]: n = int(sec * 16000) wav = torch.randn(1, n, device=device) wav = wav / wav.abs().max().clamp(min=1e-3) out = enc(wav, torch.tensor([n], device=device)) expect = int(np.ceil(sec * cfg.frame_rate_hz)) got = int(out["frame_embs"].shape[1]) good = abs(got - expect) <= 1 ok &= good report["checks"][f"{sec}s"] = {"expected~": expect, "got": got, "ok": good} assert out["frame_embs"].shape[-1] == cfg.out_dim assert out["frame_probs"].shape[-1] == cfg.n_emotions assert 0.0 <= float(out["frame_intensity"].min()) and float(out["frame_intensity"].max()) <= 1.0 report["rate_contract_ok"] = bool(ok) print("rate contract ok:", ok) # ---------- 2. fidelity on RAVDESS ---------- files = sorted(glob.glob(os.path.join(RAVDESS_DIR, "**", "*.wav"), recursive=True)) files = [f for f in files if _is_speech(f) and _actor_of(f) in TEST_ACTORS] if not files: files = [FUNASR_DIR + "/example/test.wav"] random.seed(0) sel = random.sample(files, min(60, len(files))) agree, jss, corrs = [], [], [] per_emo = {} plot_sample = None with torch.no_grad(): for f in sel: wav = load_wav(f)[: 12 * 16000] if wav.numel() < 8000: continue wav = wav / wav.abs().max().clamp(min=1e-4) wav = wav[None].to(device) out = enc(wav, torch.tensor([wav.size(1)], device=device)) states, lengths = enc.extract_teacher_states(wav, torch.tensor([wav.size(1)], device=device)) t_probs, t_inten = enc.teacher_frame_probs(states, lengths) tp = torch.softmax(enc.teacher.proj(states.mean(1)), dim=-1)[0] sp = out["utt_probs"][0] agree.append(int(sp.argmax().item() == tp.argmax().item())) jss.append(float(js_div(sp[None], tp[None]))) si = out["frame_intensity"][0, :, 0].cpu() # composite target: confidence-gated expression energy (same as training) T_frames = int(lengths[0].item()) rms = rms_envelope_50hz_t(wav.cpu().squeeze(0), T_frames) tmax = t_inten[0, :, 0].cpu().clamp(min=0, max=1).sqrt() ti = tmax * rms ti_ds = torch.nn.functional.interpolate( ti[None, None], size=si.size(0), mode="linear")[0, 0] if si.var() > 1e-6 and ti_ds.var() > 1e-6: c = float(np.corrcoef(si.numpy(), ti_ds.numpy())[0, 1]) corrs.append(c) name = os.path.basename(f) emo = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad", "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}.get( name.split("-")[2], "?") d = per_emo.setdefault(emo, {"n": 0, "agree": 0}) d["n"] += 1 d["agree"] += agree[-1] report["files"][name] = { "raveil_label": emo, "teacher_top": int(tp.argmax()), "student_top": int(sp.argmax()), "teacher_probs": [round(x, 3) for x in tp.tolist()], "student_probs": [round(x, 3) for x in sp.tolist()], "intensity_std": round(float(si.std()), 4), } if plot_sample is None and emo in ("angry", "happy", "sad"): plot_sample = (name, out["frame_probs"][0].cpu().numpy(), si.numpy(), ti_ds.numpy(), tp.cpu().numpy(), emo) report["agreement"] = float(np.mean(agree)) report["mean_js_teacher"] = float(np.mean(jss)) report["intensity_corr"] = float(np.mean(corrs)) if corrs else None report["per_emotion"] = per_emo print(f"agreement={report['agreement']:.3f} JS={report['mean_js_teacher']:.4f} " f"intensity_corr={report['intensity_corr']}") # ---------- 3. dynamics plot ---------- try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt if plot_sample is not None: name, fp, si, ti_ds, tp_, emo = plot_sample t = np.arange(len(si)) * 0.08 fig, ax = plt.subplots(3, 1, figsize=(11, 7), sharex=True) for k, lab in enumerate(["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised"]): ax[0].plot(t, fp[:, k], label=lab, lw=1.2) ax[0].set_title(f"student frame emotion probs @12.5Hz — {name} ({emo})") ax[0].legend(fontsize=7, ncol=4); ax[0].set_ylabel("p") ax[1].plot(t, si, color="crimson", lw=1.8, label="student intensity") ax[1].plot(t, ti_ds, color="steelblue", lw=1.4, ls="--", label="target (conf-gated energy)") ax[1].set_ylabel("intensity"); ax[1].set_ylim(0, 1.05); ax[1].legend(fontsize=8) ax[2].bar(["angry", "disg", "fear", "happy", "neut", "other", "sad", "surp", "unk"], tp_[:9], color="gray") ax[2].set_ylabel("teacher utt probs"); ax[2].set_xlabel("utterance (frame-aligned time axis)") fig.tight_layout() fig.savefig(os.path.join(OUT_DIR, "dynamics_plot.png"), dpi=110) print("plot saved") except Exception as e: print("plot skipped:", e) with open(os.path.join(OUT_DIR, "eval_report.json"), "w") as f: json.dump(report, f, indent=1) print("EVAL_DONE") if __name__ == "__main__": main()