tools/build_v5_benchmark.py

"""Build v5 unified benchmark on ALL 132,530 records.

For EVERY record (not just GPT):
  1. Update action labels from annotation.json (DADA + Nexar)
     DAD + DoTA already correct in _relabeled2
  2. Update/replace belief content:
     - If annotation.json has per_frame_beliefs → use those
     - Else if record has GPT belief → keep GPT
     - Else → generate from action-appropriate bank
  3. Mark belief_source field accordingly

Input:  v4_sft_{train,val,test}_full_relabeled2.jsonl (132,530 total)
Output: v5_sft_{train,val,test}.jsonl (132,530 total, same split)
"""
from __future__ import annotations
import json, hashlib, logging
from pathlib import Path
from collections import Counter, defaultdict

ROOT = Path("PROJECT_ROOT")
COT_DIR = ROOT / "data/cot_corpus_v3"
DADA_ROOT = ROOT / "DADA-2000"
NEXAR_ROOT = ROOT / "NEXAR_COLLISION/dataset"
DOTA_ANN = ROOT / "DoTA/annotations"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("v5")

# ─── Belief banks for records without GPT or annotation beliefs ───
SILENT_BANK = [
    "clear road ahead, normal traffic flow, no hazards detected",
    "steady driving, lane markings visible, surroundings stable",
    "open road with no immediate threats, maintaining safe speed",
    "traffic moving smoothly, no sudden changes observed",
    "routine driving conditions, road surface in good condition",
    "normal lane keeping, no vehicles encroaching from adjacent lanes",
    "safe following distance maintained, lead vehicle steady",
    "no pedestrians or cyclists in the immediate vicinity",
    "driving straight ahead, visibility is clear, no obstructions",
    "surrounding traffic is predictable, no erratic behavior",
    "no signs of developing hazard, all lanes flowing freely",
    "intersection clear, no conflicting traffic approaching",
    "highway driving, vehicles spaced evenly, no sudden braking",
    "residential area, low traffic volume, no unexpected obstacles",
    "parked vehicles on roadside, path clear ahead",
    "road markings intact, lane boundaries well defined",
    "crosswalk ahead but no pedestrians waiting to cross",
    "street lighting adequate, visibility acceptable",
    "wet road surface but traction appears normal",
    "cyclist on bike lane to the right, separated by marking",
]

OBSERVE_BANK = [
    "subtle change in traffic pattern, monitoring situation closely",
    "vehicle behavior ahead appears irregular, heightened awareness",
    "potential hazard developing, increased attention to surroundings",
    "traffic flow disruption possible, watching for sudden changes",
    "lead vehicle showing unusual behavior, preparing for response",
    "gap closing with vehicle ahead, monitoring deceleration",
    "unusual movement detected, staying alert",
    "road conditions may be changing, scanning for hazards",
    "intersection dynamics evolving, watching for conflicting paths",
    "pedestrian activity near roadway, heightened awareness required",
    "braking pattern of lead vehicle suggests caution ahead",
    "merging traffic creating tighter spacing, monitoring closely",
    "vehicle in adjacent lane drifting, keeping safe distance",
    "construction zone approach, expecting lane changes",
    "emergency vehicle audible, scanning for approach direction",
]

ALERT_BANK = [
    "imminent collision risk, emergency response needed",
    "critical proximity to obstacle, immediate action required",
    "vehicle cutting across path, collision risk high",
    "rapid closure with lead vehicle, braking needed now",
    "pedestrian in path, immediate alert required",
    "hard brake or evasive maneuver needed, critical situation",
    "near-impact distance, immediate driver intervention",
    "lead vehicle suddenly braking, critical TTC",
    "vehicle entering intersection on collision course",
    "loss of control situation developing, alert driver",
]

def _pick(bank, seed_str):
    h = int(hashlib.md5(seed_str.encode()).hexdigest(), 16)
    return bank[h % len(bank)]


