#!/usr/bin/env python3 """Inject belief-token fields (clip-level action + per-frame action trajectory) into an existing Nexar CoT JSONL. POMDP per-frame label rule (user thresholds 2026-04-22): label=0 / no time_of_event -> SILENT (all frames) label=1 & 0 <= tta_t < 0.5s -> ALERT label=1 & 0.5s <= tta_t < 2.5s -> OBSERVE label=1 & tta_t >= 2.5s or tta_t < 0 -> SILENT Clip-level action (legacy, kept for backward compat): SILENT if label=0 ALERT if label=1 and tta_alert < 1.5s OBSERVE otherwise (where tta_alert = time_of_event - time_of_alert) Input: - data/vla_cot/*.jsonl (GPT-4o teacher CoT) - nexar-collision-prediction/train.csv Output fields added per record: belief.action : clip-level action token (legacy) belief.tta_sec : time_of_event - time_of_alert belief.actions_per_frame : list of T action strings (POMDP target) belief.frame_times_sec : list of T float seconds (sampled times) Usage: # per-frame POMDP mode (new default): python -m training.VLA.augment_cot_with_belief \ --in_jsonl data/vla_cot/train500_cot.jsonl \ --train_csv nexar-collision-prediction/train.csv \ --video_dir nexar-collision-prediction/train \ --out_jsonl data/vla_cot_belief/train500_perframe.jsonl \ --per_frame --n_frames 8 --alert_s 0.5 --observe_s 2.5 """ from __future__ import annotations import argparse import json from pathlib import Path from typing import List, Optional def derive_clip_action(label: int, tta_alert: Optional[float]) -> str: """Legacy clip-level action (tta = time_of_event - time_of_alert).""" if label == 0: return "SILENT" if tta_alert is None or tta_alert < 0: return "OBSERVE" if tta_alert < 1.5: return "ALERT" return "OBSERVE" def per_frame_actions( label: int, time_of_event: Optional[float], frame_times: List[float], alert_s: float = 0.5, observe_s: float = 2.5, ) -> List[str]: """Strict POMDP per-frame target. tta_t = time_of_event - frame_time_t.""" if label == 0 or time_of_event is None: return ["SILENT"] * len(frame_times) out: List[str] = [] for ft in frame_times: tta = time_of_event - ft if tta < 0: out.append("SILENT") # frame is AFTER event elif tta < alert_s: out.append("ALERT") elif tta < observe_s: out.append("OBSERVE") else: out.append("SILENT") # too early, still safe return out def _probe_video(video_dir: Path, clip_id: str) -> tuple[int, float]: """Return (total_frames, fps); falls back to (0, 30) on failure.""" try: import cv2 cap = cv2.VideoCapture(str(video_dir / f"{clip_id}.mp4")) total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = float(cap.get(cv2.CAP_PROP_FPS)) or 30.0 cap.release() except Exception: total, fps = 0, 30.0 return total, fps def _frame_indices_and_times( total: int, fps: float, time_of_event: Optional[float], n_frames: int, event_anchored: bool, lookback_s: float, post_margin_s: float, ) -> tuple[list[int], list[float]]: """Return (frame_indices, frame_times_sec) for one clip.""" import numpy as np if total <= 0 or fps <= 0: return list(range(n_frames)), [i * (1.0 / 30.0) for i in range(n_frames)] if event_anchored and time_of_event is not None and time_of_event >= 0: end_s = min(float(total - 1) / fps, time_of_event + post_margin_s) start_s = max(0.0, end_s - (lookback_s + post_margin_s)) times = np.linspace(start_s, end_s, n_frames) idx = [int(round(t * fps)) for t in times] times_out = [float(i) / fps for i in idx] return idx, times_out idx = np.linspace(0, total - 1, n_frames).round().astype(int).tolist() return idx, [float(i) / fps for i in idx] def main(): ap = argparse.ArgumentParser() ap.add_argument("--in_jsonl", default="data/vla_cot/train500_cot.jsonl") ap.add_argument("--train_csv", default="nexar-collision-prediction/train.csv") ap.add_argument("--video_dir", default="nexar-collision-prediction/train") ap.add_argument("--out_jsonl", default="data/vla_cot_belief/train500_belief.jsonl") ap.add_argument("--per_frame", action="store_true", help="Also emit belief.actions_per_frame / frame_indices / frame_times_sec") ap.add_argument("--n_frames", type=int, default=8) ap.add_argument("--alert_s", type=float, default=0.5) ap.add_argument("--observe_s", type=float, default=2.5) ap.add_argument("--event_anchored", action="store_true", default=True, help="Positives: sample T frames in [event - lookback, event + post_margin]") ap.add_argument("--no_event_anchored", dest="event_anchored", action="store_false") ap.add_argument("--lookback_s", type=float, default=3.0, help="Seconds of pre-event context for event-anchored sampling") ap.add_argument("--post_margin_s", type=float, default=0.0, help="Seconds after event to include (usually 0; a small >0 " "helps fence-post frames fall in the ALERT window)") args = ap.parse_args() import pandas as pd df = pd.read_csv(args.train_csv, dtype={"id": str}) df["id"] = df["id"].str.zfill(5) toe_map: dict[str, float] = {} tta_alert_map: dict[str, float] = {} for _, row in df.iterrows(): toe = row.get("time_of_event") toa = row.get("time_of_alert") cid = row["id"] if pd.notna(toe): toe_map[cid] = float(toe) if pd.notna(toe) and pd.notna(toa): tta_alert_map[cid] = float(toe) - float(toa) video_dir = Path(args.video_dir) out_path = Path(args.out_jsonl) out_path.parent.mkdir(parents=True, exist_ok=True) n_in = n_out = 0 clip_counts = {"ALERT": 0, "OBSERVE": 0, "SILENT": 0} frame_counts = {"ALERT": 0, "OBSERVE": 0, "SILENT": 0} with open(args.in_jsonl) as f_in, open(out_path, "w") as f_out: for line in f_in: n_in += 1 rec = json.loads(line) clip_id = str(rec["id"]).zfill(5) label = int(rec["label"]) tta_a = tta_alert_map.get(clip_id) clip_action = derive_clip_action(label, tta_a) clip_counts[clip_action] += 1 belief = { "action": clip_action, "tta_sec": round(tta_a, 3) if tta_a is not None else -1.0, } if args.per_frame: toe = toe_map.get(clip_id) total, fps = _probe_video(video_dir, clip_id) frame_idx, frame_times = _frame_indices_and_times( total=total, fps=fps, time_of_event=toe, n_frames=args.n_frames, event_anchored=args.event_anchored, lookback_s=args.lookback_s, post_margin_s=args.post_margin_s, ) actions_pf = per_frame_actions(label, toe, frame_times, alert_s=args.alert_s, observe_s=args.observe_s) for a in actions_pf: frame_counts[a] += 1 belief["actions_per_frame"] = actions_pf belief["frame_indices"] = frame_idx belief["frame_times_sec"] = [round(t, 3) for t in frame_times] belief["time_of_event_sec"] = round(toe, 3) if toe is not None else -1.0 belief["alert_s"] = args.alert_s belief["observe_s"] = args.observe_s belief["sampling"] = "event_anchored" if (args.event_anchored and toe is not None) else "uniform" belief["total_frames"] = total belief["fps"] = round(fps, 3) belief["lookback_s"] = args.lookback_s belief["post_margin_s"] = args.post_margin_s rec["belief"] = belief f_out.write(json.dumps(rec) + "\n") n_out += 1 print(f"input={n_in} output={n_out}") print(f"clip-level action histogram: {clip_counts}") if args.per_frame: total_frames = sum(frame_counts.values()) rate = {k: f"{v}/{total_frames} ({v/max(1,total_frames)*100:.1f}%)" for k, v in frame_counts.items()} print(f"per-frame action histogram: {rate}") print(f"saved: {out_path}") if __name__ == "__main__": main()