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

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_classwise_counts(cls) # store classwise counts in dict
# Draw centroid of objects
# Draw tracks of objects
self.annotator.draw_centroid_and_tracks(
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
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,
):
"""
Initializes the QueueManager with specified parameters for tracking and counting objects.
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
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.
"""
# 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):
def process_queue(self, im0):
"""
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.
im0 (ndarray): The input image that will be used for processing
Returns
im0 (ndarray): The processed image for more usage
"""
self.extract_and_process_tracks(tracks, im0) # Extract and process tracks
return im0
self.counts = 0 # Reset counts every frame
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__":
classes_names = {0: "person", 1: "car"} # example class names
queue_manager = QueueManager(classes_names)
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 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):
"""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.r_s = Polygon(self.region) if len(self.region) >= 3 else LineString(self.region)
self.l_s = LineString([(self.region[0][0], self.region[0][1]), (self.region[1][0], self.region[1][1])])
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) # region segment
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):
"""