Update queue-management solution (#16772)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-10-08 23:57:01 +05:00 committed by GitHub
parent 094faeb722
commit 6509757879
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 80 additions and 152 deletions

View file

@ -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,
)
```