File size: 6,339 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | #!/usr/bin/env python3
"""
DPO Dataset — manifest-based preference pairs for HazardHead alignment.
Each sample is a (chosen, rejected) window pair where:
chosen = window where issuing an alert is CORRECT
(ego_pos, TTA ∈ [1.5, 5.0]s → "timely_alert")
rejected = window where issuing an alert is WRONG
(too_early, too_late, safe_neg, non_ego)
The dataset returns raw PIL frames; the DPO trainer handles VLM tokenisation.
"""
from __future__ import annotations
import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional
import torch
from PIL import Image
from torch.utils.data import Dataset
logger = logging.getLogger(__name__)
MAX_FRAMES = 8
# ─────────────────────────────────────────────────────────────────────────────
# Frame loader (mirrors SFT dataset)
# ─────────────────────────────────────────────────────────────────────────────
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
def _load_frames(source_dir: str, frame_indices: List[int]) -> List[Image.Image]:
src = Path(source_dir)
imgs = []
for idx in frame_indices[:MAX_FRAMES]:
img = _load_frame(src, idx)
if img is not None:
imgs.append(img)
if not imgs:
imgs = [Image.new("RGB", (384, 384), (64, 64, 64))]
return imgs
# ─────────────────────────────────────────────────────────────────────────────
# DPODataset
# ─────────────────────────────────────────────────────────────────────────────
class DPODataset(Dataset):
"""
Loads preference pairs from DPO pair manifests.
Args
----
manifests : list of paths to JSON pair manifests (as generated by make_dpo_pairs.py)
split : "train" or "val"
debug : if True, limit to debug_samples pairs
debug_samples : number of pairs to use in debug mode
"""
def __init__(
self,
manifests: List[Path],
split: str = "train",
debug: bool = False,
debug_samples: int = 64,
):
self.split = split
self.pairs: List[dict] = []
for m in manifests:
m = Path(m)
if not m.exists():
logger.warning(f"DPO manifest not found: {m}")
continue
with open(m) as f:
data = json.load(f)
p = data.get("pairs", [])
self.pairs.extend(p)
logger.info(f"Loaded {len(p)} pairs from {m.name}")
if debug:
self.pairs = self.pairs[:debug_samples]
logger.info(
f"DPODataset [{split}]: {len(self.pairs)} pairs "
f"({sum(1 for p in self.pairs if p['pair_type']=='timing')} timing, "
f"{sum(1 for p in self.pairs if p['pair_type']=='category')} category)"
)
def __len__(self) -> int:
return len(self.pairs)
def __getitem__(self, idx: int) -> Dict[str, Any]:
pair = self.pairs[idx]
c = pair["chosen"]
r = pair["rejected"]
chosen_images = _load_frames(c["source_dir"], c["frame_indices"])
rejected_images = _load_frames(r["source_dir"], r["frame_indices"])
return {
"pair_id": pair["pair_id"],
"video_id": pair["video_id"],
"source": pair["source"],
"pair_type": pair["pair_type"],
# chosen
"chosen_images": chosen_images,
"chosen_tta": float(c["tta_true"]),
"chosen_label": c["label"],
"chosen_metadata": c.get("metadata", {}),
# rejected
"rejected_images": rejected_images,
"rejected_tta": float(r["tta_true"]),
"rejected_label": r["label"],
"rejected_metadata":r.get("metadata", {}),
}
# ─────────────────────────────────────────────────────────────────────────────
# Collate
# ─────────────────────────────────────────────────────────────────────────────
def dpo_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
return {
"pair_ids": [b["pair_id"] for b in batch],
"video_ids": [b["video_id"] for b in batch],
"sources": [b["source"] for b in batch],
"pair_types": [b["pair_type"] for b in batch],
# chosen
"chosen_images": [b["chosen_images"] for b in batch],
"chosen_ttas": torch.tensor([b["chosen_tta"] for b in batch], dtype=torch.float32),
"chosen_labels": [b["chosen_label"] for b in batch],
"chosen_metadata": [b["chosen_metadata"] for b in batch],
# rejected
"rejected_images": [b["rejected_images"] for b in batch],
"rejected_ttas": torch.tensor([b["rejected_tta"] for b in batch], dtype=torch.float32),
"rejected_labels": [b["rejected_label"] for b in batch],
"rejected_metadata": [b["rejected_metadata"] for b in batch],
}
|