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中追加即可。