From 65097578797e223ffab1735a18074500b403e45c Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Tue, 8 Oct 2024 23:57:01 +0500 Subject: [PATCH] Update `queue-management` solution (#16772) Co-authored-by: UltralyticsAssistant --- docs/en/guides/queue-management.md | 59 ++++----- tests/test_solutions.py | 4 +- ultralytics/solutions/object_counter.py | 2 +- ultralytics/solutions/queue_management.py | 149 +++++++--------------- ultralytics/solutions/solutions.py | 8 +- 5 files changed, 75 insertions(+), 147 deletions(-) diff --git a/docs/en/guides/queue-management.md b/docs/en/guides/queue-management.md index 8f6610bc9e..32cb5b8a4e 100644 --- a/docs/en/guides/queue-management.md +++ b/docs/en/guides/queue-management.md @@ -40,10 +40,9 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra ```python import cv2 - from ultralytics import YOLO, solutions + from ultralytics import solutions - model = YOLO("yolo11n.pt") - 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)) @@ -53,18 +52,15 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)] queue = solutions.QueueManager( - names=model.names, - reg_pts=queue_region, - line_thickness=3, + model="yolo11n.pt", + region=queue_region, ) while cap.isOpened(): success, im0 = cap.read() if success: - tracks = model.track(im0, persist=True) - out = queue.process_queue(im0, tracks) - + out = queue.process_queue(im0) video_writer.write(im0) if cv2.waitKey(1) & 0xFF == ord("q"): break @@ -82,10 +78,9 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra ```python import cv2 - from ultralytics import YOLO, solutions + from ultralytics import solutions - model = YOLO("yolo11n.pt") - 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)) @@ -95,18 +90,15 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)] queue = solutions.QueueManager( - names=model.names, - reg_pts=queue_region, - line_thickness=3, + model="yolo11n.pt", + classes=3, ) while cap.isOpened(): success, im0 = cap.read() if success: - tracks = model.track(im0, persist=True, classes=0) # Only person class - out = queue.process_queue(im0, tracks) - + out = queue.process_queue(im0) video_writer.write(im0) if cv2.waitKey(1) & 0xFF == ord("q"): break @@ -121,13 +113,12 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra ### Arguments `QueueManager` -| Name | Type | Default | Description | -| ---------------- | ---------------- | -------------------------- | -------------------------------------------------------------------------------- | -| `names` | `dict` | `model.names` | A dictionary mapping class IDs to class names. | -| `reg_pts` | `list of tuples` | `[(20, 400), (1260, 400)]` | Points defining the counting region polygon. Defaults to a predefined rectangle. | -| `line_thickness` | `int` | `2` | Thickness of the annotation lines. | -| `view_img` | `bool` | `False` | Whether to display the image frames. | -| `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. | +| Name | Type | Default | Description | +| ------------ | ------ | -------------------------- | ---------------------------------------------------- | +| `model` | `str` | `None` | Path to Ultralytics YOLO Model File | +| `region` | `list` | `[(20, 400), (1260, 400)]` | List of points defining the queue region. | +| `line_width` | `int` | `2` | Line thickness for bounding boxes. | +| `show` | `bool` | `False` | Flag to control whether to display the video stream. | ### Arguments `model.track` @@ -149,23 +140,21 @@ Here's a minimal example: ```python import cv2 -from ultralytics import YOLO, solutions +from ultralytics import solutions -model = YOLO("yolo11n.pt") cap = cv2.VideoCapture("path/to/video.mp4") queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)] queue = solutions.QueueManager( - names=model.names, - reg_pts=queue_region, - line_thickness=3, + model="yolo11n.pt", + region=queue_region, + line_width=3, ) while cap.isOpened(): success, im0 = cap.read() if success: - tracks = model.track(im0, show=False, persist=True, verbose=False) - out = queue.process_queue(im0, tracks) + out = queue.process_queue(im0) cv2.imshow("Queue Management", im0) if cv2.waitKey(1) & 0xFF == ord("q"): break @@ -207,9 +196,9 @@ Example for airports: ```python queue_region_airport = [(50, 600), (1200, 600), (1200, 550), (50, 550)] queue_airport = solutions.QueueManager( - names=model.names, - reg_pts=queue_region_airport, - line_thickness=3, + model="yolo11n.pt", + region=queue_region_airport, + line_width=3, ) ``` diff --git a/tests/test_solutions.py b/tests/test_solutions.py index 8ec68d260f..ef9ffe7d11 100644 --- a/tests/test_solutions.py +++ b/tests/test_solutions.py @@ -22,7 +22,7 @@ def test_major_solutions(): counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False) heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False) speed = solutions.SpeedEstimator(reg_pts=region_points, names=names, view_img=False) - queue = solutions.QueueManager(names=names, reg_pts=region_points, view_img=False) + queue = solutions.QueueManager(region=region_points, model="yolo11n.pt", show=False) while cap.isOpened(): success, im0 = cap.read() if not success: @@ -32,7 +32,7 @@ def test_major_solutions(): _ = counter.count(original_im0.copy()) _ = heatmap.generate_heatmap(original_im0.copy()) _ = speed.estimate_speed(original_im0.copy(), tracks) - _ = queue.process_queue(original_im0.copy(), tracks) + _ = queue.process_queue(original_im0.copy()) cap.release() cv2.destroyAllWindows() diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index 5fdaef258a..7d9bb8c9f4 100644 --- a/ultralytics/solutions/object_counter.py +++ b/ultralytics/solutions/object_counter.py @@ -116,7 +116,7 @@ class ObjectCounter(BaseSolution): self.store_tracking_history(track_id, box) # Store track history self.store_classwise_counts(cls) # store classwise counts in dict - # Draw centroid of objects + # Draw tracks of objects self.annotator.draw_centroid_and_tracks( self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width ) diff --git a/ultralytics/solutions/queue_management.py b/ultralytics/solutions/queue_management.py index ef60150395..287f337dc5 100644 --- a/ultralytics/solutions/queue_management.py +++ b/ultralytics/solutions/queue_management.py @@ -1,127 +1,64 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license -from collections import defaultdict +from shapely.geometry import Point -import cv2 - -from ultralytics.utils.checks import check_imshow, check_requirements +from ultralytics.solutions.solutions import BaseSolution # Import a parent class from ultralytics.utils.plotting import Annotator, colors -check_requirements("shapely>=2.0.0") - -from shapely.geometry import Point, Polygon - -class QueueManager: +class QueueManager(BaseSolution): """A class to manage the queue in a real-time video stream based on object tracks.""" - def __init__( - self, - names, - reg_pts=None, - line_thickness=2, - view_img=False, - draw_tracks=False, - ): + def __init__(self, **kwargs): + """Initializes the QueueManager with specified parameters for tracking and counting objects.""" + super().__init__(**kwargs) + self.initialize_region() + self.counts = 0 # Queue counts Information + self.rect_color = (255, 255, 255) # Rectangle color + self.region_length = len(self.region) # Store region length for further usage + + def process_queue(self, im0): """ - Initializes the QueueManager with specified parameters for tracking and counting objects. + Main function to start the queue management process. Args: - names (dict): A dictionary mapping class IDs to class names. - reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined - rectangle. - line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2. - view_img (bool, optional): Whether to display the image frames. Defaults to False. - draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False. + im0 (ndarray): The input image that will be used for processing + Returns + im0 (ndarray): The processed image for more usage """ - # Region & Line Information - self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)] - self.counting_region = ( - Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)]) - ) - - # annotation Information - self.tf = line_thickness - self.view_img = view_img - - self.names = names # Class names - - # Object counting Information - self.counts = 0 - - # Tracks info - self.track_history = defaultdict(list) - self.draw_tracks = draw_tracks - - # Check if environment supports imshow - self.env_check = check_imshow(warn=True) - - def extract_and_process_tracks(self, tracks, im0): - """Extracts and processes tracks for queue management in a video stream.""" - # Initialize annotator and draw the queue region - annotator = Annotator(im0, self.tf, self.names) self.counts = 0 # Reset counts every frame - if tracks[0].boxes.id is not None: - boxes = tracks[0].boxes.xyxy.cpu() - clss = tracks[0].boxes.cls.cpu().tolist() - track_ids = tracks[0].boxes.id.int().cpu().tolist() + self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator + self.extract_tracks(im0) # Extract tracks - # Extract tracks - for box, track_id, cls in zip(boxes, track_ids, clss): - # Draw bounding box - annotator.box_label(box, label=self.names[cls], color=colors(int(track_id), True)) + self.annotator.draw_region( + reg_pts=self.region, color=self.rect_color, thickness=self.line_width * 2 + ) # Draw region - # Update track history - track_line = self.track_history[track_id] - track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))) - if len(track_line) > 30: - track_line.pop(0) + 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.store_tracking_history(track_id, box) # Store track history - # Draw track trails if enabled - if self.draw_tracks: - annotator.draw_centroid_and_tracks( - track_line, - color=colors(int(track_id), True), - track_thickness=self.line_thickness, - ) - - prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None - - # Check if the object is inside the counting region - if len(self.reg_pts) >= 3: - is_inside = self.counting_region.contains(Point(track_line[-1])) - if prev_position is not None and is_inside: - self.counts += 1 - - # Display queue counts - label = f"Queue Counts : {str(self.counts)}" - if label is not None: - annotator.queue_counts_display( - label, - points=self.reg_pts, - region_color=(255, 0, 255), - txt_color=(104, 31, 17), + # Draw tracks of objects + self.annotator.draw_centroid_and_tracks( + self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width ) - if self.env_check and self.view_img: - annotator.draw_region(reg_pts=self.reg_pts, thickness=self.tf * 2, color=(255, 0, 255)) - cv2.imshow("Ultralytics YOLOv8 Queue Manager", im0) - # Close window on 'q' key press - if cv2.waitKey(1) & 0xFF == ord("q"): - return + # Cache frequently accessed attributes + track_history = self.track_history.get(track_id, []) - def process_queue(self, im0, tracks): - """ - Main function to start the queue management process. - - Args: - im0 (ndarray): Current frame from the video stream. - tracks (list): List of tracks obtained from the object tracking process. - """ - self.extract_and_process_tracks(tracks, im0) # Extract and process tracks - return im0 + # store previous position of track and check if the object is inside the counting region + prev_position = track_history[-2] if len(track_history) > 1 else None + if self.region_length >= 3 and prev_position and self.r_s.contains(Point(self.track_line[-1])): + self.counts += 1 + # Display queue counts + self.annotator.queue_counts_display( + f"Queue Counts : {str(self.counts)}", + points=self.region, + region_color=self.rect_color, + txt_color=(104, 31, 17), + ) + self.display_output(im0) # display output with base class function -if __name__ == "__main__": - classes_names = {0: "person", 1: "car"} # example class names - queue_manager = QueueManager(classes_names) + return im0 # return output image for more usage diff --git a/ultralytics/solutions/solutions.py b/ultralytics/solutions/solutions.py index 14b3f80a52..71a92becfd 100644 --- a/ultralytics/solutions/solutions.py +++ b/ultralytics/solutions/solutions.py @@ -76,9 +76,11 @@ class BaseSolution: def initialize_region(self): """Initialize the counting region and line segment based on config.""" - self.region = [(20, 400), (1260, 400)] if self.region is None else self.region - self.r_s = Polygon(self.region) if len(self.region) >= 3 else LineString(self.region) - self.l_s = LineString([(self.region[0][0], self.region[0][1]), (self.region[1][0], self.region[1][1])]) + self.region = [(20, 400), (1080, 404), (1080, 360), (20, 360)] if self.region is None else self.region + self.r_s = Polygon(self.region) if len(self.region) >= 3 else LineString(self.region) # region segment + self.l_s = LineString( + [(self.region[0][0], self.region[0][1]), (self.region[1][0], self.region[1][1])] + ) # line segment def display_output(self, im0): """