training/VLA/augment_cot_with_belief.py

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