Chiebot-Mirror
/
ultralytics
mirror of https://github.com/ultralytics/ultralytics.git
9
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Optimize `queue` solution (#16246)

Browse Source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
multisource-docs-fix
Muhammad Rizwan Munawar 3 months ago committed by GitHub
parent 6c3297623b
commit e309b6efab
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 76 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 33
      docs/en/guides/queue-management.md
  2. 66
      ultralytics/solutions/queue_management.py

33
docs/en/guides/queue-management.md
Unescape View File

@ -56,15 +56,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
names=model.names, names=model.names,
reg_pts=queue_region, reg_pts=queue_region,
line_thickness=3, line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
) )
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if success: if success:
tracks = model.track(im0, show=False, persist=True, verbose=False) tracks = model.track(im0, persist=True)
out = queue.process_queue(im0, tracks) out = queue.process_queue(im0, tracks)
video_writer.write(im0) video_writer.write(im0)
@ -100,15 +98,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
names=model.names, names=model.names,
reg_pts=queue_region, reg_pts=queue_region,
line_thickness=3, line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
) )
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if success: 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) out = queue.process_queue(im0, tracks)
video_writer.write(im0) video_writer.write(im0)
@ -125,20 +121,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
### Arguments `QueueManager` ### Arguments `QueueManager`
| Name | Type | Default | Description | | Name | Type | Default | Description |
| ------------------- | ---------------- | -------------------------- | ----------------------------------------------------------------------------------- | | ---------------- | ---------------- | -------------------------- | -------------------------------------------------------------------------------- |
| `names` | `dict` | `model.names` | A dictionary mapping class IDs to class names. | | `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. | | `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. | | `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. |
| `view_img` | `bool` | `False` | Whether to display the image frames. | | `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. |
| `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. |
### Arguments `model.track` ### Arguments `model.track`
@ -170,8 +159,6 @@ queue = solutions.QueueManager(
names=model.names, names=model.names,
reg_pts=queue_region, reg_pts=queue_region,
line_thickness=3, line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
) )
while cap.isOpened(): while cap.isOpened():
@ -223,8 +210,6 @@ queue_airport = solutions.QueueManager(
names=model.names, names=model.names,
reg_pts=queue_region_airport, reg_pts=queue_region_airport,
line_thickness=3, line_thickness=3,
fontsize=1.0,
region_color=(0, 255, 0),
) )
``` ```

66
ultralytics/solutions/queue_management.py
Unescape View File

@ -20,15 +20,8 @@ class QueueManager:
names, names,
reg_pts=None, reg_pts=None,
line_thickness=2, line_thickness=2,
track_thickness=2,
view_img=False, view_img=False,
region_color=(255, 0, 255),
view_queue_counts=True,
draw_tracks=False, 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. 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 reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
rectangle. rectangle.
line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2. 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. 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. 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 # Region & Line Information
self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)] self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
self.counting_region = ( self.counting_region = (
Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)]) 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 # annotation Information
self.im0 = None
self.tf = line_thickness self.tf = line_thickness
self.view_img = view_img self.view_img = view_img
self.view_queue_counts = view_queue_counts
self.fontsize = fontsize
self.names = names # Class names self.names = names # Class names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Queue Manager"
# Object counting Information # Object counting Information
self.counts = 0 self.counts = 0
self.count_txt_color = count_txt_color
# Tracks info # Tracks info
self.track_history = defaultdict(list) self.track_history = defaultdict(list)
self.track_thickness = track_thickness
self.draw_tracks = draw_tracks self.draw_tracks = draw_tracks
self.track_color = track_color
# Check if environment supports imshow # Check if environment supports imshow
self.env_check = check_imshow(warn=True) 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.""" """Extracts and processes tracks for queue management in a video stream."""
# Initialize annotator and draw the queue region # 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 self.counts = 0 # Reset counts every frame
if tracks[0].boxes.id is not None: if tracks[0].boxes.id is not None:
boxes = tracks[0].boxes.xyxy.cpu() boxes = tracks[0].boxes.xyxy.cpu()
@ -98,7 +69,7 @@ class QueueManager:
# Extract tracks # Extract tracks
for box, track_id, cls in zip(boxes, track_ids, clss): for box, track_id, cls in zip(boxes, track_ids, clss):
# Draw bounding box # 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 # Update track history
track_line = self.track_history[track_id] track_line = self.track_history[track_id]
@ -108,10 +79,10 @@ class QueueManager:
# Draw track trails if enabled # Draw track trails if enabled
if self.draw_tracks: if self.draw_tracks:
self.annotator.draw_centroid_and_tracks( annotator.draw_centroid_and_tracks(
track_line, track_line,
color=self.track_color or colors(int(track_id), True), color=colors(int(track_id), True),
track_thickness=self.track_thickness, track_thickness=self.line_thickness,
) )
prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None 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 # Display queue counts
label = f"Queue Counts : {str(self.counts)}" label = f"Queue Counts : {str(self.counts)}"
if label is not None: if label is not None:
self.annotator.queue_counts_display( annotator.queue_counts_display(
label, label,
points=self.reg_pts, points=self.reg_pts,
region_color=self.region_color, region_color=(255, 0, 255),
txt_color=self.count_txt_color, 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: 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) annotator.draw_region(reg_pts=self.reg_pts, thickness=self.tf * 2, color=(255, 0, 255))
cv2.namedWindow(self.window_name) cv2.imshow("Ultralytics YOLOv8 Queue Manager", im0)
cv2.imshow(self.window_name, self.im0)
# Close window on 'q' key press # Close window on 'q' key press
if cv2.waitKey(1) & 0xFF == ord("q"): if cv2.waitKey(1) & 0xFF == ord("q"):
return return
@ -152,12 +118,8 @@ class QueueManager:
im0 (ndarray): Current frame from the video stream. im0 (ndarray): Current frame from the video stream.
tracks (list): List of tracks obtained from the object tracking process. tracks (list): List of tracks obtained from the object tracking process.
""" """
self.im0 = im0 # Store the current frame self.extract_and_process_tracks(tracks, im0) # Extract and process tracks
self.extract_and_process_tracks(tracks) # Extract and process tracks return im0
if self.view_img:
self.display_frames() # Display the frame if enabled
return self.im0
if __name__ == "__main__": if __name__ == "__main__":

Write Preview
Loading…
Cancel
Save