"""Video reading wrapper around OpenCV.""" from __future__ import annotations from pathlib import Path from typing import Iterator import cv2 import numpy as np class VideoReader: """Iterate over frames from a video file.""" def __init__(self, video_path: str | Path) -> None: self.video_path = Path(video_path) self.capture = cv2.VideoCapture(str(self.video_path)) if not self.capture.isOpened(): raise FileNotFoundError(f"Could not open video: {self.video_path}") self.fps = float(self.capture.get(cv2.CAP_PROP_FPS) or 30.0) self.width = int(self.capture.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self.capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.frame_count = int(self.capture.get(cv2.CAP_PROP_FRAME_COUNT)) def __iter__(self) -> Iterator[tuple[int, np.ndarray]]: frame_index = 0 while True: ok, frame = self.capture.read() if not ok: break yield frame_index, frame frame_index += 1 def release(self) -> None: self.capture.release() def __enter__(self) -> "VideoReader": return self def __exit__(self, *_: object) -> None: self.release()