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=['label' , 'angle1', 'angle2', 'angle3','angle4','angle5','angle5_1','angle6', 'angle7','angle8','angle9','angle10','angle11'])
    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.3) 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]

            # 获取左肩、左肘和左手腕关键点的信息 new
            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]

            # 计算腿与右手的角度
            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))

            # 获取左肩、左肘和左手腕 new
            angle_tl = math.degrees(math.atan2(left_wrist.y - left_elbow.y, left_wrist.x - left_elbow.x) -
                                    math.atan2(left_shoulder.y - left_elbow.y, left_shoulder.x - left_elbow.x))

            ##############
            #获取左髋、左膝和左踝关键点的信息
            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]

            # 获取左髋、左膝和左踝的角度
            angle_lka = 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))

            # 获取腰部、左髋和左肩的角度
            waist = results.pose_landmarks.landmark[12]
            angle_hls = math.degrees(math.atan2(left_shoulder.y - waist.y, left_shoulder.x - waist.x) -
                                     math.atan2(left_hip.y - waist.y, left_hip.x - waist.x))

            # 获取右髋、右膝和右踝关键点的信息
            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_rka = 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_hrs = math.degrees(math.atan2(right_shoulder.y - waist.y, right_shoulder.x - waist.x) -
                                     math.atan2(right_hip.y - waist.y, right_hip.x - waist.x))
            ##################
            # 获取鼻子、右肩和右手腕关键点的信息
            nose = results.pose_landmarks.landmark[mp_pose.PoseLandmark.NOSE]
            right_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
            right_wrist = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST]

            # 获取鼻子、右肩和右手腕的角度
            angle_nwr = math.degrees(math.atan2(right_wrist.y - nose.y, right_wrist.x - nose.x) -
                                     math.atan2(right_shoulder.y - nose.y, right_shoulder.x - nose.x))




            # 将角度添加到角度列表中


            angles_list.append([folder_name , angle, angle1, angle_dl , angle_dr , angle_tr, angle_tr , angle_tl,angle_lka,angle_hls,angle_rka,angle_hrs,angle_nwr])


            # 在图像上绘制骨骼
            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')

算法描述

输入: 动作图片文件夹

输出: 保存动作骨骼角度的 CSV 文件

功能: 提取动作图片中的人体骨骼关键点,计算骨骼角度,并将角度保存为 CSV 文件,用于动作识别的训练数据集构建。

算法流程:

  1. 导入库: 导入必要的库,包括 mediapipecv2pandas 等库。
  2. 定义函数: 定义一个 save_angles() 函数,用于将骨骼角度列表保存为 CSV 文件。
  3. 初始化 MediaPipe: 初始化 mediapipe,导入姿势检测和绘图库。
  4. 遍历文件夹: 遍历不同动作的文件夹,依次处理每个动作的图片。
  5. 读取图片: 读取图片,并将图片转换为 RGB 格式。
  6. 提取骨骼: 使用 mediapipe 的姿势检测模型处理图片,提取骨骼关键点。
  7. 计算角度: 根据关键点的位置,计算不同骨骼之间的角度。代码中计算了 13 个角度,包括腿与右手的角度、左臀、左膝和左踝的角度、右臀、右膝和右踝的角度、右肩、右肘和右手腕的角度、左肩、左肘和左手腕的角度、左髋、左膝和左踝的角度、腰部、左髋和左肩的角度、右髋、右膝和右踝的角度、腰部、右髋和右肩的角度、鼻子、右肩和右手腕的角度。
  8. 添加角度: 将计算得到的角度添加到角度列表中。
  9. 绘制骨骼: 在图像上绘制骨骼关键点和连接线,用于可视化。
  10. 保存图片: 保存绘制好骨骼的图像。
  11. 保存数据: 将角度列表保存为 CSV 文件。
  12. 循环处理: 循环处理所有动作文件夹中的图片,直到所有图片处理完毕。

注意: 此代码假设动作图片文件夹的命名方式为 actions,并且每个动作文件夹中包含对应动作的图片。代码中使用 os.listdir() 函数遍历文件夹,使用 cv2.imread() 函数读取图片,使用 mediapipe 的姿势检测模型提取骨骼关键点,使用 math.atan2() 函数计算角度,使用 pandasDataFrame 类将角度数据保存为 CSV 文件。

代码改进建议:

  • 可以增加参数,例如动作图片文件夹的路径,CSV 文件的保存路径等。
  • 可以使用 argparse 库解析命令行参数,方便用户自定义输入参数。
  • 可以根据需要添加更多角度的计算,例如计算头部、躯干等部位的角度。
  • 可以使用更完善的姿态识别模型,例如使用 BlazePose 模型,提高识别精度。
  • 可以添加错误处理机制,例如处理图片读取失败、骨骼识别失败等情况。
  • 可以将代码封装成函数,方便复用。

使用场景:

  • 构建动作识别的数据集
  • 分析人体动作模式
  • 开发动作识别应用
  • 辅助运动训练

相关资源:

  • MediaPipe 文档:https://google.github.io/mediapipe/solutions/pose.html
  • OpenCV 文档:https://opencv-python-tutroals.readthedocs.io/en/latest/
  • Pandas 文档:https://pandas.pydata.org/docs/
人体骨骼角度提取与分析:基于 MediaPipe 的动作识别

原文地址: https://www.cveoy.top/t/topic/gQ23 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录