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>
action-recog
Muhammad Rizwan Munawar 2 months ago committed by fcakyon
parent 9353e7d371
commit 5f1f39d36c
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
Split 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,
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),
)
```

66
ultralytics/solutions/queue_management.py
Unescape View File

@ -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__":

Write Preview
Loading…
Cancel
Save