#!/usr/bin/env python3 """ MViT video dataset for Nexar collision prediction. Loads raw .mp4 clips and returns [C, T, H, W] tensors for MViT-v2-s. Design choices: - T=16 frames (MViT-v2-s default, matches Kinetics pretraining) - H=W=224 (ImageNet-normalised) - For TRAIN positive videos: sample from 10s window ending at TTE before event - For TRAIN negative videos: sample from the last 10s of the video - For TEST clips: sample from the entire clip (already ~10s) Data-centric filtering (key insight from 1st-place winner): - Use time_of_event - time_of_alert as "clarity score" - Remove positives where the warning is very sudden (< min_warning_s = 0.5s) because they look like normal driving until the very last moment - Keep all negatives and "clear" positives """ from __future__ import annotations import logging import random from pathlib import Path from typing import List, Optional, Tuple import cv2 import numpy as np import pandas as pd import torch from torch.utils.data import Dataset from torchvision.transforms import v2 as T logger = logging.getLogger("Nexar.mvit_dataset") # MViT-v2-s canonical parameters N_FRAMES = 16 IMG_SIZE = 224 CLIP_DUR_S = 10.0 # temporal window to sample from (seconds) TTE_LIST = [0.5, 1.0, 1.5] # TTE offsets for positive train clips MEAN = [0.45, 0.45, 0.45] STD = [0.225, 0.225, 0.225] def _get_video_info(path: str) -> Tuple[float, int]: """Returns (fps, n_frames).""" try: import decord decord.bridge.set_bridge("native") vr = decord.VideoReader(path) return vr.get_avg_fps(), len(vr) except Exception: cap = cv2.VideoCapture(path) fps = cap.get(cv2.CAP_PROP_FPS) n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap.release() return fps, n def _sample_frames( path: str, start_s: float, end_s: float, n_frames: int = N_FRAMES, img_size: int = IMG_SIZE, ) -> Optional[torch.Tensor]: """ Load n_frames uniformly from [start_s, end_s] of the video. Returns FloatTensor [C, T, H, W] in [0, 1], or None on failure. """ try: import decord decord.bridge.set_bridge("native") vr = decord.VideoReader(path, width=img_size, height=img_size) fps = vr.get_avg_fps() n = len(vr) ts = [start_s + (end_s - start_s) * i / (n_frames - 1) for i in range(n_frames)] indices = [max(0, min(int(t * fps), n - 1)) for t in ts] frames = vr.get_batch(indices).asnumpy() # [T, H, W, C] uint8 # → [T, C, H, W] → float [0,1] → [C, T, H, W] t_tensor = torch.from_numpy(frames).permute(0, 3, 1, 2).float() / 255.0 return t_tensor.permute(1, 0, 2, 3) # [C, T, H, W] except Exception as e: logger.warning(f"Frame sample failed for {path}: {e}") return None def _normalize_video(video: torch.Tensor) -> torch.Tensor: """Normalize [C, T, H, W] video with ImageNet stats.""" mean = torch.tensor(MEAN, dtype=video.dtype).view(3, 1, 1, 1) std = torch.tensor(STD, dtype=video.dtype).view(3, 1, 1, 1) return (video - mean) / std def _hflip_video(video: torch.Tensor) -> torch.Tensor: """Horizontally flip [C, T, H, W] video.""" return video.flip(-1) class VideoTransform: """Simple video transform that handles [C, T, H, W] tensors.""" def __init__(self, train: bool = True): self.train = train def __call__(self, video: torch.Tensor) -> torch.Tensor: video = _normalize_video(video) if self.train: if torch.rand(1).item() > 0.5: video = _hflip_video(video) # Mild brightness/contrast jitter factor = 1.0 + (torch.rand(1).item() - 0.5) * 0.2 video = (video * factor).clamp(-5, 5) return video def _make_train_transforms(img_size: int = IMG_SIZE) -> VideoTransform: return VideoTransform(train=True) def _make_val_transforms() -> VideoTransform: return VideoTransform(train=False) class NexarMViTDataset(Dataset): """ Dataset for MViT fine-tuning on Nexar collision prediction. train_mode=True → creates multiple clips per positive video (one per TTE) and augments randomly train_mode=False → creates one clip per video (test or val) """ def __init__( self, csv_path: str, video_dir: str, # root dir containing {vid_id}.mp4 train_mode: bool = True, pos_subdir: str = "", # if set, positive videos are in video_dir/pos_subdir/ neg_subdir: str = "", # if set, negative videos are in video_dir/neg_subdir/ min_warning_s: float = 0.3, # filter positives with very short warning windows tte_list: List[float] = TTE_LIST, n_frames: int = N_FRAMES, img_size: int = IMG_SIZE, clip_dur_s: float = CLIP_DUR_S, is_test: bool = False, # test mode: no labels, single clip per video ): self.train_mode = train_mode self.n_frames = n_frames self.img_size = img_size self.clip_dur_s = clip_dur_s self.is_test = is_test self.tfm = _make_train_transforms(img_size) if train_mode else _make_val_transforms() df = pd.read_csv(csv_path) video_dir = Path(video_dir) self.samples: List[dict] = [] if is_test: # Test mode: each ID = one .mp4 clip, no label for _, row in df.iterrows(): vid_id = str(int(float(row["id"]))).zfill(5) vid_path = video_dir / f"{vid_id}.mp4" if not vid_path.exists(): continue self.samples.append({ "vid_id": vid_id, "path": str(vid_path), "label": -1, "start_s": 0.0, "end_s": -1.0, # -1 = use full clip }) else: for _, row in df.iterrows(): vid_id = str(int(float(row["id"]))).zfill(5) label = int(row["target"]) t_event = float(row["time_of_event"]) if pd.notna(row.get("time_of_event")) else None t_alert = float(row["time_of_alert"]) if pd.notna(row.get("time_of_alert")) else None # Locate video file if pos_subdir and label == 1: vid_path = video_dir / pos_subdir / f"{vid_id}.mp4" elif neg_subdir and label == 0: vid_path = video_dir / neg_subdir / f"{vid_id}.mp4" else: vid_path = video_dir / f"{vid_id}.mp4" if not vid_path.exists(): continue if label == 1 and t_event is not None: # Data-centric filter: skip sudden collisions with very short warning if t_alert is not None: warning_s = t_event - t_alert if warning_s < min_warning_s: continue # too ambiguous if train_mode: # Multiple clips: one per TTE offset for tte in tte_list: end_s = t_event - tte start_s = max(0.0, end_s - clip_dur_s) self.samples.append({ "vid_id": vid_id, "path": str(vid_path), "label": 1, "start_s": start_s, "end_s": end_s, }) else: # Validation: single clip at TTE=0.5s end_s = t_event - 0.5 start_s = max(0.0, end_s - clip_dur_s) self.samples.append({ "vid_id": vid_id, "path": str(vid_path), "label": 1, "start_s": start_s, "end_s": end_s, }) else: # Negative video: sample last clip_dur_s seconds self.samples.append({ "vid_id": vid_id, "path": str(vid_path), "label": 0, "start_s": -1.0, # -1 = auto from end "end_s": -1.0, }) n_pos = sum(1 for s in self.samples if s["label"] == 1) n_neg = sum(1 for s in self.samples if s["label"] == 0) logger.info( f"NexarMViTDataset [train={train_mode}, test={is_test}]: " f"{len(self.samples)} samples pos={n_pos} neg={n_neg}" ) def __len__(self) -> int: return len(self.samples) def __getitem__(self, idx: int) -> dict: s = self.samples[idx] # Resolve start/end for "last clip" sampling start_s, end_s = s["start_s"], s["end_s"] if start_s < 0 or end_s < 0: fps, n_total = _get_video_info(s["path"]) if fps <= 0: fps = 30.0 duration = n_total / fps end_s = duration start_s = max(0.0, duration - self.clip_dur_s) frames = _sample_frames(s["path"], start_s, end_s, self.n_frames, self.img_size) if frames is None: frames = torch.zeros(3, self.n_frames, self.img_size, self.img_size) frames = self.tfm(frames) return { "video": frames, # [C, T, H, W] "label": torch.tensor(s["label"], dtype=torch.float32), "vid_id": s["vid_id"], } def make_train_val_split( full_csv: str, val_frac: float = 0.15, seed: int = 42, min_warning_s: float = 0.3, ) -> Tuple[pd.DataFrame, pd.DataFrame]: """Stratified split returning (train_df, val_df) DataFrames.""" df = pd.read_csv(full_csv) # Filter ambiguous positives from TRAINING set (keep all in validation) pos_df = df[df["target"] == 1].copy() neg_df = df[df["target"] == 0].copy() if "time_of_event" in df.columns and "time_of_alert" in df.columns: mask = pos_df["time_of_event"].notna() & pos_df["time_of_alert"].notna() pos_df.loc[mask, "warning_s"] = pos_df.loc[mask, "time_of_event"] - pos_df.loc[mask, "time_of_alert"] else: pos_df["warning_s"] = float("nan") rng = random.Random(seed) pos_idx = pos_df.index.tolist() neg_idx = neg_df.index.tolist() rng.shuffle(pos_idx) rng.shuffle(neg_idx) n_val_pos = max(1, int(len(pos_idx) * val_frac)) n_val_neg = max(1, int(len(neg_idx) * val_frac)) val_idx = pos_idx[:n_val_pos] + neg_idx[:n_val_neg] train_idx = pos_idx[n_val_pos:] + neg_idx[n_val_neg:] return df.loc[train_idx].reset_index(drop=True), df.loc[val_idx].reset_index(drop=True)