eval_voxaudio_ood.py

#!/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()