使用 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", "nose_z", "left_shoulder_x", "left_shoulder_y",
                                              "left_shoulder_z", "right_shoulder_x", "right_shoulder_y", "right_shoulder_z",
                                              "left_elbow_x", "left_elbow_y", "left_elbow_z", "right_elbow_x", "right_elbow_y",
                                              "right_elbow_z", "left_wrist_x", "left_wrist_y", "left_wrist_z", "right_wrist_x",
                                              "right_wrist_y", "right_wrist_z", "left_hip_x", "left_hip_y", "left_hip_z",
                                              "right_hip_x", "right_hip_y", "right_hip_z", "left_knee_x", "left_knee_y",
                                              "left_knee_z", "right_knee_x", "right_knee_y", "right_knee_z", "left_ankle_x",
                                              "left_ankle_y", "left_ankle_z", "right_ankle_x", "right_ankle_y", "right_ankle_z"])
    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"):
    # 创建一个空的姿势估计结果列表
    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([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.NOSE].z,
                                 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.LEFT_SHOULDER].z,
                                 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.RIGHT_SHOULDER].z,
                                 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.LEFT_ELBOW].z,
                                 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.RIGHT_ELBOW].z,
                                 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.LEFT_WRIST].z,
                                 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.RIGHT_WRIST].z,
                                 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.LEFT_HIP].z,
                                 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.RIGHT_HIP].z,
                                 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.LEFT_KNEE].z,
                                 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.RIGHT_KNEE].z,
                                 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.LEFT_ANKLE].z,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].x,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].y,
                                 results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].z])
            # 在图像上绘制姿势估计结果
            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")

输出的数据集为空的原因内容：可能是因为在处理图片时，姿势估计模型无法检测到人体姿势，导致姿势估计结果为空。可以尝试调整模型参数或者使用其他模型来提高检测准确率。另外，也可以检查图片质量或者姿势是否符合模型要求。