naveenkm13/occupancyos

1"""2Extract frames from one or more videos at a configurable FPS.3 4Usage:5    python -m backend.scripts.extract_frames --input Dataset/videos --output Dataset/frames --fps 26"""7from __future__ import annotations8 9import argparse10from pathlib import Path11 12import cv213 14 15VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}16 17 18def extract_from_video(video: Path, out_dir: Path, target_fps: float) -> int:19    cap = cv2.VideoCapture(str(video))20    if not cap.isOpened():21        print(f"[skip] cannot open {video}")22        return 023    src_fps = cap.get(cv2.CAP_PROP_FPS) or 25.024    step = max(1, int(round(src_fps / target_fps)))25    out_dir.mkdir(parents=True, exist_ok=True)26 27    idx = 028    saved = 029    while True:30        ok, frame = cap.read()31        if not ok:32            break33        if idx % step == 0:34            out_path = out_dir / f"{video.stem}_{saved:06d}.jpg"35            cv2.imwrite(str(out_path), frame)36            saved += 137        idx += 138    cap.release()39    print(f"  {video.name}: {saved} frames saved (every {step} frames)")40    return saved41 42 43def main() -> None:44    p = argparse.ArgumentParser()45    p.add_argument("--input", required=True, help="video file or directory")46    p.add_argument("--output", required=True, help="output frames directory")47    p.add_argument("--fps", type=float, default=2.0)48    args = p.parse_args()49 50    in_path = Path(args.input)51    out_dir = Path(args.output)52 53    videos: list[Path] = []54    if in_path.is_file():55        videos = [in_path]56    else:57        videos = [v for v in in_path.rglob("*") if v.suffix.lower() in VIDEO_EXTS]58 59    if not videos:60        print(f"No videos found in {in_path}")61        return62 63    total = 064    for v in videos:65        total += extract_from_video(v, out_dir, args.fps)66    print(f"[extract_frames] done. {total} frames -> {out_dir}")67 68 69if __name__ == "__main__":70    main()71