์ฝ˜ํ…์ธ ๋กœ ๊ฑด๋„ˆ๋›ฐ๊ธฐ

์ฐธ์กฐ ultralytics/solutions/ai_gym.py

์ฐธ๊ณ 

์ด ํŒŒ์ผ์€ https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/solutions/ai_gym .py์—์„œ ํ™•์ธํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค. ๋ฌธ์ œ๋ฅผ ๋ฐœ๊ฒฌํ•˜๋ฉด ํ’€ ๋ฆฌํ€˜์ŠคํŠธ (๐Ÿ› ๏ธ) ๋ฅผ ํ†ตํ•ด ๋ฌธ์ œ๋ฅผ ํ•ด๊ฒฐํ•˜๋„๋ก ๋„์™€์ฃผ์„ธ์š”. ๊ฐ์‚ฌํ•ฉ๋‹ˆ๋‹ค ๐Ÿ™!



ultralytics.solutions.ai_gym.AIGym

ํฌ์ฆˆ๋ฅผ ๊ธฐ๋ฐ˜์œผ๋กœ ์‹ค์‹œ๊ฐ„ ๋น„๋””์˜ค ์ŠคํŠธ๋ฆผ์—์„œ ์‚ฌ๋žŒ๋“ค์˜ ํ—ฌ์Šค์žฅ ์Šคํ…์„ ๊ด€๋ฆฌํ•˜๋Š” ์ˆ˜์—…์ž…๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/solutions/ai_gym.py
class AIGym:
    """A class to manage the gym steps of people in a real-time video stream based on their poses."""

    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)

    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): Current frame count.
        """

        self.im0 = im0

        # Initialize count, angle, and stage lists on the first frame
        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=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

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

์ง€์ •๋œ ํŒŒ๋ผ๋ฏธํ„ฐ๋ฅผ ์‚ฌ์šฉํ•˜์—ฌ AIGym ํด๋ž˜์Šค๋ฅผ ์ดˆ๊ธฐํ™”ํ•ฉ๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
kpts_to_check list

ํ™•์ธํ•  ํ‚คํฌ์ธํŠธ์˜ ์ธ๋ฑ์Šค์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
line_thickness int

๊ทธ๋ ค์ง€๋Š” ์„ ์˜ ๋‘๊ป˜์ž…๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ 2์ž…๋‹ˆ๋‹ค.

2
view_img bool

์ด๋ฏธ์ง€๋ฅผ ํ‘œ์‹œํ•  ํ”Œ๋ž˜๊ทธ๋ฅผ ์„ค์ •ํ•ฉ๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ False์ž…๋‹ˆ๋‹ค.

False
pose_up_angle float

'์œ„๋กœ' ํฌ์ฆˆ์— ๋Œ€ํ•œ ๊ฐ๋„ ์ž„๊ณ„๊ฐ’์ž…๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ 145.0์ž…๋‹ˆ๋‹ค.

145.0
pose_down_angle float

'์•„๋ž˜' ํฌ์ฆˆ์— ๋Œ€ํ•œ ๊ฐ๋„ ์ž„๊ณ„๊ฐ’์ž…๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ 90.0์ž…๋‹ˆ๋‹ค.

90.0
pose_type str

๊ฐ์ง€ํ•  ํฌ์ฆˆ ์œ ํ˜•('ํ’€์—…', 'ํ‘ธ์‹œ์—…', '์• ๋ธŒ์›Œํฌ์•„์›ƒ'). ๊ธฐ๋ณธ๊ฐ’์€ 'ํ’€์—…'์ž…๋‹ˆ๋‹ค.

'pullup'
์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/solutions/ai_gym.py
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)

start_counting(im0, results, frame_count)

ํ—ฌ์Šค์žฅ ๊ฑธ์Œ ์ˆ˜๋ฅผ ๊ณ„์‚ฐํ•˜๋Š” ๋ฐ ์‚ฌ์šฉ๋˜๋Š” ๊ธฐ๋Šฅ์ž…๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
im0 ndarray

๋™์˜์ƒ ์ŠคํŠธ๋ฆผ์˜ ํ˜„์žฌ ํ”„๋ ˆ์ž„์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
results list

ํฌ์ฆˆ ์ถ”์ • ๋ฐ์ดํ„ฐ.

ํ•„์ˆ˜
frame_count int

ํ˜„์žฌ ํ”„๋ ˆ์ž„ ์ˆ˜์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/solutions/ai_gym.py
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): Current frame count.
    """

    self.im0 = im0

    # Initialize count, angle, and stage lists on the first frame
    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=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-06-02
Authors: glenn-jocher (2), Burhan-Q (1), RizwanMunawar (1)