| |
| """ |
| Extract frames from DAD (Dashcam Accident Dataset) MP4 files. |
| ============================================================= |
| DAD has no timing annotations — just positive/negative video splits. |
| positive/: 455 crash videos (no exact accident timestamp) |
| negative/: 829 safe driving (no annotations needed) |
| |
| Frames are extracted at 20 fps and saved as: |
| data/pretrain_v2/dad_frames/{positive,negative}/{vid_id}/000.jpg ... |
| |
| Each video also gets a minimal annotation.json: |
| { "split": "positive"|"negative", |
| "n_frames": int, |
| "accident": bool, |
| "fps": 20 } |
| |
| Usage: |
| python training/pretrain_v2/prepare_dad_frames.py [--workers 4] |
| """ |
|
|
| import argparse |
| import json |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from pathlib import Path |
|
|
| import cv2 |
|
|
| DAD_VIDEOS_DIR = "PROJECT_ROOT/DAD/videos/training" |
| OUTPUT_BASE = "PROJECT_ROOT/data/pretrain_v2/dad_frames" |
| DAD_TEST_VIDEOS_DIR = "PROJECT_ROOT/DAD/videos/testing" |
| OUTPUT_BASE_TEST = "PROJECT_ROOT/data/pretrain_v2/dad_frames_test" |
| TARGET_FPS = 20.0 |
| MAX_FRAMES = 400 |
|
|
|
|
| def extract_video(args): |
| """Extract one video. Returns (video_id, n_frames, error|None).""" |
| vid_path, out_dir, split = args |
| vid_id = vid_path.stem |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| ann_path = out_dir / "annotation.json" |
| if ann_path.exists(): |
| |
| existing = json.loads(ann_path.read_text()) |
| return vid_id, existing.get("n_frames", 0), None |
|
|
| cap = cv2.VideoCapture(str(vid_path)) |
| if not cap.isOpened(): |
| return vid_id, 0, f"Cannot open {vid_path}" |
|
|
| orig_fps = cap.get(cv2.CAP_PROP_FPS) or TARGET_FPS |
| interval = max(1, int(round(orig_fps / TARGET_FPS))) |
|
|
| raw_idx = 0 |
| saved = 0 |
| while saved < MAX_FRAMES: |
| ret, frame = cap.read() |
| if not ret: |
| break |
| if raw_idx % interval == 0: |
| cv2.imwrite(str(out_dir / f"{saved:03d}.jpg"), frame) |
| saved += 1 |
| raw_idx += 1 |
| cap.release() |
|
|
| if saved == 0: |
| return vid_id, 0, f"No frames extracted from {vid_path}" |
|
|
| ann_path.write_text(json.dumps({ |
| "split": split, |
| "accident": (split == "positive"), |
| "n_frames": saved, |
| "fps": TARGET_FPS, |
| }, indent=2), encoding="utf-8") |
|
|
| return vid_id, saved, None |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--workers", type=int, default=4) |
| parser.add_argument("--split", choices=["training", "testing"], default="training", |
| help="DAD split to extract: training or testing") |
| args = parser.parse_args() |
|
|
| if args.split == "testing": |
| dad_root = Path(DAD_TEST_VIDEOS_DIR) |
| output_base = Path(OUTPUT_BASE_TEST) |
| else: |
| dad_root = Path(DAD_VIDEOS_DIR) |
| output_base = Path(OUTPUT_BASE) |
|
|
| tasks = [] |
| for split in ("positive", "negative"): |
| vid_dir = dad_root / split |
| if not vid_dir.exists(): |
| print(f"Skip {vid_dir} — not found") |
| continue |
| vids = sorted(vid_dir.glob("*.mp4")) |
| print(f" {split}: {len(vids)} videos") |
| for v in vids: |
| out_dir = output_base / split / v.stem |
| tasks.append((v, out_dir, split)) |
|
|
| print(f"\nExtracting {len(tasks)} videos with {args.workers} workers ...") |
|
|
| ok = 0 |
| fail = 0 |
| with ThreadPoolExecutor(max_workers=args.workers) as ex: |
| futs = {ex.submit(extract_video, t): t for t in tasks} |
| from tqdm import tqdm |
| for fut in tqdm(as_completed(futs), total=len(futs)): |
| vid_id, n, err = fut.result() |
| if err: |
| print(f" FAIL {vid_id}: {err}") |
| fail += 1 |
| else: |
| ok += 1 |
|
|
| print(f"\nDone: {ok} extracted, {fail} failed") |
| print(f"Frames saved to: {output_base}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
