training/Nexar/nexar_train_extractor.py

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