diff --git a/docs/en/guides/distance-calculation.md b/docs/en/guides/distance-calculation.md index dbc15a4a2a..761d4ea9cf 100644 --- a/docs/en/guides/distance-calculation.md +++ b/docs/en/guides/distance-calculation.md @@ -30,8 +30,7 @@ Measuring the gap between two objects is known as distance calculation within a ## Advantages of Distance Calculation? - **Localization Precision:** Enhances accurate spatial positioning in computer vision tasks. -- **Size Estimation:** Allows estimation of physical sizes for better contextual understanding. -- **Scene Understanding:** Contributes to a 3D understanding of the environment for improved decision-making. +- **Size Estimation:** Allows estimation of object size for better contextual understanding. ???+ tip "Distance Calculation" @@ -85,14 +84,13 @@ Measuring the gap between two objects is known as distance calculation within a ### Arguments `DistanceCalculation()` -| `Name` | `Type` | `Default` | Description | -| ------------------ | ------- | --------------- | --------------------------------------------------------- | -| `names` | `dict` | `None` | Dictionary of classes names. | -| `pixels_per_meter` | `int` | `10` | Conversion factor from pixels to meters. | -| `view_img` | `bool` | `False` | Flag to indicate if the video stream should be displayed. | -| `line_thickness` | `int` | `2` | Thickness of the lines drawn on the image. | -| `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). | +| `Name` | `Type` | `Default` | Description | +| ---------------- | ------- | --------------- | --------------------------------------------------------- | +| `names` | `dict` | `None` | Dictionary of classes names. | +| `view_img` | `bool` | `False` | Flag to indicate if the video stream should be displayed. | +| `line_thickness` | `int` | `2` | Thickness of the lines drawn on the image. | +| `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` @@ -133,7 +131,6 @@ To delete points drawn during distance calculation with Ultralytics YOLOv8, you The key arguments for initializing the `DistanceCalculation` class in Ultralytics YOLOv8 include: - `names`: Dictionary mapping class indices to class names. -- `pixels_per_meter`: Conversion factor from pixels to meters. - `view_img`: Flag to indicate if the video stream should be displayed. - `line_thickness`: Thickness of the lines drawn on the image. - `line_color`: Color of the lines drawn on the image (BGR format). diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 37925ac9e7..c5d3a82d51 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license -__version__ = "8.2.87" +__version__ = "8.2.88" import os diff --git a/ultralytics/solutions/distance_calculation.py b/ultralytics/solutions/distance_calculation.py index 35a791ca5f..dccd1687c6 100644 --- a/ultralytics/solutions/distance_calculation.py +++ b/ultralytics/solutions/distance_calculation.py @@ -14,18 +14,16 @@ class DistanceCalculation: def __init__( self, names, - pixels_per_meter=10, view_img=False, line_thickness=2, - line_color=(255, 255, 0), - centroid_color=(255, 0, 255), + 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. - pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10. view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False. line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2. line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0). @@ -39,7 +37,6 @@ class DistanceCalculation: self.centroid_color = centroid_color # Prediction & tracking information - self.clss = None self.names = names self.boxes = None self.line_thickness = line_thickness @@ -47,7 +44,6 @@ class DistanceCalculation: # Distance calculation information self.centroids = [] - self.pixel_per_meter = pixels_per_meter # Mouse event information self.left_mouse_count = 0 @@ -55,6 +51,7 @@ class DistanceCalculation: # 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): """ @@ -78,46 +75,6 @@ class DistanceCalculation: self.selected_boxes = {} self.left_mouse_count = 0 - def extract_tracks(self, tracks): - """ - Extracts tracking results from the provided data. - - Args: - tracks (list): List of tracks obtained from the object tracking process. - """ - self.boxes = tracks[0].boxes.xyxy.cpu() - self.clss = tracks[0].boxes.cls.cpu().tolist() - self.trk_ids = tracks[0].boxes.id.int().cpu().tolist() - - @staticmethod - def calculate_centroid(box): - """ - Calculates the centroid of a bounding box. - - Args: - box (list): Bounding box coordinates [x1, y1, x2, y2]. - - Returns: - (tuple): Centroid coordinates (x, y). - """ - return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2) - - def calculate_distance(self, centroid1, centroid2): - """ - Calculates the distance between two centroids. - - Args: - centroid1 (tuple): Coordinates of the first centroid (x, y). - centroid2 (tuple): Coordinates of the second centroid (x, y). - - Returns: - (tuple): Distance in meters and millimeters. - """ - pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2) - distance_m = pixel_distance / self.pixel_per_meter - distance_mm = distance_m * 1000 - return distance_m, distance_mm - def start_process(self, im0, tracks): """ Processes the video frame and calculates the distance between two bounding boxes. @@ -135,10 +92,13 @@ class DistanceCalculation: self.display_frames() return im0 - self.extract_tracks(tracks) + self.boxes = tracks[0].boxes.xyxy.cpu() + clss = tracks[0].boxes.cls.cpu().tolist() + 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, self.clss, self.trk_ids): + 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: @@ -147,12 +107,15 @@ class DistanceCalculation: self.selected_boxes[track_id] = box if len(self.selected_boxes) == 2: - self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes] - - distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1]) - self.annotator.plot_distance_and_line( - distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color + # Store user selected boxes in centroids list + 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 + 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.centroids = [] @@ -163,9 +126,9 @@ class DistanceCalculation: def display_frames(self): """Displays the current frame with annotations.""" - cv2.namedWindow("Ultralytics Distance Estimation") - cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance) - cv2.imshow("Ultralytics Distance Estimation", self.im0) + 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 diff --git a/ultralytics/utils/plotting.py b/ultralytics/utils/plotting.py index 5fc7d0e741..dfaa21490a 100644 --- a/ultralytics/utils/plotting.py +++ b/ultralytics/utils/plotting.py @@ -756,39 +756,35 @@ 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, distance_m, distance_mm, centroids, line_color, centroid_color): + def plot_distance_and_line(self, pixels_distance, centroids, line_color, centroid_color): """ Plot the distance and line on frame. Args: - distance_m (float): Distance between two bbox centroids in meters. - distance_mm (float): Distance between two bbox centroids in millimeters. + pixels_distance (float): Pixels distance between two bbox centroids. centroids (list): Bounding box centroids data. line_color (RGB): Distance line color. centroid_color (RGB): Bounding box centroid color. """ - (text_width_m, text_height_m), _ = cv2.getTextSize(f"Distance M: {distance_m:.2f}m", 0, self.sf, self.tf) - cv2.rectangle(self.im, (15, 25), (15 + text_width_m + 10, 25 + text_height_m + 20), line_color, -1) - cv2.putText( - self.im, - f"Distance M: {distance_m:.2f}m", - (20, 50), - 0, - self.sf, - centroid_color, - self.tf, - cv2.LINE_AA, + # Get the text size + (text_width_m, text_height_m), _ = cv2.getTextSize( + f"Pixels Distance: {pixels_distance:.2f}", 0, self.sf, self.tf ) - (text_width_mm, text_height_mm), _ = cv2.getTextSize(f"Distance MM: {distance_mm:.2f}mm", 0, self.sf, self.tf) - cv2.rectangle(self.im, (15, 75), (15 + text_width_mm + 10, 75 + text_height_mm + 20), line_color, -1) + # Define corners with 10-pixel margin and draw rectangle + top_left = (15, 25) + 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) cv2.putText( self.im, - f"Distance MM: {distance_mm:.2f}mm", - (20, 100), + f"Pixels Distance: {pixels_distance:.2f}", + text_position, 0, self.sf, - centroid_color, + (255, 255, 255), self.tf, cv2.LINE_AA, )