feat(pose): add joint similarity visualization in skeleton drawing
This commit is contained in:
parent
1d94e77557
commit
389a7fd16b
@ -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)
|
||||
|
@ -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."""
|
||||
|
Loading…
Reference in New Issue
Block a user