training/SFT/dataset.py

#!/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']}")