98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
"""视频读取模块"""
|
||
|
||
import cv2
|
||
from pathlib import Path
|
||
from typing import Iterator, Optional
|
||
import numpy as np
|
||
|
||
|
||
class VideoReader:
|
||
"""读取视频文件，按帧提取"""
|
||
|
||
def __init__(self, path: str | Path):
|
||
self.path = Path(path)
|
||
self._cap: Optional[cv2.VideoCapture] = None
|
||
self._frame_count: Optional[int] = None
|
||
self._fps: Optional[float] = None
|
||
self._width: Optional[int] = None
|
||
self._height: Optional[int] = None
|
||
|
||
@property
|
||
def frame_count(self) -> int:
|
||
if self._frame_count is None:
|
||
self._open()
|
||
return self._frame_count
|
||
|
||
@property
|
||
def fps(self) -> float:
|
||
if self._fps is None:
|
||
self._open()
|
||
return self._fps
|
||
|
||
@property
|
||
def width(self) -> int:
|
||
if self._width is None:
|
||
self._open()
|
||
return self._width
|
||
|
||
@property
|
||
def height(self) -> int:
|
||
if self._height is None:
|
||
self._open()
|
||
return self._height
|
||
|
||
def _open(self) -> None:
|
||
self._cap = cv2.VideoCapture(str(self.path))
|
||
if not self._cap.isOpened():
|
||
raise FileNotFoundError(f"Cannot open video: {self.path}")
|
||
self._frame_count = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
self._fps = self._cap.get(cv2.CAP_PROP_FPS)
|
||
self._width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
self._height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
|
||
def read_frame(self, index: int) -> Optional[np.ndarray]:
|
||
"""读取指定索引的一帧 (BGR) — 随机访问，可能慢"""
|
||
if self._cap is None:
|
||
self._open()
|
||
self._cap.set(cv2.CAP_PROP_POS_FRAMES, index)
|
||
ret, frame = self._cap.read()
|
||
if not ret:
|
||
return None
|
||
return frame
|
||
|
||
def reset(self) -> None:
|
||
"""重开视频，回到第 0 帧，用于顺序读取"""
|
||
self.close()
|
||
self._open()
|
||
|
||
def read_next(self) -> Optional[np.ndarray]:
|
||
"""顺序读取下一帧（不 seek，比 read_frame 可靠）"""
|
||
if self._cap is None:
|
||
self._open()
|
||
ret, frame = self._cap.read()
|
||
if not ret:
|
||
return None
|
||
return frame
|
||
|
||
def iter_frames(self) -> Iterator[tuple[int, np.ndarray]]:
|
||
"""逐帧迭代，返回 (index, frame)"""
|
||
self.reset()
|
||
idx = 0
|
||
while True:
|
||
frame = self.read_next()
|
||
if frame is None:
|
||
break
|
||
yield idx, frame
|
||
idx += 1
|
||
|
||
def close(self) -> None:
|
||
if self._cap is not None:
|
||
self._cap.release()
|
||
self._cap = None
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *args):
|
||
self.close()
|