ultralytics 8.3.86__py3-none-any.whl → 8.3.88__py3-none-any.whl

ultralytics/engine/trainer.py

@@ -452,7 +452,8 @@ class BaseTrainer:
                 self.scheduler.last_epoch = self.epoch  # do not move
                 self.stop |= epoch >= self.epochs  # stop if exceeded epochs
             self.run_callbacks("on_fit_epoch_end")
-            self._clear_memory()
+            if self._get_memory(fraction=True) > 0.9:
+                self._clear_memory()  # clear if memory utilization > 90%
 
             # Early Stopping
             if RANK != -1:  # if DDP training
@@ -485,15 +486,20 @@ class BaseTrainer:
             max_num_obj=max_num_obj,
         )  # returns batch size
 
-    def _get_memory(self):
-        """Get accelerator memory utilization in GB."""
+    def _get_memory(self, fraction=False):
+        """Get accelerator memory utilization in GB or fraction."""
+        memory, total = 0, 0
         if self.device.type == "mps":
             memory = torch.mps.driver_allocated_memory()
+            if fraction:
+                total = torch.mps.get_mem_info()[0]
         elif self.device.type == "cpu":
-            memory = 0
+            pass
         else:
             memory = torch.cuda.memory_reserved()
-        return memory / (2**30)
+            if fraction:
+                total = torch.cuda.get_device_properties(self.device).total_memory
+        return ((memory / total) if total > 0 else 0) if fraction else (memory / 2**30)
 
     def _clear_memory(self):
         """Clear accelerator memory on different platforms."""

ultralytics/engine/tuner.py

@@ -191,8 +191,9 @@ class Tuner:
             weights_dir = save_dir / "weights"
             try:
                 # Train YOLO model with mutated hyperparameters (run in subprocess to avoid dataloader hang)
-                cmd = ["yolo", "train", *(f"{k}={v}" for k, v in train_args.items())]
-                return_code = subprocess.run(" ".join(cmd), check=True, shell=True).returncode
+                launch = [__import__("sys").executable, "-m", "ultralytics.cfg.__init__"]  # workaround yolo not found
+                cmd = [*launch, "train", *(f"{k}={v}" for k, v in train_args.items())]
+                return_code = subprocess.run(cmd, check=True).returncode
                 ckpt_file = weights_dir / ("best.pt" if (weights_dir / "best.pt").exists() else "last.pt")
                 metrics = torch.load(ckpt_file)["train_metrics"]
                 assert return_code == 0, "training failed"

ultralytics/nn/autobackend.py

@@ -244,7 +244,7 @@ class AutoBackend(nn.Module):
         # OpenVINO
         elif xml:
             LOGGER.info(f"Loading {w} for OpenVINO inference...")
-            check_requirements("openvino>=2024.0.0,<2025.0.0")
+            check_requirements("openvino>=2024.0.0,!=2025.0.0")
             import openvino as ov
 
             core = ov.Core()

ultralytics/nn/tasks.py

@@ -1119,7 +1119,7 @@ def guess_model_scale(model_path):
         (str): The size character of the model's scale, which can be n, s, m, l, or x.
     """
     try:
-        return re.search(r"yolo[v]?\d+([nslmx])", Path(model_path).stem).group(1)  # noqa, returns n, s, m, l, or x
+        return re.search(r"yolo[v]?\d+([nslmx])", Path(model_path).stem).group(1)  # returns n, s, m, l, or x
     except AttributeError:
         return ""
 

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/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/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)]

ultralytics/solutions/distance_calculation.py

@@ -4,8 +4,8 @@ import math
 
 import cv2
 
-from ultralytics.solutions.solutions import BaseSolution
-from ultralytics.utils.plotting import Annotator, colors
+from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
+from ultralytics.utils.plotting import colors
 
 
 class DistanceCalculation(BaseSolution):
@@ -18,22 +18,17 @@ class DistanceCalculation(BaseSolution):
     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.
+        process: 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)
+        >>> results = distance_calc.process(frame)
+        >>> cv2.imshow("Distance Calculation", results.plot_im)
         >>> cv2.waitKey(0)
     """
 
