Update workouts_monitoring solution (#16706)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-10-05 18:08:37 +05:00 committed by GitHub
parent c17ddcdf70
commit 73e6861d95
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 162 additions and 245 deletions

View file

@ -697,14 +697,13 @@ class Annotator:
angle = 360 - angle
return angle
def draw_specific_points(self, keypoints, indices=None, shape=(640, 640), radius=2, conf_thres=0.25):
def draw_specific_points(self, keypoints, indices=None, radius=2, conf_thres=0.25):
"""
Draw specific keypoints for gym steps counting.
Args:
keypoints (list): Keypoints data to be plotted.
indices (list, optional): Keypoint indices to be plotted. Defaults to [2, 5, 7].
shape (tuple, optional): Image size for model inference. Defaults to (640, 640).
radius (int, optional): Keypoint radius. Defaults to 2.
conf_thres (float, optional): Confidence threshold for keypoints. Defaults to 0.25.
@ -715,90 +714,71 @@ class Annotator:
Keypoint format: [x, y] or [x, y, confidence].
Modifies self.im in-place.
"""
if indices is None:
indices = [2, 5, 7]
for i, k in enumerate(keypoints):
if i in indices:
x_coord, y_coord = k[0], k[1]
if x_coord % shape[1] != 0 and y_coord % shape[0] != 0:
if len(k) == 3:
conf = k[2]
if conf < conf_thres:
continue
cv2.circle(self.im, (int(x_coord), int(y_coord)), radius, (0, 255, 0), -1, lineType=cv2.LINE_AA)
indices = indices or [2, 5, 7]
points = [(int(k[0]), int(k[1])) for i, k in enumerate(keypoints) if i in indices and k[2] >= conf_thres]
# Draw lines between consecutive points
for start, end in zip(points[:-1], points[1:]):
cv2.line(self.im, start, end, (0, 255, 0), 2, lineType=cv2.LINE_AA)
# Draw circles for keypoints
for pt in points:
cv2.circle(self.im, pt, radius, (0, 0, 255), -1, lineType=cv2.LINE_AA)
return self.im
def plot_workout_information(self, display_text, position, color=(104, 31, 17), txt_color=(255, 255, 255)):
"""
Draw text with a background on the image.
Args:
display_text (str): The text to be displayed.
position (tuple): Coordinates (x, y) on the image where the text will be placed.
color (tuple, optional): Text background color
txt_color (tuple, optional): Text foreground color
"""
(text_width, text_height), _ = cv2.getTextSize(display_text, 0, self.sf, self.tf)
# Draw background rectangle
cv2.rectangle(
self.im,
(position[0], position[1] - text_height - 5),
(position[0] + text_width + 10, position[1] - text_height - 5 + text_height + 10 + self.tf),
color,
-1,
)
# Draw text
cv2.putText(self.im, display_text, position, 0, self.sf, txt_color, self.tf)
return text_height
def plot_angle_and_count_and_stage(
self, angle_text, count_text, stage_text, center_kpt, color=(104, 31, 17), txt_color=(255, 255, 255)
):
"""
Plot the pose angle, count value and step stage.
Plot the pose angle, count value, and step stage.
Args:
angle_text (str): angle value for workout monitoring
count_text (str): counts value for workout monitoring
stage_text (str): stage decision for workout monitoring
center_kpt (list): centroid pose index for workout monitoring
color (tuple): text background color for workout monitoring
txt_color (tuple): text foreground color for workout monitoring
angle_text (str): Angle value for workout monitoring
count_text (str): Counts value for workout monitoring
stage_text (str): Stage decision for workout monitoring
center_kpt (list): Centroid pose index for workout monitoring
color (tuple, optional): Text background color
txt_color (tuple, optional): Text foreground color
"""
angle_text, count_text, stage_text = (f" {angle_text:.2f}", f"Steps : {count_text}", f" {stage_text}")
# Format text
angle_text, count_text, stage_text = f" {angle_text:.2f}", f"Steps : {count_text}", f" {stage_text}"
# Draw angle
(angle_text_width, angle_text_height), _ = cv2.getTextSize(angle_text, 0, self.sf, self.tf)
angle_text_position = (int(center_kpt[0]), int(center_kpt[1]))
angle_background_position = (angle_text_position[0], angle_text_position[1] - angle_text_height - 5)
angle_background_size = (angle_text_width + 2 * 5, angle_text_height + 2 * 5 + (self.tf * 2))
cv2.rectangle(
self.im,
angle_background_position,
(
angle_background_position[0] + angle_background_size[0],
angle_background_position[1] + angle_background_size[1],
),
color,
-1,
# Draw angle, count and stage text
angle_height = self.plot_workout_information(
angle_text, (int(center_kpt[0]), int(center_kpt[1])), color, txt_color
)
cv2.putText(self.im, angle_text, angle_text_position, 0, self.sf, txt_color, self.tf)
# Draw Counts
(count_text_width, count_text_height), _ = cv2.getTextSize(count_text, 0, self.sf, self.tf)
count_text_position = (angle_text_position[0], angle_text_position[1] + angle_text_height + 20)
count_background_position = (
angle_background_position[0],
angle_background_position[1] + angle_background_size[1] + 5,
count_height = self.plot_workout_information(
count_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + 20), color, txt_color
)
count_background_size = (count_text_width + 10, count_text_height + 10 + self.tf)
cv2.rectangle(
self.im,
count_background_position,
(
count_background_position[0] + count_background_size[0],
count_background_position[1] + count_background_size[1],
),
color,
-1,
self.plot_workout_information(
stage_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + count_height + 40), color, txt_color
)
cv2.putText(self.im, count_text, count_text_position, 0, self.sf, txt_color, self.tf)
# Draw Stage
(stage_text_width, stage_text_height), _ = cv2.getTextSize(stage_text, 0, self.sf, self.tf)
stage_text_position = (int(center_kpt[0]), int(center_kpt[1]) + angle_text_height + count_text_height + 40)
stage_background_position = (stage_text_position[0], stage_text_position[1] - stage_text_height - 5)
stage_background_size = (stage_text_width + 10, stage_text_height + 10)
cv2.rectangle(
self.im,
stage_background_position,
(
stage_background_position[0] + stage_background_size[0],
stage_background_position[1] + stage_background_size[1],
),
color,
-1,
)
cv2.putText(self.im, stage_text, stage_text_position, 0, self.sf, txt_color, self.tf)
def seg_bbox(self, mask, mask_color=(255, 0, 255), label=None, txt_color=(255, 255, 255)):
"""