3. 视频特征提取(I3D)
任务说明
本案例使用PyTorch实现基于预训练I3D模型的视频时空特征提取。输入为任意长度视频片段,输出为对应的512维特征向量,并保存为.npy文件。I3D模型在Kinetics-400上预训练,可捕获视频中的时空信息,是视频理解任务中常用的特征提取器。
环境依赖
需要安装以下依赖:
pip install torch torchvision opencv-python-headless numpy decord
其中decord用于高效视频解码,opencv-python-headless用于备用解码方案。
完整实现
import torch
import numpy as np
import decord
from decord import VideoReader, cpu
from pathlib import Path
# I3D模型定义(基于Kinetics-400预训练)
class I3DEncoder(torch.nn.Module):
def __init__(self):
super().__init__()
# 使用torchvision的r3d_18作为I3D的替代实现
# 实际I3D结构更复杂,此处用3D ResNet保证可复现性
from torchvision.models.video import r3d_18, R3D_18_Weights
self.backbone = r3d_18(weights=R3D_18_Weights.KINETICS400_V1)
# 移除分类头,保留特征层
self.backbone.fc = torch.nn.Identity()
self._freeze()
def _freeze(self):
for param in self.backbone.parameters():
param.requires_grad = False
def forward(self, x):
# x: [B, C, T, H, W] = [B, 3, 16, 224, 224]
return self.backbone(x) # [B, 512]
def load_video_frames(video_path, num_frames=16, size=224):
"""读取视频并采样16帧,resize到224x224"""
vr = VideoReader(str(video_path), ctx=cpu(0))
total_frames = len(vr)
if total_frames == 0:
raise ValueError(f"视频为空: {video_path}")
# 均匀采样16帧
indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
frames = vr.get_batch(indices).asnumpy() # [16, H, W, C]
# 转换为tensor并归一化
frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float() # [16, C, H, W]
frames = frames / 255.0
# ImageNet归一化
mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
frames = (frames - mean) / std
# resize到224x224
frames = torch.nn.functional.interpolate(
frames, size=(size, size), mode='bilinear', align_corners=False
)
# 调整维度为 [C, T, H, W]
return frames.permute(1, 0, 2, 3).unsqueeze(0) # [1, 3, 16, 224, 224]
def extract_features(video_path, model, device):
"""提取视频特征"""
frames = load_video_frames(video_path)
frames = frames.to(device)
with torch.no_grad():
features = model(frames) # [1, 512]
return features.cpu().numpy().squeeze() # [512]
def main():
# 配置
video_path = "sample_video.mp4" # 输入视频路径
output_path = "video_features.npy" # 输出特征路径
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 初始化模型
model = I3DEncoder().to(device)
model.eval()
print(f"使用设备: {device}")
# 提取特征
features = extract_features(video_path, model, device)
print(f"特征维度: {features.shape}")
# 保存为numpy数组
np.save(output_path, features)
print(f"特征已保存至: {output_path}")
if __name__ == "__main__":
main()
运行方式
- 将待处理的视频文件命名为
sample_video.mp4并放在当前目录(或修改代码中的video_path变量); - 直接运行脚本:
python extract_i3d_features.py; - 脚本将输出特征维度并保存
video_features.npy文件。
验证方法
运行成功后,通过以下方式验证结果正确性:
import numpy as np
features = np.load("video_features.npy")
print(f"特征形状: {features.shape}") # 应为 (512,)
print(f"特征类型: {features.dtype}") # 应为 float32 或 float64
print(f"特征范数: {np.linalg.norm(features):.4f}") # 非零值
判定通过条件:
- 输出特征形状为
(512,); - 特征值不全为零(预训练模型对有效视频输入不会输出全零向量);
- 不同视频输入产生不同特征向量(可对两个不同视频分别提取并比较余弦相似度,应明显小于1)。
说明
本实现使用torchvision中的r3d_18(3D ResNet-18)作为I3D的替代,两者均为视频时空特征提取器,输出维度一致(512维)。如需使用官方I3D(含Inception模块),可替换为pytorch_i3d库,但核心调用逻辑不变。本案例聚焦于特征提取流程的完整闭环,模型结构差异不影响特征保存与验证方式。