@@ -44,8 +39,7 @@ class DistanceCalculation(BaseSolution):
         # Mouse event information
         self.left_mouse_count = 0
         self.selected_boxes = {}
-
-        self.centroids = []  # Initialize empty list to store centroids
+        self.centroids = []  # Store centroids of selected objects
 
     def mouse_event_for_distance(self, event, x, y, flags, param):
         """
@@ -56,7 +50,7 @@ class DistanceCalculation(BaseSolution):
             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).
-            param (Dict): Additional parameters passed to the function.
+            param (Any): Additional parameters passed to the function.
 
         Examples:
             >>> # Assuming 'dc' is an instance of DistanceCalculation
@@ -73,7 +67,7 @@ class DistanceCalculation(BaseSolution):
             self.selected_boxes = {}
             self.left_mouse_count = 0
 
-    def calculate(self, im0):
+    def process(self, im0):
         """
         Processes a video frame and calculates the distance between two selected bounding boxes.
 
@@ -84,41 +78,47 @@ class DistanceCalculation(BaseSolution):
             im0 (numpy.ndarray): The input image frame to process.
 
         Returns:
-            (numpy.ndarray): The processed image frame with annotations and distance calculations.
+            (SolutionResults): Contains processed image `plot_im`, `total_tracks` (int) representing the total number
+                of tracked objects, and `pixels_distance` (float) representing the distance between selected objects
+                in pixels.
 
         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)
+            >>> results = dc.process(frame)
+            >>> print(f"Distance: {results.pixels_distance:.2f} pixels")
         """
-        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
         self.extract_tracks(im0)  # Extract tracks
+        annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator
 
+        pixels_distance = 0
         # Iterate over bounding boxes, track ids and classes index
         for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
-            self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
+            annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
 
+            # Update selected boxes if they're being tracked
             if len(self.selected_boxes) == 2:
                 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:
-            # Store user selected boxes in centroids list
+            # Calculate centroids of selected boxes
             self.centroids.extend(
                 [[int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)] for box in self.selected_boxes.values()]
             )
-            # Calculate pixels distance
+            # Calculate Euclidean distance between centroids
             pixels_distance = math.sqrt(
                 (self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
             )
-            self.annotator.plot_distance_and_line(pixels_distance, self.centroids)
-
-        self.centroids = []
+            annotator.plot_distance_and_line(pixels_distance, self.centroids)
 
-        self.display_output(im0)  # display output with base class function
+        self.centroids = []  # Reset centroids for next frame
+        plot_im = annotator.result()
+        self.display_output(plot_im)  # Display output with base class function
         cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)
 
-        return im0  # return output image for more usage
+        # Return SolutionResults with processed image and calculated metrics
+        return SolutionResults(plot_im=plot_im, pixels_distance=pixels_distance, total_tracks=len(self.track_ids))

ultralytics/solutions/heatmap.py

@@ -4,7 +4,7 @@ import cv2
 import numpy as np
 
 from ultralytics.solutions.object_counter import ObjectCounter
-from ultralytics.utils.plotting import Annotator
+from ultralytics.solutions.solutions import SolutionAnnotator, SolutionResults
 
 
 class Heatmap(ObjectCounter):
@@ -18,28 +18,33 @@ class Heatmap(ObjectCounter):
         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.
+        annotator (SolutionAnnotator): 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.
+        process: Generates and applies the heatmap effect to each frame.
 
     Examples:
         >>> from ultralytics.solutions import Heatmap
         >>> heatmap = Heatmap(model="yolo11n.pt", colormap=cv2.COLORMAP_JET)
         >>> frame = cv2.imread("frame.jpg")
-        >>> processed_frame = heatmap.generate_heatmap(frame)
+        >>> processed_frame = heatmap.process(frame)
     """
 
     def __init__(self, **kwargs):
