ultralytics 8.3.87__tar.gz → 8.3.88__tar.gz

{ultralytics-8.3.87/ultralytics.egg-info → ultralytics-8.3.88}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ultralytics
-Version: 8.3.87
+Version: 8.3.88
 Summary: Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{ultralytics-8.3.87 → ultralytics-8.3.88}/tests/test_solutions.py

@@ -1,11 +1,13 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+# This is file for Ultralytics Solutions tests: https://docs.ultralytics.com/solutions/,
+# It includes every solution excluding DistanceCalculation and Security Alarm System.
 
 import cv2
 import pytest
 
 from tests import TMP
-from ultralytics import YOLO, solutions
-from ultralytics.utils import ASSETS_URL, WEIGHTS_DIR
+from ultralytics import solutions
+from ultralytics.utils import ASSETS_URL
 from ultralytics.utils.downloads import safe_download
 
 DEMO_VIDEO = "solutions_ci_demo.mp4"  # for all the solutions, except workout and parking
@@ -24,16 +26,23 @@ def test_major_solutions():
     region_points = [(20, 400), (1080, 400), (1080, 360), (20, 360)]
     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
-    heatmap_count = solutions.Heatmap(
+    heatmapcounter = solutions.Heatmap(
         colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False, region=region_points
     )  # Test heatmaps with object counting
     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
-    trackzone = solutions.TrackZone(region=region_points, model="yolo11n.pt", show=False)  # Test trackzone
+    lineanalytics = solutions.Analytics(analytics_type="line", model="yolo11n.pt", show=False)  # line analytics
+    pieanalytics = solutions.Analytics(analytics_type="pie", model="yolo11n.pt", show=False)  # line analytics
+    baranalytics = solutions.Analytics(analytics_type="bar", model="yolo11n.pt", show=False)  # line analytics
+    areaanalytics = solutions.Analytics(analytics_type="area", model="yolo11n.pt", show=False)  # line analytics
+    trackzone = solutions.TrackZone(region=region_points, model="yolo11n.pt", show=False)  # trackzone
+    objectcropper = solutions.ObjectCropper(
+        model="yolo11n.pt", show=False, crop_dir=str(TMP / "cropped-detections")
+    )  # object cropping
+    objectblurrer = solutions.ObjectBlurrer(blur_ratio=0.5, model="yolo11n.pt", show=False)  # Object blurring
+    isegment = solutions.InstanceSegmentation(model="yolo11n-seg.pt", show=False)  # Instance segmentation
+    visioneye = solutions.VisionEye(model="yolo11n.pt", show=False)  # Visioneye
+    regioncounter = solutions.RegionCounter(region=region_points, model="yolo11n.pt", show=False)  # Region counter
     frame_count = 0  # Required for analytics
     while cap.isOpened():
         success, im0 = cap.read()
@@ -41,16 +50,21 @@ def test_major_solutions():
             break
         frame_count += 1
         original_im0 = im0.copy()
-        _ = counter.count(original_im0.copy())
-        _ = heatmap.generate_heatmap(original_im0.copy())
-        _ = heatmap_count.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)
-        _ = trackzone.trackzone(original_im0.copy())
+        _ = counter(original_im0.copy())
+        _ = heatmap(original_im0.copy())
+        _ = heatmapcounter(original_im0.copy())
+        _ = speed(original_im0.copy())
+        _ = queue(original_im0.copy())
+        _ = lineanalytics(original_im0.copy(), frame_count)
+        _ = pieanalytics(original_im0.copy(), frame_count)
+        _ = baranalytics(original_im0.copy(), frame_count)
+        _ = areaanalytics(original_im0.copy(), frame_count)
+        _ = trackzone(original_im0.copy())
+        _ = objectcropper(original_im0.copy())
+        _ = isegment(original_im0.copy())
+        _ = objectblurrer(original_im0.copy())
+        _ = visioneye(original_im0.copy())
+        _ = regioncounter(original_im0.copy())
     cap.release()
 
     # Test workouts monitoring
@@ -62,7 +76,7 @@ def test_major_solutions():
         success, im0 = cap.read()
         if not success:
             break
-        _ = gym.monitor(im0)
+        _ = gym(im0)
     cap.release()
 
     # Test parking management
@@ -78,35 +92,10 @@ def test_major_solutions():
         success, im0 = cap.read()
         if not success:
             break
