انتقل إلى المحتوى

مرجع ل ultralytics/solutions/ai_gym.py

ملاحظه

هذا الملف متاح في https://github.com/ultralytics/ultralytics/ نقطة / الرئيسية /ultralytics/الحلول/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):
        """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__()

تهيئة AIGym مع القيم الافتراضية للمعلمات المرئية والصورة.

شفرة المصدر في ultralytics/solutions/ai_gym.py
12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 2829 3031 3233 34 3536 37
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')

يقوم بتكوين line_thickness AIGym وحفظ الصورة وعرض معلمات الصورة.

البارامترات:

اسم نوع وصف افتراضي
kpts_to_check list

3 نقاط رئيسية للعد

مطلوب
line_thickness int

سمك الخط للمربعات المحيطة.

2
view_img bool

عرض IM0

False
pose_up_angle float

زاوية لضبط موضع الوضع لأعلى

145.0
pose_down_angle float

زاوية لضبط موضع الوضع لأسفل

90.0
pose_type str

"تمرين الضغط" أو "السحب" أو "التمرين"

'pullup'
شفرة المصدر في ultralytics/solutions/ai_gym.py
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 565758596061626364
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)

وظيفة تستخدم لحساب خطوات الصالة الرياضية.

البارامترات:

اسم نوع وصف افتراضي
im0 ndarray

الإطار الحالي من دفق الفيديو.

مطلوب
results list

طرح بيانات التقدير

مطلوب
frame_count int

تخزين عدد الإطارات الحالي

مطلوب
شفرة المصدر في ultralytics/solutions/ai_gym.py
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 130131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146147148149 150151 152153 154
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





تم إنشاء 2023-12-02, اخر تحديث 2023-12-02
المؤلفون: رضوان منور (1)