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()

运行方式

  1. 将待处理的视频文件命名为sample_video.mp4并放在当前目录(或修改代码中的video_path变量);
  2. 直接运行脚本:python extract_i3d_features.py;
  3. 脚本将输出特征维度并保存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库,但核心调用逻辑不变。本案例聚焦于特征提取流程的完整闭环,模型结构差异不影响特征保存与验证方式。