feat(ui): support multi-person skeleton drawing
This commit is contained in:
parent
389a7fd16b
commit
12d9c59a53
@ -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,8 +42,16 @@ 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, :, :]
|
||||
|
||||
# 遍历所有人
|
||||
num_instances = keypoints.shape[0]
|
||||
for person_idx in range(num_instances):
|
||||
person_kpts = keypoints[person_idx]
|
||||
person_scores = scores[person_idx]
|
||||
|
||||
# 绘制骨骼
|
||||
for limb_id, limb in enumerate(skeleton):
|
||||
|
Loading…
Reference in New Issue
Block a user