From 389a7fd16b55091e04f2cb1b5fbc72bf74373ed5 Mon Sep 17 00:00:00 2001 From: gameloader Date: Mon, 30 Jun 2025 18:53:51 +0800 Subject: [PATCH] feat(pose): add joint similarity visualization in skeleton drawing --- motion_app.py | 85 +++++++++++++++++++++++++++++++++++++++++++++--- pose_analyzer.py | 13 ++++++++ 2 files changed, 94 insertions(+), 4 deletions(-) diff --git a/motion_app.py b/motion_app.py index 6c3db64..6255d72 100644 --- a/motion_app.py +++ b/motion_app.py @@ -12,6 +12,75 @@ 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): + """ + 自定义骨骼绘制函数，根据关节相似度设置颜色 + 相似度 > 90: 绿色 + 相似度 80-90: 黄色 + 相似度 < 80: 红色 + """ + if keypoints is None or len(keypoints) == 0: + return img + + # OpenPose连接关系 + skeleton = [ + [1, 0], [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], + [1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], + [0, 14], [0, 15], [14, 16], [15, 17] + ] + + # 关节与关节角度的映射 + joint_to_angle_mapping = { + (2, 3): 'left_shoulder', (3, 4): 'left_elbow', + (5, 6): 'right_shoulder', (6, 7): 'right_elbow', + (8, 9): 'left_hip', (9, 10): 'left_knee', + (11, 12): 'right_hip', (12, 13): 'right_knee' + } + + # 骨骼的默认颜色 + 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 + + # 绘制骨骼 + for limb_id, limb in enumerate(skeleton): + joint_a, joint_b = limb + + 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] + + # 确定线条颜色 + 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) + + return img + if REALSENSE_AVAILABLE: import pyrealsense2 as rs @@ -119,7 +188,7 @@ class MotionComparisonApp: if self.body_detector: try: keypoints, scores = self.body_detector(frame) - frame = draw_skeleton(frame.copy(), keypoints, scores, openpose_skeleton=True, kpt_thr=0.43) + frame = draw_skeleton_with_similarity(frame.copy(), keypoints, scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=1) except Exception: pass return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) @@ -341,9 +410,17 @@ class MotionComparisonApp: # --- 绘制与显示 (可以保持原来的逻辑，只修改UI更新部分) --- try: - # 绘制骨骼 - standard_with_keypoints = draw_skeleton(standard_frame.copy(), standard_keypoints, standard_scores, openpose_skeleton=True, kpt_thr=0.43) - webcam_with_keypoints = draw_skeleton(webcam_frame.copy(), webcam_keypoints, webcam_scores, openpose_skeleton=True, kpt_thr=0.43) + # 计算关节相似度 + joint_similarities = None + if standard_keypoints is not None and webcam_keypoints is not None: + standard_angles = self.similarity_analyzer.extract_joint_angles(standard_keypoints, standard_scores) + webcam_angles = self.similarity_analyzer.extract_joint_angles(webcam_keypoints, webcam_scores) + if standard_angles and webcam_angles: + joint_similarities = self.similarity_analyzer.calculate_joint_similarities(standard_angles, webcam_angles) + + # 绘制骨骼 (线条更窄，根据相似度设置颜色) + standard_with_keypoints = draw_skeleton_with_similarity(standard_frame.copy(), standard_keypoints, standard_scores, joint_similarities=None, openpose_skeleton=True, kpt_thr=0.43, line_width=1) + webcam_with_keypoints = draw_skeleton_with_similarity(webcam_frame.copy(), webcam_keypoints, webcam_scores, joint_similarities=joint_similarities, openpose_skeleton=True, kpt_thr=0.43, line_width=1) # 更新视频画面 (主线程专注于实时显示) standard_placeholder.image(cv2.cvtColor(standard_with_keypoints, cv2.COLOR_BGR2RGB), use_container_width=True) diff --git a/pose_analyzer.py b/pose_analyzer.py index 996fd4a..0d84e3e 100644 --- a/pose_analyzer.py +++ b/pose_analyzer.py @@ -93,6 +93,19 @@ class PoseSimilarityAnalyzer: final_similarity = (weighted_similarity / total_weight) * 100 if total_weight > 0 else 0 return min(max(final_similarity, 0), 100) + + def calculate_joint_similarities(self, angles1, angles2): + """计算每个关节的相似度""" + joint_similarities = {} + if not angles1 or not angles2: return joint_similarities + + common_joints = set(angles1.keys()) & set(angles2.keys()) + for joint in common_joints: + angle_diff = abs(angles1[joint] - angles2[joint]) + similarity = math.exp(-(angle_diff ** 2) / (2 * (30 ** 2))) + joint_similarities[joint] = min(max(similarity * 100, 0), 100) + + return joint_similarities def add_similarity_score(self, score, timestamp=None): """Adds a similarity score to the history."""