utils/video_utils.py

"""Frame extraction and video writing utilities."""

import cv2
import numpy as np
from typing import List

__all__ = [
    "extract_frames_from_video",
    "save_frames_to_video",
]


def extract_frames_from_video(video_path: str, target_fps: int = 30) -> List[np.ndarray]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise IOError(f"Cannot open video: {video_path}")

    original_fps = cap.get(cv2.CAP_PROP_FPS) or target_fps
    frame_interval = max(1, round(original_fps / target_fps))

    frames, count = [], 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if count % frame_interval == 0:
            frames.append(frame)
        count += 1
    cap.release()
    return frames


def save_frames_to_video(frames: List[np.ndarray], output_path: str, fps: int = 30):
    if not frames:
        raise ValueError("No frames to save")

    height, width = frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    for frm in frames:
        writer.write(frm)
    writer.release()