HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Update queue-management solution (#16772)

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

59
docs/en/guides/queue-management.md
View file

@ -40,10 +40,9 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra
```python ```python
import cv2 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" 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)) 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_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager( queue = solutions.QueueManager(
names=model.names, model="yolo11n.pt",
reg_pts=queue_region, region=queue_region,
line_thickness=3,
) )
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if success: if success:
tracks = model.track(im0, persist=True) out = queue.process_queue(im0)
out = queue.process_queue(im0, tracks)
video_writer.write(im0) video_writer.write(im0)
if cv2.waitKey(1) & 0xFF == ord("q"): if cv2.waitKey(1) & 0xFF == ord("q"):
break break
@ -82,10 +78,9 @@ Queue management using [Ultralytics YOLO11](https://github.com/ultralytics/ultra
```python ```python
import cv2 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" 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)) 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_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager( queue = solutions.QueueManager(
names=model.names, model="yolo11n.pt",
reg_pts=queue_region, classes=3,
line_thickness=3,
) )
while cap.isOpened(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if success: if success:
tracks = model.track(im0, persist=True, classes=0) # Only person class out = queue.process_queue(im0)
out = queue.process_queue(im0, tracks)
video_writer.write(im0) video_writer.write(im0)
if cv2.waitKey(1) & 0xFF == ord("q"): if cv2.waitKey(1) & 0xFF == ord("q"):
break break
@ -121,13 +113,12 @@ Queue management using [Ultralytics YOLO11](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. | | `model` | `str` | `None` | Path to Ultralytics YOLO Model File |
| `reg_pts` | `list of tuples` | `[(20, 400), (1260, 400)]` | Points defining the counting region polygon. Defaults to a predefined rectangle. | | `region` | `list` | `[(20, 400), (1260, 400)]` | List of points defining the queue region. |
| `line_thickness` | `int` | `2` | Thickness of the annotation lines. | | `line_width` | `int` | `2` | Line thickness for bounding boxes. |
| `view_img` | `bool` | `False` | Whether to display the image frames. | | `show` | `bool` | `False` | Flag to control whether to display the video stream. |
| `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. |
### Arguments `model.track` ### Arguments `model.track`
@ -149,23 +140,21 @@ Here's a minimal example:
```python ```python
import cv2 import cv2
from ultralytics import YOLO, solutions from ultralytics import solutions
model = YOLO("yolo11n.pt")
cap = cv2.VideoCapture("path/to/video.mp4") cap = cv2.VideoCapture("path/to/video.mp4")
queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)] queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager( queue = solutions.QueueManager(
names=model.names, model="yolo11n.pt",
reg_pts=queue_region, region=queue_region,
line_thickness=3, line_width=3,
) )
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) out = queue.process_queue(im0)
out = queue.process_queue(im0, tracks)
cv2.imshow("Queue Management", im0) cv2.imshow("Queue Management", im0)
if cv2.waitKey(1) & 0xFF == ord("q"): if cv2.waitKey(1) & 0xFF == ord("q"):
break break
@ -207,9 +196,9 @@ Example for airports:
```python ```python
queue_region_airport = [(50, 600), (1200, 600), (1200, 550), (50, 550)] queue_region_airport = [(50, 600), (1200, 600), (1200, 550), (50, 550)]
queue_airport = solutions.QueueManager( queue_airport = solutions.QueueManager(
names=model.names, model="yolo11n.pt",
reg_pts=queue_region_airport, region=queue_region_airport,
line_thickness=3, line_width=3,
) )
``` ```

4
tests/test_solutions.py
View file

@ -22,7 +22,7 @@ def test_major_solutions():
counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False) counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False)
heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, 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) 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(): while cap.isOpened():
success, im0 = cap.read() success, im0 = cap.read()
if not success: if not success:
@ -32,7 +32,7 @@ def test_major_solutions():
_ = counter.count(original_im0.copy()) _ = counter.count(original_im0.copy())
_ = heatmap.generate_heatmap(original_im0.copy()) _ = heatmap.generate_heatmap(original_im0.copy())
_ = speed.estimate_speed(original_im0.copy(), tracks) _ = speed.estimate_speed(original_im0.copy(), tracks)
_ = queue.process_queue(original_im0.copy(), tracks) _ = queue.process_queue(original_im0.copy())
cap.release() cap.release()
cv2.destroyAllWindows() cv2.destroyAllWindows()

2
ultralytics/solutions/object_counter.py
View file

@ -116,7 +116,7 @@ class ObjectCounter(BaseSolution):
self.store_tracking_history(track_id, box) # Store track history self.store_tracking_history(track_id, box) # Store track history
self.store_classwise_counts(cls) # store classwise counts in dict 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.annotator.draw_centroid_and_tracks(
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
) )

159
ultralytics/solutions/queue_management.py
View file

@ -1,127 +1,64 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license # Ultralytics YOLO 🚀, AGPL-3.0 license
from collections import defaultdict from shapely.geometry import Point
import cv2 from ultralytics.solutions.solutions import BaseSolution # Import a parent class
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator, colors from ultralytics.utils.plotting import Annotator, colors
check_requirements("shapely>=2.0.0")
from shapely.geometry import Point, Polygon class QueueManager(BaseSolution):
class QueueManager:
"""A class to manage the queue in a real-time video stream based on object tracks.""" """A class to manage the queue in a real-time video stream based on object tracks."""
def __init__( def __init__(self, **kwargs):
self, """Initializes the QueueManager with specified parameters for tracking and counting objects."""
names, super().__init__(**kwargs)
reg_pts=None, self.initialize_region()
line_thickness=2, self.counts = 0 # Queue counts Information
view_img=False, self.rect_color = (255, 255, 255) # Rectangle color
draw_tracks=False, self.region_length = len(self.region) # Store region length for further usage
):
"""
Initializes the QueueManager with specified parameters for tracking and counting objects.
Args: def process_queue(self, im0):
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.
"""
# 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()
# 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))
# 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)
# 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),
)
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
def process_queue(self, im0, tracks):
""" """
Main function to start the queue management process. Main function to start the queue management process.
Args: Args:
im0 (ndarray): Current frame from the video stream. im0 (ndarray): The input image that will be used for processing
tracks (list): List of tracks obtained from the object tracking process. Returns
im0 (ndarray): The processed image for more usage
""" """
self.extract_and_process_tracks(tracks, im0) # Extract and process tracks self.counts = 0 # Reset counts every frame
return im0 self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator
self.extract_tracks(im0) # Extract tracks
self.annotator.draw_region(
reg_pts=self.region, color=self.rect_color, thickness=self.line_width * 2
) # Draw region
if __name__ == "__main__": for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
classes_names = {0: "person", 1: "car"} # example class names # Draw bounding box and counting region
queue_manager = QueueManager(classes_names) 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 tracks of objects
self.annotator.draw_centroid_and_tracks(
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
)
# Cache frequently accessed attributes
track_history = self.track_history.get(track_id, [])
# 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
return im0 # return output image for more usage

8
ultralytics/solutions/solutions.py
View file

@ -76,9 +76,11 @@ class BaseSolution:
def initialize_region(self): def initialize_region(self):
"""Initialize the counting region and line segment based on config.""" """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.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) 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])]) 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): def display_output(self, im0):
""" """