model/sampling.py

import cv2, numpy as np
from typing import List, Tuple, Dict

def read_video_meta(video_path: str):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f'无法打开视频: {video_path}')
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
    cap.release()
    return {'fps': fps, 'total': total, 'w': w, 'h': h}

def sample_indices_uniform(total: int, n: int) -> List[int]:
    if n<=1 or total<=1:
        return [0]
    idxs = np.linspace(0, total-1, n, dtype=int).tolist()
    return sorted(set(int(i) for i in idxs))

def sample_indices_fps(total: int, native_fps: float, target_fps: float) -> List[int]:
    if target_fps <= 0:
        target_fps = 1.0
    step = max(int(round(native_fps / target_fps)), 1)
    idxs = list(range(0, total, step))
    return idxs

def grab_frames(video_path: str, indices: List[int]) -> Tuple[List[np.ndarray], Dict[int,int]]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f'无法打开视频: {video_path}')
    frames = []
    mapping = {}  # t -> global index
    for i, gi in enumerate(indices):
        cap.set(cv2.CAP_PROP_POS_FRAMES, gi)
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(frame[:, :, ::-1])  # BGR->RGB for model
        mapping[i] = gi
    cap.release()
    return frames, mapping