#!/usr/bin/env python3 """ SFT Dataset — manifest-based, dual-head (hazard + TTA). Sample categories ----------------- ego_positive : ego-vehicle crash; hazard_label=1, TTA supervised non_ego : accident in scene, not ego-relevant; hazard_label=0 (soft, weight=0.35), no TTA supervision; near-accident windows oversampled safe_neg : no accident; hazard_label=0, no TTA supervision Pre-risky windows from ego_positive videos are tagged as safe_neg with hazard_weight=0.8 (slightly soft — annotation boundary may be imprecise). All timestamps are 20 Hz frame indices (0.05 s / frame). Folder assignment is the source of truth; accident boolean is ignored. """ from __future__ import annotations import json import logging import random from collections import defaultdict from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch from PIL import Image from torch.utils.data import Dataset, DataLoader, Sampler logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) # ── constants ───────────────────────────────────────────────────────────────── FRAME_RATE = 20 FRAME_INTERVAL = 1.0 / FRAME_RATE # 0.05 s WINDOW_STD = 40 # 2.0 s WINDOW_EXT = 60 # 3.0 s MAX_FRAMES_PER_SAMPLE = 8 FRAME_SAMPLE_RATE = 4 # pool every 4th frame inside window MAX_TTA = 10.0 MIN_TTA = 0.1 # Hazard supervision weights W_EGO_POS = 1.0 W_SAFE_NEG = 1.0 W_PRE_RISKY = 0.8 # pre-risky window: from a crash video but before risk onset W_NON_EGO = 0.35 # genuinely ambiguous; push softly toward no-alert # Positive sampling POS_STRIDE_NORMAL = 10 # frames POS_STRIDE_CLOSE = 5 # frames (when TTA ≤ POS_CLOSE_TTA_S) POS_CLOSE_TTA_S = 3.0 # seconds PRE_RISKY_BUFFER_S = 3.0 # seconds before risky_time that positives can start # Negative / non-ego sampling NEG_STRIDE = 10 NEG_NUM_OFFSETS = 3 # staggered starts for negative videos NONEEGO_STRIDE = 10 NONEEGO_NEAR_STRIDE = 5 # denser near accident NONEEGO_NEAR_PRE_S = 3.0 # seconds before accident_frame to oversample NONEEGO_NEAR_POST_S = 1.0 # seconds after accident_frame to oversample # Balance NEG_POS_RATIO = 2.0 # cap: total_neg ≤ ratio × total_pos NONEEGO_NEG_FLOOR = 0.30 # non-ego ≥ 30% of all negative samples # ── data structures ─────────────────────────────────────────────────────────── @dataclass class VideoInfo: video_id: str source: str category: str # "ego_positive" | "non_ego" | "safe_neg" source_dir: Path num_frames: int accident_frame: Optional[int] # 20Hz; non_ego: sampling density only risky_frame: Optional[int] # 20Hz; non_ego: sampling density only metadata: Dict[str, Any] = field(default_factory=dict) @property def is_ego_positive(self) -> bool: return self.category == "ego_positive" @property def is_non_ego(self) -> bool: return self.category == "non_ego" @dataclass class TTASample: video_id: str source: str category: str # same as VideoInfo.category source_dir: str frame_indices: List[int] # Supervision signals hazard_label: float # 1.0 (ego_pos) or 0.0 (others) hazard_weight: float # see W_* constants tta_label: float # valid only when is_ego_positive and not is_censored is_ego_positive: bool is_non_ego: bool is_censored: bool # tta_raw > MAX_TTA (ego_pos only) # Window metadata accident_frame: Optional[int] risky_frame: Optional[int] window_end: int # exclusive window_len: int window_type: str # "standard" | "extended" tta_raw: float tta_cap: float difficulty: str metadata: Dict[str, Any] = field(default_factory=dict) # ── helpers ─────────────────────────────────────────────────────────────────── def _safe_int(x: Any) -> Optional[int]: if x is None: return None try: return int(float(str(x).strip())) except Exception: return None def _classify_difficulty(tta: float, category: str) -> str: if category == "safe_neg": return "easy" if category == "non_ego": return "hard" # ego_positive if tta <= 2.0: return "easy" if tta <= 5.0: return "hard" return "medium" def _load_manifest(path: Path) -> List[Dict[str, Any]]: with open(path, "r") as f: obj = json.load(f) return obj.get("videos", []) # ── dataset ─────────────────────────────────────────────────────────────────── class SFTDataset(Dataset): """ Args ---- manifests : list of Path / str pointing to manifest JSON files. Each file's "videos" list is loaded; split is already encoded in the manifest (train vs val). split : "train" or "val" — controls stochastic frame sampling. """ def __init__( self, manifests: List[Any], split: str = "train", seed: int = 42, debug: bool = False, debug_samples: int = 100, # sampling overrides pos_stride: int = POS_STRIDE_NORMAL, neg_stride: int = NEG_STRIDE, max_frames: int = MAX_FRAMES_PER_SAMPLE, frame_sample_rate: int = FRAME_SAMPLE_RATE, multi_window: bool = True, neg_pos_ratio: float = NEG_POS_RATIO, ): self.split = split self.seed = seed self.debug = debug self.debug_samples = debug_samples self.pos_stride = pos_stride self.neg_stride = neg_stride self.max_frames = max_frames self.frame_sample_rate = frame_sample_rate self.multi_window = multi_window self.neg_pos_ratio = neg_pos_ratio self.stochastic = (split == "train") random.seed(seed) np.random.seed(seed) self.videos: List[VideoInfo] = [] self.samples: List[TTASample] = [] for m in manifests: self._load_manifest(Path(m)) self._balance() if debug and len(self.samples) > debug_samples: self.samples = random.sample(self.samples, debug_samples) if split == "train": random.shuffle(self.samples) self._log_stats() # ── loading ────────────────────────────────────────────────────────────── def _load_manifest(self, path: Path) -> None: if not path.exists(): logger.warning(f"Manifest not found: {path}") return entries = _load_manifest(path) for e in entries: vi = VideoInfo( video_id = e["video_id"], source = e["source"], category = e["category"], source_dir = Path(e["source_dir"]), num_frames = int(e["num_frames"]), accident_frame= _safe_int(e.get("accident_frame")), risky_frame = _safe_int(e.get("risky_frame")), metadata = dict(e.get("metadata", {})), ) self.videos.append(vi) self._generate_samples(vi) # ── sample generation ───────────────────────────────────────────────────── def _generate_samples(self, vi: VideoInfo) -> None: if vi.is_ego_positive: self._gen_ego_positive(vi) elif vi.is_non_ego: self._gen_non_ego(vi) else: self._gen_safe_neg(vi) # ── ego_positive ────────────────────────────────────────────────────────── def _gen_ego_positive(self, vi: VideoInfo) -> None: n = vi.num_frames acc = vi.accident_frame # guaranteed not None and < n (manifest filtered) rsk = vi.risky_frame # may be None or 0 base_win = WINDOW_EXT if self.multi_window else WINDOW_STD # 1) Pre-risky windows (safe_neg from this video) if rsk is not None: safe_end = max(base_win, rsk) # windows must end at or before risky_frame pre_risky_buffer = int(PRE_RISKY_BUFFER_S / FRAME_INTERVAL) start = max(base_win, rsk - pre_risky_buffer - base_win) if start < safe_end and safe_end > base_win: # sample a few pre-risky windows for we in range(start + base_win, safe_end, self.neg_stride): self._add_safe_neg(vi, we, WINDOW_STD, neg_tag="pre_risky", weight=W_PRE_RISKY) if self.multi_window: self._add_safe_neg(vi, we, WINDOW_EXT, neg_tag="pre_risky", weight=W_PRE_RISKY) # 2) Positive windows: from pos_start_frame onward to accident_frame if rsk is not None: buf = int(PRE_RISKY_BUFFER_S / FRAME_INTERVAL) pos_start = max(0, rsk - buf) else: pos_start = 0 seen: set = set() def add_pos(we: int) -> None: if we in seen: return seen.add(we) cur = we - 1 tta_raw = (acc - cur) * FRAME_INTERVAL self._add_ego_pos(vi, we, WINDOW_STD, tta_raw) if self.multi_window: self._add_ego_pos(vi, we, WINDOW_EXT, tta_raw) # TTA anchor sampling (biased toward 2–7s) if self.split == "train": targets_s = [2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 7.0, 8.0] repeats = 2 jitter = 3 else: targets_s = [2.0, 3.0, 4.0, 5.0, 6.0, 7.0] repeats = 1 jitter = 0 for tta_t in targets_s: off = int(round(tta_t / FRAME_INTERVAL)) we_base = (acc - off) + 1 for _ in range(repeats): j = random.randint(-jitter, jitter) if jitter else 0 we = we_base + j we = max(base_win, min(we, acc + 1)) if we - 1 >= pos_start: add_pos(we) # Stride-based sweep we = max(base_win, pos_start + base_win) while we <= acc + 1: if we - 1 >= pos_start: add_pos(we) cur = we - 1 tta = (acc - cur) * FRAME_INTERVAL we += POS_STRIDE_CLOSE if tta <= POS_CLOSE_TTA_S else self.pos_stride def _add_ego_pos(self, vi: VideoInfo, window_end: int, window_len: int, tta_raw: float) -> None: n = vi.num_frames acc = vi.accident_frame window_end = max(window_len, min(window_end, n)) if window_end <= 0: return if tta_raw < MIN_TTA: return is_censored = tta_raw > MAX_TTA tta_label = min(tta_raw, MAX_TTA) if not is_censored else MAX_TTA fi = self._sample_frames(window_end - window_len, window_end) meta = dict(vi.metadata) meta["neg_tag"] = "ego_pos" self.samples.append(TTASample( video_id = vi.video_id, source = vi.source, category = "ego_positive", source_dir = str(vi.source_dir), frame_indices = fi, hazard_label = 1.0, hazard_weight = W_EGO_POS, tta_label = tta_label, is_ego_positive = True, is_non_ego = False, is_censored = is_censored, accident_frame = acc, risky_frame = vi.risky_frame, window_end = window_end, window_len = window_len, window_type = "extended" if window_len == WINDOW_EXT else "standard", tta_raw = tta_raw, tta_cap = MAX_TTA, difficulty = _classify_difficulty(tta_label, "ego_positive"), metadata = meta, )) def _add_safe_neg(self, vi: VideoInfo, window_end: int, window_len: int, neg_tag: str = "neg_video", weight: float = W_SAFE_NEG) -> None: n = vi.num_frames window_end = max(window_len, min(window_end, n)) if window_end <= 0: return fi = self._sample_frames(window_end - window_len, window_end) meta = dict(vi.metadata) meta["neg_tag"] = neg_tag self.samples.append(TTASample( video_id = vi.video_id, source = vi.source, category = "safe_neg", source_dir = str(vi.source_dir), frame_indices = fi, hazard_label = 0.0, hazard_weight = weight, tta_label = MAX_TTA, is_ego_positive = False, is_non_ego = False, is_censored = True, accident_frame = None, risky_frame = None, window_end = window_end, window_len = window_len, window_type = "standard", tta_raw = float("inf"), tta_cap = MAX_TTA, difficulty = _classify_difficulty(MAX_TTA, "safe_neg"), metadata = meta, )) # ── safe_neg (negative video) ───────────────────────────────────────────── def _gen_safe_neg(self, vi: VideoInfo) -> None: n = vi.num_frames offsets = [int(i * self.neg_stride / NEG_NUM_OFFSETS) for i in range(NEG_NUM_OFFSETS)] for off in offsets: we = WINDOW_STD + off while we <= n: self._add_safe_neg(vi, we, WINDOW_STD, neg_tag="neg_video", weight=W_SAFE_NEG) we += self.neg_stride # ── non_ego ─────────────────────────────────────────────────────────────── def _gen_non_ego(self, vi: VideoInfo) -> None: n = vi.num_frames acc = vi.accident_frame # used for density only, NOT as label # Cap sampling to accident_frame: never include post-accident content. # At inference time the system predicts before accidents occur, so # post-accident frames (debris, aftermath) are out-of-distribution. video_cap = acc if acc is not None else n # exclusive upper bound for window_end # Define near-accident zone (pre-accident only) if acc is not None: near_pre = int(NONEEGO_NEAR_PRE_S / FRAME_INTERVAL) near_start = max(0, acc - near_pre) near_end = acc # cap at accident_frame (no post-accident) else: near_start = near_end = -1 # Normal stride — only up to accident_frame offsets = [int(i * NONEEGO_STRIDE / max(1, NEG_NUM_OFFSETS - 1)) for i in range(NEG_NUM_OFFSETS)] for off in offsets: we = WINDOW_STD + off while we <= video_cap: self._add_non_ego(vi, we, WINDOW_STD, near=(near_start <= we <= near_end)) we += NONEEGO_STRIDE # Dense near-accident sampling (pre-accident only) if acc is not None and near_end > near_start: we = near_start + WINDOW_STD while we <= near_end: self._add_non_ego(vi, we, WINDOW_STD, near=True) we += NONEEGO_NEAR_STRIDE def _add_non_ego(self, vi: VideoInfo, window_end: int, window_len: int, near: bool = False) -> None: n = vi.num_frames window_end = max(window_len, min(window_end, n)) if window_end <= 0: return fi = self._sample_frames(window_end - window_len, window_end) meta = dict(vi.metadata) meta["neg_tag"] = "non_ego_near" if near else "non_ego" self.samples.append(TTASample( video_id = vi.video_id, source = vi.source, category = "non_ego", source_dir = str(vi.source_dir), frame_indices = fi, hazard_label = 0.0, hazard_weight = W_NON_EGO, tta_label = MAX_TTA, is_ego_positive = False, is_non_ego = True, is_censored = True, accident_frame = None, # NOT passed as label risky_frame = None, window_end = window_end, window_len = window_len, window_type = "standard", tta_raw = float("inf"), tta_cap = MAX_TTA, difficulty = "hard", # always hard — visually accident-like metadata = meta, )) # ── frame sampling ──────────────────────────────────────────────────────── def _sample_frames(self, start: int, end: int) -> List[int]: start = max(0, start) end = max(start + 1, end) pool = list(range(start, end, max(1, self.frame_sample_rate))) if len(pool) <= self.max_frames: return pool if self.stochastic and self.max_frames >= 3: first = pool[0] last = pool[-1] mid = pool[1:-1] k = self.max_frames - 2 chosen = sorted(random.sample(mid, k=min(k, len(mid)))) return [first] + chosen + [last] else: idx = np.linspace(0, len(pool) - 1, self.max_frames, dtype=int) return [pool[i] for i in idx] # ── balancing ───────────────────────────────────────────────────────────── def _balance(self) -> None: pos = [s for s in self.samples if s.is_ego_positive] non_ego = [s for s in self.samples if s.is_non_ego] neg = [s for s in self.samples if not s.is_ego_positive and not s.is_non_ego] n_pos = len(pos) if n_pos == 0: return # Total negatives target target_total_neg = int(n_pos * self.neg_pos_ratio) # Floor: non_ego ≥ 30% of all negatives target_ne = max(len(non_ego), int(target_total_neg * NONEEGO_NEG_FLOOR)) target_neg = max(0, target_total_neg - target_ne) # Cap non_ego if len(non_ego) > target_ne: non_ego = random.sample(non_ego, target_ne) # Cap safe_neg if len(neg) > target_neg: neg = random.sample(neg, target_neg) self.samples = pos + non_ego + neg random.shuffle(self.samples) # ── stats ───────────────────────────────────────────────────────────────── def _log_stats(self) -> None: n = len(self.samples) cats = defaultdict(int) tags = defaultdict(int) for s in self.samples: cats[s.category] += 1 tags[str(s.metadata.get("neg_tag", ""))] += 1 n_cens = sum(1 for s in self.samples if s.is_censored) logger.info("=" * 55) logger.info(f"SFTDataset [{self.split}] total={n}") for cat, cnt in sorted(cats.items()): logger.info(f" {cat:20s}: {cnt:6d} ({100*cnt/max(1,n):.1f}%)") logger.info(f" {'censored':20s}: {n_cens:6d}") pos = [s for s in self.samples if s.is_ego_positive and not s.is_censored] if pos: ttas = [s.tta_label for s in pos] logger.info(f" TTA pos: mean={np.mean(ttas):.2f} " f"min={min(ttas):.2f} max={max(ttas):.2f} " f"std={np.std(ttas):.2f}") logger.info(f" neg_tags: { {k: v for k,v in sorted(tags.items()) if k} }") logger.info("=" * 55) # ── Dataset protocol ────────────────────────────────────────────────────── def __len__(self) -> int: return len(self.samples) def __getitem__(self, idx: int) -> Dict[str, Any]: s = self.samples[idx] src = Path(s.source_dir) images: List[Any] = [] for fi in s.frame_indices: img = self._load_frame(src, fi) if img is not None: images.append(img) if not images: logger.warning(f"No frames loaded for {s.video_id} idx={idx}; using blank.") images = [Image.new("RGB", (384, 384), (64, 64, 64))] return { "video_id": s.video_id, "source": s.source, "category": s.category, "images": images, "frame_indices": s.frame_indices, "hazard_label": float(s.hazard_label), "hazard_weight": float(s.hazard_weight), "tta_label": float(s.tta_label), "is_ego_positive":bool(s.is_ego_positive), "is_non_ego": bool(s.is_non_ego), "is_censored": bool(s.is_censored), "tta_raw": float(s.tta_raw) if np.isfinite(s.tta_raw) else MAX_TTA, "tta_cap": float(s.tta_cap), "window_type": s.window_type, "difficulty": s.difficulty, "metadata": s.metadata, } @staticmethod def _load_frame(src_dir: Path, frame_idx: int) -> Optional[Image.Image]: for fmt in ["{:03d}", "{:04d}", "{:05d}", "{:06d}", "{}"]: for ext in [".jpg", ".jpeg", ".png"]: p = src_dir / (fmt.format(frame_idx) + ext) if p.exists(): try: return Image.open(p).convert("RGB") except Exception: pass return None # ── collate ─────────────────────────────────────────────────────────────────── def sft_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]: return { "video_ids": [b["video_id"] for b in batch], "sources": [b["source"] for b in batch], "categories": [b["category"] for b in batch], "images": [b["images"] for b in batch], "frame_indices": [b["frame_indices"] for b in batch], "metadata": [b["metadata"] for b in batch], "window_types": [b["window_type"] for b in batch], "difficulties": [b["difficulty"] for b in batch], "hazard_labels": torch.tensor([b["hazard_label"] for b in batch], dtype=torch.float32), "hazard_weights": torch.tensor([b["hazard_weight"] for b in batch], dtype=torch.float32), "tta_labels": torch.tensor([b["tta_label"] for b in batch], dtype=torch.float32), "is_ego_positive":torch.tensor([b["is_ego_positive"]for b in batch], dtype=torch.bool), "is_non_ego": torch.tensor([b["is_non_ego"] for b in batch], dtype=torch.bool), "is_censored": torch.tensor([b["is_censored"] for b in batch], dtype=torch.bool), "tta_caps": torch.tensor([b["tta_cap"] for b in batch], dtype=torch.float32), "tta_raws": torch.tensor([b["tta_raw"] for b in batch], dtype=torch.float32), } # ── quick smoke ─────────────────────────────────────────────────────────────── if __name__ == "__main__": import sys manifest_dir = Path("PROJECT_ROOT/data/sft_manifests") train_manifests = [ manifest_dir / "nexar_train.json", manifest_dir / "dada_pos_train.json", manifest_dir / "dada_noneego_train.json", manifest_dir / "dada_neg_train.json", ] ds = SFTDataset(train_manifests, split="train", debug=True, debug_samples=40) print(f"\nSize: {len(ds)}") item = ds[0] print(f" video_id={item['video_id']} category={item['category']}") print(f" hazard_label={item['hazard_label']} hazard_weight={item['hazard_weight']}") print(f" tta_label={item['tta_label']:.2f} is_censored={item['is_censored']}") print(f" n_images={len(item['images'])}") loader = DataLoader(ds, batch_size=4, collate_fn=sft_collate_fn, num_workers=0) b = next(iter(loader)) print(f"\nBatch hazard_labels: {b['hazard_labels']}") print(f"Batch tta_labels: {b['tta_labels']}") print(f"Batch is_ego_pos: {b['is_ego_positive']}") print(f"Batch is_censored: {b['is_censored']}")