diff --git a/motion_app.py b/motion_app.py index 6255d72..fee0763 100644 --- a/motion_app.py +++ b/motion_app.py @@ -12,12 +12,14 @@ from audio_player import AudioPlayer from pose_analyzer import PoseSimilarityAnalyzer from config import REALSENSE_AVAILABLE -def draw_skeleton_with_similarity(img, keypoints, scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=1): +def draw_skeleton_with_similarity(img, keypoints, scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=2): """ 自定义骨骼绘制函数,根据关节相似度设置颜色 相似度 > 90: 绿色 相似度 80-90: 黄色 相似度 < 80: 红色 + + 支持多人检测,将绘制所有检测到的人体骨骼 """ if keypoints is None or len(keypoints) == 0: return img @@ -40,44 +42,52 @@ def draw_skeleton_with_similarity(img, keypoints, scores, joint_similarities=Non # 骨骼的默认颜色 default_color = (0, 255, 255) # 黄色 - person_kpts = keypoints[0] if len(keypoints.shape) > 2 else keypoints - person_scores = scores[0] if len(scores.shape) > 1 else scores + # 确保keypoints和scores的形状正确,处理多人情况 + if len(keypoints.shape) == 2: + keypoints = keypoints[None, :, :] + scores = scores[None, :, :] - # 绘制骨骼 - for limb_id, limb in enumerate(skeleton): - joint_a, joint_b = limb + # 遍历所有人 + num_instances = keypoints.shape[0] + for person_idx in range(num_instances): + person_kpts = keypoints[person_idx] + person_scores = scores[person_idx] - if joint_a >= len(person_scores) or joint_b >= len(person_scores): - continue + # 绘制骨骼 + for limb_id, limb in enumerate(skeleton): + joint_a, joint_b = limb - if person_scores[joint_a] < kpt_thr or person_scores[joint_b] < kpt_thr: - continue + if joint_a >= len(person_scores) or joint_b >= len(person_scores): + continue + + if person_scores[joint_a] < kpt_thr or person_scores[joint_b] < kpt_thr: + continue + + x_a, y_a = person_kpts[joint_a] + x_b, y_b = person_kpts[joint_b] - x_a, y_a = person_kpts[joint_a] - x_b, y_b = person_kpts[joint_b] + # 确定线条颜色 + color = default_color + if joint_similarities is not None: + # 检查这个连接是否有对应的关节角度 + if (joint_a, joint_b) in joint_to_angle_mapping: + angle_name = joint_to_angle_mapping[(joint_a, joint_b)] + if angle_name in joint_similarities: + similarity = joint_similarities[angle_name] + if similarity > 90: + color = (0, 255, 0) # 绿色 + elif similarity > 80: + color = (0, 255, 255) # 黄色 + else: + color = (0, 0, 255) # 红色 + + cv2.line(img, (int(x_a), int(y_a)), (int(x_b), int(y_b)), color, thickness=line_width) - # 确定线条颜色 - color = default_color - if joint_similarities is not None: - # 检查这个连接是否有对应的关节角度 - if (joint_a, joint_b) in joint_to_angle_mapping: - angle_name = joint_to_angle_mapping[(joint_a, joint_b)] - if angle_name in joint_similarities: - similarity = joint_similarities[angle_name] - if similarity > 90: - color = (0, 255, 0) # 绿色 - elif similarity > 80: - color = (0, 255, 255) # 黄色 - else: - color = (0, 0, 255) # 红色 - - cv2.line(img, (int(x_a), int(y_a)), (int(x_b), int(y_b)), color, thickness=line_width) - - # 绘制关键点 - for kpt_id, (x, y) in enumerate(person_kpts): - if person_scores[kpt_id] < kpt_thr: - continue - cv2.circle(img, (int(x), int(y)), 3, (255, 0, 255), -1) + # 绘制关键点 + for kpt_id, (x, y) in enumerate(person_kpts): + if person_scores[kpt_id] < kpt_thr: + continue + cv2.circle(img, (int(x), int(y)), 3, (255, 0, 255), -1) return img