File size: 1,256 Bytes
c94751c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """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()
|