eval_v11a_ov1_sim.py

#!/usr/bin/env python3
"""Evaluate v11a_real_balanced_10hz ckpt on **sim ov1 test split only**.

The v11a / v9 chain uses supervision_mode='local_spatial_track' and
readout_scheme='local_spatial_track', i.e. K=4 per-frame track queries with
frame-level Hungarian matching.  There is no mono_ast clip token, so
visualize_spatial_latents.py does not apply.  This script feeds test batches
through the model and reports:

  classification
    - oracle_class_acc       (GT-active frames, matcher without activity cost)
    - activity_precision     (mean sigmoid(pred_act) on supposed-active frames)
    - activity_recall        (mean sigmoid(pred_act) on supposed-inactive)
    - (DCASE) F20, LR_CD     (official class-gated detection metrics)

  spatial
    - oracle_azi_mae_deg     (GT-active frames)
    - oracle_ele_mae_deg
    - oracle_dist_mae
    - (DCASE) LE_CD, ER20, SELD_score

Usage:
    python eval_v11a_ov1_sim.py \
        --checkpoint checkpoints/spatial_beats_ov1_local_spatial_v11a_real_balanced_10hz_exp/03_ov123_top4/best.pt \
        --preset ov1_local_spatial_v11a_real_balanced_10hz \
        --batch-size 8 --num-workers 8 --amp bf16
"""
from __future__ import annotations

import argparse
import contextlib
import copy
import dataclasses
import functools
import json
from pathlib import Path
from types import SimpleNamespace
from typing import Dict, List, Optional

import torch
from tqdm.auto import tqdm

