training/Nexar/nexar_extractor.py

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