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