Skip to content

Référence pour ultralytics/solutions/ai_gym.py

Note

Ce fichier est disponible à l'adresse https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/solutions/ai_gym .py. Si tu repères un problème, aide à le corriger en contribuant à une Pull Request 🛠️. Merci 🙏 !



ultralytics.solutions.ai_gym.AIGym

Une classe pour gérer les pas de gym des personnes dans un flux vidéo en temps réel en fonction de leurs poses.

Code source dans 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):
        """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)

        num_keypoints = len(results[0])

        # Resize self.angle, self.count, and self.stage if the number of keypoints has changed
        if len(self.angle) != num_keypoints:
            self.angle = [0] * num_keypoints
            self.count = [0] * num_keypoints
            self.stage = ["-" for _ in range(num_keypoints)]

        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

__init__()

Initialise l'AIGym avec les valeurs par défaut des paramètres visuels et d'image.

Code source dans ultralytics/solutions/ai_gym.py
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)

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

Configure les paramètres d'épaisseur de ligne, d'enregistrement de l'image et d'affichage de l'image de l'AIGym.

Paramètres :

Nom Type Description DĂ©faut
kpts_to_check list

3 points clés pour le comptage

requis
line_thickness int

Épaisseur du trait pour les boîtes de délimitation.

2
view_img bool

affiche l'im0

False
pose_up_angle float

Angle pour poser la position vers le haut

145.0
pose_down_angle float

Angle pour fixer la position de la pose vers le bas

90.0
pose_type str

"pushup", "pullup" ou "abworkout".

'pullup'
Code source dans ultralytics/solutions/ai_gym.py
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

start_counting(im0, results, frame_count)

Fonction utilisée pour compter les pas de gymnastique.

Paramètres :

Nom Type Description DĂ©faut
im0 ndarray

Image actuelle du flux vidéo.

requis
results list

Données d'estimation de la pose

requis
frame_count int

enregistrer le nombre d'images en cours

requis
Code source dans 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): 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)

    num_keypoints = len(results[0])

    # Resize self.angle, self.count, and self.stage if the number of keypoints has changed
    if len(self.angle) != num_keypoints:
        self.angle = [0] * num_keypoints
        self.count = [0] * num_keypoints
        self.stage = ["-" for _ in range(num_keypoints)]

    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





Créé le 2023-12-02, Mis à jour le 2023-12-02
Auteurs : RizwanMunawar (1)