Update distance-calculation solution (#16907)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>

docs/en/guides/distance-calculation.md

@@ -43,12 +43,9 @@ Measuring the gap between two objects is known as distance calculation within a
         ```python
         import cv2
 
-        from ultralytics import YOLO, solutions
+        from ultralytics import solutions
 
-        model = YOLO("yolo11n.pt")
+        cap = cv2.VideoCapture("Path/to/video/file.mp4")
-        names = model.model.names
-
-        cap = cv2.VideoCapture("path/to/video/file.mp4")
         assert cap.isOpened(), "Error reading video file"
         w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
 
@@ -56,16 +53,14 @@ Measuring the gap between two objects is known as distance calculation within a
         video_writer = cv2.VideoWriter("distance_calculation.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
 
         # Init distance-calculation obj
-        dist_obj = solutions.DistanceCalculation(names=names, view_img=True)
+        distance = solutions.DistanceCalculation(model="yolo11n.pt", show=True)
 
         while cap.isOpened():
             success, im0 = cap.read()
             if not success:
                 print("Video frame is empty or video processing has been successfully completed.")
                 break
-
+            im0 = distance.calculate(im0)
-            tracks = model.track(im0, persist=True, show=False)
-            im0 = dist_obj.start_process(im0, tracks)
             video_writer.write(im0)
 
         cap.release()
@@ -84,13 +79,11 @@ Measuring the gap between two objects is known as distance calculation within a
 
 ### Arguments `DistanceCalculation()`
 
-| `Name`           | `Type`  | `Default`       | Description                                               |
+| `Name`       | `Type` | `Default` | Description                                          |
-| ---------------- | ------- | --------------- | --------------------------------------------------------- |
+| ------------ | ------ | --------- | ---------------------------------------------------- |
-| `names`          | `dict`  | `None`          | Dictionary of classes names.                              |
+| `model`      | `str`  | `None`    | Path to Ultralytics YOLO Model File                  |
-| `view_img`       | `bool`  | `False`         | Flag to indicate if the video stream should be displayed. |
+| `line_width` | `int`  | `2`       | Line thickness for bounding boxes.                   |
-| `line_thickness` | `int`   | `2`             | Thickness of the lines drawn on the image.                |
+| `show`       | `bool` | `False`   | Flag to control whether to display the video stream. |
-| `line_color`     | `tuple` | `(255, 255, 0)` | Color of the lines drawn on the image (BGR format).       |
-| `centroid_color` | `tuple` | `(255, 0, 255)` | Color of the centroids drawn (BGR format).                |
 
 ### Arguments `model.track`
 
@@ -122,10 +115,8 @@ To delete points drawn during distance calculation with Ultralytics YOLO11, you
 
 The key arguments for initializing the `DistanceCalculation` class in Ultralytics YOLO11 include:
 
-- `names`: Dictionary mapping class indices to class names.
+- `model`: Model file path.
-- `view_img`: Flag to indicate if the video stream should be displayed.
+- `show`: Flag to indicate if the video stream should be displayed.
-- `line_thickness`: Thickness of the lines drawn on the image.
+- `line_width`: Thickness of bounding box and the lines drawn on the image.
-- `line_color`: Color of the lines drawn on the image (BGR format).
-- `centroid_color`: Color of the centroids (BGR format).
 
 For an exhaustive list and default values, see the [arguments of DistanceCalculation](#arguments-distancecalculation).

docs/en/guides/heatmaps.md

@@ -222,6 +222,7 @@ A heatmap generated with [Ultralytics YOLO11](https://github.com/ultralytics/ult
 
 | Name         | Type   | Default            | Description                                                       |
 | ------------ | ------ | ------------------ | ----------------------------------------------------------------- |
+| `model`      | `str`  | `None`             | Path to Ultralytics YOLO Model File                               |
 | `colormap`   | `int`  | `cv2.COLORMAP_JET` | Colormap to use for the heatmap.                                  |
 | `show`       | `bool` | `False`            | Whether to display the image with the heatmap overlay.            |
 | `show_in`    | `bool` | `True`             | Whether to display the count of objects entering the region.      |

docs/en/guides/workouts-monitoring.md

@@ -106,6 +106,7 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
 | `show`       | `bool`  | `False` | Flag to display the image.                                                             |
 | `up_angle`   | `float` | `145.0` | Angle threshold for the 'up' pose.                                                     |
 | `down_angle` | `float` | `90.0`  | Angle threshold for the 'down' pose.                                                   |
+| `model`      | `str`   | `None`  | Path to Ultralytics YOLO Pose Model File                                               |
 
 ### Arguments `model.predict`
 

ultralytics/solutions/distance_calculation.py

@@ -4,55 +4,21 @@ import math
 
 import cv2
 
-from ultralytics.utils.checks import check_imshow
+from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
 from ultralytics.utils.plotting import Annotator, colors
 
 
-class DistanceCalculation:
+class DistanceCalculation(BaseSolution):
     """A class to calculate distance between two objects in a real-time video stream based on their tracks."""
 
-    def __init__(
+    def __init__(self, **kwargs):
-        self,
+        """Initializes the DistanceCalculation class with the given parameters."""
-        names,
+        super().__init__(**kwargs)
-        view_img=False,
-        line_thickness=2,
-        line_color=(255, 0, 255),
-        centroid_color=(104, 31, 17),
-    ):
-        """
-        Initializes the DistanceCalculation class with the given parameters.
-
-        Args:
-            names (dict): Dictionary of classes names.
-            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 = view_img
-        self.line_color = line_color
-        self.centroid_color = centroid_color
-
-        # Prediction & tracking information
-        self.names = names
-        self.boxes = None
-        self.line_thickness = line_thickness
-        self.trk_ids = None
-
-        # Distance calculation information
-        self.centroids = []
 
         # Mouse event information
         self.left_mouse_count = 0
         self.selected_boxes = {}
 
-        # Check if environment supports imshow
-        self.env_check = check_imshow(warn=True)
-        self.window_name = "Ultralytics Solutions"
-
     def mouse_event_for_distance(self, event, x, y, flags, param):
         """
         Handles mouse events to select regions in a real-time video stream.
@@ -67,7 +33,7 @@ class DistanceCalculation:
         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):
+                for box, track_id in zip(self.boxes, self.track_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] = box
 
@@ -75,30 +41,21 @@ class DistanceCalculation:
             self.selected_boxes = {}
             self.left_mouse_count = 0
 
-    def start_process(self, im0, tracks):
+    def calculate(self, im0):
         """
         Processes the video frame and calculates the distance between two bounding boxes.
 
         Args:
             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
+        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
-        if tracks[0].boxes.id is None:
+        self.extract_tracks(im0)  # Extract tracks
-            if self.view_img:
-                self.display_frames()
-            return im0
 
-        self.boxes = tracks[0].boxes.xyxy.cpu()
+        # Iterate over bounding boxes, track ids and classes index
-        clss = tracks[0].boxes.cls.cpu().tolist()
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
-        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
-
-        self.annotator = Annotator(self.im0, line_width=self.line_thickness)
-
-        for box, cls, track_id in zip(self.boxes, 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:
@@ -115,25 +72,11 @@ class DistanceCalculation:
             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.line_color, self.centroid_color)
+            self.annotator.plot_distance_and_line(pixels_distance, self.centroids)
 
         self.centroids = []
 
-        if self.view_img and self.env_check:
+        self.display_output(im0)  # display output with base class function
-            self.display_frames()
+        cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)
 
-        return im0
+        return im0  # return output image for more usage
-
-    def display_frames(self):
-        """Displays the current frame with annotations."""
-        cv2.namedWindow(self.window_name)
-        cv2.setMouseCallback(self.window_name, self.mouse_event_for_distance)
-        cv2.imshow(self.window_name, self.im0)
-
-        if cv2.waitKey(1) & 0xFF == ord("q"):
-            return
-
-
-if __name__ == "__main__":
-    names = {0: "person", 1: "car"}  # example class names
-    distance_calculation = DistanceCalculation(names)

ultralytics/solutions/object_counter.py

@@ -112,13 +112,13 @@ class ObjectCounter(BaseSolution):
         # 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
-            self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
+            self.annotator.box_label(box, label=self.names[cls], color=colors(cls, True))
             self.store_tracking_history(track_id, box)  # Store track history
             self.store_classwise_counts(cls)  # store classwise counts in dict
 
             # Draw tracks of objects
             self.annotator.draw_centroid_and_tracks(
-                self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
+                self.track_line, color=colors(int(cls), True), track_thickness=self.line_width
             )
 
             # store previous position of track for object counting

ultralytics/utils/plotting.py

@@ -804,31 +804,30 @@ class Annotator:
                 self.im, label, (int(mask[0][0]) - text_size[0] // 2, int(mask[0][1])), 0, self.sf, txt_color, self.tf
             )
 
-    def plot_distance_and_line(self, pixels_distance, centroids, line_color, centroid_color):
+    def plot_distance_and_line(
+        self, pixels_distance, centroids, line_color=(104, 31, 17), centroid_color=(255, 0, 255)
+    ):
         """
         Plot the distance and line on frame.
 
         Args:
             pixels_distance (float): Pixels distance between two bbox centroids.
             centroids (list): Bounding box centroids data.
-            line_color (tuple): RGB distance line color.
+            line_color (tuple, optional): Distance line color.
-            centroid_color (tuple): RGB bounding box centroid color.
+            centroid_color (tuple, optional): Bounding box centroid color.
         """
         # Get the text size
-        (text_width_m, text_height_m), _ = cv2.getTextSize(
+        text = f"Pixels Distance: {pixels_distance:.2f}"
-            f"Pixels Distance: {pixels_distance:.2f}", 0, self.sf, self.tf
+        (text_width_m, text_height_m), _ = cv2.getTextSize(text, 0, self.sf, self.tf)
-        )
 
         # Define corners with 10-pixel margin and draw rectangle
-        top_left = (15, 25)
+        cv2.rectangle(self.im, (15, 25), (15 + text_width_m + 20, 25 + text_height_m + 20), line_color, -1)
-        bottom_right = (15 + text_width_m + 20, 25 + text_height_m + 20)
-        cv2.rectangle(self.im, top_left, bottom_right, centroid_color, -1)
 
         # Calculate the position for the text with a 10-pixel margin and draw text
-        text_position = (top_left[0] + 10, top_left[1] + text_height_m + 10)
+        text_position = (25, 25 + text_height_m + 10)
         cv2.putText(
             self.im,
-            f"Pixels Distance: {pixels_distance:.2f}",
+            text,
             text_position,
             0,
             self.sf,