1. 视频帧采样与预处理

任务目标

本案例实现一个视频帧预处理流水线,完成三个核心操作:均匀采样(从视频中按时间均匀抽取N帧)、缩放(统一尺寸为224×224)、归一化(像素值映射到[0,1]并标准化)。输出结果为可直接送入视觉编码器的张量批次。

环境依赖

需要 opencv-python、numpy、torch 和 torchvision。若使用GPU加速张量操作,需安装对应CUDA版本的PyTorch。

完整实现

import cv2
import numpy as np
import torch
from torchvision import transforms

class VideoFramePreprocessor:
    """视频帧均匀采样 + 缩放 + 归一化流水线"""
    
    def __init__(self, num_frames=8, target_size=(224, 224), 
                 mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        self.num_frames = num_frames
        self.target_size = target_size
        self.mean = mean
        self.std = std
        
    def _uniform_sample(self, video_path):
        """从视频中均匀抽取num_frames帧,返回BGR数组列表"""
        cap = cv2.VideoCapture(video_path)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total_frames <= 0:
            cap.release()
            raise ValueError(f"无法读取视频或视频无帧: {video_path}")
            
        # 计算采样索引(均匀分布,包含首尾)
        indices = np.linspace(0, total_frames - 1, self.num_frames, dtype=int)
        frames = []
        
        for idx in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
            ret, frame = cap.read()
            if not ret:
                # 若读取失败,用最后成功帧填充
                if frames:
                    frames.append(frames[-1].copy())
                continue
            frames.append(frame)
            
        cap.release()
        
        # 若实际帧数不足,循环填充至num_frames
        while len(frames) < self.num_frames:
            frames.append(frames[-1].copy() if frames else np.zeros((224,224,3), dtype=np.uint8))
        return frames[:self.num_frames]
    
    def _preprocess_frames(self, frames_bgr):
        """将BGR帧列表转为归一化后的tensor [num_frames, 3, H, W]"""
        transform = transforms.Compose([
            transforms.ToPILImage(),          # 从numpy数组(H,W,C)转为PIL
            transforms.Resize(self.target_size),
            transforms.ToTensor(),            # [0,255] -> [0,1] 且 HWC -> CHW
            transforms.Normalize(mean=self.mean, std=self.std)
        ])
        
        tensor_list = []
        for frame in frames_bgr:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            tensor_list.append(transform(rgb_frame))
            
        return torch.stack(tensor_list)  # [num_frames, 3, H, W]
    
    def process(self, video_path):
        """完整流水线:采样 -> 缩放 -> 归一化,返回tensor"""
        sampled_frames = self._uniform_sample(video_path)
        return self._preprocess_frames(sampled_frames)

# 最小调用入口
if __name__ == "__main__":
    # 使用示例(需替换为实际视频路径)
    preprocessor = VideoFramePreprocessor(num_frames=8)
    # output_tensor = preprocessor.process("sample_video.mp4")
    # print(f"输出张量形状: {output_tensor.shape}")  # 预期: torch.Size([8, 3, 224, 224])

运行与验证

运行方式:将上述代码保存为 video_preprocess.py,在Python环境中执行以下命令验证(需准备一个测试视频文件):

python -c "
from video_preprocess import VideoFramePreprocessor
import torch

proc = VideoFramePreprocessor(num_frames=8)
tensor = proc.process('your_test_video.mp4')
print('输出形状:', tensor.shape)
print('数值范围:', tensor.min().item(), '~', tensor.max().item())
"

通过条件:

  • 输出张量形状为 [8, 3, 224, 224](对应8帧、RGB三通道、224×224分辨率);
  • 归一化后张量数值大致落在 [-2.5, 2.5] 区间(因ImageNet均值和标准差的作用);
  • 采样帧在时间轴上均匀分布——可通过打印采样索引或逐帧保存对比确认。

关键说明

  • 均匀采样逻辑:使用 np.linspace 生成包含首尾帧的等间隔索引,避免使用 range 可能导致的边界偏移。
  • 帧读取失败处理:当视频编码异常导致某帧读取失败时,用前一帧填充,保证输出帧数恒定。
  • 归一化参数:默认采用ImageNet统计值(mean/std),若用于自监督训练可替换为 (0.5,0.5,0.5) 或数据集统计值。
  • 扩展性:如需增加随机裁剪、色彩抖动等增强,在 _preprocess_frames 的 transforms.Compose 中追加即可。