#!/usr/bin/env python """Score val videos with VLAlert-v3 + BADAS, find 5 where VLAlert >> BADAS. Uses pre-computed belief caches (no VLM needed). Outputs selected videos to demo/C/selected_videos.json. """ import json, sys, logging, torch from pathlib import Path from collections import defaultdict from tqdm import tqdm ROOT = Path("PROJECT_ROOT") sys.path.insert(0, str(ROOT)) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") logger = logging.getLogger("select") device = "cuda" if torch.cuda.is_available() else "cpu" def load_val_gt(): """Load v5 val benchmark ground truth, grouped by video.""" lines = Path(ROOT / "data/cot_corpus_v3/v5_sft_val.jsonl").read_text().strip().split("\n") videos = {} tick_to_vid = {} for i, l in enumerate(lines): d = json.loads(l) vid = d["video_id"] actions = d.get("actions_per_frame", []) gt_action = actions[-1] if actions else "SILENT" cat = d.get("category", "") src = d.get("source", "") if vid not in videos: videos[vid] = {"ticks": [], "category": cat, "source": src} videos[vid]["ticks"].append({"idx": i, "gt": gt_action}) tick_to_vid[i] = vid return videos, tick_to_vid, len(lines) def load_badas_scores(n_ticks): """Load BADAS per-sample p_alert.""" d = json.load(open(ROOT / "eval_results/benchmark_v1_val/badas_per_sample.json")) scores = [] for i in range(n_ticks): p = d[str(i)]["p_alert"] if p > 0.5: action = "ALERT" elif p > 0.07: action = "OBSERVE" else: action = "SILENT" scores.append({"p_alert": p, "action": action}) return scores def load_vlalert_v3_scores(n_ticks, videos): """Run DangerHead + PolicyHead on v3 cache, return per-tick predictions.""" logger.info("Loading v3 cache + heads...") cache = torch.load(ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val_narrow.pt", weights_only=False, map_location="cpu") cache_ids = cache["ids"] cache_vid = cache.get("video_id", cache_ids) val_vids = set(videos.keys()) val_lines = Path(ROOT / "data/cot_corpus_v3/v5_sft_val.jsonl").read_text().strip().split("\n") vid_tick_counter = defaultdict(int) cache_idx_for_val = [] cache_vid_tick = defaultdict(list) for ci, vid in enumerate(cache_vid): cache_vid_tick[vid].append(ci) for i, l in enumerate(val_lines): d = json.loads(l) vid = d["video_id"] tick_num = vid_tick_counter[vid] vid_tick_counter[vid] += 1 if vid in cache_vid_tick and tick_num < len(cache_vid_tick[vid]): cache_idx_for_val.append(cache_vid_tick[vid][tick_num]) else: cache_idx_for_val.append(-1) matched = sum(1 for x in cache_idx_for_val if x >= 0) logger.info(f"Matched {matched}/{n_ticks} val ticks to v3 cache") from lkalert.models.danger_head import DangerHead from lkalert.models.policy_head_v2 import PolicyHeadV2 ck = torch.load(ROOT / "checkpoints/danger_v3_hazard/best.pt", weights_only=False, map_location="cpu") danger = DangerHead(in_dim=ck["in_dim"], n_hazards=int(ck.get("n_hazards", 0) or 0)).to(device).eval() danger.load_state_dict(ck["model"]) pk = torch.load(ROOT / "checkpoints/policy_v3_strong/best.pt", weights_only=False, map_location="cpu") sd = pk["model"] mapped = {k.replace("fuse.0.", "fuse_pre.0.").replace("fuse.3.", "cls_head."): v for k, v in sd.items()} policy = PolicyHeadV2( policy_dim=pk.get("policy_dim", 2560), perception_dim_per_query=pk.get("perception_dim_per_query", 512), k_queries=pk.get("k_queries", 4), ).to(device).eval() policy.load_state_dict(mapped, strict=False) belief_all = cache["belief_content"] policy_all = cache["policy_position"] valid_all = cache["valid_frames"] results = [] BS = 128 logger.info("Running DangerHead + PolicyHead on val ticks...") for start in tqdm(range(0, n_ticks, BS), desc="v3 heads"): end = min(start + BS, n_ticks) idxs = cache_idx_for_val[start:end] valid_idxs = [x for x in idxs if x >= 0] if not valid_idxs: for _ in range(end - start): results.append({"action": "SILENT", "p_alert": 0.0}) continue b = belief_all[valid_idxs].to(device, dtype=torch.float32) pp = policy_all[valid_idxs].to(device, dtype=torch.float32) v = valid_all[valid_idxs].to(device) prev = torch.full((len(valid_idxs),), 3, device=device, dtype=torch.long) with torch.no_grad(): d_out = danger(b, valid_frames=v) logits = policy(pp, d_out["perception_summary"], d_out["per_frame"], prev, valid_frames=v) probs = torch.softmax(logits, dim=-1) j = 0 for i_rel in range(end - start): ci = idxs[i_rel] if ci < 0: results.append({"action": "SILENT", "p_alert": 0.0}) else: p_alert = float(probs[j, 2].cpu()) p_obs = float(probs[j, 1].cpu()) act_idx = int(probs[j].argmax().cpu()) action = ["SILENT", "OBSERVE", "ALERT"][act_idx] results.append({"action": action, "p_alert": p_alert, "p_observe": p_obs}) j += 1 return results def select_top_videos(videos, badas_scores, vlalert_scores, n=5): """Select videos where VLAlert >> BADAS.""" scores = [] for vid, info in videos.items(): if info["category"] not in ("ego_positive",): continue n_alert_gt = sum(1 for t in info["ticks"] if t["gt"] == "ALERT") if n_alert_gt == 0: continue badas_correct_alert = 0 vlalert_correct_alert = 0 badas_false_alert = 0 vlalert_false_alert = 0 badas_miss = 0 vlalert_miss = 0 for t in info["ticks"]: idx = t["idx"] gt = t["gt"] ba = badas_scores[idx]["action"] va = vlalert_scores[idx]["action"] if gt == "ALERT": if ba == "ALERT": badas_correct_alert += 1 else: badas_miss += 1 if va == "ALERT": vlalert_correct_alert += 1 else: vlalert_miss += 1 elif gt == "SILENT": if ba == "ALERT": badas_false_alert += 1 if va == "ALERT": vlalert_false_alert += 1 advantage = (vlalert_correct_alert - badas_correct_alert) - 0.5 * (vlalert_false_alert - badas_false_alert) if advantage > 0: scores.append({ "video_id": vid, "source": info["source"], "category": info["category"], "n_ticks": len(info["ticks"]), "n_alert_gt": n_alert_gt, "vlalert_correct": vlalert_correct_alert, "badas_correct": badas_correct_alert, "vlalert_miss": vlalert_miss, "badas_miss": badas_miss, "vlalert_fa": vlalert_false_alert, "badas_fa": badas_false_alert, "advantage": advantage, }) scores.sort(key=lambda x: x["advantage"], reverse=True) selected = [] sources_used = set() for s in scores: if len(selected) >= n: break if len(selected) >= 3 and s["source"] in sources_used: continue selected.append(s) sources_used.add(s["source"]) if len(selected) < n: for s in scores: if len(selected) >= n: break if s not in selected: selected.append(s) return selected def main(): out_dir = ROOT / "demo/C" out_dir.mkdir(exist_ok=True) videos, tick_to_vid, n_ticks = load_val_gt() logger.info(f"Val: {n_ticks} ticks, {len(videos)} videos") badas_scores = load_badas_scores(n_ticks) logger.info(f"BADAS: {n_ticks} scores loaded") vlalert_scores = load_vlalert_v3_scores(n_ticks, videos) logger.info(f"VLAlert-v3: {len(vlalert_scores)} scores") selected = select_top_videos(videos, badas_scores, vlalert_scores, n=5) logger.info(f"\n{'='*60}") logger.info(f" Top 5 videos where VLAlert >> BADAS") logger.info(f"{'='*60}") for i, s in enumerate(selected): logger.info(f" #{i+1}: {s['video_id']} ({s['source']}/{s['category']})") logger.info(f" {s['n_ticks']} ticks, {s['n_alert_gt']} GT ALERT") logger.info(f" VLAlert: {s['vlalert_correct']}/{s['n_alert_gt']} correct, {s['vlalert_fa']} FA") logger.info(f" BADAS: {s['badas_correct']}/{s['n_alert_gt']} correct, {s['badas_fa']} FA") logger.info(f" Advantage: {s['advantage']:.1f}") # Save per-tick predictions for selected videos for s in selected: vid = s["video_id"] info = videos[vid] ticks = [] for t in info["ticks"]: idx = t["idx"] ticks.append({ "tick_idx": idx, "gt": t["gt"], "badas": badas_scores[idx], "vlalert_v3": vlalert_scores[idx], }) s["ticks"] = ticks json.dump(selected, open(out_dir / "selected_videos.json", "w"), indent=2) logger.info(f"\nSaved → {out_dir / 'selected_videos.json'}") if __name__ == "__main__": main()