student-v2 / code /eval_surgical.py
amir0907's picture
Surgical emotion module v2: 2160-utt distillation (2xT4), 2000 steps
1419b86 verified
Raw History Blame Contribute Delete
8.25 kB
"""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()