ultralytics 8.2.17__py3-none-any.whl → 8.2.18__py3-none-any.whl

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.2.17"
+__version__ = "8.2.18"
 
 from ultralytics.data.explorer.explorer import Explorer
 from ultralytics.models import RTDETR, SAM, YOLO, YOLOWorld

ultralytics/solutions/__init__.py

@@ -1 +1,19 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from .ai_gym import AIGym
+from .distance_calculation import DistanceCalculation
+from .heatmap import Heatmap
+from .object_counter import ObjectCounter
+from .parking_management import ParkingManagement
+from .queue_management import QueueManager
+from .speed_estimation import SpeedEstimator
+
+__all__ = (
+    "AIGym",
+    "DistanceCalculation",
+    "Heatmap",
+    "ObjectCounter",
+    "ParkingManagement",
+    "QueueManager",
+    "SpeedEstimator",
+)

ultralytics/solutions/ai_gym.py

@@ -9,79 +9,75 @@ 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."""
+    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 = None
+        self.tf = line_thickness
 
         # Keypoints and count information
         self.keypoints = None
-        self.poseup_angle = None
-        self.posedown_angle = 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 = "pushup"
-        self.kpts_to_check = None
+        self.pose_type = pose_type
+        self.kpts_to_check = kpts_to_check
 
         # Visual Information
-        self.view_img = False
+        self.view_img = view_img
         self.annotator = None
 
-        # Check if environment support imshow
+        # Check if environment supports 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
+            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=2)
 
         for ind, k in enumerate(reversed(self.keypoints)):
-            if self.pose_type in {"pushup", "pullup"}:
+            # Estimate angle and draw specific points based on pose type
+            if self.pose_type in {"pushup", "pullup", "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(),
@@ -89,55 +85,32 @@ class AIGym:
                 )
                 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,
-                )
+                # 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 == "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
 
-            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,
                 )
 
+            # 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"):
@@ -147,4 +120,5 @@ class AIGym:
 
 
 if __name__ == "__main__":
-    AIGym()
+    kpts_to_check = [0, 1, 2]  # example keypoints
+    aigym = AIGym(kpts_to_check)

ultralytics/solutions/distance_calculation.py

@@ -9,94 +9,78 @@ from ultralytics.utils.plotting import Annotator, colors
 
 
 class DistanceCalculation:
-    """A class to calculate distance between two objects in real-time video stream based on their tracks."""
+    """A class to calculate distance between two objects in a real-time video stream based on their tracks."""
 
-    def __init__(self):
-        """Initializes the distance calculation class with default values for Visual, Image, track and distance
-        parameters.
+    def __init__(
+        self,
+        names,
+        pixels_per_meter=10,
+        view_img=False,
+        line_thickness=2,
+        line_color=(255, 255, 0),
+        centroid_color=(255, 0, 255),
+    ):
         """
+        Initializes the DistanceCalculation class with the given parameters.
 
-        # Visual & im0 information
+        Args:
+            names (dict): Dictionary mapping class indices to class names.
+            pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
+            view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
+            line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
+            line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
+            centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
+        """
+        # Visual & image information
         self.im0 = None
         self.annotator = None
-        self.view_img = False
-        self.line_color = (255, 255, 0)
-        self.centroid_color = (255, 0, 255)
+        self.view_img = view_img
+        self.line_color = line_color
+        self.centroid_color = centroid_color
 
-        # Predict/track information
+        # Prediction & tracking information
         self.clss = None
-        self.names = None
+        self.names = names
         self.boxes = None
-        self.line_thickness = 2
+        self.line_thickness = line_thickness
         self.trk_ids = None
 
         # Distance calculation information
         self.centroids = []
-        self.pixel_per_meter = 10
+        self.pixel_per_meter = pixels_per_meter
 
-        # Mouse event
+        # Mouse event information
         self.left_mouse_count = 0
         self.selected_boxes = {}
 
-        # Check if environment support imshow
+        # Check if environment supports imshow
         self.env_check = check_imshow(warn=True)
 
