HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Update workouts_monitoring solution (#16706)

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

4
docs/en/guides/object-counting.md
View file

@ -286,7 +286,7 @@ def count_objects_in_region(video_path, output_video_path, model_path):
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
im0 = counter.start_counting(im0) im0 = counter.count(im0)
video_writer.write(im0) video_writer.write(im0)
cap.release() cap.release()
@ -334,7 +334,7 @@ def count_specific_classes(video_path, output_video_path, model_path, classes_to
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
im0 = counter.start_counting(im0) im0 = counter.count(im0)
video_writer.write(im0) video_writer.write(im0)
cap.release() cap.release()

89
docs/en/guides/workouts-monitoring.md
View file

@ -41,18 +41,16 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
```python ```python
import cv2 import cv2
from ultralytics import YOLO, solutions from ultralytics import solutions
model = YOLO("yolo11n-pose.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4") cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
gym_object = solutions.AIGym( gym = solutions.AIGym(
line_thickness=2, model="yolo11n-pose.pt",
view_img=True, show=True,
pose_type="pushup", kpts=[6, 8, 10],
kpts_to_check=[6, 8, 10],
) )
while cap.isOpened(): while cap.isOpened():
@ -60,9 +58,7 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
results = model.track(im0, verbose=False) # Tracking recommended im0 = gym.monitor(im0)
# results = model.predict(im0) # Prediction also supported
im0 = gym_object.start_counting(im0, results)
cv2.destroyAllWindows() cv2.destroyAllWindows()
``` ```
@ -72,20 +68,17 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
```python ```python
import cv2 import cv2
from ultralytics import YOLO, solutions from ultralytics import solutions
model = YOLO("yolo11n-pose.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4") cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
video_writer = cv2.VideoWriter("workouts.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) video_writer = cv2.VideoWriter("workouts.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
gym_object = solutions.AIGym( gym = solutions.AIGym(
line_thickness=2, show=True,
view_img=True, kpts=[6, 8, 10],
pose_type="pushup",
kpts_to_check=[6, 8, 10],
) )
while cap.isOpened(): while cap.isOpened():
@ -93,33 +86,26 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
results = model.track(im0, verbose=False) # Tracking recommended im0 = gym.monitor(im0)
# results = model.predict(im0) # Prediction also supported
im0 = gym_object.start_counting(im0, results)
video_writer.write(im0) video_writer.write(im0)
cv2.destroyAllWindows() cv2.destroyAllWindows()
video_writer.release() video_writer.release()
``` ```
???+ tip "Support"
"pushup", "pullup" and "abworkout" supported
### KeyPoints Map ### KeyPoints Map
![keyPoints Order Ultralytics YOLO11 Pose](https://github.com/ultralytics/docs/releases/download/0/keypoints-order-ultralytics-yolov8-pose.avif) ![keyPoints Order Ultralytics YOLO11 Pose](https://github.com/ultralytics/docs/releases/download/0/keypoints-order-ultralytics-yolov8-pose.avif)
### Arguments `AIGym` ### Arguments `AIGym`
| Name | Type | Default | Description | | Name | Type | Default | Description |
| ----------------- | ------- | -------- | -------------------------------------------------------------------------------------- | | ------------ | ------- | ------- | -------------------------------------------------------------------------------------- |
| `kpts_to_check` | `list` | `None` | List of three keypoints index, for counting specific workout, followed by keypoint Map | | `kpts` | `list` | `None` | List of three keypoints index, for counting specific workout, followed by keypoint Map |
| `line_thickness` | `int` | `2` | Thickness of the lines drawn. | | `line_width` | `int` | `2` | Thickness of the lines drawn. |
| `view_img` | `bool` | `False` | Flag to display the image. | | `show` | `bool` | `False` | Flag to display the image. |
| `pose_up_angle` | `float` | `145.0` | Angle threshold for the 'up' pose. | | `up_angle` | `float` | `145.0` | Angle threshold for the 'up' pose. |
| `pose_down_angle` | `float` | `90.0` | Angle threshold for the 'down' pose. | | `down_angle` | `float` | `90.0` | Angle threshold for the 'down' pose. |
| `pose_type` | `str` | `pullup` | Type of pose to detect (`'pullup`', `pushup`, `abworkout`, `squat`). |
### Arguments `model.predict` ### Arguments `model.predict`
@ -138,18 +124,16 @@ To monitor your workouts using Ultralytics YOLO11, you can utilize the pose esti
```python ```python
import cv2 import cv2
from ultralytics import YOLO, solutions from ultralytics import solutions
model = YOLO("yolo11n-pose.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4") cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
gym_object = solutions.AIGym( gym = solutions.AIGym(
line_thickness=2, line_width=2,
view_img=True, show=True,
pose_type="pushup", kpts=[6, 8, 10],
kpts_to_check=[6, 8, 10],
) )
while cap.isOpened(): while cap.isOpened():
@ -157,8 +141,7 @@ while cap.isOpened():
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
results = model.track(im0, verbose=False) im0 = gym.monitor(im0)
im0 = gym_object.start_counting(im0, results)
cv2.destroyAllWindows() cv2.destroyAllWindows()
``` ```
@ -188,11 +171,10 @@ Yes, Ultralytics YOLO11 can be adapted for custom workout routines. The `AIGym`
```python ```python
from ultralytics import solutions from ultralytics import solutions
gym_object = solutions.AIGym( gym = solutions.AIGym(
line_thickness=2, line_width=2,
view_img=True, show=True,
pose_type="squat", kpts=[6, 8, 10],
kpts_to_check=[6, 8, 10],
) )
``` ```
@ -205,20 +187,18 @@ To save the workout monitoring output, you can modify the code to include a vide
```python ```python
import cv2 import cv2
from ultralytics import YOLO, solutions from ultralytics import solutions
model = YOLO("yolo11n-pose.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4") cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
video_writer = cv2.VideoWriter("workouts.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) video_writer = cv2.VideoWriter("workouts.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
gym_object = solutions.AIGym( gym = solutions.AIGym(
line_thickness=2, line_width=2,
view_img=True, show=True,
pose_type="pushup", kpts=[6, 8, 10],
kpts_to_check=[6, 8, 10],
) )
while cap.isOpened(): while cap.isOpened():
@ -226,8 +206,7 @@ while cap.isOpened():
if not success: if not success:
print("Video frame is empty or video processing has been successfully completed.") print("Video frame is empty or video processing has been successfully completed.")
break break
results = model.track(im0, verbose=False) im0 = gym.monitor(im0)
im0 = gym_object.start_counting(im0, results)
video_writer.write(im0) video_writer.write(im0)
cv2.destroyAllWindows() cv2.destroyAllWindows()

6
tests/test_solutions.py
View file

@ -41,16 +41,14 @@ def test_major_solutions():
def test_aigym(): def test_aigym():
"""Test the workouts monitoring solution.""" """Test the workouts monitoring solution."""
safe_download(url=WORKOUTS_SOLUTION_DEMO) safe_download(url=WORKOUTS_SOLUTION_DEMO)
model = YOLO("yolo11n-pose.pt")
cap = cv2.VideoCapture("solution_ci_pose_demo.mp4") cap = cv2.VideoCapture("solution_ci_pose_demo.mp4")
assert cap.isOpened(), "Error reading video file" assert cap.isOpened(), "Error reading video file"
gym_object = solutions.AIGym(line_thickness=2, pose_type="squat", kpts_to_check=[5, 11, 13]) gym = solutions.AIGym(line_width=2, kpts=[5, 11, 13])
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if not success: if not success:
break break
results = model.track(im0, verbose=False) _ = gym.monitor(im0)
_ = gym_object.start_counting(im0, results)
cap.release() cap.release()
cv2.destroyAllWindows() cv2.destroyAllWindows()

4
ultralytics/cfg/solutions/default.yaml
View file

@ -10,3 +10,7 @@ show: True # Flag to control whether to display output image or not
show_in: True # Flag to display objects moving *into* the defined region show_in: True # Flag to display objects moving *into* the defined region
show_out: True # Flag to display objects moving *out of* the defined region show_out: True # Flag to display objects moving *out of* the defined region
classes: # To count specific classes classes: # To count specific classes
up_angle: 145.0 # workouts up_angle for counts, 145.0 is default value
down_angle: 90 # workouts down_angle for counts, 90 is default value
kpts: [6, 8, 10] # keypoints for workouts monitoring

166
ultralytics/solutions/ai_gym.py
View file

@ -1,127 +1,79 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license # Ultralytics YOLO 🚀, AGPL-3.0 license
import cv2 from ultralytics.solutions.solutions import BaseSolution # Import a parent class
from ultralytics.utils.checks import check_imshow
from ultralytics.utils.plotting import Annotator from ultralytics.utils.plotting import Annotator
class AIGym: class AIGym(BaseSolution):
"""A class to manage the gym steps of people in a real-time video stream based on their poses.""" """A class to manage the gym steps of people in a real-time video stream based on their poses."""
def __init__( def __init__(self, **kwargs):
self, """Initialization function for AiGYM class, a child class of BaseSolution class, can be used for workouts
kpts_to_check, monitoring.
line_thickness=2,
view_img=False,
pose_up_angle=145.0,
pose_down_angle=90.0,
pose_type="pullup",
):
""" """
Initializes the AIGym class with the specified parameters. # Check if the model name ends with '-pose'
if "model" in kwargs and "-pose" not in kwargs["model"]:
kwargs["model"] = "yolo11n-pose.pt"
elif "model" not in kwargs:
kwargs["model"] = "yolo11n-pose.pt"
super().__init__(**kwargs)
self.count = [] # List for counts, necessary where there are multiple objects in frame
self.angle = [] # List for angle, necessary where there are multiple objects in frame
self.stage = [] # List for stage, necessary where there are multiple objects in frame
# Extract details from CFG single time for usage later
self.initial_stage = None
self.up_angle = float(self.CFG["up_angle"]) # Pose up predefined angle to consider up pose
self.down_angle = float(self.CFG["down_angle"]) # Pose down predefined angle to consider down pose
self.kpts = self.CFG["kpts"] # User selected kpts of workouts storage for further usage
self.lw = self.CFG["line_width"] # Store line_width for usage
def monitor(self, im0):
"""
Monitor the workouts using Ultralytics YOLOv8 Pose Model: https://docs.ultralytics.com/tasks/pose/.
Args: Args:
kpts_to_check (list): Indices of keypoints to check. im0 (ndarray): The input image that will be used for processing
line_thickness (int, optional): Thickness of the lines drawn. Defaults to 2. Returns
view_img (bool, optional): Flag to display the image. Defaults to False. im0 (ndarray): The processed image for more usage
pose_up_angle (float, optional): Angle threshold for the 'up' pose. Defaults to 145.0.
pose_down_angle (float, optional): Angle threshold for the 'down' pose. Defaults to 90.0.
pose_type (str, optional): Type of pose to detect ('pullup', 'pushup', 'abworkout'). Defaults to "pullup".
""" """
# Image and line thickness # Extract tracks
self.im0 = None tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"])[0]
self.tf = line_thickness
# Keypoints and count information if tracks.boxes.id is not None:
self.keypoints = None # Extract and check keypoints
self.poseup_angle = pose_up_angle if len(tracks) > len(self.count):
self.posedown_angle = pose_down_angle new_human = len(tracks) - len(self.count)
self.threshold = 0.001 self.angle += [0] * new_human
self.count += [0] * new_human
self.stage += ["-"] * new_human
# Store stage, count and angle information # Initialize annotator
self.angle = None self.annotator = Annotator(im0, line_width=self.lw)
self.count = None
self.stage = None
self.pose_type = pose_type
self.kpts_to_check = kpts_to_check
# Visual Information # Enumerate over keypoints
self.view_img = view_img for ind, k in enumerate(reversed(tracks.keypoints.data)):
self.annotator = None # Get keypoints and estimate the angle
kpts = [k[int(self.kpts[i])].cpu() for i in range(3)]
self.angle[ind] = self.annotator.estimate_pose_angle(*kpts)
im0 = self.annotator.draw_specific_points(k, self.kpts, radius=self.lw * 3)
# Check if environment supports imshow # Determine stage and count logic based on angle thresholds
self.env_check = check_imshow(warn=True) if self.angle[ind] < self.down_angle:
self.count = [] if self.stage[ind] == "up":
self.angle = []
self.stage = []
def start_counting(self, im0, results):
"""
Function used to count the gym steps.
Args:
im0 (ndarray): Current frame from the video stream.
results (list): Pose estimation data.
"""
self.im0 = im0
if not len(results[0]):
return self.im0
if len(results[0]) > len(self.count):
new_human = len(results[0]) - len(self.count)
self.count += [0] * new_human
self.angle += [0] * new_human
self.stage += ["-"] * new_human
self.keypoints = results[0].keypoints.data
self.annotator = Annotator(im0, line_width=self.tf)
for ind, k in enumerate(reversed(self.keypoints)):
# Estimate angle and draw specific points based on pose type
if self.pose_type in {"pushup", "pullup", "abworkout", "squat"}:
self.angle[ind] = self.annotator.estimate_pose_angle(
k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu(),
)
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
# Check and update pose stages and counts based on angle
if self.pose_type in {"abworkout", "pullup"}:
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
elif self.pose_type in {"pushup", "squat"}:
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "up"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
self.stage[ind] = "down"
self.count[ind] += 1 self.count[ind] += 1
self.stage[ind] = "down"
elif self.angle[ind] > self.up_angle:
self.stage[ind] = "up"
# Display angle, count, and stage text
self.annotator.plot_angle_and_count_and_stage( self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind], angle_text=self.angle[ind], # angle text for display
count_text=self.count[ind], count_text=self.count[ind], # count text for workouts
stage_text=self.stage[ind], stage_text=self.stage[ind], # stage position text
center_kpt=k[int(self.kpts_to_check[1])], center_kpt=k[int(self.kpts[1])], # center keypoint for display
) )
# Draw keypoints self.display_output(im0) # Display output image, if environment support display
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True) return im0 # return an image for writing or further usage
# Display the image if environment supports it and view_img is True
if self.env_check and self.view_img:
cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
return self.im0
if __name__ == "__main__":
kpts_to_check = [0, 1, 2] # example keypoints
aigym = AIGym(kpts_to_check)

12
ultralytics/solutions/solutions.py
View file

@ -4,11 +4,13 @@ from collections import defaultdict
from pathlib import Path from pathlib import Path
import cv2 import cv2
from shapely.geometry import LineString, Polygon
from ultralytics import YOLO from ultralytics import YOLO
from ultralytics.utils import yaml_load from ultralytics.utils import LOGGER, yaml_load
from ultralytics.utils.checks import check_imshow from ultralytics.utils.checks import check_imshow, check_requirements
check_requirements("shapely>=2.0.0")
from shapely.geometry import LineString, Polygon
DEFAULT_SOL_CFG_PATH = Path(__file__).resolve().parents[1] / "cfg/solutions/default.yaml" DEFAULT_SOL_CFG_PATH = Path(__file__).resolve().parents[1] / "cfg/solutions/default.yaml"
@ -25,7 +27,7 @@ class BaseSolution:
# Load config and update with args # Load config and update with args
self.CFG = yaml_load(DEFAULT_SOL_CFG_PATH) self.CFG = yaml_load(DEFAULT_SOL_CFG_PATH)
self.CFG.update(kwargs) self.CFG.update(kwargs)
print("Ultralytics Solutions: ✅", self.CFG) LOGGER.info(f"Ultralytics Solutions: ✅ {self.CFG}")
self.region = self.CFG["region"] # Store region data for other classes usage self.region = self.CFG["region"] # Store region data for other classes usage
self.line_width = self.CFG["line_width"] # Store line_width for usage self.line_width = self.CFG["line_width"] # Store line_width for usage
@ -54,6 +56,8 @@ class BaseSolution:
self.boxes = self.track_data.xyxy.cpu() self.boxes = self.track_data.xyxy.cpu()
self.clss = self.track_data.cls.cpu().tolist() self.clss = self.track_data.cls.cpu().tolist()
self.track_ids = self.track_data.id.int().cpu().tolist() self.track_ids = self.track_data.id.int().cpu().tolist()
else:
LOGGER.warning("WARNING ⚠️ tracks none, no keypoints will be considered.")
def store_tracking_history(self, track_id, box): def store_tracking_history(self, track_id, box):
""" """

126
ultralytics/utils/plotting.py
View file

@ -697,14 +697,13 @@ class Annotator:
angle = 360 - angle angle = 360 - angle
return 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. Draw specific keypoints for gym steps counting.
Args: Args:
keypoints (list): Keypoints data to be plotted. keypoints (list): Keypoints data to be plotted.
indices (list, optional): Keypoint indices to be plotted. Defaults to [2, 5, 7]. 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. radius (int, optional): Keypoint radius. Defaults to 2.
conf_thres (float, optional): Confidence threshold for keypoints. Defaults to 0.25. 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]. Keypoint format: [x, y] or [x, y, confidence].
Modifies self.im in-place. Modifies self.im in-place.
""" """
if indices is None: indices = indices or [2, 5, 7]
indices = [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]
for i, k in enumerate(keypoints):
if i in indices: # Draw lines between consecutive points
x_coord, y_coord = k[0], k[1] for start, end in zip(points[:-1], points[1:]):
if x_coord % shape[1] != 0 and y_coord % shape[0] != 0: cv2.line(self.im, start, end, (0, 255, 0), 2, lineType=cv2.LINE_AA)
if len(k) == 3:
conf = k[2] # Draw circles for keypoints
if conf < conf_thres: for pt in points:
continue cv2.circle(self.im, pt, radius, (0, 0, 255), -1, lineType=cv2.LINE_AA)
cv2.circle(self.im, (int(x_coord), int(y_coord)), radius, (0, 255, 0), -1, lineType=cv2.LINE_AA)
return self.im 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( 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) 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: Args:
angle_text (str): angle value for workout monitoring angle_text (str): Angle value for workout monitoring
count_text (str): counts value for workout monitoring count_text (str): Counts value for workout monitoring
stage_text (str): stage decision for workout monitoring stage_text (str): Stage decision for workout monitoring
center_kpt (list): centroid pose index for workout monitoring center_kpt (list): Centroid pose index for workout monitoring
color (tuple): text background color for workout monitoring color (tuple, optional): Text background color
txt_color (tuple): text foreground color for workout monitoring 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 # Draw angle, count and stage text
(angle_text_width, angle_text_height), _ = cv2.getTextSize(angle_text, 0, self.sf, self.tf) angle_height = self.plot_workout_information(
angle_text_position = (int(center_kpt[0]), int(center_kpt[1])) angle_text, (int(center_kpt[0]), int(center_kpt[1])), color, txt_color
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,
) )
cv2.putText(self.im, angle_text, angle_text_position, 0, self.sf, txt_color, self.tf) count_height = self.plot_workout_information(
count_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + 20), color, txt_color
# 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_background_size = (count_text_width + 10, count_text_height + 10 + self.tf) self.plot_workout_information(
stage_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + count_height + 40), color, txt_color
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,
) )
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)): def seg_bbox(self, mask, mask_color=(255, 0, 255), label=None, txt_color=(255, 255, 255)):
""" """