File size: 5,696 Bytes
69a5bd9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
# Ultralytics YOLO 🚀, AGPL-3.0 license
import cv2
from ultralytics.utils.checks import check_imshow
from ultralytics.utils.plotting import Annotator
class AIGym:
"""A class to manage the gym steps of people in a real-time video stream based on their poses."""
def __init__(self):
"""Initializes the AIGym with default values for Visual and Image parameters."""
# Image and line thickness
self.im0 = None
self.tf = None
# Keypoints and count information
self.keypoints = None
self.poseup_angle = None
self.posedown_angle = None
self.threshold = 0.001
# Store stage, count and angle information
self.angle = None
self.count = None
self.stage = None
self.pose_type = "pushup"
self.kpts_to_check = None
# Visual Information
self.view_img = False
self.annotator = None
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
self,
kpts_to_check,
line_thickness=2,
view_img=False,
pose_up_angle=145.0,
pose_down_angle=90.0,
pose_type="pullup",
):
"""
Configures the AIGym line_thickness, save image and view image parameters.
Args:
kpts_to_check (list): 3 keypoints for counting
line_thickness (int): Line thickness for bounding boxes.
view_img (bool): display the im0
pose_up_angle (float): Angle to set pose position up
pose_down_angle (float): Angle to set pose position down
pose_type (str): "pushup", "pullup" or "abworkout"
"""
self.kpts_to_check = kpts_to_check
self.tf = line_thickness
self.view_img = view_img
self.poseup_angle = pose_up_angle
self.posedown_angle = pose_down_angle
self.pose_type = pose_type
def start_counting(self, im0, results, frame_count):
"""
Function used to count the gym steps.
Args:
im0 (ndarray): Current frame from the video stream.
results (list): Pose estimation data
frame_count (int): store current frame count
"""
self.im0 = im0
if frame_count == 1:
self.count = [0] * len(results[0])
self.angle = [0] * len(results[0])
self.stage = ["-" for _ in results[0]]
self.keypoints = results[0].keypoints.data
self.annotator = Annotator(im0, line_width=2)
for ind, k in enumerate(reversed(self.keypoints)):
if self.pose_type in ["pushup", "pullup"]:
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)
if self.pose_type == "abworkout":
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)
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
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == "pushup":
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.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == "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
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=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__":
AIGym()
|