eval_v12_per_subset.py

#!/usr/bin/env python3
"""Per-subset evaluation for v12 (ov1_unified_v12).

对 v12 训练得到的 checkpoint（unified_spatial_foa_fsd63_all 数据集上训练）
按子集独立评估，输出每个子集的 F20/ER20/LE_CD/LR_CD/SELD_score
以及 oracle_class_acc / oracle_azi_mae_deg / oracle_ele_mae_deg / oracle_dist_mae
等诊断指标，用于分析：
  1) cls vs 空间 谁是瓶颈
  2) 真实 vs 仿真 的 gap
  3) 多源（ov2/ov3）vs 单源 的差距

支持的子集：
  - ov1_sim / ov2_sim / ov3_sim   （旧仿真 OV 数据，static，所有帧同一 DOA）
  - ov1_real / ov2_real / ov3_real（DCASE real static mapped，更接近真实混响）
  - dcase_starss                  （DCASE STARSS23 / 22 / TAU 混合，real dynamic）
  - unified_valid / unified_test  （新 unified 数据，混合 sim_static + qa_sim + dcase_real）

用法：
    python eval_v12_per_subset.py \\
        --checkpoint checkpoints/spatial_beats_ov1_unified_v12_exp/03_ov123_top4/best.pt \\
        --preset ov1_unified_v12 \\
        --split valid \\
        --batch-size 8 --num-workers 8 --amp bf16

    # 跑 test 集合
    python eval_v12_per_subset.py \\
        --checkpoint checkpoints/spatial_beats_ov1_unified_v12_exp/03_ov123_top4/best.pt \\
        --preset ov1_unified_v12 \\
        --split test \\
        --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, Tuple

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_DCASE_STARSS_VALID_MANIFEST,
    DEFAULT_OV1_MANIFEST,
    DEFAULT_OV2_MANIFEST,
    DEFAULT_OV3_MANIFEST,
    DEFAULT_OV1_REAL_MANIFEST,
    DEFAULT_OV2_REAL_MANIFEST,
    DEFAULT_OV3_REAL_MANIFEST,
    DEFAULT_UNIFIED_TRAIN_MANIFEST,
    DEFAULT_UNIFIED_VALID_MANIFEST,
    TrainSpatialBEATsConfig,
    build_dataset_config,
    build_model_config,
    build_train_config_from_args,
)


DEFAULT_DCASE_STARSS_TEST_MANIFEST = (
    "/apdcephfs_cq10/share_1603164/user/schmittzhu/data/metadata/"
    "dcase_starss_foa.test.jsonl"
)
DEFAULT_UNIFIED_TEST_MANIFEST = (
    "/apdcephfs_cq12/share_302080740/user/schmittzhu/data/"
    "unified_spatial_foa_fsd63_all/test.jsonl"
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--preset", required=True)
    p.add_argument(
        "--split",
        choices=("valid", "test"),
        default="valid",
        help="valid uses each dataset's valid split; test uses the test split",
    )
    # sim manifests (ov1/2/3)
    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)
    # real (DCASE static mapped) manifests
    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)
    # dcase_starss (full real, dynamic)
    p.add_argument(
        "--dcase-starss-valid-manifest",
        default=DEFAULT_DCASE_STARSS_VALID_MANIFEST,
    )
    p.add_argument(
        "--dcase-starss-test-manifest",
        default=DEFAULT_DCASE_STARSS_TEST_MANIFEST,
    )
    # unified
    p.add_argument("--unified-valid-manifest", default=DEFAULT_UNIFIED_VALID_MANIFEST)
    p.add_argument("--unified-test-manifest", default=DEFAULT_UNIFIED_TEST_MANIFEST)
    # runtime
    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)
    p.add_argument(
        "--max-samples-per-subset",
        type=int,
        default=-1,
        help="-1 = no cap (evaluate everything)",
    )
    p.add_argument(
        "--only-subsets",
        default="",
        help="comma-separated subset names to evaluate (empty=all). "
        "Options: ov1_sim,ov2_sim,ov3_sim,ov1_real,ov2_real,ov3_real,"
        "dcase_starss,unified",
    )
    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,
        dcase_starss_valid_manifest=args.dcase_starss_valid_manifest,
        unified_train_manifest=DEFAULT_UNIFIED_TRAIN_MANIFEST,
        unified_valid_manifest=args.unified_valid_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
    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 unexpected:
        print(f"[Eval] WARN unexpected({len(unexpected)}): {unexpected[:6]}")
    model.to(device).eval()
    return model


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 build_subset_plan(args: argparse.Namespace) -> List[Tuple[str, str, Tuple[str, ...]]]:
    """返回 [(subset_name, manifest_path, allowed_splits), ...]."""
    if args.split == "valid":
        plan = [
            ("ov1_sim", args.ov1_manifest, ("valid",)),
            ("ov2_sim", args.ov2_manifest, ("valid",)),
            ("ov3_sim", args.ov3_manifest, ("valid",)),
            ("ov1_real", args.ov1_real_manifest, ("valid",)),
            ("ov2_real", args.ov2_real_manifest, ("valid",)),
            ("ov3_real", args.ov3_real_manifest, ("valid",)),
            ("dcase_starss", args.dcase_starss_valid_manifest, ("valid",)),
            ("unified", args.unified_valid_manifest, ("valid",)),
        ]
    else:
        # test：旧 ov 数据集有 test split，dcase 也有
        plan = [
            ("ov1_sim", args.ov1_manifest, ("test",)),
            ("ov2_sim", args.ov2_manifest, ("test",)),
            ("ov3_sim", args.ov3_manifest, ("test",)),
            ("ov1_real", args.ov1_real_manifest, ("test",)),
            ("ov2_real", args.ov2_real_manifest, ("test",)),
            ("ov3_real", args.ov3_real_manifest, ("test",)),
            ("dcase_starss", args.dcase_starss_test_manifest, ("test",)),
            ("unified", args.unified_test_manifest, ("test",)),
        ]

    filter_set = {s.strip() for s in args.only_subsets.split(",") if s.strip()}
    if filter_set:
        plan = [entry for entry in plan if entry[0] in filter_set]
    # drop entries whose manifest file does not exist
    filtered = []
    for name, path, splits in plan:
        if not Path(path).exists():
            print(f"[Eval] SKIP subset={name} (manifest not found: {path})")
            continue
        filtered.append((name, path, splits))
    return filtered


def _empty_running() -> Dict[str, float]:
    return {
        "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,
        "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,
    }


def eval_one_subset(
    model: SpatialBEATs,
    cfg: TrainSpatialBEATsConfig,
    subset_name: str,
    manifest_path: str,
    allowed_splits: Tuple[str, ...],
    device: torch.device,
    activity_threshold: float,
    max_samples: int,
) -> Dict[str, float]:
    ds_cfg = copy.deepcopy(build_dataset_config(cfg))
    ds_cfg.allowed_splits = allowed_splits
    dataset = SpatialDataset(manifest_path=manifest_path, config=ds_cfg)
    total = len(dataset)
    print(f"\n[Eval][{subset_name}] manifest={manifest_path}")
    print(f"[Eval][{subset_name}] split={allowed_splits}  size={total}")
    if total == 0:
        print(f"[Eval][{subset_name}] empty — skip")
        return {"subset": subset_name, "manifest": manifest_path, "size": 0}
    if max_samples > 0 and total > max_samples:
        indices = list(range(max_samples))
        dataset = torch.utils.data.Subset(dataset, indices)
        print(f"[Eval][{subset_name}] capped to first {max_samples}")

    collate = functools.partial(collate_spatial_batch, config=ds_cfg)
    loader = 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,
    )

    running = _empty_running()
    num_batches = 0
    seld_acc = OfficialDCASEMetricsAccumulator()
    with torch.no_grad():
        for batch in tqdm(loader, desc=f"Eval {subset_name}", leave=False):
            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")
            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=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()}
    metrics.update(seld_acc.compute())
    metrics["subset"] = subset_name
    metrics["manifest"] = manifest_path
    metrics["size"] = total
    return metrics


def _fmt(v, digits=4):
    try:
        return f"{float(v):.{digits}f}"
    except Exception:
        return "nan"


def print_summary(all_metrics: List[Dict[str, float]], split: str) -> None:
    if not all_metrics:
        print("[Eval] no metrics")
        return
    print("\n" + "=" * 96)
    print(f"  v12 per-subset ({split} split)")
    print("=" * 96)
    hdr = (
        f"{'subset':<16} {'N':>6} "
        f"{'F20':>6} {'ER20':>6} {'LE_CD':>7} {'LR_CD':>6} {'SELD':>6}  "
        f"{'o_cls':>6} {'o_azi':>6} {'o_ele':>6} {'o_dst':>6}  "
        f"{'a_P':>5} {'a_R':>5}"
    )
    print(hdr)
    print("-" * len(hdr))
    for m in all_metrics:
        if m.get("size", 0) == 0:
            print(f"{m['subset']:<16} {'EMPTY':>6}")
            continue
        print(
            f"{m['subset']:<16} "
            f"{m['size']:>6} "
            f"{_fmt(m.get('F20'), 4):>6} "
            f"{_fmt(m.get('ER20'), 4):>6} "
            f"{_fmt(m.get('LE_CD'), 2):>7} "
            f"{_fmt(m.get('LR_CD'), 4):>6} "
            f"{_fmt(m.get('SELD_score'), 4):>6}  "
            f"{_fmt(m.get('oracle_class_acc'), 4):>6} "
            f"{_fmt(m.get('oracle_azi_mae_deg'), 2):>6} "
            f"{_fmt(m.get('oracle_ele_mae_deg'), 2):>6} "
            f"{_fmt(m.get('oracle_dist_mae'), 4):>6}  "
            f"{_fmt(m.get('activity_precision'), 3):>5} "
            f"{_fmt(m.get('activity_recall'), 3):>5}"
        )
    print("=" * 96)
    print("  legend:  F20↑ ER20↓ LE_CD↓ LR_CD↑ SELD↓  "
          "o_cls=oracle_class_acc  o_azi/ele=oracle doa MAE (deg)  "
          "a_P/a_R=activity precision/recall")
    print("=" * 96)


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}")
    print(f"[Eval] Split: {args.split}")

    cfg = build_cfg(args)
    assert cfg.loss.supervision_mode == "local_spatial_track", (
        f"Expected local_spatial_track, got {cfg.loss.supervision_mode}."
    )
    if device.type != "cuda":
        cfg.amp_dtype = "fp32"

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

    plan = build_subset_plan(args)
    all_metrics: List[Dict[str, float]] = []
    for subset_name, manifest_path, allowed_splits in plan:
        m = eval_one_subset(
            model=model,
            cfg=cfg,
            subset_name=subset_name,
            manifest_path=manifest_path,
            allowed_splits=allowed_splits,
            device=device,
            activity_threshold=args.activity_threshold,
            max_samples=args.max_samples_per_subset,
        )
        all_metrics.append(m)
        # 打印单子集即时结果
        if m.get("size", 0) > 0:
            print(
                f"[Eval][{subset_name}] F20={_fmt(m.get('F20'))}  "
                f"ER20={_fmt(m.get('ER20'))}  LE_CD={_fmt(m.get('LE_CD'), 2)}  "
                f"LR_CD={_fmt(m.get('LR_CD'))}  SELD={_fmt(m.get('SELD_score'))}  "
                f"o_cls={_fmt(m.get('oracle_class_acc'))}  "
                f"o_azi={_fmt(m.get('oracle_azi_mae_deg'), 2)}  "
                f"o_ele={_fmt(m.get('oracle_ele_mae_deg'), 2)}  "
                f"o_dst={_fmt(m.get('oracle_dist_mae'))}  "
                f"a_P={_fmt(m.get('activity_precision'))}  "
                f"a_R={_fmt(m.get('activity_recall'))}"
            )

    print_summary(all_metrics, args.split)

    out_path = args.output_json
    if out_path is None:
        out_path = str(
            Path(args.checkpoint).parent
            / f"eval_v12_per_subset_{args.split}.json"
        )
    with open(out_path, "w") as f:
        json.dump(
            {
                "checkpoint": args.checkpoint,
                "preset": args.preset,
                "split": args.split,
                "activity_threshold": args.activity_threshold,
                "per_subset": all_metrics,
            },
            f,
            indent=2,
            ensure_ascii=True,
        )
    print(f"\n[Eval] Summary saved to {out_path}")


if __name__ == "__main__":
    main()