-        """Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks."""
+        """
+        Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed to the parent ObjectCounter class.
+        """
         super().__init__(**kwargs)
 
-        self.initialized = False  # bool variable for heatmap initialization
-        if self.region is not None:  # check if user provided the region coordinates
+        self.initialized = False  # Flag for heatmap initialization
+        if self.region is not None:  # Check if user provided the region coordinates
             self.initialize_region()
 
-        # store colormap
+        # Store colormap
         self.colormap = cv2.COLORMAP_PARULA if self.CFG["colormap"] is None else self.CFG["colormap"]
         self.heatmap = None
 
@@ -49,11 +54,6 @@ class Heatmap(ObjectCounter):
 
         Args:
             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
@@ -70,7 +70,7 @@ class Heatmap(ObjectCounter):
         # Update only the values within the bounding box in a single vectorized operation
         self.heatmap[y0:y1, x0:x1][within_radius] += 2
 
-    def generate_heatmap(self, im0):
+    def process(self, im0):
         """
         Generate heatmap for each frame using Ultralytics.
 
@@ -78,50 +78,52 @@ class Heatmap(ObjectCounter):
             im0 (np.ndarray): Input image array for processing.
 
         Returns:
-            (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)
+            (SolutionResults): Contains processed image `plot_im`,
+                'in_count' (int, count of objects entering the region),
+                'out_count' (int, count of objects exiting the region),
+                'classwise_count' (dict, per-class object count), and
+                'total_tracks' (int, total number of tracked objects).
         """
         if not self.initialized:
             self.heatmap = np.zeros_like(im0, dtype=np.float32) * 0.99
-        self.initialized = True  # Initialize heatmap only once
+            self.initialized = True  # Initialize heatmap only once
 
-        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
         self.extract_tracks(im0)  # Extract tracks
+        self.annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator
 
         # Iterate over bounding boxes, track ids and classes index
         for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
-            # Draw bounding box and counting region
+            # Apply heatmap effect for the bounding box
             self.heatmap_effect(box)
 
             if self.region is not None:
                 self.annotator.draw_region(reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2)
                 self.store_tracking_history(track_id, box)  # Store track history
-                self.store_classwise_counts(cls)  # store classwise counts in dict
+                self.store_classwise_counts(cls)  # Store classwise counts in dict
                 current_centroid = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
-                # Store tracking previous position and perform object counting
+                # Get previous position if available
                 prev_position = None
                 if len(self.track_history[track_id]) > 1:
                     prev_position = self.track_history[track_id][-2]
                 self.count_objects(current_centroid, track_id, prev_position, cls)  # Perform object counting
 
+        plot_im = self.annotator.result()
         if self.region is not None:
-            self.display_counts(im0)  # Display the counts on the frame
+            self.display_counts(plot_im)  # Display the counts on the frame
 
         # Normalize, apply colormap to heatmap and combine with original image
         if self.track_data.id is not None:
-            im0 = cv2.addWeighted(
-                im0,
-                0.5,
-                cv2.applyColorMap(
-                    cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8), self.colormap
-                ),
-                0.5,
-                0,
-            )
-
-        self.display_output(im0)  # display output with base class function
-        return im0  # return output image for more usage
+            normalized_heatmap = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+            colored_heatmap = cv2.applyColorMap(normalized_heatmap, self.colormap)
+            plot_im = cv2.addWeighted(plot_im, 0.5, colored_heatmap, 0.5, 0)
+
+        self.display_output(plot_im)  # Display output with base class function
+
+        # Return SolutionResults
+        return SolutionResults(
+            plot_im=plot_im,
+            in_count=self.in_count,
+            out_count=self.out_count,
+            classwise_count=self.classwise_counts,
+            total_tracks=len(self.track_ids),
+        )