tools/score_val_select_demos.py

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