training/Nexar/mvit_dataset.py

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