clip-annotator/src/river_annotation_tool/video_loader.py

import os
import tempfile
import zipfile
from pathlib import Path

import cv2


def load_frames(zip_path: Path, max_frames: int, display_max: int, fps_fallback: int):
    video_bytes = zipfile.ZipFile(zip_path).read("left.mp4")

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
        f.write(video_bytes)
        tmp_path = f.name

    cap = cv2.VideoCapture(tmp_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or fps_fallback

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    step = max(1, total // max_frames)

    frames = []
    i = 0
    while len(frames) < max_frames:
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        i += step

    cap.release()
    os.unlink(tmp_path)

    if not frames:
        raise RuntimeError(f"No frames found in {zip_path}")

    h, w = frames[0].shape[:2]
    scale = display_max / max(h, w)
    dh, dw = int(h * scale), int(w * scale)

    frames = [cv2.resize(f, (dw, dh)) for f in frames]

    return frames, fps, dh, dw, h, w