xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-07-07 22:04:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Resize angle, count, and stage on keypoint number change (#7598)

Browse Source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mike Gvozdev 2024-01-15 15:04:24 -05:00 committed by GitHub
parent b093dce415
commit 22651d01cf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
ultralytics/solutions/ai_gym.py
View File

@ -78,6 +78,14 @@ class AIGym:
self.keypoints = results[0].keypoints.data
self.annotator = Annotator(im0, line_width=2)
num_keypoints = len(results[0])
# Resize self.angle, self.count, and self.stage if the number of keypoints has changed
if len(self.angle) != num_keypoints:
self.angle = [0] * num_keypoints
self.count = [0] * num_keypoints
self.stage = ["-" for _ in range(num_keypoints)]
for ind, k in enumerate(reversed(self.keypoints)):
if self.pose_type in ["pushup", "pullup"]:
self.angle[ind] = self.annotator.estimate_pose_angle(