#!/usr/bin/env python3 """ Extract Nexar TRAIN features with proper TTE-alignment. For positive videos: - We know time_of_event from train.csv - We create N synthetic clips per video, each ending at TTE=[0.5, 1.0, 1.5]s before event - For each clip, we extract n_windows temporal windows from the last 9s of the clip - The clip length mirrors the test clips (~10s) For negative videos: - No event time — extract windows from the LAST 10s of the video - (Negative videos contain normal driving; the last portion is most similar to test) This ensures train/test feature distributions are aligned. Usage: python -m training.Nexar.nexar_train_extractor \ --sft_checkpoint checkpoints/SFT/sft_v2/best \ --policy_checkpoint checkpoints/Policy/policy_warmstart_v2/best \ --train_csv nexar-collision-prediction/train.csv \ --train_pos_dir NEXAR_COLLISION/train/positive \ --train_neg_dir NEXAR_COLLISION/train/negative \ --out_dir data/nexar_cache \ --n_windows 3 \ --batch_size 8 """ from __future__ import annotations import argparse import logging from pathlib import Path from typing import Dict, List, Optional, Tuple import pandas as pd import torch from torch.amp import autocast from tqdm import tqdm import torch.nn.functional as F import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from training.Policy.policy_model import PolicyModel from training.Nexar.video_utils import sample_multi_windows, get_video_info logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.train_extractor") FRAME_W = 640 FRAME_H = 360 # TTE offsets for synthetic positive clips TTE_OFFSETS = [0.5, 1.0, 1.5] # seconds before event CLIP_DURATION = 9.0 # seconds per synthetic clip N_WINDOWS_DEFAULT = 3 WINDOW_DUR_DEFAULT = 3.0 # each window covers 3s N_FRAMES_DEFAULT = 8 def build_positive_tasks( train_csv: str, pos_dir: str, n_windows: int, window_dur: float, n_frames: int, max_clips: int = 0, ) -> List[Tuple[str, str, float, int]]: """ Returns list of (video_path, clip_id, end_time_s, label) for positive training clips. clip_id = f"{vid_id}_tte{tte_ms}" """ df = pd.read_csv(train_csv) pos_df = df[df["target"] == 1].dropna(subset=["time_of_event"]) tasks = [] for _, row in pos_df.iterrows(): vid_id = str(row["id"]).zfill(5) vid_path = Path(pos_dir) / f"{vid_id}.mp4" if not vid_path.exists(): logger.warning(f"Missing positive video: {vid_path}") continue t_event = float(row["time_of_event"]) for tte in TTE_OFFSETS: end_time = t_event - tte # clip ends this many seconds before event if end_time < CLIP_DURATION: end_time = CLIP_DURATION # ensure at least 9s of context clip_id = f"{vid_id}_tte{int(tte*1000)}" tasks.append((str(vid_path), clip_id, end_time, 1)) if max_clips > 0: tasks = tasks[:max_clips] logger.info(f"Positive training clips: {len(tasks)} (from {len(pos_df)} videos × {len(TTE_OFFSETS)} TTEs)") return tasks def build_negative_tasks( train_csv: str, neg_dir: str, n_per_video: int = 1, max_clips: int = 0, ) -> List[Tuple[str, str, float, int]]: """ Returns list of (video_path, clip_id, end_time_s, label) for negative training clips. end_time_s = duration of the video (sample from the end) """ df = pd.read_csv(train_csv) neg_df = df[df["target"] == 0] tasks = [] for _, row in neg_df.iterrows(): vid_id = str(row["id"]).zfill(5) vid_path = Path(neg_dir) / f"{vid_id}.mp4" if not vid_path.exists(): logger.warning(f"Missing negative video: {vid_path}") continue # Use video end (negative videos: last 9s = representative sample) for i in range(n_per_video): clip_id = f"{vid_id}_neg{i}" tasks.append((str(vid_path), clip_id, -1.0, 0)) # -1 = use video end if max_clips > 0: tasks = tasks[:max_clips] logger.info(f"Negative training clips: {len(tasks)} (from {len(neg_df)} videos × {n_per_video} clips)") return tasks @torch.no_grad() def extract_tasks( model: PolicyModel, tasks: List[Tuple[str, str, float, int]], # (path, clip_id, end_time, label) n_windows: int, window_dur: float, n_frames: int, batch_size: int, ) -> Tuple[List[str], Dict[str, dict], List[int]]: """Process all tasks through SFT backbone.""" model.eval() # Load all frames first logger.info(f"Loading frames for {len(tasks)} tasks ...") flat: List[Tuple[str, int, List]] = [] # (clip_id, window_idx, frames) for vid_path, clip_id, end_time, label in tqdm(tasks, desc="Loading frames"): fps, n_total = get_video_info(vid_path) if fps <= 0: fps = 30.0 duration = n_total / fps if end_time < 0: end_t = duration else: end_t = min(end_time, duration) clip_start = max(0.0, end_t - n_windows * window_dur) try: from PIL import Image import numpy as np import decord decord.bridge.set_bridge("native") vr = decord.VideoReader(vid_path, width=FRAME_W, height=FRAME_H) n_vid = len(vr) for w_idx in range(n_windows): ws = clip_start + w_idx * window_dur we = ws + window_dur we = min(we, end_t) times = [ws + (we - ws) * k / (n_frames - 1) for k in range(n_frames)] indices = [max(0, min(int(t * fps), n_vid - 1)) for t in times] frame_arr = vr.get_batch(indices).asnumpy() frames = [Image.fromarray(f) for f in frame_arr] flat.append((clip_id, w_idx, frames)) except Exception as e: logger.warning(f"Frame load failed for {clip_id}: {e}") from PIL import Image dummy = [Image.new("RGB", (FRAME_W, FRAME_H), (64, 64, 64))] * n_frames for w_idx in range(n_windows): flat.append((clip_id, w_idx, dummy)) logger.info(f"Total VLM passes: {len(flat)} (batch={batch_size})") # Process in batches results: Dict[str, dict] = {} for i in tqdm(range(0, len(flat), batch_size), desc="VLM encode"): batch_tasks = flat[i : i + batch_size] batch_imgs = [t[2] for t in batch_tasks] batch_meta = [{} for _ in batch_tasks] try: inputs = model._build_inputs(batch_imgs, batch_meta) inputs = {k: v.to(model.device) for k, v in inputs.items() if hasattr(v, "to")} with autocast(device_type="cuda", dtype=model._amp_dtype, enabled=True): beliefs_b = model.sft.encode_observation(inputs) tta_mean_b, tta_lv_b = model.sft.tta_head(beliefs_b) tta_var_b = torch.exp(tta_lv_b.float().clamp(-20, 20)) bel_f = beliefs_b.float() tmu_f = tta_mean_b.float() B = bel_f.shape[0] prev = torch.zeros(B, dtype=torch.long, device=model.device) logits = model.policy_head(bel_f, tmu_f, tta_var_b, prev) probs = F.softmax(logits, dim=-1) except Exception as e: logger.warning(f"VLM batch i={i} failed: {e}") B = len(batch_tasks) bel_f = torch.zeros(B, model.hidden_dim) tmu_f = torch.full((B,), 10.0) tta_var_b = torch.ones(B) probs = torch.full((B, 3), 1/3) for j, (clip_id, w_idx, _) in enumerate(batch_tasks): if clip_id not in results: results[clip_id] = { "beliefs": [], "tta_means": [], "tta_vars": [], "p_silent": [], "p_obs": [], "p_alert": [], } r = results[clip_id] r["beliefs"].append(bel_f[j].cpu()) r["tta_means"].append(tmu_f[j].item()) r["tta_vars"].append(tta_var_b[j].item()) r["p_silent"].append(probs[j][0].item()) r["p_obs"].append(probs[j][1].item()) r["p_alert"].append(probs[j][2].item()) for clip_id, r in results.items(): r["beliefs"] = torch.stack(r["beliefs"]) r["tta_means"] = torch.tensor(r["tta_means"]) r["tta_vars"] = torch.tensor(r["tta_vars"]) r["p_silent"] = torch.tensor(r["p_silent"]) r["p_obs"] = torch.tensor(r["p_obs"]) r["p_alert"] = torch.tensor(r["p_alert"]) clip_ids = [t[1] for t in tasks] labels = [t[3] for t in tasks] return clip_ids, results, labels def main(): parser = argparse.ArgumentParser("nexar_train_extractor") parser.add_argument("--sft_checkpoint", required=True) parser.add_argument("--policy_checkpoint", default=None) parser.add_argument("--train_csv", default="nexar-collision-prediction/train.csv") parser.add_argument("--train_pos_dir", default="NEXAR_COLLISION/train/positive") parser.add_argument("--train_neg_dir", default="NEXAR_COLLISION/train/negative") parser.add_argument("--out_dir", default="data/nexar_cache") parser.add_argument("--n_windows", type=int, default=N_WINDOWS_DEFAULT) parser.add_argument("--window_dur", type=float, default=WINDOW_DUR_DEFAULT) parser.add_argument("--n_frames", type=int, default=N_FRAMES_DEFAULT) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--max_clips", type=int, default=0, help="Debug: limit number of positive clips (0=all)") args = parser.parse_args() out_pos = Path(args.out_dir) / "train_positive.pt" out_neg = Path(args.out_dir) / "train_negative.pt" # Check if caches exist pos_exists = out_pos.exists() neg_exists = out_neg.exists() if pos_exists and neg_exists: logger.info("Both train caches already exist — skipping extraction.") return Path(args.out_dir).mkdir(parents=True, exist_ok=True) model = PolicyModel(args.sft_checkpoint, use_bf16=True) if args.policy_checkpoint: model.load_policy_checkpoint(args.policy_checkpoint) # ── Positive ────────────────────────────────────────────────────────────── if not pos_exists: logger.info("Building positive train cache ...") pos_tasks = build_positive_tasks( args.train_csv, args.train_pos_dir, args.n_windows, args.window_dur, args.n_frames, args.max_clips, ) clip_ids, results, labels = extract_tasks( model, pos_tasks, args.n_windows, args.window_dur, args.n_frames, args.batch_size, ) torch.save({ "video_ids": clip_ids, "labels": labels, "features": results, }, out_pos) logger.info(f"Saved → {out_pos}") else: logger.info(f"Positive cache exists: {out_pos}") # ── Negative ────────────────────────────────────────────────────────────── if not neg_exists: logger.info("Building negative train cache ...") neg_tasks = build_negative_tasks( args.train_csv, args.train_neg_dir, n_per_video=1, max_clips=args.max_clips * 3 if args.max_clips > 0 else 0, ) clip_ids, results, labels = extract_tasks( model, neg_tasks, args.n_windows, args.window_dur, args.n_frames, args.batch_size, ) torch.save({ "video_ids": clip_ids, "labels": labels, "features": results, }, out_neg) logger.info(f"Saved → {out_neg}") else: logger.info(f"Negative cache exists: {out_neg}") logger.info("\n✅ Train feature extraction complete.") if __name__ == "__main__": main()