Skip to content

Reference for ultralytics/solutions/ai_gym.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/ai_gym.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.solutions.ai_gym.AIGym 

AIGym(kpts_to_check, line_thickness=2, view_img=False, pose_up_angle=145.0, pose_down_angle=90.0, pose_type='pullup')

A class to manage the gym steps of people in a real-time video stream based on their poses.

Parameters:

Name Type Description Default
kpts_to_check list

Indices of keypoints to check.

 required
line_thickness int

Thickness of the lines drawn. Defaults to 2.

 2
view_img bool

Flag to display the image. Defaults to False.

 False
pose_up_angle float

Angle threshold for the 'up' pose. Defaults to 145.0.

 145.0
pose_down_angle float

Angle threshold for the 'down' pose. Defaults to 90.0.

 90.0
pose_type str

Type of pose to detect ('pullup', 'pushup', 'abworkout'). Defaults to "pullup".

 'pullup'
Source code in ultralytics/solutions/ai_gym.py 
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
def __init__(
    self,
    kpts_to_check,
    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.

    Args:
        kpts_to_check (list): Indices of keypoints to check.
        line_thickness (int, optional): Thickness of the lines drawn. Defaults to 2.
        view_img (bool, optional): Flag to display the image. Defaults to False.
        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
    self.im0 = None
    self.tf = line_thickness

    # Keypoints and count information
    self.keypoints = None
    self.poseup_angle = pose_up_angle
    self.posedown_angle = pose_down_angle
    self.threshold = 0.001

    # Store stage, count and angle information
    self.angle = None
    self.count = None
    self.stage = None
    self.pose_type = pose_type
    self.kpts_to_check = kpts_to_check

    # Visual Information
    self.view_img = view_img
    self.annotator = None

    # Check if environment supports imshow
    self.env_check = check_imshow(warn=True)
    self.count = []
    self.angle = []
    self.stage = []

start_counting 

start_counting(im0, results)

Function used to count the gym steps.

Parameters:

Name Type Description Default
im0 ndarray

Current frame from the video stream.

 required
results list

Pose estimation data.

 required
Source code in ultralytics/solutions/ai_gym.py 
 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
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.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])],
            )

        # Draw keypoints
        self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)

    # 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





Created 2023-12-02, Updated 2024-07-21
Authors: glenn-jocher (3), Burhan-Q (1), RizwanMunawar (1)