35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
"""
|
||
视频帧提取（供 modules.segment 的 segment_person_from_video_frame 使用）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
|
||
class VideoExtractor:
|
||
"""从视频中按时间戳抽取一帧（RGB numpy 数组）。"""
|
||
|
||
def __init__(self, video_path: str) -> None:
|
||
self.video_path = video_path
|
||
self._cap = cv2.VideoCapture(video_path)
|
||
if not self._cap.isOpened():
|
||
raise RuntimeError(f"无法打开视频: {video_path}")
|
||
|
||
def extract_frame_at_time(self, time_seconds: float) -> np.ndarray:
|
||
self._cap.set(cv2.CAP_PROP_POS_MSEC, max(0.0, float(time_seconds)) * 1000.0)
|
||
ret, frame = self._cap.read()
|
||
if not ret or frame is None:
|
||
self._cap.release()
|
||
raise RuntimeError(f"无法在 {time_seconds}s 处提取视频帧: {self.video_path}")
|
||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
self._cap.release()
|
||
return rgb
|
||
|
||
def __del__(self) -> None:
|
||
try:
|
||
if hasattr(self, "_cap") and self._cap is not None:
|
||
self._cap.release()
|
||
except Exception:
|
||
pass
|