"""VLAlert-X v2 Phase 5 — score a benchmark via cached features + heads. Load a dual-stream cache (from `tools/make_cache_x_v2.py`) and the trained DangerHead + PolicyHeadV2 checkpoints. Forward to produce per-tick action probabilities, then save in the standard `per_tick.pt` schema that the existing `tools/compute_daus_*.py` utilities consume. Output schema: { "ids": list[str] (N,) "indices": LongTensor [N] "scores_binary": FloatTensor [N, 1] # P(ALERT) "scores_3class": FloatTensor [N, 1, 3] # P(S), P(O), P(A) "tta_per_tick": FloatTensor [N, 1] "frame_indices": LongTensor [N, 8] "category": list[str] "source": list[str] "tta_raw": FloatTensor [N] "n_ticks": int = 1 "method": "VLAlert-X-v2" "danger_ckpt": str "policy_ckpt": str } Usage: python tools/score_vlalert_x_v2.py \ --cache data/belief_cache_v2/sft_x_v2__multisrc_val_full.pt \ --manifest data/cot_corpus_v2/multisrc_val_full_perframe_v2.jsonl \ --danger_ckpt checkpoints/danger_v2/seed2/best.pt \ --policy_ckpt checkpoints/policy_v2/seed2/best.pt \ --out eval_results/aus_metric/multisrc_per_tick/vlalert_x_v2.pt """ from __future__ import annotations import argparse import json import logging import sys from pathlib import Path from typing import Dict import torch import torch.nn.functional as F from tqdm import tqdm ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from lkalert.models.danger_head import DangerHead from lkalert.models.policy_head_v2 import PolicyHeadV2 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("score_vlalert_x_v2") @torch.no_grad() def main(): ap = argparse.ArgumentParser() ap.add_argument("--cache", type=Path, required=True, help="Dual-stream cache .pt from tools/make_cache_x_v2.py") ap.add_argument("--manifest", type=Path, required=True, help="Perframe-v2 jsonl that was used to build the cache " "(needed for category/source/tta_raw metadata)") ap.add_argument("--danger_ckpt", type=Path, required=True) ap.add_argument("--policy_ckpt", type=Path, required=True) ap.add_argument("--out", type=Path, required=True) ap.add_argument("--batch_size", type=int, default=64) ap.add_argument("--prev_action", type=int, default=3, help="prev_action embedding index; 3=BOS (no temporal context)") args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" # ── load cache ── logger.info(f"[load] cache: {args.cache}") d = torch.load(args.cache, weights_only=False, map_location="cpu") belief = d["belief_content"].float() # [N, 8, D_belief] policy = d["policy_position"].float() # [N, 8, D_policy] valid = d["valid_frames"] # [N, 8] bool ids_cache = list(d["ids"]) N = belief.shape[0] logger.info(f" N={N} belief={tuple(belief.shape)} policy={tuple(policy.shape)}") # ── load manifest for metadata (category, source, tta_raw, frame_indices) ── logger.info(f"[load] manifest: {args.manifest}") meta_by_id: Dict[str, Dict] = {} with open(args.manifest) as f: for ln in f: ln = ln.strip() if not ln: continue r = json.loads(ln) mid = r.get("id") or r.get("video_id") if mid: meta_by_id[mid] = r logger.info(f" manifest records: {len(meta_by_id)}") # ── load heads ── logger.info(f"[load] DangerHead: {args.danger_ckpt}") ck_d = torch.load(args.danger_ckpt, weights_only=False, map_location="cpu") danger = DangerHead(in_dim=ck_d["in_dim"]).to(device) danger.load_state_dict(ck_d["model"]) danger.eval() logger.info(f"[load] PolicyHeadV2: {args.policy_ckpt}") ck_p = torch.load(args.policy_ckpt, weights_only=False, map_location="cpu") policy_head = PolicyHeadV2( policy_dim=ck_p["policy_dim"], perception_dim_per_query=ck_p["perception_dim_per_query"], k_queries=ck_p["k_queries"], ).to(device) policy_head.load_state_dict(ck_p["model"]) policy_head.eval() logger.info(f" Phase 3 val: per_frame_auc={ck_d['val_metrics'].get('per_frame_auc',0):.4f}") logger.info(f" Phase 4 val: bal_acc={ck_p['val_metrics']['balanced_acc']:.4f} " f"per_class_recall={ck_p['val_metrics']['per_class_recall']}") # ── infer per-tick scores ── scores_3class = torch.zeros(N, 1, 3, dtype=torch.float32) n_failed = 0 prev_act_tensor = torch.full((args.batch_size,), args.prev_action, dtype=torch.long, device=device) bs = args.batch_size for i in tqdm(range(0, N, bs), ncols=80, desc="infer"): end = min(N, i + bs) b_belief = belief[i:end].to(device, non_blocking=True) b_policy = policy[i:end].to(device, non_blocking=True) b_valid = valid[i:end].to(device, non_blocking=True) cur_bs = end - i # Danger forward → perception_summary + danger_per_frame d_out = danger(b_belief, valid_frames=b_valid) perc = d_out["perception_summary"] # [B, K, hidden] danger_pf = d_out["per_frame"] # [B, 8] # Policy forward prev = prev_act_tensor[:cur_bs] logits = policy_head(b_policy, perc, danger_pf, prev, valid_frames=b_valid) # [B, 3] probs = F.softmax(logits, dim=-1).cpu() scores_3class[i:end, 0] = probs scores_binary = scores_3class[:, :, 2].clone() # P(ALERT) # ── assemble per_tick.pt metadata ── ids_out: list = [] cat_out: list = [] src_out: list = [] tta_raw_out = torch.zeros(N, dtype=torch.float32) tta_per_tick_out = torch.zeros(N, 1, dtype=torch.float32) frame_indices_out = torch.zeros(N, 8, dtype=torch.long) indices_out = torch.arange(N, dtype=torch.long) # IMPORTANT: cache stores per-tick category/source already (correctly # tied to each specific tick's TTA). Manifest meta_by_id dedups on # video_id and clobbers earlier ticks' category — DON'T use it for # category/source. Only use manifest for `frame_indices` lookup. cache_category = list(d.get("category", [""] * N)) cache_source = list(d.get("source", [""] * N)) cache_tick_tta = d.get("tick_tta_raw", torch.full((N,), -1.0)) n_missing_meta = 0 for i, vid in enumerate(ids_cache): m = meta_by_id.get(vid, {}) if not m: n_missing_meta += 1 ids_out.append(vid) cat_out.append(cache_category[i] if i < len(cache_category) else "") src_out.append(cache_source[i] if i < len(cache_source) else "") tta_v = (cache_tick_tta[i].item() if hasattr(cache_tick_tta[i], "item") else float(cache_tick_tta[i])) tta_raw_out[i] = tta_v tta_per_tick_out[i, 0] = tta_v fi = m.get("frame_indices", [0]*8) frame_indices_out[i] = torch.tensor(fi[:8], dtype=torch.long) if n_missing_meta: logger.warning(f" {n_missing_meta} cache ids had no matching manifest record " f"(only frame_indices lost; category/source still correct from cache)") out_dict = { "ids": ids_out, "indices": indices_out, "scores_binary": scores_binary, "scores_3class": scores_3class, "tta_per_tick": tta_per_tick_out, "frame_indices": frame_indices_out, "category": cat_out, "source": src_out, "tta_raw": tta_raw_out, "n_ticks": 1, "method": "VLAlert-X-v2", "danger_ckpt": str(args.danger_ckpt), "policy_ckpt": str(args.policy_ckpt), } args.out.parent.mkdir(parents=True, exist_ok=True) torch.save(out_dict, args.out) logger.info(f"[save] {args.out}") logger.info(f" N={len(ids_out)} " f"P(ALERT) range=[{scores_binary.min():.4f}, {scores_binary.max():.4f}]") # distribution by category from collections import Counter cc = Counter(cat_out) logger.info(f" category dist: {dict(cc)}") if __name__ == "__main__": main()