#!/usr/bin/env python3 """ Extract belief vectors from Nexar video clips using the frozen SFT backbone. For each clip, we sample one or more temporal windows, run through the SFT model, and cache [belief, tta_mean, tta_var, p_alert] per window. Usage (feature extraction): python -m training.Nexar.nexar_extractor \ --sft_checkpoint checkpoints/SFT/sft_v2/best \ --policy_checkpoint checkpoints/Policy/policy_warmstart_v2/best \ --video_dir nexar-collision-prediction/test \ --out_file data/nexar_cache/test.pt \ --n_windows 3 \ --batch_size 8 python -m training.Nexar.nexar_extractor \ --sft_checkpoint checkpoints/SFT/sft_v2/best \ --policy_checkpoint checkpoints/Policy/policy_warmstart_v2/best \ --video_dir NEXAR_COLLISION/train/positive \ --out_file data/nexar_cache/train_positive.pt \ --n_windows 3 \ --batch_size 8 """ from __future__ import annotations import argparse import logging import os from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import torch import torch.nn.functional as F from tqdm import tqdm 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 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.extractor") # Frame resolution passed to the VLM (lower = faster, still captures content) FRAME_W = 640 FRAME_H = 360 @torch.no_grad() def extract_features_for_clips( model: PolicyModel, video_paths: List[str], video_ids: List[str], n_windows: int = 3, window_dur_s: float = 3.0, n_frames: int = 8, batch_size: int = 4, end_offset_s: float = 0.0, ) -> Dict[str, dict]: """ Process each clip through the frozen SFT backbone. For each clip, extract n_windows temporal windows and compute: belief [n_windows, hidden_dim] tta_mean [n_windows] tta_var [n_windows] p_alert [n_windows] P(ALERT) from PolicyHead p_obs [n_windows] P(OBSERVE) p_silent [n_windows] P(SILENT) Returns: dict keyed by video_id → {beliefs, tta_means, tta_vars, p_alerts, ...} """ from torch.amp import autocast model.eval() results: Dict[str, dict] = {} # Build flat list of (video_id, window_idx, frames) tasks # We batch across windows AND clips to maximise GPU utilisation flat_tasks: List[Tuple[str, int, List]] = [] logger.info(f"Loading frames from {len(video_paths)} clips ({n_windows} windows each) ...") for vid_path, vid_id in zip(tqdm(video_paths, desc="Loading frames"), video_ids): try: windows = sample_multi_windows( vid_path, n_windows, window_dur_s, n_frames, FRAME_W, FRAME_H, end_offset_s, ) except Exception as e: logger.warning(f" Frame extract failed for {vid_id}: {e}") # Use dummy frames from PIL import Image dummy = [Image.new("RGB", (FRAME_W, FRAME_H), (64, 64, 64))] * n_frames windows = [dummy] * n_windows for w_idx, frames in enumerate(windows): flat_tasks.append((vid_id, w_idx, frames)) logger.info(f"Total VLM forward passes: {len(flat_tasks)} (batch={batch_size})") # Process in batches for i in tqdm(range(0, len(flat_tasks), batch_size), desc="VLM encode"): batch_tasks = flat_tasks[i : i + batch_size] batch_images = [t[2] for t in batch_tasks] # list of List[PIL] batch_metadata = [{} for _ in batch_tasks] # no metadata → "Urban driving" try: inputs = model._build_inputs(batch_images, batch_metadata) 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 = model.sft.encode_observation(inputs) tta_mean, tta_logvar = model.sft.tta_head(beliefs) tta_var = torch.exp(tta_logvar.float().clamp(-20, 20)) tta_mean_f = tta_mean.float() beliefs_f = beliefs.float() B = beliefs_f.shape[0] prev_action = torch.zeros(B, dtype=torch.long, device=model.device) logits = model.policy_head(beliefs_f, tta_mean_f, tta_var, prev_action) probs = F.softmax(logits, dim=-1) # [B, 3] except Exception as e: logger.warning(f" VLM batch failed (i={i}): {e}") B = len(batch_tasks) beliefs_f = torch.zeros(B, model.hidden_dim) tta_mean_f = torch.full((B,), 10.0) tta_var = torch.ones(B) probs = torch.full((B, 3), 1/3) for j, (vid_id, w_idx, _) in enumerate(batch_tasks): if vid_id not in results: results[vid_id] = { "beliefs": [], "tta_means": [], "tta_vars": [], "p_silent": [], "p_obs": [], "p_alert": [], } r = results[vid_id] r["beliefs"].append(beliefs_f[j].cpu()) r["tta_means"].append(tta_mean_f[j].item()) r["tta_vars"].append(tta_var[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()) # Stack per-video tensors for vid_id, r in results.items(): r["beliefs"] = torch.stack(r["beliefs"]) # [n_windows, hidden_dim] r["tta_means"] = torch.tensor(r["tta_means"]) # [n_windows] r["tta_vars"] = torch.tensor(r["tta_vars"]) # [n_windows] r["p_silent"] = torch.tensor(r["p_silent"]) r["p_obs"] = torch.tensor(r["p_obs"]) r["p_alert"] = torch.tensor(r["p_alert"]) return results def cache_to_pt(results: dict, video_ids: list, out_file: str): """Save extraction results to a .pt cache file.""" out = Path(out_file) out.parent.mkdir(parents=True, exist_ok=True) torch.save({"video_ids": video_ids, "features": results}, out) logger.info(f"Saved cache → {out} ({len(results)} clips)") def load_cache(cache_file: str) -> Tuple[List[str], Dict[str, dict]]: """Load .pt cache produced by cache_to_pt.""" d = torch.load(cache_file, map_location="cpu", weights_only=False) return d["video_ids"], d["features"] def main(): parser = argparse.ArgumentParser("nexar_extractor") parser.add_argument("--sft_checkpoint", required=True) parser.add_argument("--policy_checkpoint", default=None) parser.add_argument("--video_dir", required=True, help="Directory of .mp4 clips to process") parser.add_argument("--out_file", required=True, help="Output .pt cache file") parser.add_argument("--n_windows", type=int, default=3) parser.add_argument("--window_dur", type=float, default=3.0) parser.add_argument("--n_frames", type=int, default=8) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--end_offset_s", type=float, default=0.0, help="Skip last N seconds of clip (e.g. for train videos with known TTE)") parser.add_argument("--max_clips", type=int, default=0, help="Limit number of clips (0=all); for debugging") args = parser.parse_args() if Path(args.out_file).exists(): logger.info(f"Cache already exists: {args.out_file} — skipping.") return model = PolicyModel(args.sft_checkpoint, use_bf16=True) if args.policy_checkpoint: model.load_policy_checkpoint(args.policy_checkpoint) video_dir = Path(args.video_dir) video_files = sorted(video_dir.glob("*.mp4")) if args.max_clips > 0: video_files = video_files[:args.max_clips] video_paths = [str(v) for v in video_files] video_ids = [v.stem for v in video_files] logger.info(f"Processing {len(video_paths)} clips from {video_dir}") results = extract_features_for_clips( model, video_paths, video_ids, n_windows = args.n_windows, window_dur_s = args.window_dur, n_frames = args.n_frames, batch_size = args.batch_size, end_offset_s = args.end_offset_s, ) cache_to_pt(results, video_ids, args.out_file) if __name__ == "__main__": main()