使用 mediapipe 提取人物动作角度并生成数据集

本文介绍了使用 mediapipe 库提取人物动作角度并生成数据集的代码。代码使用 mediapipe 库分析图像，提取关键点信息，并计算角度，最终生成 CSV 文件数据集。

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]
            right_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST]

            # 获取左臀、左膝和左踝关键点的信息
            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]
            right_knee = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE]
            right_ankle = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE]

            # 计算腿与右手的角度
            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.y) -
                                    math.atan2(right_hip.y - right_knee.y, right_hip.x - right_knee.y))

            # 获取右肩、右肘和右手腕
            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')

代码说明:

使用 mediapipe 库提取人物关键点信息，包括左肩、左肘、左手腕、右肩、右肘、右手腕、左臀、左膝、左踝、右臀、右膝、右踝。
计算关键点之间的角度，包括腿与右手的角度、左肩左肘左手腕的角度、左臀左膝左踝的角度、右臀右膝右踝的角度、右肩右肘右手腕的角度。
将计算得到的角度保存到 CSV 文件中，方便后续使用。

注意:

代码中使用了 os.listdir() 函数遍历文件夹，需要确保 actions 文件夹存在，并且包含不同动作的图片。
代码中使用了 cv2.imwrite() 函数保存图片，需要确保 annotated_images 文件夹存在，以便保存绘制好骨骼的图片。

使用场景:

用于生成动作识别的训练数据。
用于分析人物动作变化，例如运动员的动作分析。
用于开发动作捕捉系统。

未来改进:

可以增加更多关键点和角度的计算，例如头部角度、躯干角度等。
可以使用更先进的算法，例如深度学习，对动作进行更精确的分析和识别。
可以开发更友好的界面，方便用户进行动作采集和分析。

希望这篇文章对您有所帮助！如果您有任何问题，请随时提出。