| """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() |
|
|
|
|
