11

scripts/cover/modules/extractor.py

@@ -0,0 +1,34 @@
+"""
+视频帧提取（供 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