-    def set_args(
-        self,
-        names,
-        pixels_per_meter=10,
-        view_img=False,
-        line_thickness=2,
-        line_color=(255, 255, 0),
-        centroid_color=(255, 0, 255),
-    ):
-        """
-        Configures the distance calculation and display parameters.
-
-        Args:
-            names (dict): object detection classes names
-            pixels_per_meter (int): Number of pixels in meter
-            view_img (bool): Flag indicating frame display
-            line_thickness (int): Line thickness for bounding boxes.
-            line_color (RGB): color of centroids line
-            centroid_color (RGB): colors of bbox centroids
-        """
-        self.names = names
-        self.pixel_per_meter = pixels_per_meter
-        self.view_img = view_img
-        self.line_thickness = line_thickness
-        self.line_color = line_color
-        self.centroid_color = centroid_color
-
     def mouse_event_for_distance(self, event, x, y, flags, param):
         """
-        This function is designed to move region with mouse events in a real-time video stream.
+        Handles mouse events to select regions in a real-time video stream.
 
         Args:
-            event (int): The type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
-            x (int): The x-coordinate of the mouse pointer.
-            y (int): The y-coordinate of the mouse pointer.
-            flags (int): Any flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY,
-                cv2.EVENT_FLAG_SHIFTKEY, etc.).
-            param (dict): Additional parameters you may want to pass to the function.
+            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
+            x (int): X-coordinate of the mouse pointer.
+            y (int): Y-coordinate of the mouse pointer.
+            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).
+            param (dict): Additional parameters passed to the function.
         """
-        global selected_boxes
-        global left_mouse_count
         if event == cv2.EVENT_LBUTTONDOWN:
             self.left_mouse_count += 1
             if self.left_mouse_count <= 2:
                 for box, track_id in zip(self.boxes, self.trk_ids):
                     if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
-                        self.selected_boxes[track_id] = []
                         self.selected_boxes[track_id] = box
 
-        if event == cv2.EVENT_RBUTTONDOWN:
+        elif event == cv2.EVENT_RBUTTONDOWN:
             self.selected_boxes = {}
             self.left_mouse_count = 0
 
     def extract_tracks(self, tracks):
         """
-        Extracts results from the provided data.
+        Extracts tracking results from the provided data.
 
         Args:
             tracks (list): List of tracks obtained from the object tracking process.
@@ -105,55 +89,65 @@ class DistanceCalculation:
         self.clss = tracks[0].boxes.cls.cpu().tolist()
         self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
 
-    def calculate_centroid(self, box):
+    @staticmethod
+    def calculate_centroid(box):
         """
-        Calculate the centroid of bounding box.
+        Calculates the centroid of a bounding box.
 
         Args:
-            box (list): Bounding box data
+            box (list): Bounding box coordinates [x1, y1, x2, y2].
+
+        Returns:
+            (tuple): Centroid coordinates (x, y).
         """
         return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
 
     def calculate_distance(self, centroid1, centroid2):
         """
-        Calculate distance between two centroids.
+        Calculates the distance between two centroids.
 
         Args:
-            centroid1 (point): First bounding box data
-            centroid2 (point): Second bounding box data
+            centroid1 (tuple): Coordinates of the first centroid (x, y).
+            centroid2 (tuple): Coordinates of the second centroid (x, y).
+
+        Returns:
+            (tuple): Distance in meters and millimeters.
         """
         pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
-        return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000
+        distance_m = pixel_distance / self.pixel_per_meter
+        distance_mm = distance_m * 1000
+        return distance_m, distance_mm
 
     def start_process(self, im0, tracks):
         """
-        Calculate distance between two bounding boxes based on tracking data.
+        Processes the video frame and calculates the distance between two bounding boxes.
 
         Args:
-            im0 (nd array): Image
+            im0 (ndarray): The image frame.
             tracks (list): List of tracks obtained from the object tracking process.
+
+        Returns:
+            (ndarray): The processed image frame.
         """
         self.im0 = im0
         if tracks[0].boxes.id is None:
             if self.view_img:
                 self.display_frames()
-            return
-        self.extract_tracks(tracks)
+            return im0
 
-        self.annotator = Annotator(self.im0, line_width=2)
+        self.extract_tracks(tracks)
+        self.annotator = Annotator(self.im0, line_width=self.line_thickness)
 
         for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
             self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
 
             if len(self.selected_boxes) == 2:
-                for trk_id, _ in self.selected_boxes.items():
+                for trk_id in self.selected_boxes.keys():
                     if trk_id == track_id:
                         self.selected_boxes[track_id] = box
 
         if len(self.selected_boxes) == 2:
-            for trk_id, box in self.selected_boxes.items():
-                centroid = self.calculate_centroid(self.selected_boxes[trk_id])
-                self.centroids.append(centroid)
+            self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes]
 
             distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
             self.annotator.plot_distance_and_line(
@@ -168,7 +162,7 @@ class DistanceCalculation:
         return im0
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame with annotations."""
         cv2.namedWindow("Ultralytics Distance Estimation")
         cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
         cv2.imshow("Ultralytics Distance Estimation", self.im0)
@@ -178,4 +172,5 @@ class DistanceCalculation:
 
 
 if __name__ == "__main__":
-    DistanceCalculation()
+    names = {0: "person", 1: "car"}  # example class names
+    distance_calculation = DistanceCalculation(names)