使用 PyTorch 和 Streamlit 实现视频动作识

近年来，视频动作识别因其在体育分析、安全和娱乐等领域的广泛应用而受到关注。在这篇博文中，我将向您介绍如何使用 PyTorch 和 Streamlit 开发视频动作识别应用程序。此应用程序允许用户上传视频并让深度学习模型预测正在执行的作。

应用程序概述

该应用程序实现了以下功能：

动作识别模型：在自定义数据集上微调的预先训练的视频动作识别模型（r2plus1d_18来自 PyTorch 的模型动物园）。

视频上传和转换：用户可以上传各种格式的视频文件。该应用程序将它们转换为 MP4 以确保兼容性和最佳播放效果。

预测：该应用程序使用深度学习模型从上传的视频中预测作，并显示预测的类别和置信度分数。

Confidence Threshold （置信度阈值）：如果模型的置信度低于某个阈值，则该作将被归类为 “Unknown”（未知）。

应用程序结构

1. 加载模型和标签映射

第一步是加载推理所需的标签映射和模型。我们使用 JSON 文件来存储标签映射（类名映射到数字标签）和 idx_to_class（将数字标签映射回类名）。

@st.cache_resourcedef load_label_mappings(label_map_path, idx_to_class_path):    with open(label_map_path, 'r') as f:        label_map = json.load(f)    with open(idx_to_class_path, 'r') as f:        idx_to_class = json.load(f)    idx_to_class = {int(k): v for k, v in idx_to_class.items()}    return label_map, idx_to_class

加载和修改模型本身以匹配数据集中的类数。然后加载训练模型的权重。

@st.cache_resourcedef load_model(num_classes, device, model_path):    model = models.video.r2plus1d_18(pretrained=False)    for param in model.parameters():        param.requires_grad = False    model.layer4.requires_grad = True    model.fc = nn.Linear(model.fc.in_features, num_classes)    model.load_state_dict(torch.load(model_path, map_location=device))    model.eval()    return model

2. 视频推理数据集

由于我们处理的是视频，因此输入数据格式与标准图像分类不同。为此，我们定义了一个自定义 VideoInferenceDataset 类，该类从视频加载帧，应用转换（例如调整大小和规范化），并为模型输入做好准备。

class VideoInferenceDataset(Dataset):    def __init__(self, video_path, transform=None, frames_per_clip=32):        self.video_path = video_path        self.transform = transform        self.frames_per_clip = frames_per_clip
    def __getitem__(self, idx):        frames = self._load_video_frames(self.video_path)        frames = [self.transform(frame) for frame in frames]        return torch.stack(frames).permute(1, 0, 2, 3)

3. 视频处理和转换

由于视频可以有多种格式，因此我们使用 MoviePy 将上传的视频转换为 MP4 格式，然后再将其输入到模型中。该应用程序还会显示视频以便于预览。

def convert_to_mp4(input_video_path):    output_video_path = input_video_path.replace(os.path.splitext(input_video_path)[1], ".mp4")    clip = VideoFileClip(input_video_path)    clip.write_videofile(output_video_path, codec="libx264", audio_codec="aac")    clip.close()    return output_video_path

4. 预测

视频经过预处理和转换后，会将其传递给模型进行推理。模型输出每个类的概率，并显示预测的类以及相应的置信度分数。如果置信度低于预定义的阈值，则预测被视为 “未知”。

def predict(model, video_path, transform, device, idx_to_class, threshold=0.5):    dataset = VideoInferenceDataset(video_path, transform=transform, frames_per_clip=32)    loader = DataLoader(dataset, batch_size=1, shuffle=False)    with torch.no_grad():        for inputs in loader:            inputs = inputs.to(device)            outputs = model(inputs)            probabilities = torch.nn.functional.softmax(outputs, dim=1)            confidence, preds = torch.max(probabilities, 1)            if confidence.item() < threshold:                return "Unknown", confidence.item()            return idx_to_class[preds.item()], confidence.item()

5. Streamlit 应用程序

最后，整个管道被包装到 Streamlit 接口中。该应用程序允许用户上传视频文件、对其进行处理、运行推理并显示预测。