-        _ = parkingmanager.process_data(im0)
+        _ = parkingmanager(im0)
     cap.release()
 
 
-@pytest.mark.slow
-def test_instance_segmentation():
-    """Test the instance segmentation solution."""
-    from ultralytics.utils.plotting import Annotator, colors
-
-    model = YOLO(WEIGHTS_DIR / "yolo11n-seg.pt")
-    names = model.names
-    cap = cv2.VideoCapture(TMP / DEMO_VIDEO)
-    assert cap.isOpened(), "Error reading video file"
-    while cap.isOpened():
-        success, im0 = cap.read()
-        if not success:
-            break
-        results = model.predict(im0)
-        annotator = Annotator(im0, line_width=2)
-        if results[0].masks is not None:
-            clss = results[0].boxes.cls.cpu().tolist()
-            masks = results[0].masks.xy
-            for mask, cls in zip(masks, clss):
-                color = colors(int(cls), True)
-                annotator.seg_bbox(mask=mask, mask_color=color, label=names[int(cls)])
-    cap.release()
-    cv2.destroyAllWindows()
-
-
 @pytest.mark.slow
 def test_streamlit_predict():
     """Test streamlit predict live inference solution."""

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.87"
+__version__ = "8.3.88"
 
 import os
 

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/cfg/__init__.py

@@ -5,7 +5,7 @@ import subprocess
 import sys
 from pathlib import Path
 from types import SimpleNamespace
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Union
 
 import cv2
 
