broadleaf-weed-detector / scripts /video_to_frames.py
rgthelen's picture
Broadleaf weed detector: yolo11n/s + Hailo-10H HEFs + full dataset + repro scripts
33475e9 verified
Raw
History Blame Contribute Delete
1.91 kB
#!/usr/bin/env python3
"""Extract frames from real drone/robot footage at a fixed sample rate.
Writes {out}/frames/{video}_{idx:05d}.jpg + a manifest.jsonl (one entry per
frame) matching the iNat downloader schema so the SAME Gemini labeler can box
them as ground truth.
"""
import argparse
import json
from pathlib import Path
import cv2
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--video", required=True, help="Path to a video file")
ap.add_argument("--out", required=True)
ap.add_argument("--fps", type=float, default=2.5, help="frames per second to sample")
ap.add_argument("--max-frames", type=int, default=0, help="0 = no cap")
args = ap.parse_args()
out = Path(args.out)
(out / "frames").mkdir(parents=True, exist_ok=True)
cap = cv2.VideoCapture(args.video)
if not cap.isOpened():
raise SystemExit(f"cannot open {args.video}")
src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
step = max(1, int(round(src_fps / args.fps)))
stem = Path(args.video).stem.replace(" ", "_")
manifest = (out / "manifest.jsonl").open("a")
fi = kept = 0
while True:
ok, frame = cap.read()
if not ok:
break
if fi % step == 0:
name = f"{stem}_{kept:05d}.jpg"
dst = out / "frames" / name
cv2.imwrite(str(dst), frame)
manifest.write(json.dumps({
"common": "weed", "sci": "field footage",
"obs_id": stem, "photo_id": kept,
"src_frame": fi, "file": str(dst),
}) + "\n")
kept += 1
if args.max_frames and kept >= args.max_frames:
break
fi += 1
cap.release()
manifest.close()
print(f"src_fps={src_fps:.1f} step={step}{kept} frames from {fi} "
f"(sampled ~{args.fps} fps) → {out/'frames'}")
if __name__ == "__main__":
main()