使用MediaPipe提取人体姿态角度数据并生成CSV数据集
使用MediaPipe提取人体姿态角度数据并生成CSV数据集
本文介绍如何使用MediaPipe库从不同文件夹的图片中提取人体骨骼角度数据,并将数据与对应的动作标签一起保存到CSV文件中,方便机器学习模型训练。
实现思路:
- 遍历不同文件夹中的图片,使用MediaPipe库获取每张图片中人体骨骼的角度信息。2. 将每张图片中的角度信息与对应的动作标签存储到一个字典中。3. 将字典中的数据转换成CSV格式,并存储到本地文件中。
**代码实现:**pythonimport osimport csvimport mediapipe as mp
定义MediaPipe的姿态估计类mp_pose = mp.solutions.posemp_drawing = mp.solutions.drawing_utils
定义动作名称和对应的文件夹路径actions = { 'squat': 'squat_images', 'pushup': 'pushup_images', 'stand': 'stand_images'}
定义一个空字典,用于存储每个动作对应的角度信息angles = {}
遍历每个动作对应的文件夹for action, folder in actions.items(): # 获取文件夹中的所有图片路径 image_files = os.listdir(folder) # 创建一个空列表,用于存储每张图片的角度信息 action_angles = [] # 遍历每张图片 for image_file in image_files: # 使用MediaPipe获取图片中的角度信息 with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose: image_path = os.path.join(folder, image_file) image = mp.solutions.pose.Pose.load_image_file(image_path) results = pose.process(image) # 将角度信息存储到列表中 if results.pose_landmarks is not None: landmarks = results.pose_landmarks.landmark # 此处需要根据MediaPipe的骨骼点定义提取所需的骨骼角度 # 例如,提取左肘关节角度: # left_elbow_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value], # landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value], # landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value]) # 将提取的骨骼角度添加到angle_list中 # angle_list.append(left_elbow_angle) action_angles.append(angle_list) # 将动作名称和角度信息存储到字典中 angles[action] = action_angles
将字典中的数据转换成CSV格式,并存储到本地文件中with open('angles.csv', mode='w', newline='') as file: writer = csv.writer(file) writer.writerow(['action', 'angles']) for action, action_angles in angles.items(): for angle_list in action_angles: writer.writerow([action, angle_list])
代码说明:
- 代码首先定义了动作名称和对应图片文件夹的字典
actions。2. 然后遍历每个动作,读取对应文件夹下的所有图片。3. 使用MediaPipe库检测图片中的人体姿态,并提取骨骼点坐标。4. 根据骨骼点坐标计算所需的骨骼角度。5. 将提取到的角度信息与动作标签一起存储到CSV文件中。
注意:
- 以上代码仅提供一个框架,需要根据实际需求修改骨骼角度计算方法和存储格式。* 可以使用OpenCV等库读取和处理图片数据。* 可以使用pandas等库更方便地处理和存储数据。
原文地址: https://www.cveoy.top/t/topic/gnU3 著作权归作者所有。请勿转载和采集!