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 pose_analyzer import PoseSimilarityAnalyzer
|
||||||
from config import REALSENSE_AVAILABLE
|
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: 绿色
|
相似度 > 90: 绿色
|
||||||
相似度 80-90: 黄色
|
相似度 80-90: 黄色
|
||||||
相似度 < 80: 红色
|
相似度 < 80: 红色
|
||||||
|
|
||||||
|
支持多人检测,将绘制所有检测到的人体骨骼
|
||||||
"""
|
"""
|
||||||
if keypoints is None or len(keypoints) == 0:
|
if keypoints is None or len(keypoints) == 0:
|
||||||
return img
|
return img
|
||||||
@ -40,44 +42,52 @@ def draw_skeleton_with_similarity(img, keypoints, scores, joint_similarities=Non
|
|||||||
# 骨骼的默认颜色
|
# 骨骼的默认颜色
|
||||||
default_color = (0, 255, 255) # 黄色
|
default_color = (0, 255, 255) # 黄色
|
||||||
|
|
||||||
person_kpts = keypoints[0] if len(keypoints.shape) > 2 else keypoints
|
# 确保keypoints和scores的形状正确,处理多人情况
|
||||||
person_scores = scores[0] if len(scores.shape) > 1 else scores
|
if len(keypoints.shape) == 2:
|
||||||
|
keypoints = keypoints[None, :, :]
|
||||||
|
scores = scores[None, :, :]
|
||||||
|
|
||||||
# 绘制骨骼
|
# 遍历所有人
|
||||||
for limb_id, limb in enumerate(skeleton):
|
num_instances = keypoints.shape[0]
|
||||||
joint_a, joint_b = limb
|
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:
|
if joint_a >= len(person_scores) or joint_b >= len(person_scores):
|
||||||
continue
|
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
|
for kpt_id, (x, y) in enumerate(person_kpts):
|
||||||
if joint_similarities is not None:
|
if person_scores[kpt_id] < kpt_thr:
|
||||||
# 检查这个连接是否有对应的关节角度
|
continue
|
||||||
if (joint_a, joint_b) in joint_to_angle_mapping:
|
cv2.circle(img, (int(x), int(y)), 3, (255, 0, 255), -1)
|
||||||
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
|
return img
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user