def load_dada_annotations():
    lookup = {}
    for cat in ["positive", "non-ego", "negative"]:
        cat_dir = DADA_ROOT / cat
        if not cat_dir.exists(): continue
        for clip_dir in cat_dir.iterdir():
            ann_path = clip_dir / "annotation.json"
            if not ann_path.exists(): continue
            ann = json.load(open(ann_path))
            lookup[f"dada_{clip_dir.name}"] = ann
    return lookup


def load_nexar_annotations():
    lookup = {}
    for split in ["train", "test-public", "test-private"]:
        for pol in ["positive", "negative"]:
            parent = NEXAR_ROOT / split / pol
            if not parent.exists(): continue
            for clip_dir in parent.iterdir():
                if not clip_dir.is_dir(): continue
                ann_path = clip_dir / "annotation.json"
                if not ann_path.exists(): continue
                ann = json.load(open(ann_path))
                lookup[f"nexar_{clip_dir.name}"] = ann
    return lookup


def load_dota_annotations():
    lookup = {}
    for p in sorted(DOTA_ANN.glob("*.json")):
        d = json.load(open(p))
        vname = d.get("video_name", p.stem)
        lookup[vname] = d
    return lookup


def map_labels(frame_indices, per_frame_labels):
    n = len(per_frame_labels) if per_frame_labels else 0
    return [per_frame_labels[fi] if 0 <= fi < n else "SILENT" for fi in frame_indices]


def map_beliefs(frame_indices, per_frame_beliefs):
    if not per_frame_beliefs: return [None] * len(frame_indices)
    n = len(per_frame_beliefs)
    return [per_frame_beliefs[fi] if 0 <= fi < n and per_frame_beliefs[fi] else None
            for fi in frame_indices]


def fill_missing_beliefs(actions, beliefs, vid, frame_indices):
    """For any frame where belief is None, generate from the appropriate bank."""
    result = list(beliefs) if beliefs else [None] * 8
    for i in range(len(actions)):
        if result[i] is None or result[i] == "":
            fi = frame_indices[i] if i < len(frame_indices) else i
            seed = f"{vid}_{fi}"
            act = actions[i] if i < len(actions) else "SILENT"
            if act == "ALERT":
                result[i] = _pick(ALERT_BANK, seed)
            elif act == "OBSERVE":
                result[i] = _pick(OBSERVE_BANK, seed)
            else:
                result[i] = _pick(SILENT_BANK, seed)
    return result