@@ -35,14 +35,18 @@ from ultralytics.utils import (
 
 # Define valid solutions
 SOLUTION_MAP = {
-    "count": ("ObjectCounter", "count"),
-    "heatmap": ("Heatmap", "generate_heatmap"),
-    "queue": ("QueueManager", "process_queue"),
-    "speed": ("SpeedEstimator", "estimate_speed"),
-    "workout": ("AIGym", "monitor"),
-    "analytics": ("Analytics", "process_data"),
-    "trackzone": ("TrackZone", "trackzone"),
-    "inference": ("Inference", "inference"),
+    "count": "ObjectCounter",
+    "crop": "ObjectCropper",
+    "blur": "ObjectBlurrer",
+    "workout": "AIGym",
+    "heatmap": "Heatmap",
+    "isegment": "InstanceSegmentation",
+    "visioneye": "VisionEye",
+    "speed": "SpeedEstimator",
+    "queue": "QueueManager",
+    "analytics": "Analytics",
+    "inference": "Inference",
+    "trackzone": "TrackZone",
     "help": None,
 }
 
@@ -238,7 +242,7 @@ CFG_BOOL_KEYS = frozenset(
 )
 
 
-def cfg2dict(cfg):
+def cfg2dict(cfg: Union[str, Path, Dict, SimpleNamespace]) -> Dict:
     """
     Converts a configuration object to a dictionary.
 
@@ -273,7 +277,7 @@ def cfg2dict(cfg):
     return cfg
 
 
-def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None):
+def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None) -> SimpleNamespace:
     """
     Load and merge configuration data from a file or dictionary, with optional overrides.
 
@@ -321,7 +325,7 @@ def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, ove
     return IterableSimpleNamespace(**cfg)
 
 
-def check_cfg(cfg, hard=True):
+def check_cfg(cfg: Dict, hard: bool = True) -> None:
     """
     Checks configuration argument types and values for the Ultralytics library.
 
@@ -383,7 +387,7 @@ def check_cfg(cfg, hard=True):
                 cfg[k] = bool(v)
 
 
-def get_save_dir(args, name=None):
+def get_save_dir(args: SimpleNamespace, name: str = None) -> Path:
     """
     Returns the directory path for saving outputs, derived from arguments or default settings.
 
@@ -415,7 +419,7 @@ def get_save_dir(args, name=None):
     return Path(save_dir)
 
 
-def _handle_deprecation(custom):
+def _handle_deprecation(custom: Dict) -> Dict:
     """
     Handles deprecated configuration keys by mapping them to current equivalents with deprecation warnings.
 
@@ -453,7 +457,7 @@ def _handle_deprecation(custom):
     return custom
 
 
-def check_dict_alignment(base: Dict, custom: Dict, e=None):
+def check_dict_alignment(base: Dict, custom: Dict, e: Exception = None) -> None:
     """
     Checks alignment between custom and base configuration dictionaries, handling deprecated keys and providing error
     messages for mismatched keys.
@@ -507,7 +511,7 @@ def merge_equals_args(args: List[str]) -> List[str]:
         args (List[str]): A list of strings where each element represents an argument or fragment.
 
     Returns:
-        List[str]: A list of strings where the arguments around isolated '=' are merged and fragments with brackets are joined.
+        (List[str]): A list of strings where the arguments around isolated '=' are merged and fragments with brackets are joined.
 
     Examples:
         >>> args = ["arg1", "=", "value", "arg2=", "value2", "arg3", "=value3", "imgsz=[3,", "640,", "640]"]
@@ -634,9 +638,6 @@ def handle_yolo_solutions(args: List[str]) -> None:
             solutions: https://docs.ultralytics.com/solutions/, It can include solution name, source,
             and other configuration parameters.
 
-    Returns:
-        None: The function processes video frames and saves the output but doesn't return any value.
-
     Examples:
         Run people counting solution with default settings:
         >>> handle_yolo_solutions(["count"])
@@ -660,7 +661,13 @@ def handle_yolo_solutions(args: List[str]) -> None:
         - The inference solution will be launched using the 'streamlit run' command.
         - The Streamlit app file is located in the Ultralytics package directory.
     """
-    full_args_dict = {**DEFAULT_SOL_DICT, **DEFAULT_CFG_DICT}  # arguments dictionary
+    full_args_dict = {
+        **DEFAULT_SOL_DICT,
+        **DEFAULT_CFG_DICT,
+        "blur_ratio": 0.5,
+        "vision_point": (20, 20),
+        "crop_dir": "cropped-detections",
+    }  # arguments dictionary
     overrides = {}
 
     # check dictionary alignment
@@ -705,21 +712,19 @@ def handle_yolo_solutions(args: List[str]) -> None:
     else:
         from ultralytics import solutions
 
-        cls, method = SOLUTION_MAP[solution_name]  # solution class name, method name and default source
-        solution = getattr(solutions, cls)(IS_CLI=True, **overrides)  # get solution class i.e ObjectCounter
-        process = getattr(
-            solution, method
-        )  # get specific function of class for processing i.e, count from ObjectCounter
+        solution = getattr(solutions, SOLUTION_MAP[solution_name])(is_cli=True, **overrides)  # class i.e ObjectCounter
 
         cap = cv2.VideoCapture(solution.CFG["source"])  # read the video file
-
-        # extract width, height and fps of the video file, create save directory and initialize video writer
-        w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
-        if solution_name == "analytics":  # analytical graphs follow fixed shape for output i.e w=1920, h=1080
-            w, h = 1920, 1080
-        save_dir = get_save_dir(SimpleNamespace(project="runs/solutions", name="exp", exist_ok=False))
-        save_dir.mkdir(parents=True)  # create the output directory i.e. runs/solutions/exp
-        vw = cv2.VideoWriter(str(save_dir / f"{solution_name}.avi"), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+        if solution_name != "crop":
+            # extract width, height and fps of the video file, create save directory and initialize video writer
+            w, h, fps = (
+                int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)
+            )
+            if solution_name == "analytics":  # analytical graphs follow fixed shape for output i.e w=1920, h=1080
+                w, h = 1280, 720
+            save_dir = get_save_dir(SimpleNamespace(project="runs/solutions", name="exp", exist_ok=False))
+            save_dir.mkdir(parents=True)  # create the output directory i.e. runs/solutions/exp
+            vw = cv2.VideoWriter(str(save_dir / f"{solution_name}.avi"), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
 
         try:  # Process video frames
             f_n = 0  # frame number, required for analytical graphs
@@ -727,15 +732,17 @@ def handle_yolo_solutions(args: List[str]) -> None:
                 success, frame = cap.read()
                 if not success:
                     break
-                frame = process(frame, f_n := f_n + 1) if solution_name == "analytics" else process(frame)
-                vw.write(frame)
+                results = solution(frame, f_n := f_n + 1) if solution_name == "analytics" else solution(frame)
+                LOGGER.info(f"🚀 Results: {results}")
+                if solution_name != "crop":
+                    vw.write(results.plot_im)
                 if cv2.waitKey(1) & 0xFF == ord("q"):
                     break
         finally:
             cap.release()
 
 
-def parse_key_value_pair(pair: str = "key=value"):
+def parse_key_value_pair(pair: str = "key=value") -> tuple:
     """
     Parses a key-value pair string into separate key and value components.
 
@@ -769,7 +776,7 @@ def parse_key_value_pair(pair: str = "key=value"):
     return k, smart_value(v)
 
 
-def smart_value(v):
+def smart_value(v: str) -> Any:
     """
     Converts a string representation of a value to its appropriate Python type.
 
@@ -814,7 +821,7 @@ def smart_value(v):
             return v
 
 
-def entrypoint(debug=""):
+def entrypoint(debug: str = "") -> None:
     """
     Ultralytics entrypoint function for parsing and executing command-line arguments.
 
@@ -986,7 +993,7 @@ def entrypoint(debug=""):
 
 
 # Special modes --------------------------------------------------------------------------------------------------------
-def copy_default_cfg():
+def copy_default_cfg() -> None:
     """
     Copies the default configuration file and creates a new one with '_copy' appended to its name.
 

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/data/augment.py

@@ -1193,8 +1193,8 @@ class RandomPerspective:
         Args:
             labels (Dict): A dictionary containing image data and annotations.
                 Must include:
-                    'img' (ndarray): The input image.
-                    'cls' (ndarray): Class labels.
+                    'img' (np.ndarray): The input image.
+                    'cls' (np.ndarray): Class labels.
                     'instances' (Instances): Object instances with bounding boxes, segments, and keypoints.
                 May include:
                     'mosaic_border' (Tuple[int, int]): Border size for mosaic augmentation.

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/solutions/__init__.py

@@ -4,7 +4,10 @@ from .ai_gym import AIGym
 from .analytics import Analytics
 from .distance_calculation import DistanceCalculation
 from .heatmap import Heatmap
+from .instance_segmentation import InstanceSegmentation
+from .object_blurrer import ObjectBlurrer
 from .object_counter import ObjectCounter
+from .object_cropper import ObjectCropper
 from .parking_management import ParkingManagement, ParkingPtsSelection
 from .queue_management import QueueManager
 from .region_counter import RegionCounter
@@ -12,19 +15,24 @@ from .security_alarm import SecurityAlarm
 from .speed_estimation import SpeedEstimator
 from .streamlit_inference import Inference
 from .trackzone import TrackZone
+from .vision_eye import VisionEye
 
 __all__ = (
+    "ObjectCounter",
+    "ObjectCropper",
+    "ObjectBlurrer",
     "AIGym",
-    "DistanceCalculation",
+    "RegionCounter",
+    "SecurityAlarm",
     "Heatmap",
-    "ObjectCounter",
+    "InstanceSegmentation",
+    "VisionEye",
+    "SpeedEstimator",
+    "DistanceCalculation",
+    "QueueManager",
     "ParkingManagement",
     "ParkingPtsSelection",
-    "QueueManager",
-    "SpeedEstimator",
     "Analytics",
     "Inference",
-    "RegionCounter",
     "TrackZone",
-    "SecurityAlarm",
 )

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/solutions/ai_gym.py

@@ -1,7 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from ultralytics.solutions.solutions import BaseSolution
-from ultralytics.utils.plotting import Annotator
+from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
 
 
 class AIGym(BaseSolution):
@@ -19,27 +18,28 @@ class AIGym(BaseSolution):
         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.
-        annotator (Annotator): Object for drawing annotations on the image.
 
     Methods:
-        monitor: Processes a frame to detect poses, calculate angles, and count repetitions.
+        process: Processes a frame to detect poses, calculate angles, and count repetitions.
 
     Examples:
         >>> gym = AIGym(model="yolo11n-pose.pt")
         >>> image = cv2.imread("gym_scene.jpg")
-        >>> processed_image = gym.monitor(image)
+        >>> results = gym.process(image)
+        >>> processed_image = results.plot_im
         >>> cv2.imshow("Processed Image", processed_image)
         >>> cv2.waitKey(0)
     """
 
     def __init__(self, **kwargs):
-        """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"
-        elif "model" not in kwargs:
-            kwargs["model"] = "yolo11n-pose.pt"
+        """
+        Initializes AIGym for workout monitoring using pose estimation and predefined angles.
 
+        Args:
+            **kwargs (Any): Keyword arguments passed to the parent class constructor.
+                model (str): Model name or path, defaults to "yolo11n-pose.pt".
+        """
+        kwargs["model"] = kwargs.get("model", "yolo11n-pose.pt")
         super().__init__(**kwargs)
         self.count = []  # List for counts, necessary where there are multiple objects in frame
         self.angle = []  # List for angle, necessary where there are multiple objects in frame
@@ -51,7 +51,7 @@ class AIGym(BaseSolution):
         self.down_angle = float(self.CFG["down_angle"])  # Pose down predefined angle to consider down pose
         self.kpts = self.CFG["kpts"]  # User selected kpts of workouts storage for further usage
 
-    def monitor(self, im0):
+    def process(self, im0):
         """
         Monitors workouts using Ultralytics YOLO Pose Model.
 
@@ -60,36 +60,39 @@ class AIGym(BaseSolution):
         angle thresholds.
 
         Args:
-            im0 (ndarray): Input image for processing.
+            im0 (np.ndarray): Input image for processing.
 
         Returns:
-            (ndarray): Processed image with annotations for workout monitoring.
+            (SolutionResults): Contains processed image `plot_im`,
+                'workout_count' (list of completed reps),
+                'workout_stage' (list of current stages),
+                'workout_angle' (list of angles), and
+                'total_tracks' (total number of tracked individuals).
 
         Examples:
             >>> gym = AIGym()
             >>> image = cv2.imread("workout.jpg")
-            >>> processed_image = gym.monitor(image)
+            >>> results = gym.process(image)
+            >>> processed_image = results.plot_im
         """
-        # Extract tracks
-        tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"], **self.track_add_args)[0]
+        annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator
+
+        self.extract_tracks(im0)  # Extract tracks (bounding boxes, classes, and masks)
+        tracks = self.tracks[0]
 
         if tracks.boxes.id is not None:
-            # Extract and check keypoints
-            if len(tracks) > len(self.count):
+            if len(tracks) > len(self.count):  # Add new entries for newly detected people
                 new_human = len(tracks) - len(self.count)
                 self.angle += [0] * new_human
                 self.count += [0] * new_human
                 self.stage += ["-"] * new_human
 
-            # Initialize annotator
-            self.annotator = Annotator(im0, line_width=self.line_width)
-
             # Enumerate over keypoints
             for ind, k in enumerate(reversed(tracks.keypoints.data)):
                 # Get keypoints and estimate the angle
                 kpts = [k[int(self.kpts[i])].cpu() for i in range(3)]
-                self.angle[ind] = self.annotator.estimate_pose_angle(*kpts)
-                im0 = self.annotator.draw_specific_points(k, self.kpts, radius=self.line_width * 3)
+                self.angle[ind] = annotator.estimate_pose_angle(*kpts)
+                annotator.draw_specific_kpts(k, self.kpts, radius=self.line_width * 3)
 
                 # Determine stage and count logic based on angle thresholds
                 if self.angle[ind] < self.down_angle:
@@ -100,12 +103,20 @@ class AIGym(BaseSolution):
                     self.stage[ind] = "up"
 
                 # Display angle, count, and stage text
-                self.annotator.plot_angle_and_count_and_stage(
+                annotator.plot_angle_and_count_and_stage(
                     angle_text=self.angle[ind],  # angle text for display
                     count_text=self.count[ind],  # count text for workouts
                     stage_text=self.stage[ind],  # stage position text
                     center_kpt=k[int(self.kpts[1])],  # center keypoint for display
                 )
-
-        self.display_output(im0)  # Display output image, if environment support display
-        return im0  # return an image for writing or further usage
+        plot_im = annotator.result()
+        self.display_output(plot_im)  # Display output image, if environment support display
+
+        # Return SolutionResults
+        return SolutionResults(
+            plot_im=plot_im,
+            workout_count=self.count,
+            workout_stage=self.stage,
+            workout_angle=self.angle,
+            total_tracks=len(self.track_ids),
+        )

{ultralytics-8.3.87 → ultralytics-8.3.88}/ultralytics/solutions/analytics.py

@@ -8,7 +8,7 @@ import numpy as np
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 from matplotlib.figure import Figure
 
-from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
+from ultralytics.solutions.solutions import BaseSolution, SolutionResults  # Import a parent class
 
 
 class Analytics(BaseSolution):
@@ -33,16 +33,18 @@ class Analytics(BaseSolution):
         fig (Figure): Matplotlib figure object for the chart.
         ax (Axes): Matplotlib axes object for the chart.
         canvas (FigureCanvas): Canvas for rendering the chart.
+        lines (Dict): Dictionary to store line objects for area charts.
+        color_mapping (Dict[str, str]): Dictionary mapping class labels to colors for consistent visualization.
 
     Methods:
-        process_data: Processes image data and updates the chart.
+        process: 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)
+        >>> results = analytics.process(frame, frame_number=1)
+        >>> cv2.imshow("Analytics", results.plot_im)
     """
 
     def __init__(self, **kwargs):
@@ -59,7 +61,7 @@ class Analytics(BaseSolution):
         self.title = "Ultralytics Solutions"  # window name
         self.max_points = 45  # maximum points to be drawn on window
         self.fontsize = 25  # text font size for display
-        figsize = (19.2, 10.8)  # Set output image size 1920 * 1080
+        figsize = (12.8, 7.2)  # Set output image size 1280 * 720
         self.color_cycle = cycle(["#DD00BA", "#042AFF", "#FF4447", "#7D24FF", "#BD00FF"])
 
         self.total_counts = 0  # count variable for storing total counts i.e. for line
@@ -83,7 +85,7 @@ class Analytics(BaseSolution):
             if self.type == "pie":  # Ensure pie chart is circular
                 self.ax.axis("equal")
 
-    def process_data(self, im0, frame_number):
+    def process(self, im0, frame_number):
         """
         Processes image data and runs object tracking to update analytics charts.
 
@@ -92,7 +94,8 @@ class Analytics(BaseSolution):
             frame_number (int): Video frame number for plotting the data.
 
         Returns:
-            (np.ndarray): Processed image with updated analytics chart.
+            (SolutionResults): Contains processed image `plot_im`, 'total_tracks' (int, total number of tracked objects)
+                and 'classwise_count' (dict, per-class object count).
 
         Raises:
             ModuleNotFoundError: If an unsupported chart type is specified.
@@ -100,26 +103,27 @@ class Analytics(BaseSolution):
         Examples:
             >>> analytics = Analytics(analytics_type="line")
             >>> frame = np.zeros((480, 640, 3), dtype=np.uint8)
-            >>> processed_frame = analytics.process_data(frame, frame_number=1)
+            >>> results = analytics.process(frame, frame_number=1)
         """
         self.extract_tracks(im0)  # Extract tracks
-
         if self.type == "line":
             for _ in self.boxes:
                 self.total_counts += 1
-            im0 = self.update_graph(frame_number=frame_number)
+            plot_im = self.update_graph(frame_number=frame_number)
             self.total_counts = 0
         elif self.type in {"pie", "bar", "area"}:
             self.clswise_count = {}
-            for box, cls in zip(self.boxes, self.clss):
+            for cls in self.clss:
                 if self.names[int(cls)] in self.clswise_count:
                     self.clswise_count[self.names[int(cls)]] += 1
                 else:
                     self.clswise_count[self.names[int(cls)]] = 1
-            im0 = self.update_graph(frame_number=frame_number, count_dict=self.clswise_count, plot=self.type)
+            plot_im = self.update_graph(frame_number=frame_number, count_dict=self.clswise_count, plot=self.type)
         else:
             raise ModuleNotFoundError(f"{self.type} chart is not supported ❌")
-        return im0
+
+        # return output dictionary with summary for more usage
+        return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids), classwise_count=self.clswise_count)
 
     def update_graph(self, frame_number, count_dict=None, plot="line"):
         """
@@ -135,10 +139,10 @@ class Analytics(BaseSolution):
             (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")
+            >>> analytics = Analytics(analytics_type="bar")
+            >>> frame_num = 10
+            >>> results_dict = {"person": 5, "car": 3}
+            >>> updated_image = analytics.update_graph(frame_num, results_dict, plot="bar")
         """
         if count_dict is None:
             # Single line update
@@ -216,7 +220,7 @@ class Analytics(BaseSolution):
                 self.ax.clear()
 
                 # Create pie chart and create legend labels with percentages
-                wedges, autotexts = self.ax.pie(
+                wedges, _ = self.ax.pie(
                     counts, labels=labels, startangle=start_angle, textprops={"color": self.fg_color}, autopct=None
                 )
                 legend_labels = [f"{label} ({percentage:.1f}%)" for label, percentage in zip(labels, percentages)]