from spatial_beats import SpatialBEATs
from spatial_dataset import SpatialDataset, collate_spatial_batch
from spatial_loss import (
    OfficialDCASEMetricsAccumulator,
    accumulate_frame_track_seld,
    compute_frame_track_validation_metrics,
)
from train_spatial_beats import (
    DEFAULT_OV1_MANIFEST,
    DEFAULT_OV2_MANIFEST,
    DEFAULT_OV3_MANIFEST,
    DEFAULT_OV1_REAL_MANIFEST,
    DEFAULT_OV2_REAL_MANIFEST,
    DEFAULT_OV3_REAL_MANIFEST,
    TrainSpatialBEATsConfig,
    build_dataset_config,
    build_model_config,
    build_train_config_from_args,
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--preset", required=True)
    p.add_argument("--ov1-manifest", default=DEFAULT_OV1_MANIFEST)
    p.add_argument("--ov2-manifest", default=DEFAULT_OV2_MANIFEST)
    p.add_argument("--ov3-manifest", default=DEFAULT_OV3_MANIFEST)
    p.add_argument("--ov1-real-manifest", default=DEFAULT_OV1_REAL_MANIFEST)
    p.add_argument("--ov2-real-manifest", default=DEFAULT_OV2_REAL_MANIFEST)
    p.add_argument("--ov3-real-manifest", default=DEFAULT_OV3_REAL_MANIFEST)
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--num-workers", type=int, default=8)
    p.add_argument("--amp", choices=("fp32", "bf16", "fp16"), default="bf16")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--output-json", default=None)
    p.add_argument("--activity-threshold", type=float, default=0.5)
    return p.parse_args()


def build_cfg(args: argparse.Namespace) -> TrainSpatialBEATsConfig:
    ns = SimpleNamespace(
        preset=args.preset,
        ov1_manifest=args.ov1_manifest,
        ov2_manifest=args.ov2_manifest,
        ov3_manifest=args.ov3_manifest,
        ov1_real_manifest=args.ov1_real_manifest,
        ov2_real_manifest=args.ov2_real_manifest,
        ov3_real_manifest=args.ov3_real_manifest,
        batch_size=None,
        num_workers=None,
        amp=None,
        num_epochs=None,
        learning_rate=None,
        weight_decay=None,
        output_dir=None,
        class_finetuned_ckpt=None,
        init_from_spatial_ckpt=None,
        resume=None,
        no_resume_optimizer=False,
        reset_epoch_on_resume=False,
        reset_best_on_resume=False,
        crop_mode=None,
        max_clip_duration_seconds=None,
        save_every_n_epochs=None,
        train_projector_in_stage1=False,
        freeze_trunk=False,
        no_progress=False,
        distributed=False,
        local_rank=None,
        distributed_backend=None,
        ddp_find_unused_parameters=False,
    )
    cfg = build_train_config_from_args(ns)
    cfg.batch_size = int(args.batch_size)
    cfg.num_workers = int(args.num_workers)
    cfg.amp_dtype = args.amp
    cfg.distributed = False
    cfg.show_progress_bars = True
    cfg.dump_val_predictions = False
    cfg.num_val_prediction_examples = 0
    # Force evaluation on sim ov1 test split only, no matter what the preset said.
    cfg.test_splits = ("test",)
    cfg.test_manifest_paths = (args.ov1_manifest,)
    cfg.train_splits = ()
    cfg.val_splits = ()
    return cfg


def load_model(ckpt_path: str, cfg: TrainSpatialBEATsConfig, device: torch.device) -> SpatialBEATs:
    model_cfg = build_model_config(cfg)
    model = SpatialBEATs(model_cfg)
    sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state_dict = sd["model_state_dict"] if "model_state_dict" in sd else sd.get("model", sd)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing:
        print(f"[Eval] WARN missing({len(missing)}): {missing[:6]}{'...' if len(missing) > 6 else ''}")
    if unexpected:
        print(f"[Eval] WARN unexpected({len(unexpected)}): {unexpected[:6]}{'...' if len(unexpected) > 6 else ''}")
    model.to(device).eval()
    return model


def build_loader(cfg: TrainSpatialBEATsConfig) -> torch.utils.data.DataLoader:
    ds_cfg = copy.deepcopy(build_dataset_config(cfg))
    ds_cfg.allowed_splits = cfg.test_splits
    path = cfg.test_manifest_paths[0]
    dataset = SpatialDataset(manifest_path=path, config=ds_cfg)
    print(f"[Eval] Test manifest: {path}")
    print(f"[Eval] Test size: {len(dataset)}")
    collate = functools.partial(collate_spatial_batch, config=ds_cfg)
    return torch.utils.data.DataLoader(
        dataset,
        batch_size=cfg.batch_size,
        shuffle=False,
        num_workers=cfg.num_workers,
        collate_fn=collate,
        pin_memory=True,
        drop_last=False,
        persistent_workers=cfg.num_workers > 0,
        prefetch_factor=4 if cfg.num_workers > 0 else None,
    )


def _amp_ctx(dtype: str):
    if not torch.cuda.is_available():
        return contextlib.nullcontext()
    if dtype == "bf16":
        return torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
    if dtype == "fp16":
        return torch.amp.autocast(device_type="cuda", dtype=torch.float16)
    return contextlib.nullcontext()


def _move_to_device(batch, device):
    field_vals = {}
    for f in dataclasses.fields(batch):
        v = getattr(batch, f.name)
        field_vals[f.name] = v.to(device) if isinstance(v, torch.Tensor) else v
    return type(batch)(**field_vals)


def main() -> None:
    args = parse_args()
    device = torch.device(args.device)
    print(f"[Eval] Device: {device}")
    print(f"[Eval] Checkpoint: {args.checkpoint}")
    print(f"[Eval] Preset: {args.preset}")

    cfg = build_cfg(args)
    assert cfg.loss.supervision_mode == "local_spatial_track", (
        f"Expected local_spatial_track, got {cfg.loss.supervision_mode}. "
        "This script is for track-supervised ckpts (v7f chain and descendants)."
    )
    if device.type != "cuda":
        cfg.amp_dtype = "fp32"

    model = load_model(args.checkpoint, cfg, device)
    loader = build_loader(cfg)

    running = {
        "oracle_class_acc": 0.0,
        "oracle_azi_mae_deg": 0.0,
        "oracle_ele_mae_deg": 0.0,
        "oracle_dist_mae": 0.0,
        "class_acc": 0.0,            # tier-1, activity-gated via training matcher
        "azi_mae_deg": 0.0,
        "ele_mae_deg": 0.0,
        "dist_mae": 0.0,
        "activity_precision": 0.0,
        "activity_recall": 0.0,
        "activity_acc": 0.0,
        "matched_count": 0.0,
    }
    num_batches = 0
    seld_acc = OfficialDCASEMetricsAccumulator()

    with torch.no_grad():
        for batch in tqdm(loader, desc="Eval sim ov1 test", leave=True):
            batch = _move_to_device(batch, device)
            with _amp_ctx(cfg.amp_dtype):
                model_output = model(
                    waveform=batch.waveform,
                    padding_mask=batch.waveform_padding_mask,
                    clip_duration_seconds=batch.clip_duration_seconds,
                    mono_window_mask=None,
                )
            pred_out = model_output.frame_track_prediction_output
            if pred_out is None:
                raise RuntimeError(
                    "frame_track_prediction_output is None — the loaded model does not "
                    "expose the track head.  Check readout_scheme / preset."
                )
            metric_output = compute_frame_track_validation_metrics(
                prediction_output=pred_out,
                batch=batch,
                temporal_padding_mask=model_output.temporal_padding_mask,
                config=cfg.loss,
            )
            accumulate_frame_track_seld(
                prediction_output=pred_out,
                batch=batch,
                temporal_padding_mask=model_output.temporal_padding_mask,
                accumulator=seld_acc,
                activity_threshold=args.activity_threshold,
            )
            for key in running:
                v = getattr(metric_output, key, None)
                if v is None:
                    continue
                running[key] += float(v.item())
            num_batches += 1

    metrics = {k: v / max(num_batches, 1) for k, v in running.items()}
    dcase = seld_acc.compute()
    metrics.update(dcase)

    print("\n" + "=" * 60)
    print("  v11a @ sim ov1 test split")
    print("=" * 60)
    print("  [classification]")
    print(f"    oracle_class_acc         : {metrics['oracle_class_acc']:.4f}")
    print(f"    class_acc (gated)        : {metrics['class_acc']:.4f}")
    print(f"    activity_precision       : {metrics['activity_precision']:.4f}")
    print(f"    activity_recall          : {metrics['activity_recall']:.4f}")
    print(f"    activity_acc (P-R)       : {metrics['activity_acc']:.4f}")
    print(f"    F20 (DCASE)              : {metrics['F20']:.4f}")
    print(f"    LR_CD (class-dep recall) : {metrics['LR_CD']:.4f}")
    print("  [spatial]")
    print(f"    oracle_azi_mae_deg       : {metrics['oracle_azi_mae_deg']:.2f}")
    print(f"    oracle_ele_mae_deg       : {metrics['oracle_ele_mae_deg']:.2f}")
    print(f"    oracle_dist_mae          : {metrics['oracle_dist_mae']:.4f}")
    print(f"    azi_mae_deg (gated)      : {metrics['azi_mae_deg']:.2f}")
    print(f"    ele_mae_deg (gated)      : {metrics['ele_mae_deg']:.2f}")
    print(f"    dist_mae (gated)         : {metrics['dist_mae']:.4f}")
    print(f"    LE_CD (DCASE, deg)       : {metrics['LE_CD']:.2f}")
    print(f"    ER20                     : {metrics['ER20']:.4f}")
    print(f"    SELD_score (lower=better): {metrics['SELD_score']:.4f}")
    print("=" * 60)

    out_path = args.output_json
    if out_path is None:
        out_path = str(Path(args.checkpoint).parent / "eval_ov1_sim_summary.json")
    with open(out_path, "w") as f:
        json.dump(
            {
                "checkpoint": args.checkpoint,
                "preset": args.preset,
                "manifest": cfg.test_manifest_paths[0],
                "split": list(cfg.test_splits),
                "activity_threshold": args.activity_threshold,
                "metrics": metrics,
            },
            f,
            indent=2,
            ensure_ascii=True,
        )
    print(f"[Eval] Summary saved to {out_path}")


if __name__ == "__main__":
    main()