From 5f1f39d36c6f5fb1f0ecc3fd28868bcd784399d1 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Fri, 13 Sep 2024 22:36:23 +0500 Subject: [PATCH] Optimize `queue` solution (#16246) Co-authored-by: UltralyticsAssistant --- docs/en/guides/queue-management.md | 33 ++++-------- ultralytics/solutions/queue_management.py | 66 +++++------------------ 2 files changed, 23 insertions(+), 76 deletions(-) diff --git a/docs/en/guides/queue-management.md b/docs/en/guides/queue-management.md index bf9beba4fa..993d414102 100644 --- a/docs/en/guides/queue-management.md +++ b/docs/en/guides/queue-management.md @@ -56,15 +56,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra names=model.names, reg_pts=queue_region, line_thickness=3, - fontsize=1.0, - region_color=(255, 144, 31), ) while cap.isOpened(): success, im0 = cap.read() if success: - tracks = model.track(im0, show=False, persist=True, verbose=False) + tracks = model.track(im0, persist=True) out = queue.process_queue(im0, tracks) video_writer.write(im0) @@ -100,15 +98,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra names=model.names, reg_pts=queue_region, line_thickness=3, - fontsize=1.0, - region_color=(255, 144, 31), ) while cap.isOpened(): success, im0 = cap.read() if success: - tracks = model.track(im0, show=False, persist=True, verbose=False, classes=0) # Only person class + tracks = model.track(im0, persist=True, classes=0) # Only person class out = queue.process_queue(im0, tracks) video_writer.write(im0) @@ -125,20 +121,13 @@ Queue management using [Ultralytics YOLOv8](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. | -| `track_thickness` | `int` | `2` | Thickness of the track lines. | -| `view_img` | `bool` | `False` | Whether to display the image frames. | -| `region_color` | `tuple` | `(255, 0, 255)` | Color of the counting region lines (BGR). | -| `view_queue_counts` | `bool` | `True` | Whether to display the queue counts. | -| `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. | -| `count_txt_color` | `tuple` | `(255, 255, 255)` | Color of the count text (BGR). | -| `track_color` | `tuple` | `None` | Color of the tracks. If `None`, different colors will be used for different tracks. | -| `region_thickness` | `int` | `5` | Thickness of the counting region lines. | -| `fontsize` | `float` | `0.7` | Font size for the text annotations. | +| 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. | ### Arguments `model.track` @@ -170,8 +159,6 @@ queue = solutions.QueueManager( names=model.names, reg_pts=queue_region, line_thickness=3, - fontsize=1.0, - region_color=(255, 144, 31), ) while cap.isOpened(): @@ -223,8 +210,6 @@ queue_airport = solutions.QueueManager( names=model.names, reg_pts=queue_region_airport, line_thickness=3, - fontsize=1.0, - region_color=(0, 255, 0), ) ``` diff --git a/ultralytics/solutions/queue_management.py b/ultralytics/solutions/queue_management.py index 16df28b1bb..ef60150395 100644 --- a/ultralytics/solutions/queue_management.py +++ b/ultralytics/solutions/queue_management.py @@ -20,15 +20,8 @@ class QueueManager: names, reg_pts=None, line_thickness=2, - track_thickness=2, view_img=False, - region_color=(255, 0, 255), - view_queue_counts=True, draw_tracks=False, - count_txt_color=(255, 255, 255), - track_color=None, - region_thickness=5, - fontsize=0.7, ): """ Initializes the QueueManager with specified parameters for tracking and counting objects. @@ -38,57 +31,35 @@ class QueueManager: 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. - track_thickness (int, optional): Thickness of the track lines. Defaults to 2. view_img (bool, optional): Whether to display the image frames. Defaults to False. - region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255). - view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True. draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False. - count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255). - track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different - tracks. Defaults to None. - region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5. - fontsize (float, optional): Font size for the text annotations. Defaults to 0.7. """ - # Mouse events state - self.is_drawing = False - self.selected_point = None - # 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)]) ) - self.region_color = region_color - self.region_thickness = region_thickness - # Image and annotation Information - self.im0 = None + # annotation Information self.tf = line_thickness self.view_img = view_img - self.view_queue_counts = view_queue_counts - self.fontsize = fontsize self.names = names # Class names - self.annotator = None # Annotator - self.window_name = "Ultralytics YOLOv8 Queue Manager" # Object counting Information self.counts = 0 - self.count_txt_color = count_txt_color # Tracks info self.track_history = defaultdict(list) - self.track_thickness = track_thickness self.draw_tracks = draw_tracks - self.track_color = track_color # Check if environment supports imshow self.env_check = check_imshow(warn=True) - def extract_and_process_tracks(self, tracks): + 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 - self.annotator = Annotator(self.im0, self.tf, self.names) + 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() @@ -98,7 +69,7 @@ class QueueManager: # Extract tracks for box, track_id, cls in zip(boxes, track_ids, clss): # Draw bounding box - self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True)) + annotator.box_label(box, label=self.names[cls], color=colors(int(track_id), True)) # Update track history track_line = self.track_history[track_id] @@ -108,10 +79,10 @@ class QueueManager: # Draw track trails if enabled if self.draw_tracks: - self.annotator.draw_centroid_and_tracks( + annotator.draw_centroid_and_tracks( track_line, - color=self.track_color or colors(int(track_id), True), - track_thickness=self.track_thickness, + 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 @@ -125,21 +96,16 @@ class QueueManager: # Display queue counts label = f"Queue Counts : {str(self.counts)}" if label is not None: - self.annotator.queue_counts_display( + annotator.queue_counts_display( label, points=self.reg_pts, - region_color=self.region_color, - txt_color=self.count_txt_color, + region_color=(255, 0, 255), + txt_color=(104, 31, 17), ) - self.display_frames() - - def display_frames(self): - """Displays the current frame with annotations.""" if self.env_check and self.view_img: - self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color) - cv2.namedWindow(self.window_name) - cv2.imshow(self.window_name, self.im0) + 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 @@ -152,12 +118,8 @@ class QueueManager: im0 (ndarray): Current frame from the video stream. tracks (list): List of tracks obtained from the object tracking process. """ - self.im0 = im0 # Store the current frame - self.extract_and_process_tracks(tracks) # Extract and process tracks - - if self.view_img: - self.display_frames() # Display the frame if enabled - return self.im0 + self.extract_and_process_tracks(tracks, im0) # Extract and process tracks + return im0 if __name__ == "__main__":