#!/usr/bin/env python3 """OOD inference + pairwise comparison on voxaudio reconstruction data. For each of the 4 reconstruction model directories under ``/apdcephfs_cq12/share_302080740/user/schmittzhu/data/voxaudio/data`` and each sample sub-directory, we run the Spatial-BEATs v13_D ``best.pt`` checkpoint on both the GT FOA clip and the reconstructed FOA clip, and compare the two sets of model predictions (events + DOA + distance). Notes / conventions ------------------- * The raw 4-ch WAV files store FOA in DCASE waveform order ``[W, Y, Z, X]``. ``SpatialBEATsPreprocessor`` does the internal ``[0,3,1,2]`` permutation back to ``[W, X, Y, Z]``. We therefore feed the 4-ch waveform *as-is*. * Source sample rate is 44.1 kHz (or 24 kHz for ``mono_vae``); we resample to 16 kHz first. * The checkpoint uses ``readout_scheme='local_spatial_track'`` with K=4 track queries at 10 Hz. We decode each frame with an activity threshold of 0.5 and take the argmax class per active (track, frame). Outputs ------- Per-sample JSON with track-level event lists for both GT and Recon, and an aggregated ``summary.json`` with mean angular / distance error, class agreement, and activity Jaccard across all samples per model. """ from __future__ import annotations import argparse import csv import json import math import os from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import soundfile as sf import torch import torch.nn.functional as F from tqdm import tqdm # Local imports — must run from beats/ directory or have it on PYTHONPATH. from spatial_beats import SpatialBEATs from train_spatial_beats import make_ov1_unified_v13d_config VOXAUDIO_ROOT = "/apdcephfs_cq12/share_302080740/user/schmittzhu/data/voxaudio/data" CKPT_PATH = ( "/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/" "checkpoints/spatial_beats_ov1_unified_v13d_exp/03_ov123_top4/best.pt" ) TARGET_SR = 16000 ACTIVITY_THRESHOLD = 0.5 # ---------------------------------------------------------------------------- # Utilities # ---------------------------------------------------------------------------- def load_class_names(vocab_path: str) -> List[str]: rows = [] with open(vocab_path, "r", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: rows.append(row) rows.sort(key=lambda r: int(r["label_id"])) return [r["final_label"] for r in rows] def resample_numpy(x: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray: """Resample multi-channel numpy array ``x`` of shape (T, C) from ``src_sr`` to ``dst_sr`` using torchaudio if available, else scipy. """ if src_sr == dst_sr: return x try: import torchaudio wav = torch.from_numpy(x.T.astype(np.float32)) # [C, T] out = torchaudio.functional.resample(wav, src_sr, dst_sr) return out.numpy().T except Exception: import scipy.signal as sps g = math.gcd(src_sr, dst_sr) up = dst_sr // g down = src_sr // g return sps.resample_poly(x, up, down, axis=0).astype(np.float32) def list_sample_dirs(model_dir: Path) -> List[Path]: return sorted(p for p in model_dir.iterdir() if p.is_dir()) def find_foa_files(sample_dir: Path) -> Optional[Tuple[Path, Path, int]]: """Return (gt_path, recon_path, expected_sr) or None.""" # Standard (dacvae / flow2gan / stable_audio_vae): gt_foa4ch.wav, recon_foa4ch.wav gt = sample_dir / "gt_foa4ch.wav" rc = sample_dir / "recon_foa4ch.wav" if gt.exists() and rc.exists(): return gt, rc, 0 # 0 → detect from file header # mono_vae variants for sr_tag, sr in [("16k", 16000), ("24k", 24000), ("44k", 44100), ("48k", 48000)]: gt = sample_dir / f"gt_foa4ch_{sr_tag}.wav" rc = sample_dir / f"recon_foa4ch_{sr_tag}.wav" if gt.exists() and rc.exists(): return gt, rc, sr # Fall back: reconstruct from per-channel files (W/X/Y/Z) per_ch_gt = [sample_dir / f"gt_{c}.wav" for c in ("W", "Y", "Z", "X")] per_ch_rc = [sample_dir / f"recon_{c}.wav" for c in ("W", "Y", "Z", "X")] if all(p.exists() for p in per_ch_gt) and all(p.exists() for p in per_ch_rc): return sample_dir, sample_dir, -1 # sentinel: load per channel return None def load_foa_4ch(path_or_dir: Path, special_sr: int) -> Tuple[np.ndarray, int]: """Load a 4-ch FOA clip in channel order matching the .wav file. Returns (waveform [T, 4], sample_rate). """ if special_sr == -1: # per-channel fallback, assemble WYZX wavs = [] sr_ref = None for c in ("W", "Y", "Z", "X"): p = (path_or_dir if path_or_dir.is_dir() else path_or_dir.parent) / f"{c}.wav" w, sr = sf.read(p) if sr_ref is None: sr_ref = sr wavs.append(w.astype(np.float32)) length = min(len(w) for w in wavs) arr = np.stack([w[:length] for w in wavs], axis=1) return arr, sr_ref w, sr = sf.read(path_or_dir) return w.astype(np.float32), sr def load_and_prepare(path: Path, special_sr: int) -> torch.Tensor: """Load a FOA wav, resample to 16 kHz, return [4, T] float tensor in WYZX order.""" x, sr = load_foa_4ch(path, special_sr) if x.ndim == 1: raise ValueError(f"{path}: expected multi-channel audio, got mono") if x.shape[1] != 4: raise ValueError(f"{path}: expected 4 channels, got shape {x.shape}") x = resample_numpy(x, sr, TARGET_SR) # The files contain WYZX order (per user note); SpatialBEATsPreprocessor # will permute [0,3,1,2] → [W,X,Y,Z] internally. return torch.from_numpy(x.T).float().contiguous() # ---------------------------------------------------------------------------- # Prediction decoding # ---------------------------------------------------------------------------- def decode_frame_track( pred, target_num_steps: int, activity_threshold: float, class_names: List[str], ) -> Dict: """Decode a FrameTrackPredictionOutput (B=1) into a list of active per-frame per-track detections plus a clip-level event summary. """ # Shapes: [1, K, T_s], [1, K, T_s, C], [1, K, T_s, 3], [1, K, T_s] act = torch.sigmoid(pred.pred_activity[0]).cpu() # [K, T_s] cls = pred.pred_class_logits[0].cpu() # [K, T_s, C] direc = pred.pred_direction[0].cpu() # [K, T_s, 3] dist = pred.pred_distance[0].cpu() # [K, T_s] K, T_s = act.shape T_s = min(T_s, target_num_steps) act = act[:, :T_s] cls = cls[:, :T_s] direc = direc[:, :T_s] dist = dist[:, :T_s] direc_n = F.normalize(direc, dim=-1) azi_deg = torch.rad2deg(torch.atan2(direc_n[..., 1], direc_n[..., 0])) # y, x ele_deg = torch.rad2deg(torch.asin(direc_n[..., 2].clamp(-1, 1))) cls_prob = cls.softmax(dim=-1) cls_idx = cls_prob.argmax(dim=-1) cls_conf = cls_prob.amax(dim=-1) # Per-frame per-track detections frames = [] # list of lists — frames[t] is list of detected tracks for t in range(T_s): frame_list = [] for k in range(K): a = float(act[k, t]) if a >= activity_threshold: frame_list.append({ "track": k, "activity": round(a, 3), "class_idx": int(cls_idx[k, t]), "class_name": class_names[int(cls_idx[k, t])], "class_conf": round(float(cls_conf[k, t]), 3), "azi_deg": round(float(azi_deg[k, t]), 2), "ele_deg": round(float(ele_deg[k, t]), 2), "dist_m": round(float(dist[k, t]), 3), }) frames.append(frame_list) # Clip-level event = class most frequently predicted among active frames class_votes: Dict[int, float] = {} for t in range(T_s): for d in frames[t]: class_votes[d["class_idx"]] = class_votes.get(d["class_idx"], 0.0) + d["activity"] if class_votes: top_class = max(class_votes, key=class_votes.get) else: # fall back to most confident class regardless of activity flat_idx = cls_conf.reshape(-1).argmax().item() top_class = int(cls_idx.reshape(-1)[flat_idx]) return { "frames": frames, "top_class_idx": int(top_class), "top_class_name": class_names[int(top_class)], "T_s": T_s, # Raw tensors for downstream pairwise comparison. "_act": act.numpy(), "_cls_idx": cls_idx.numpy(), "_cls_conf": cls_conf.numpy(), "_direction": direc_n.numpy(), "_azi_deg": azi_deg.numpy(), "_ele_deg": ele_deg.numpy(), "_dist": dist.numpy(), } def angular_error_deg(a: np.ndarray, b: np.ndarray) -> float: """Great-circle angular error in degrees between two unit 3-vectors.""" dot = float(np.clip(np.dot(a, b), -1.0, 1.0)) return math.degrees(math.acos(dot)) def compare_predictions(gt_dec: Dict, rc_dec: Dict, activity_threshold: float) -> Dict: """Compare two decoded outputs with identical (K, T_s) shapes.""" T_s = min(gt_dec["T_s"], rc_dec["T_s"]) gt_act = gt_dec["_act"][:, :T_s] rc_act = rc_dec["_act"][:, :T_s] gt_cls = gt_dec["_cls_idx"][:, :T_s] rc_cls = rc_dec["_cls_idx"][:, :T_s] gt_dir = gt_dec["_direction"][:, :T_s] rc_dir = rc_dec["_direction"][:, :T_s] gt_dist = gt_dec["_dist"][:, :T_s] rc_dist = rc_dec["_dist"][:, :T_s] gt_on = gt_act >= activity_threshold rc_on = rc_act >= activity_threshold both_on = gt_on & rc_on # Activity agreement activity_jaccard = float((gt_on & rc_on).sum()) / max(1, int((gt_on | rc_on).sum())) activity_f1_tp = float((gt_on & rc_on).sum()) activity_f1_fp = float((~gt_on & rc_on).sum()) activity_f1_fn = float((gt_on & ~rc_on).sum()) prec = activity_f1_tp / max(1e-8, activity_f1_tp + activity_f1_fp) rec = activity_f1_tp / max(1e-8, activity_f1_tp + activity_f1_fn) f1 = 2 * prec * rec / max(1e-8, prec + rec) # Class agreement on both-on cells if both_on.any(): class_match = float((gt_cls[both_on] == rc_cls[both_on]).mean()) else: class_match = float("nan") # DOA angular error on both-on cells ang_errs = [] for k in range(gt_dir.shape[0]): for t in range(T_s): if both_on[k, t]: ang_errs.append(angular_error_deg(gt_dir[k, t], rc_dir[k, t])) doa_mae_deg = float(np.mean(ang_errs)) if ang_errs else float("nan") doa_median_deg = float(np.median(ang_errs)) if ang_errs else float("nan") # Distance MAE on both-on cells if both_on.any(): dist_mae = float(np.mean(np.abs(gt_dist[both_on] - rc_dist[both_on]))) else: dist_mae = float("nan") # Top-class agreement top_match = int(gt_dec["top_class_idx"] == rc_dec["top_class_idx"]) return { "T_s": T_s, "activity_gt_frac": float(gt_on.mean()), "activity_rc_frac": float(rc_on.mean()), "activity_jaccard": activity_jaccard, "activity_precision_rc_vs_gt": prec, "activity_recall_rc_vs_gt": rec, "activity_f1_rc_vs_gt": f1, "class_match_rate": class_match, "doa_angular_error_deg_mean": doa_mae_deg, "doa_angular_error_deg_median": doa_median_deg, "distance_mae_m": dist_mae, "top_class_agreement": top_match, "gt_top_class": gt_dec["top_class_name"], "rc_top_class": rc_dec["top_class_name"], } # ---------------------------------------------------------------------------- # Model loading # ---------------------------------------------------------------------------- def load_model(device: torch.device) -> Tuple[SpatialBEATs, List[str], object]: ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False) # Use the in-code factory to reconstruct a compatible TrainSpatialBEATsConfig, # then overlay the checkpoint's stored model config to guarantee exact match # with the weights. train_cfg = make_ov1_unified_v13d_config() model_cfg = ckpt["train_cfg"]["model"] model = SpatialBEATs(model_cfg) state = ckpt["model_state_dict"] missing, unexpected = model.load_state_dict(state, strict=False) if missing: print(f"[WARN] Missing keys ({len(missing)}): {missing[:3]}...") if unexpected: print(f"[WARN] Unexpected keys ({len(unexpected)}): {unexpected[:3]}...") model = model.to(device).eval() class_names = load_class_names(model_cfg.source_vocab_path) return model, class_names, model_cfg # ---------------------------------------------------------------------------- # Main # ---------------------------------------------------------------------------- def run_sample( model: SpatialBEATs, class_names: List[str], model_cfg, gt_path: Path, rc_path: Path, special_sr: int, device: torch.device, ) -> Dict: gt_wav = load_and_prepare(gt_path, special_sr).unsqueeze(0).to(device) # [1, 4, T] rc_wav = load_and_prepare(rc_path, special_sr).unsqueeze(0).to(device) # clip duration tensor (seconds) dur_gt = torch.tensor([gt_wav.shape[-1] / TARGET_SR], device=device, dtype=torch.float32) dur_rc = torch.tensor([rc_wav.shape[-1] / TARGET_SR], device=device, dtype=torch.float32) T_s_gt = int(round(float(dur_gt.item()) * model_cfg.target_token_rate)) T_s_rc = int(round(float(dur_rc.item()) * model_cfg.target_token_rate)) with torch.no_grad(): gt_out = model(waveform=gt_wav, padding_mask=None, clip_duration_seconds=dur_gt) rc_out = model(waveform=rc_wav, padding_mask=None, clip_duration_seconds=dur_rc) gt_dec = decode_frame_track(gt_out.frame_track_prediction_output, T_s_gt, ACTIVITY_THRESHOLD, class_names) rc_dec = decode_frame_track(rc_out.frame_track_prediction_output, T_s_rc, ACTIVITY_THRESHOLD, class_names) cmp = compare_predictions(gt_dec, rc_dec, ACTIVITY_THRESHOLD) return { "gt_top_class": gt_dec["top_class_name"], "rc_top_class": rc_dec["top_class_name"], "gt_frames_preview": gt_dec["frames"][:5], "rc_frames_preview": rc_dec["frames"][:5], "comparison": cmp, } def aggregate(sample_results: List[Dict]) -> Dict: keys_mean = [ "activity_jaccard", "activity_precision_rc_vs_gt", "activity_recall_rc_vs_gt", "activity_f1_rc_vs_gt", "class_match_rate", "doa_angular_error_deg_mean", "doa_angular_error_deg_median", "distance_mae_m", "top_class_agreement", "activity_gt_frac", "activity_rc_frac", ] out: Dict[str, float] = {} for k in keys_mean: vals = [s["comparison"][k] for s in sample_results if s["comparison"][k] is not None and not (isinstance(s["comparison"][k], float) and math.isnan(s["comparison"][k]))] out[f"mean_{k}"] = float(np.mean(vals)) if vals else float("nan") out[f"n_valid_{k}"] = len(vals) out["n_samples"] = len(sample_results) return out def main(): parser = argparse.ArgumentParser() parser.add_argument("--root", default=VOXAUDIO_ROOT) parser.add_argument("--output-dir", default="eval_voxaudio_ood_results") parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") parser.add_argument("--models", nargs="+", default=["dacvae", "flow2gan", "mono_vae", "stable_audio_vae", "foa_vae"]) parser.add_argument("--max-per-model", type=int, default=0, help="Debug limit; 0 = all.") args = parser.parse_args() out_root = Path(args.output_dir) out_root.mkdir(parents=True, exist_ok=True) device = torch.device(args.device) print(f"[Load] checkpoint: {CKPT_PATH}") model, class_names, model_cfg = load_model(device) print(f"[Load] {len(class_names)} classes, K={model_cfg.frame_track_num_queries}, " f"token_rate={model_cfg.target_token_rate} Hz") all_summary: Dict[str, Dict] = {} for m in args.models: model_dir = Path(args.root) / m if not model_dir.is_dir(): print(f"[Skip] {m}: dir not found") continue samples = list_sample_dirs(model_dir) if args.max_per_model: samples = samples[: args.max_per_model] print(f"\n=== {m}: {len(samples)} samples ===") results: List[Dict] = [] per_sample_detail = {} for s_dir in tqdm(samples, desc=m): paths = find_foa_files(s_dir) if paths is None: continue gt_path, rc_path, special_sr = paths try: res = run_sample(model, class_names, model_cfg, gt_path, rc_path, special_sr, device) except Exception as e: print(f"[Err] {s_dir.name}: {e}") continue res["sample"] = s_dir.name results.append(res) per_sample_detail[s_dir.name] = res # Persist per-model details + summary model_out_dir = out_root / m model_out_dir.mkdir(parents=True, exist_ok=True) with open(model_out_dir / "per_sample.json", "w") as f: json.dump(per_sample_detail, f, indent=2, ensure_ascii=False) summary = aggregate(results) all_summary[m] = summary with open(model_out_dir / "summary.json", "w") as f: json.dump(summary, f, indent=2) print(f"[{m}] summary: {json.dumps(summary, indent=2)}") with open(out_root / "summary_all.json", "w") as f: json.dump(all_summary, f, indent=2) # Pretty print comparison across recon models print("\n" + "=" * 80) print(" OOD recon-vs-gt (model self-consistency) summary") print("=" * 80) metric_keys = [ "mean_top_class_agreement", "mean_class_match_rate", "mean_activity_f1_rc_vs_gt", "mean_activity_jaccard", "mean_doa_angular_error_deg_mean", "mean_doa_angular_error_deg_median", "mean_distance_mae_m", "mean_activity_gt_frac", "mean_activity_rc_frac", ] header = f"{'metric':45s} " + " ".join(f"{m:>18s}" for m in all_summary.keys()) print(header) for k in metric_keys: row = f"{k:45s} " + " ".join( f"{all_summary[m].get(k, float('nan')):>18.4f}" for m in all_summary.keys() ) print(row) print("=" * 80) print(f"[Done] detailed results under: {out_root.resolve()}") if __name__ == "__main__": main()