ultralytics 8.3.14__py3-none-any.whl → 8.3.16__py3-none-any.whl

tests/test_solutions.py

@@ -17,10 +17,15 @@ def test_major_solutions():
     cap = cv2.VideoCapture("solutions_ci_demo.mp4")
     assert cap.isOpened(), "Error reading video file"
     region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
-    counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False)
-    heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False)
-    speed = solutions.SpeedEstimator(region=region_points, model="yolo11n.pt", show=False)
-    queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False)
+    counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False)  # Test object counter
+    heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False)  # Test heatmaps
+    speed = solutions.SpeedEstimator(region=region_points, model="yolo11n.pt", show=False)  # Test queue manager
+    queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False)  # Test speed estimation
+    line_analytics = solutions.Analytics(analytics_type="line", model="yolo11n.pt", show=False)  # line analytics
+    pie_analytics = solutions.Analytics(analytics_type="pie", model="yolo11n.pt", show=False)  # line analytics
+    bar_analytics = solutions.Analytics(analytics_type="bar", model="yolo11n.pt", show=False)  # line analytics
+    area_analytics = solutions.Analytics(analytics_type="area", model="yolo11n.pt", show=False)  # line analytics
+    frame_count = 0  # Required for analytics
     while cap.isOpened():
         success, im0 = cap.read()
         if not success:
@@ -30,24 +35,23 @@ def test_major_solutions():
         _ = heatmap.generate_heatmap(original_im0.copy())
         _ = speed.estimate_speed(original_im0.copy())
         _ = queue.process_queue(original_im0.copy())
+        _ = line_analytics.process_data(original_im0.copy(), frame_count)
+        _ = pie_analytics.process_data(original_im0.copy(), frame_count)
+        _ = bar_analytics.process_data(original_im0.copy(), frame_count)
+        _ = area_analytics.process_data(original_im0.copy(), frame_count)
     cap.release()
-    cv2.destroyAllWindows()
-
 
-@pytest.mark.slow
-def test_aigym():
-    """Test the workouts monitoring solution."""
+    # Test workouts monitoring
     safe_download(url=WORKOUTS_SOLUTION_DEMO)
-    cap = cv2.VideoCapture("solution_ci_pose_demo.mp4")
-    assert cap.isOpened(), "Error reading video file"
-    gym = solutions.AIGym(line_width=2, kpts=[5, 11, 13])
-    while cap.isOpened():
-        success, im0 = cap.read()
+    cap1 = cv2.VideoCapture("solution_ci_pose_demo.mp4")
+    assert cap1.isOpened(), "Error reading video file"
+    gym = solutions.AIGym(line_width=2, kpts=[5, 11, 13], show=False)
+    while cap1.isOpened():
+        success, im0 = cap1.read()
         if not success:
             break
         _ = gym.monitor(im0)
-    cap.release()
-    cv2.destroyAllWindows()
+    cap1.release()
 
 
 @pytest.mark.slow

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.3.14"
+__version__ = "8.3.16"
 
 import os
 

ultralytics/cfg/solutions/default.yaml

@@ -15,3 +15,4 @@ down_angle: 90 # Workouts down_angle for counts, 90 is default value. You can ch
 kpts: [6, 8, 10] # Keypoints for workouts monitoring, i.e. If you want to consider keypoints for pushups that have mostly values of [6, 8, 10].
 colormap: # Colormap for heatmap, Only OPENCV supported colormaps can be used. By default COLORMAP_PARULA will be used for visualization.
 analytics_type: "line" # Analytics type i.e "line", "pie", "bar" or "area" charts. By default, "line" analytics will be used for processing.
+json_file: # parking system regions file path.

ultralytics/data/split_dota.py

@@ -13,9 +13,6 @@ from tqdm import tqdm
 from ultralytics.data.utils import exif_size, img2label_paths
 from ultralytics.utils.checks import check_requirements
 
-check_requirements("shapely")
-from shapely.geometry import Polygon
-
 
 def bbox_iof(polygon1, bbox2, eps=1e-6):
     """
@@ -33,6 +30,9 @@ def bbox_iof(polygon1, bbox2, eps=1e-6):
         Polygon format: [x1, y1, x2, y2, x3, y3, x4, y4].
         Bounding box format: [x_min, y_min, x_max, y_max].
     """
+    check_requirements("shapely")
+    from shapely.geometry import Polygon
+
     polygon1 = polygon1.reshape(-1, 4, 2)
     lt_point = np.min(polygon1, axis=-2)  # left-top
     rb_point = np.max(polygon1, axis=-2)  # right-bottom

