From 8e3846d377ad36c3a1f79d024231544bdd39e237 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Mon, 14 Oct 2024 19:44:36 +0500 Subject: [PATCH] Update `distance-calculation` solution (#16907) Co-authored-by: UltralyticsAssistant --- docs/en/guides/distance-calculation.md | 33 +++---- docs/en/guides/heatmaps.md | 1 + docs/en/guides/workouts-monitoring.md | 1 + ultralytics/solutions/distance_calculation.py | 87 ++++--------------- ultralytics/solutions/object_counter.py | 4 +- ultralytics/utils/plotting.py | 21 +++-- 6 files changed, 41 insertions(+), 106 deletions(-) diff --git a/docs/en/guides/distance-calculation.md b/docs/en/guides/distance-calculation.md index b0b12f919b..009899ae3c 100644 --- a/docs/en/guides/distance-calculation.md +++ b/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") - names = model.model.names - - cap = cv2.VideoCapture("path/to/video/file.mp4") + 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 - - tracks = model.track(im0, persist=True, show=False) - im0 = dist_obj.start_process(im0, tracks) + im0 = distance.calculate(im0) 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 | -| ---------------- | ------- | --------------- | --------------------------------------------------------- | -| `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). | +| `Name` | `Type` | `Default` | Description | +| ------------ | ------ | --------- | ---------------------------------------------------- | +| `model` | `str` | `None` | Path to Ultralytics YOLO Model File | +| `line_width` | `int` | `2` | Line thickness for bounding boxes. | +| `show` | `bool` | `False` | Flag to control whether to display the video stream. | ### 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. -- `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). -- `centroid_color`: Color of the centroids (BGR format). +- `model`: Model file path. +- `show`: Flag to indicate if the video stream should be displayed. +- `line_width`: Thickness of bounding box and the lines drawn on the image. For an exhaustive list and default values, see the [arguments of DistanceCalculation](#arguments-distancecalculation). diff --git a/docs/en/guides/heatmaps.md b/docs/en/guides/heatmaps.md index f33993134f..7919bc7d94 100644 --- a/docs/en/guides/heatmaps.md +++ b/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. | diff --git a/docs/en/guides/workouts-monitoring.md b/docs/en/guides/workouts-monitoring.md index 78d894e81d..34056da3fb 100644 --- a/docs/en/guides/workouts-monitoring.md +++ b/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` diff --git a/ultralytics/solutions/distance_calculation.py b/ultralytics/solutions/distance_calculation.py index dccd1687c6..773b6086da 100644 --- a/ultralytics/solutions/distance_calculation.py +++ b/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__( - self, - names, - 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 = [] + def __init__(self, **kwargs): + """Initializes the DistanceCalculation class with the given parameters.""" + super().__init__(**kwargs) # 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 - if tracks[0].boxes.id is None: - if self.view_img: - self.display_frames() - return im0 + self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator + self.extract_tracks(im0) # Extract 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, clss, self.trk_ids): + # 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)]) 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_frames() - - return im0 - - 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 - + self.display_output(im0) # display output with base class function + cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance) -if __name__ == "__main__": - names = {0: "person", 1: "car"} # example class names - distance_calculation = DistanceCalculation(names) + return im0 # return output image for more usage diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index 7d9bb8c9f4..d576746421 100644 --- a/ultralytics/solutions/object_counter.py +++ b/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 diff --git a/ultralytics/utils/plotting.py b/ultralytics/utils/plotting.py index 1fda09d917..8295f77f2f 100644 --- a/ultralytics/utils/plotting.py +++ b/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. - centroid_color (tuple): RGB bounding box centroid color. + line_color (tuple, optional): Distance line color. + centroid_color (tuple, optional): Bounding box centroid color. """ # Get the text size - (text_width_m, text_height_m), _ = cv2.getTextSize( - f"Pixels Distance: {pixels_distance:.2f}", 0, self.sf, self.tf - ) + text = f"Pixels Distance: {pixels_distance:.2f}" + (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) - bottom_right = (15 + text_width_m + 20, 25 + text_height_m + 20) - cv2.rectangle(self.im, top_left, bottom_right, centroid_color, -1) + cv2.rectangle(self.im, (15, 25), (15 + text_width_m + 20, 25 + text_height_m + 20), line_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,