scripts/predict_video.py

"""Headless batch video annotation.

Reads a video, draws detection boxes above the confidence threshold, and writes
the annotated copy next to it as <name>_out.mp4. Run from the repo root:

    uv run python scripts/predict_video.py
    uv run python scripts/predict_video.py --video videos/3.mp4 --conf 0.5 --device mps
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

import cv2

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from src import inference  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("predict_video")


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--video", default="videos/2.mp4")
    p.add_argument("--weights", default=inference.DEFAULT_WEIGHTS)
    p.add_argument("--conf", type=float, default=0.7)
    p.add_argument("--device", default=None, help="cpu / mps / cuda; auto-picks when omitted")
    p.add_argument("--output", default=None, help="defaults to <video>_out.mp4")
    args = p.parse_args()

    out_path = args.output or f"{args.video}_out.mp4"
    model = inference.load_model(args.weights, args.device)

    cap = cv2.VideoCapture(args.video)
    ret, frame = cap.read()
    if not ret:
        sys.exit(f"Cannot read {args.video}")

    h, w = frame.shape[:2]
    writer = cv2.VideoWriter(
        out_path, cv2.VideoWriter_fourcc(*"MP4V"), int(cap.get(cv2.CAP_PROP_FPS)), (w, h)
    )

    n = 0
    while ret:
        writer.write(inference.detect_and_annotate(model, frame, conf=args.conf))
        n += 1
        ret, frame = cap.read()

    cap.release()
    writer.release()
    logger.info("Annotated %d frames -> %s", n, out_path)


if __name__ == "__main__":
    main()