基于Mediapipe和OpenCV的人物动作识别程序

本程序使用Python、Mediapipe和OpenCV库来识别视频中的人物动作。

数据集生成

首先，我们需要生成包含不同动作的骨骼角度数据集。

import mediapipe as mp
import cv2
import os
import pandas as pd
import math


# 定义保存骨骼角度的函数
def save_angles(angles_list, action_name, folder_name):
    filename = f'{folder_name}_{action_name}.csv'
    df = pd.DataFrame(angles_list, columns=['angle1', 'angle2', 'angle3', 'angle4', 'angle5'])
    df.to_csv(filename, index=False)
    print(f'{filename} saved successfully')

# 初始化mediapipe
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

# 遍历不同动作的文件夹
for folder_name in os.listdir('actions'):
    # 创建一个空的角度列表
    angles_list = []
    # 遍历文件夹中的所有图片
    for filename in os.listdir(f'actions/{folder_name}'):
        # 读取图片
        image = cv2.imread(f'actions/{folder_name}/{filename}')
        # 将图片转换为RGB格式
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 处理图片，提取骨骼
        with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
            results = pose.process(image)
            if results.pose_landmarks is None:
                continue
            # 获取关键点信息
            right_knee = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE]
            right_ankle = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE]
            right_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST]
            left_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER]
            left_elbow = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW]
            left_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST]
            right_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
            right_elbow = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW]
            left_hip = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP]
            left_knee = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE]
            left_ankle = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE]
            right_hip = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP]

            # 计算角度
            angle = math.degrees(math.atan2(right_wrist.y - right_ankle.y, right_wrist.x - right_ankle.x) -
                                 math.atan2(right_knee.y - right_ankle.y, right_knee.x - right_ankle.x))
            angle1 = math.degrees(math.atan2(right_wrist.y - right_ankle.y, right_wrist.x - right_ankle.x) -
                                 math.atan2(right_knee.y - right_ankle.y, right_knee.x - right_ankle.x))
            angle_dl = math.degrees(math.atan2(left_ankle.y - left_knee.y, left_ankle.x - left_knee.x) -
                                    math.atan2(left_hip.y - left_knee.y, left_hip.x - left_knee.x))
            angle_dr = math.degrees(math.atan2(right_ankle.y - right_knee.y, right_ankle.x - right_knee.x) -
                                    math.atan2(right_hip.y - right_knee.y, right_hip.x - right_knee.x))
            angle_tr = math.degrees(math.atan2(right_wrist.y - right_elbow.y, right_wrist.x - right_elbow.x) -
                                    math.atan2(right_shoulder.y - right_elbow.y, right_shoulder.x - right_elbow.x))

            # 将角度添加到角度列表中
            angles_list.append([angle, angle1, angle_dl,angle_dr,angle_tr])
            # 在图像上绘制骨骼
            annotated_image = image.copy()
            mp_drawing.draw_landmarks(annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
        # 保存绘制好骨骼的图像
        cv2.imwrite(f'annotated_images/{filename}', cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR))
    # 将角度列表保存为csv文件
    save_angles(angles_list, folder_name, 'dataset')

动作识别程序

以下代码展示了如何使用生成的数据集来识别视频'1.mp4'中的人物动作。

# 导入必要的库
import mediapipe as mp
import cv2
import pandas as pd
import math

# 定义计算角度的函数
def calculate_angles(image):
    # 处理图片，提取骨骼
    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
        results = pose.process(image)
        if results.pose_landmarks is None:
            return None
        # 获取关键点信息 (代码同上)
        # ...
        # 计算角度 (代码同上)
        # ...
        angles = [angle, angle1, angle_dl, angle_dr, angle_tr]
        return angles

# 初始化mediapipe
mp_drawing = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose

# 加载模型
# 此处需要加载您训练好的动作识别模型

# 打开视频文件
cap = cv2.VideoCapture('1.mp4')
# 定义一个空的角度列表
angles_list = []
# 循环读取视频帧
while cap.isOpened():
    # 读取视频帧
    success, image = cap.read()
    if not success:
        break
    # 将图片转换为RGB格式
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 处理图片，提取骨骼
    angles = calculate_angles(image)
    if angles is not None:
        angles_list.append(angles)
        # 使用加载的模型进行动作识别
        # ...
        # 在图像上绘制骨骼
        annotated_image = image.copy()
        mp_drawing.draw_landmarks(annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
        # 显示绘制好骨骼的图像
        cv2.imshow('Annotated Image', cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR))
    # 按下q键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
# 释放资源
cap.release()
cv2.destroyAllWindows()
# 将角度列表保存为csv文件
df = pd.DataFrame(angles_list, columns=['angle1', 'angle2', 'angle3', 'angle4', 'angle5'])
df.to_csv('1.csv', index=False)
print('1.csv saved successfully')

注意

• 您需要将代码中的'actions'文件夹替换为您自己存储动作图片的文件夹。
• 您需要训练自己的机器学习模型并将其加载到动作识别程序中。

结论

本程序提供了一个使用Mediapipe和OpenCV进行人物动作识别的基本框架。您可以根据自己的需要修改和扩展此程序。