以下是一个基于ST-Transformer（Spatio-Temporal Transformer）的人体动作识别程序的实现示例。该程序使用Python和PyTorch框架，从指定视频文件中识别人体的动作。 ### 实现步骤 1. **安装依赖库**：确保安装了必要的库，如`opencv-python`用于视频处理，`torch`和`torchvision`用于深度学习模型。 2. **加载预训练模型**：假设你已经有了一个预训练好的ST-Transformer模型。 3. **视频预处理**：将视频转换为帧，并提取人体关键点（例如通过OpenPose或AlphaPose）。 4. **推理阶段**：将提取的关键点输入到ST-Transformer模型中进行预测。 以下是完整的代码实现： ```python import cv2 import torch from torchvision import transforms from st_transformer_model import ST_Transformer # 假设这是你的ST-Transformer模型定义文件 # 加载预训练的ST-Transformer模型 def load_model(model_path): model = ST_Transformer(num_classes=60, num_frames=32, num_joints=25) # 示例参数 model.load_state_dict(torch.load(model_path)) model.eval() return model # 视频预处理函数 def preprocess_video(video_path, num_frames=32, num_joints=25): cap = cv2.VideoCapture(video_path) frames = [] while True: ret, frame = cap.read() if not ret: break # 将帧调整为固定大小并归一化 frame = cv2.resize(frame, (224, 224)) frame = transforms.ToTensor()(frame) frames.append(frame) # 提取固定数量的帧 if len(frames) > num_frames: frames = [frames[i] for i in range(0, len(frames), len(frames) // num_frames)[:num_frames]] else: frames += [frames[-1]] * (num_frames - len(frames)) # 如果帧数不足，则重复最后一帧 # 假设这里有一个函数 extract_keypoints_from_frames 来提取关键点 keypoints = extract_keypoints_from_frames(frames, num_joints) return torch.tensor(keypoints).unsqueeze(0) # 添加批次维度 # 假设的函数：从帧中提取关键点 def extract_keypoints_from_frames(frames, num_joints): # 这里可以调用 OpenPose 或 AlphaPose 等工具来提取关键点 # 返回形状为 (num_frames, num_joints, 3) 的张量 # 其中每个关键点包含 (x, y, confidence) dummy_keypoints = torch.rand(len(frames), num_joints, 3) # 示例数据 return dummy_keypoints # 动作识别函数 def recognize_action(video_path, model): keypoints = preprocess_video(video_path) with torch.no_grad(): output = model(keypoints) _, predicted_class = torch.max(output, 1) return predicted_class.item() # 主函数 if __name__ == "__main__": video_path = "path_to_your_video.mp4" # 替换为你的视频路径 model_path = "path_to_pretrained_st_transformer.pth" # 替换为你的模型路径 model = load_model(model_path) action_class = recognize_action(video_path, model) print(f"Predicted Action Class: {action_class}") ``` ### 代码解释 1. **模型加载**： - `load_model` 函数加载预训练的ST-Transformer模型，并将其设置为评估模式。 - 模型的输入参数包括类别数 (`num_classes`)、帧数 (`num_frames`) 和关节数 (`num_joints`)。 2. **视频预处理**： - `preprocess_video` 函数读取视频文件，并将其转换为固定数量的帧。 - 使用 `cv2` 对帧进行缩放和归一化。 - 调用 `extract_keypoints_from_frames` 函数从帧中提取人体关键点。 3. **动作识别**： - `recognize_action` 函数将提取的关键点输入到ST-Transformer模型中进行推理。 - 输出结果为预测的动作类别。 4. **主函数**： - 在主函数中，指定视频路径和模型路径，调用上述函数完成动作识别。 ### 注意事项 - 需要提前准备好ST-Transformer模型的定义文件 `st_transformer_model.py`。 - 提取关键点的部分需要结合具体工具（如OpenPose或AlphaPose）实现。 - 视频帧数和关节数需要根据实际情况调整。