def main():
    logger.info("Loading annotations...")
    dada_ann = load_dada_annotations()
    nexar_ann = load_nexar_annotations()
    dota_ann = load_dota_annotations()
    logger.info(f"  DADA: {len(dada_ann)}  Nexar: {len(nexar_ann)}  DoTA: {len(dota_ann)}")

    for split in ["v4_sft_train_full", "v4_sft_val_full", "v4_sft_test_full"]:
        in_path = COT_DIR / f"{split}_relabeled2.jsonl"
        out_tag = split.replace("v4_sft_", "v5_sft_").replace("_full", "")
        out_path = COT_DIR / f"{out_tag}.jsonl"
        if not in_path.exists():
            logger.warning(f"skip {in_path}"); continue

        stats = Counter()
        src_action = defaultdict(Counter)

        with in_path.open() as fin, out_path.open("w") as fout:
            for ln in fin:
                ln = ln.strip()
                if not ln: continue
                rec = json.loads(ln)
                src = rec.get("source", "?")
                vid = rec.get("video_id", "")
                fi = rec.get("frame_indices", [])
                old_beliefs = rec.get("beliefs_per_frame", [None]*8)

                # ── 1. Update action labels ──
                if src == "dada" and vid in dada_ann:
                    ann = dada_ann[vid]
                    pfl = ann.get("per_frame_labels", [])
                    if pfl and fi:
                        new_acts = map_labels(fi, pfl)
                        rec["actions_per_frame"] = new_acts
                        rec["tick_action"] = new_acts[-1]
                    stats["dada_action_updated"] += 1

                elif src == "nexar" and vid in nexar_ann:
                    ann = nexar_ann[vid]
                    pfl = ann.get("per_frame_labels", [])
                    if pfl and fi:
                        new_acts = map_labels(fi, pfl)
                        rec["actions_per_frame"] = new_acts
                        rec["tick_action"] = new_acts[-1]
                    stats["nexar_action_updated"] += 1

                # DAD + DoTA: already correct in _relabeled2

                # ── 2. Update belief content ──
                acts = rec.get("actions_per_frame", ["SILENT"]*8)
                ann_beliefs = None

                if src == "dada" and vid in dada_ann:
                    pfb = dada_ann[vid].get("per_frame_beliefs")
                    if pfb:
                        ann_beliefs = map_beliefs(fi, pfb)

                elif src == "dota":
                    vid_key = vid.replace("dota_", "", 1) if vid.startswith("dota_") else vid
                    if vid_key in dota_ann:
                        pfb = dota_ann[vid_key].get("per_frame_beliefs")
                        if pfb:
                            ann_beliefs = map_beliefs(fi, pfb)

                # Merge: annotation > GPT > bank-generated
                merged = [None] * 8
                for i in range(8):
                    ab = ann_beliefs[i] if ann_beliefs and i < len(ann_beliefs) else None
                    gb = old_beliefs[i] if i < len(old_beliefs) and old_beliefs[i] else None
                    merged[i] = ab if ab else gb  # prefer annotation over GPT

                # Fill remaining Nones from bank
                merged = fill_missing_beliefs(acts, merged, vid, fi)
                rec["beliefs_per_frame"] = merged

                # Update belief_source
                has_gpt = rec.get("belief_source") in ("gpt4o",)
                has_ann = ann_beliefs and any(b is not None for b in ann_beliefs)
                if has_ann and has_gpt:
                    rec["belief_source"] = "annotation+gpt4o"
                elif has_ann:
                    rec["belief_source"] = "annotation"
                elif has_gpt:
                    rec["belief_source"] = "gpt4o"
                else:
                    rec["belief_source"] = "auto_generated"

                src_action[src][rec.get("tick_action", "?")] += 1
                stats[f"{src}_total"] += 1
                fout.write(json.dumps(rec) + "\n")

        total = sum(v for k, v in stats.items() if k.endswith("_total"))
        logger.info(f"[{out_tag}] {total} records written → {out_path}")
        for src in ['dad', 'dada', 'dota', 'nexar']:
            sa = src_action.get(src, {})
            s = sa.get('SILENT',0); o = sa.get('OBSERVE',0); a = sa.get('ALERT',0)
            t = s+o+a
            if t > 0:
                logger.info(f"  {src:>8s}: S={s:>6d} O={o:>5d} A={a:>5d}  total={t}")

    # Summary
    print("\n" + "=" * 80)
    print("  v5 Benchmark — ALL 132,530 records")
    print("=" * 80)
    for tag in ["v5_sft_train", "v5_sft_val", "v5_sft_test"]:
        path = COT_DIR / f"{tag}.jsonl"
        if not path.exists(): continue
        acts = Counter(); srcs = Counter(); bsrcs = Counter()
        with open(path) as f:
            for ln in f:
                d = json.loads(ln)
                acts[d.get("tick_action","?")] += 1
                srcs[d.get("source","?")] += 1
                bsrcs[d.get("belief_source","?")] += 1
        n = sum(acts.values())
        s,o,a = acts.get("SILENT",0), acts.get("OBSERVE",0), acts.get("ALERT",0)
        print(f"\n  {tag}: {n:,} records")
        print(f"    sources: {dict(srcs)}")
        print(f"    actions: SILENT={s:,} ({100*s/n:.1f}%)  OBSERVE={o:,} ({100*o/n:.1f}%)  ALERT={a:,} ({100*a/n:.1f}%)")
        print(f"    belief:  {dict(bsrcs)}")


if __name__ == "__main__":
    main()