使用 MediaPipe 进行人体姿态估计并保存结果

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


# 定义保存姿势估计结果的函数
def save_pose_results(pose_results, action_name, folder_name):
    filename = f'{folder_name}_{action_name}.csv'
    df = pd.DataFrame(pose_results, columns=['label', 'nose_x', 'nose_y', 'left_shoulder_x', 'left_shoulder_y',
                                              'right_shoulder_x', 'right_shoulder_y', 'left_elbow_x', 'left_elbow_y',
                                              'right_elbow_x', 'right_elbow_y', 'left_wrist_x', 'left_wrist_y',
                                              'right_wrist_x', 'right_wrist_y', 'left_hip_x', 'left_hip_y',
                                              'right_hip_x', 'right_hip_y', 'left_knee_x', 'left_knee_y',
                                              'right_knee_x', 'right_knee_y', 'left_ankle_x', 'left_ankle_y',
                                              'right_ankle_x', 'right_ankle_y'])
    df.to_csv(filename, index=False)
    print(f'{filename} saved successfully')

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

# 定义动作名称和数字之间的映射关系
action_map = {'warm-up': 0, 'jumping-jacks': 1, 'squats': 2, 'lunges': 3, 'push-ups': 4, 'crunches': 5}

# 遍历不同动作的文件夹
for folder_name in os.listdir('actions'):
    # 创建一个空的姿势估计结果列表
    pose_results = []
    # 遍历文件夹中的所有图片
    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
            # 将姿势估计结果添加到列表中
            pose_results.append([action_map[folder_name], results.pose_landmarks.landmark[mp_pose.PoseLandmark.NOSE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.NOSE].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].y])
            # 在图像上绘制姿势估计结果
            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_pose_results(pose_results, folder_name, 'dataset')

代码说明：

导入库: 导入必要的库，包括 mediapipe、cv2、os 和 pandas。
定义保存函数: 定义 save_pose_results 函数，用于将姿势估计结果保存为 CSV 文件。
初始化 MediaPipe: 初始化 MediaPipe 的 drawing_utils 和 pose 模块。
定义动作映射: 定义一个字典 action_map，将动作名称映射为数字。
遍历文件夹: 遍历 actions 文件夹中的所有动作文件夹。
遍历图片: 遍历每个动作文件夹中的所有图片。
读取图片: 读取图片并将其转换为 RGB 格式。
进行姿势估计: 使用 MediaPipe 的 Pose 模型进行姿势估计。
保存结果: 将姿势估计结果添加到列表中，并将其保存为 CSV 文件。
绘制结果: 在图像上绘制姿势估计结果，并保存绘制好的图像。

解决错误：

代码中使用 action_map 将动作名称转换为数字，避免了将字符串保存到 CSV 文件中导致的错误。

注意：

请确保 actions 文件夹中包含不同动作的图片文件夹。
annotated_images 文件夹用于保存绘制好姿势估计结果的图像。
可以根据需要修改 action_map 字典，添加或删除动作名称。

运行代码：

安装必要的库：pip install mediapipe opencv-python pandas
将代码保存为 Python 文件，例如 pose_estimation.py
运行代码：python pose_estimation.py

代码运行后，将在 dataset 文件夹中生成包含姿势估计结果的 CSV 文件。