ultralytics/engine/exporter.py

@@ -398,7 +398,7 @@ class Exporter:
         """YOLO ONNX export."""
         requirements = ["onnx>=1.12.0"]
         if self.args.simplify:
-            requirements += ["onnxslim==0.1.34", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
+            requirements += ["onnxslim", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
         check_requirements(requirements)
         import onnx  # noqa
 

ultralytics/nn/autobackend.py

@@ -126,7 +126,7 @@ class AutoBackend(nn.Module):
         fp16 &= pt or jit or onnx or xml or engine or nn_module or triton  # FP16
         nhwc = coreml or saved_model or pb or tflite or edgetpu  # BHWC formats (vs torch BCWH)
         stride = 32  # default stride
-        model, metadata = None, None
+        model, metadata, task = None, None, None
 
         # Set device
         cuda = torch.cuda.is_available() and device.type != "cpu"  # use CUDA
@@ -336,11 +336,15 @@ class AutoBackend(nn.Module):
 
                 Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate
             if edgetpu:  # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
-                LOGGER.info(f"Loading {w} for TensorFlow Lite Edge TPU inference...")
+                device = device[3:] if str(device).startswith("tpu") else ":0"
+                LOGGER.info(f"Loading {w} on device {device[1:]} for TensorFlow Lite Edge TPU inference...")
                 delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
                     platform.system()
                 ]
-                interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
+                interpreter = Interpreter(
+                    model_path=w,
+                    experimental_delegates=[load_delegate(delegate, options={"device": device})],
+                )
             else:  # TFLite
                 LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
                 interpreter = Interpreter(model_path=w)  # load TFLite model

ultralytics/solutions/ai_gym.py

@@ -1,16 +1,40 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
+from ultralytics.solutions.solutions import BaseSolution
 from ultralytics.utils.plotting import Annotator
 
 
 class AIGym(BaseSolution):
-    """A class to manage the gym steps of people in a real-time video stream based on their poses."""
+    """
+    A class to manage gym steps of people in a real-time video stream based on their poses.
+
+    This class extends BaseSolution to monitor workouts using YOLO pose estimation models. It tracks and counts
+    repetitions of exercises based on predefined angle thresholds for up and down positions.
+
+    Attributes:
+        count (List[int]): Repetition counts for each detected person.
+        angle (List[float]): Current angle of the tracked body part for each person.
+        stage (List[str]): Current exercise stage ('up', 'down', or '-') for each person.
+        initial_stage (str | None): Initial stage of the exercise.
+        up_angle (float): Angle threshold for considering the 'up' position of an exercise.
+        down_angle (float): Angle threshold for considering the 'down' position of an exercise.
+        kpts (List[int]): Indices of keypoints used for angle calculation.
+        lw (int): Line width for drawing annotations.
+        annotator (Annotator): Object for drawing annotations on the image.
+
+    Methods:
+        monitor: Processes a frame to detect poses, calculate angles, and count repetitions.
+
+    Examples:
+        >>> gym = AIGym(model="yolov8n-pose.pt")
+        >>> image = cv2.imread("gym_scene.jpg")
+        >>> processed_image = gym.monitor(image)
+        >>> cv2.imshow("Processed Image", processed_image)
+        >>> cv2.waitKey(0)
+    """
 
     def __init__(self, **kwargs):
-        """Initialization function for AiGYM class, a child class of BaseSolution class, can be used for workouts
-        monitoring.
-        """
+        """Initializes AIGym for workout monitoring using pose estimation and predefined angles."""
         # Check if the model name ends with '-pose'
         if "model" in kwargs and "-pose" not in kwargs["model"]:
             kwargs["model"] = "yolo11n-pose.pt"
@@ -31,12 +55,22 @@ class AIGym(BaseSolution):
 
     def monitor(self, im0):
         """
-        Monitor the workouts using Ultralytics YOLOv8 Pose Model: https://docs.ultralytics.com/tasks/pose/.
+        Monitors workouts using Ultralytics YOLO Pose Model.
+
+        This function processes an input image to track and analyze human poses for workout monitoring. It uses
+        the YOLO Pose model to detect keypoints, estimate angles, and count repetitions based on predefined
+        angle thresholds.
 
         Args:
-            im0 (ndarray): The input image that will be used for processing
-        Returns
-            im0 (ndarray): The processed image for more usage
+            im0 (ndarray): Input image for processing.
+
+        Returns:
+            (ndarray): Processed image with annotations for workout monitoring.
+
+        Examples:
+            >>> gym = AIGym()
+            >>> image = cv2.imread("workout.jpg")
+            >>> processed_image = gym.monitor(image)
         """
         # Extract tracks
         tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"])[0]

ultralytics/solutions/analytics.py

@@ -12,10 +12,41 @@ from ultralytics.solutions.solutions import BaseSolution  # Import a parent clas
 
 
 class Analytics(BaseSolution):
-    """A class to create and update various types of charts (line, bar, pie, area) for visual analytics."""
+    """
+    A class for creating and updating various types of charts for visual analytics.
+
+    This class extends BaseSolution to provide functionality for generating line, bar, pie, and area charts
+    based on object detection and tracking data.
+
+    Attributes:
+        type (str): The type of analytics chart to generate ('line', 'bar', 'pie', or 'area').
+        x_label (str): Label for the x-axis.
+        y_label (str): Label for the y-axis.
+        bg_color (str): Background color of the chart frame.
+        fg_color (str): Foreground color of the chart frame.
+        title (str): Title of the chart window.
+        max_points (int): Maximum number of data points to display on the chart.
+        fontsize (int): Font size for text display.
+        color_cycle (cycle): Cyclic iterator for chart colors.
+        total_counts (int): Total count of detected objects (used for line charts).
+        clswise_count (Dict[str, int]): Dictionary for class-wise object counts.
+        fig (Figure): Matplotlib figure object for the chart.
+        ax (Axes): Matplotlib axes object for the chart.
+        canvas (FigureCanvas): Canvas for rendering the chart.
+
+    Methods:
+        process_data: Processes image data and updates the chart.
+        update_graph: Updates the chart with new data points.
+
+    Examples:
+        >>> analytics = Analytics(analytics_type="line")
+        >>> frame = cv2.imread("image.jpg")
+        >>> processed_frame = analytics.process_data(frame, frame_number=1)
+        >>> cv2.imshow("Analytics", processed_frame)
+    """
 
     def __init__(self, **kwargs):
-        """Initialize the Analytics class with various chart types."""
+        """Initialize Analytics class with various chart types for visual data representation."""
         super().__init__(**kwargs)
 
         self.type = self.CFG["analytics_type"]  # extract type of analytics
@@ -31,8 +62,8 @@ class Analytics(BaseSolution):
         figsize = (19.2, 10.8)  # Set output image size 1920 * 1080
         self.color_cycle = cycle(["#DD00BA", "#042AFF", "#FF4447", "#7D24FF", "#BD00FF"])
 
-        self.total_counts = 0  # count variable for storing total counts i.e for line
-        self.clswise_count = {}  # dictionary for classwise counts
+        self.total_counts = 0  # count variable for storing total counts i.e. for line
+        self.clswise_count = {}  # dictionary for class-wise counts
 
         # Ensure line and area chart
         if self.type in {"line", "area"}:
@@ -48,15 +79,28 @@ class Analytics(BaseSolution):
             self.canvas = FigureCanvas(self.fig)  # Set common axis properties
             self.ax.set_facecolor(self.bg_color)
             self.color_mapping = {}
-            self.ax.axis("equal") if self.type == "pie" else None  # Ensure pie chart is circular
+
+            if self.type == "pie":  # Ensure pie chart is circular
+                self.ax.axis("equal")
 
     def process_data(self, im0, frame_number):
         """
-        Process the image data, run object tracking.
+        Processes image data and runs object tracking to update analytics charts.
 
         Args:
-            im0 (ndarray): Input image for processing.
-            frame_number (int): Video frame # for plotting the data.
+            im0 (np.ndarray): Input image for processing.
+            frame_number (int): Video frame number for plotting the data.
+
+        Returns:
+            (np.ndarray): Processed image with updated analytics chart.
+
+        Raises:
+            ModuleNotFoundError: If an unsupported chart type is specified.
+
+        Examples:
+            >>> analytics = Analytics(analytics_type="line")
+            >>> frame = np.zeros((480, 640, 3), dtype=np.uint8)
+            >>> processed_frame = analytics.process_data(frame, frame_number=1)
         """
         self.extract_tracks(im0)  # Extract tracks
 
@@ -79,13 +123,22 @@ class Analytics(BaseSolution):
 
     def update_graph(self, frame_number, count_dict=None, plot="line"):
         """
-        Update the graph (line or area) with new data for single or multiple classes.
+        Updates the graph with new data for single or multiple classes.
 
         Args:
             frame_number (int): The current frame number.
-            count_dict (dict, optional): Dictionary with class names as keys and counts as values for multiple classes.
-                                          If None, updates a single line graph.
-            plot (str): Type of the plot i.e. line, bar or area.
+            count_dict (Dict[str, int] | None): Dictionary with class names as keys and counts as values for multiple
+                classes. If None, updates a single line graph.
+            plot (str): Type of the plot. Options are 'line', 'bar', 'pie', or 'area'.
+
+        Returns:
+            (np.ndarray): Updated image containing the graph.
+
+        Examples:
+            >>> analytics = Analytics()
+            >>> frame_number = 10
+            >>> count_dict = {"person": 5, "car": 3}
+            >>> updated_image = analytics.update_graph(frame_number, count_dict, plot="bar")
         """
         if count_dict is None:
             # Single line update

ultralytics/solutions/distance_calculation.py

@@ -4,15 +4,41 @@ import math
 
 import cv2
 
-from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
+from ultralytics.solutions.solutions import BaseSolution
 from ultralytics.utils.plotting import Annotator, colors
 
 
 class DistanceCalculation(BaseSolution):
-    """A class to calculate distance between two objects in a 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.
+
+    This class extends BaseSolution to provide functionality for selecting objects and calculating the distance
+    between them in a video stream using YOLO object detection and tracking.
+
+    Attributes:
+        left_mouse_count (int): Counter for left mouse button clicks.
+        selected_boxes (Dict[int, List[float]]): Dictionary to store selected bounding boxes and their track IDs.
+        annotator (Annotator): An instance of the Annotator class for drawing on the image.
+        boxes (List[List[float]]): List of bounding boxes for detected objects.
+        track_ids (List[int]): List of track IDs for detected objects.
+        clss (List[int]): List of class indices for detected objects.
+        names (List[str]): List of class names that the model can detect.
+        centroids (List[List[int]]): List to store centroids of selected bounding boxes.
+
+    Methods:
+        mouse_event_for_distance: Handles mouse events for selecting objects in the video stream.
+        calculate: Processes video frames and calculates the distance between selected objects.
+
+    Examples:
+        >>> distance_calc = DistanceCalculation()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_frame = distance_calc.calculate(frame)
+        >>> cv2.imshow("Distance Calculation", processed_frame)
+        >>> cv2.waitKey(0)
+    """
 
     def __init__(self, **kwargs):
-        """Initializes the DistanceCalculation class with the given parameters."""
+        """Initializes the DistanceCalculation class for measuring object distances in video streams."""
         super().__init__(**kwargs)
 
         # Mouse event information
@@ -21,14 +47,18 @@ class DistanceCalculation(BaseSolution):
 
     def mouse_event_for_distance(self, event, x, y, flags, param):
         """
-        Handles mouse events to select regions in a real-time video stream.
+        Handles mouse events to select regions in a real-time video stream for distance calculation.
 
         Args:
-            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
+            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN).
             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.
+            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY).
+            param (Dict): Additional parameters passed to the function.
+
+        Examples:
+            >>> # Assuming 'dc' is an instance of DistanceCalculation
+            >>> cv2.setMouseCallback("window_name", dc.mouse_event_for_distance)
         """
         if event == cv2.EVENT_LBUTTONDOWN:
             self.left_mouse_count += 1
@@ -43,13 +73,23 @@ class DistanceCalculation(BaseSolution):
 
     def calculate(self, im0):
         """
-        Processes the video frame and calculates the distance between two bounding boxes.
+        Processes a video frame and calculates the distance between two selected bounding boxes.
+
+        This method extracts tracks from the input frame, annotates bounding boxes, and calculates the distance
+        between two user-selected objects if they have been chosen.
 
         Args:
-            im0 (ndarray): The image frame.
+            im0 (numpy.ndarray): The input image frame to process.
 
         Returns:
-            (ndarray): The processed image frame.
+            (numpy.ndarray): The processed image frame with annotations and distance calculations.
+
+        Examples:
+            >>> import numpy as np
+            >>> from ultralytics.solutions import DistanceCalculation
+            >>> dc = DistanceCalculation()
+            >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+            >>> processed_frame = dc.calculate(frame)
         """
         self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
         self.extract_tracks(im0)  # Extract tracks

ultralytics/solutions/heatmap.py

@@ -3,15 +3,40 @@
 import cv2
 import numpy as np
 
-from ultralytics.solutions.object_counter import ObjectCounter  # Import object counter class
+from ultralytics.solutions.object_counter import ObjectCounter
 from ultralytics.utils.plotting import Annotator
 
 
 class Heatmap(ObjectCounter):
-    """A class to draw heatmaps in real-time video stream based on their tracks."""
+    """
+    A class to draw heatmaps in real-time video streams based on object tracks.
+
+    This class extends the ObjectCounter class to generate and visualize heatmaps of object movements in video
+    streams. It uses tracked object positions to create a cumulative heatmap effect over time.
+
+    Attributes:
+        initialized (bool): Flag indicating whether the heatmap has been initialized.
+        colormap (int): OpenCV colormap used for heatmap visualization.
+        heatmap (np.ndarray): Array storing the cumulative heatmap data.
+        annotator (Annotator): Object for drawing annotations on the image.
+
+    Methods:
+        heatmap_effect: Calculates and updates the heatmap effect for a given bounding box.
+        generate_heatmap: Generates and applies the heatmap effect to each frame.
+
+    Examples:
+        >>> from ultralytics.solutions import Heatmap
+        >>> heatmap = Heatmap(model="yolov8n.pt", colormap=cv2.COLORMAP_JET)
+        >>> results = heatmap("path/to/video.mp4")
+        >>> for result in results:
+        ...     print(result.speed)  # Print inference speed
+        ...     cv2.imshow("Heatmap", result.plot())
+        ...     if cv2.waitKey(1) & 0xFF == ord("q"):
+        ...         break
+    """
 
     def __init__(self, **kwargs):
-        """Initializes function for heatmap class with default values."""
+        """Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks."""
         super().__init__(**kwargs)
 
         self.initialized = False  # bool variable for heatmap initialization
@@ -23,10 +48,15 @@ class Heatmap(ObjectCounter):
 
     def heatmap_effect(self, box):
         """
-        Efficient calculation of heatmap area and effect location for applying colormap.
+        Efficiently calculates heatmap area and effect location for applying colormap.
 
         Args:
-            box (list): Bounding Box coordinates data [x0, y0, x1, y1]
+            box (List[float]): Bounding box coordinates [x0, y0, x1, y1].
+
+        Examples:
+            >>> heatmap = Heatmap()
+            >>> box = [100, 100, 200, 200]
+            >>> heatmap.heatmap_effect(box)
         """
         x0, y0, x1, y1 = map(int, box)
         radius_squared = (min(x1 - x0, y1 - y0) // 2) ** 2
@@ -48,9 +78,15 @@ class Heatmap(ObjectCounter):
         Generate heatmap for each frame using Ultralytics.
 
         Args:
-            im0 (ndarray): Input image array for processing
+            im0 (np.ndarray): Input image array for processing.
+
         Returns:
-            im0 (ndarray): Processed image for further usage
+            (np.ndarray): Processed image with heatmap overlay and object counts (if region is specified).
+
+        Examples:
+            >>> heatmap = Heatmap()
+            >>> im0 = cv2.imread("image.jpg")
+            >>> result = heatmap.generate_heatmap(im0)
         """
         if not self.initialized:
             self.heatmap = np.zeros_like(im0, dtype=np.float32) * 0.99
@@ -70,16 +106,17 @@ class Heatmap(ObjectCounter):
                 self.store_classwise_counts(cls)  # store classwise counts in dict
 
                 # Store tracking previous position and perform object counting
-                prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
+                prev_position = None
+                if len(self.track_history[track_id]) > 1:
+                    prev_position = self.track_history[track_id][-2]
                 self.count_objects(self.track_line, box, track_id, prev_position, cls)  # Perform object counting
 
-        self.display_counts(im0) if self.region is not None else None  # Display the counts on the frame
+        if self.region is not None:
+            self.display_counts(im0)  # Display the counts on the frame
 
         # Normalize, apply colormap to heatmap and combine with original image
-        im0 = (
-            im0
-            if self.track_data.id is None
-            else cv2.addWeighted(
+        if self.track_data.id is not None:
+            im0 = cv2.addWeighted(
                 im0,
                 0.5,
                 cv2.applyColorMap(
@@ -88,7 +125,6 @@ class Heatmap(ObjectCounter):
                 0.5,
                 0,
             )
-        )
 
         self.display_output(im0)  # display output with base class function
         return im0  # return output image for more usage