Francesco Mattioli
2 months ago
committed by
GitHub
commit
cc80f984c9
38 changed files with 397 additions and 435 deletions
@ -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: |
class QueueManager(BaseSolution): |
||||||
"""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 |
||||||
): |
|
||||||
|
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: |
Args: |
||||||
names (dict): A dictionary mapping class IDs to class names. |
im0 (ndarray): The input image that will be used for processing |
||||||
reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined |
Returns |
||||||
rectangle. |
im0 (ndarray): The processed image for more usage |
||||||
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 |
self.counts = 0 # Reset counts every frame |
||||||
if tracks[0].boxes.id is not None: |
self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator |
||||||
boxes = tracks[0].boxes.xyxy.cpu() |
self.extract_tracks(im0) # Extract tracks |
||||||
clss = tracks[0].boxes.cls.cpu().tolist() |
|
||||||
track_ids = tracks[0].boxes.id.int().cpu().tolist() |
|
||||||
|
|
||||||
# Extract tracks |
self.annotator.draw_region( |
||||||
for box, track_id, cls in zip(boxes, track_ids, clss): |
reg_pts=self.region, color=self.rect_color, thickness=self.line_width * 2 |
||||||
# Draw bounding box |
) # Draw region |
||||||
annotator.box_label(box, label=self.names[cls], color=colors(int(track_id), True)) |
|
||||||
|
|
||||||
# Update track history |
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss): |
||||||
track_line = self.track_history[track_id] |
# Draw bounding box and counting region |
||||||
track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))) |
self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True)) |
||||||
if len(track_line) > 30: |
self.store_tracking_history(track_id, box) # Store track history |
||||||
track_line.pop(0) |
|
||||||
|
|
||||||
# Draw track trails if enabled |
# Draw tracks of objects |
||||||
if self.draw_tracks: |
self.annotator.draw_centroid_and_tracks( |
||||||
annotator.draw_centroid_and_tracks( |
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width |
||||||
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: |
# Cache frequently accessed attributes |
||||||
annotator.draw_region(reg_pts=self.reg_pts, thickness=self.tf * 2, color=(255, 0, 255)) |
track_history = self.track_history.get(track_id, []) |
||||||
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): |
# 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 |
||||||
Main function to start the queue management process. |
if self.region_length >= 3 and prev_position and self.r_s.contains(Point(self.track_line[-1])): |
||||||
|
self.counts += 1 |
||||||
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 |
|
||||||
|
|
||||||
|
# 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__": |
return im0 # return output image for more usage |
||||||
classes_names = {0: "person", 1: "car"} # example class names |
|
||||||
queue_manager = QueueManager(classes_names) |
|
||||||
|
|||||||
@ -1,116 +1,76 @@ |
|||||||
# Ultralytics YOLO 🚀, AGPL-3.0 license |
# Ultralytics YOLO 🚀, AGPL-3.0 license |
||||||
|
|
||||||
from collections import defaultdict |
|
||||||
from time import time |
from time import time |
||||||
|
|
||||||
import cv2 |
|
||||||
import numpy as np |
import numpy as np |
||||||
|
|
||||||
from ultralytics.utils.checks import check_imshow |
from ultralytics.solutions.solutions import BaseSolution, LineString |
||||||
from ultralytics.utils.plotting import Annotator, colors |
from ultralytics.utils.plotting import Annotator, colors |
||||||
|
|
||||||
|
|
||||||
class SpeedEstimator: |
class SpeedEstimator(BaseSolution): |
||||||
"""A class to estimate the speed of objects in a real-time video stream based on their tracks.""" |
"""A class to estimate the speed of objects in a real-time video stream based on their tracks.""" |
||||||
|
|
||||||
def __init__(self, names, reg_pts=None, view_img=False, line_thickness=2, spdl_dist_thresh=10): |
def __init__(self, **kwargs): |
||||||
""" |
"""Initializes the SpeedEstimator with the given parameters.""" |
||||||
Initializes the SpeedEstimator with the given parameters. |
super().__init__(**kwargs) |
||||||
|
|
||||||
Args: |
|
||||||
names (dict): Dictionary of class names. |
|
||||||
reg_pts (list, optional): List of region points for speed estimation. Defaults to [(20, 400), (1260, 400)]. |
|
||||||
view_img (bool, optional): Whether to display the image with annotations. Defaults to False. |
|
||||||
line_thickness (int, optional): Thickness of the lines for drawing boxes and tracks. Defaults to 2. |
|
||||||
spdl_dist_thresh (int, optional): Distance threshold for speed calculation. Defaults to 10. |
|
||||||
""" |
|
||||||
# Region information |
|
||||||
self.reg_pts = reg_pts if reg_pts is not None else [(20, 400), (1260, 400)] |
|
||||||
|
|
||||||
self.names = names # Classes names |
self.initialize_region() # Initialize speed region |
||||||
|
|
||||||
# Tracking information |
|
||||||
self.trk_history = defaultdict(list) |
|
||||||
|
|
||||||
self.view_img = view_img # bool for displaying inference |
|
||||||
self.tf = line_thickness # line thickness for annotator |
|
||||||
self.spd = {} # set for speed data |
self.spd = {} # set for speed data |
||||||
self.trkd_ids = [] # list for already speed_estimated and tracked ID's |
self.trkd_ids = [] # list for already speed_estimated and tracked ID's |
||||||
self.spdl = spdl_dist_thresh # Speed line distance threshold |
|
||||||
self.trk_pt = {} # set for tracks previous time |
self.trk_pt = {} # set for tracks previous time |
||||||
self.trk_pp = {} # set for tracks previous point |
self.trk_pp = {} # set for tracks previous point |
||||||
|
|
||||||
# Check if the environment supports imshow |
def estimate_speed(self, im0): |
||||||
self.env_check = check_imshow(warn=True) |
|
||||||
|
|
||||||
def estimate_speed(self, im0, tracks): |
|
||||||
""" |
""" |
||||||
Estimates the speed of objects based on tracking data. |
Estimates the speed of objects based on tracking data. |
||||||
|
|
||||||
Args: |
Args: |
||||||
im0 (ndarray): Image. |
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 |
||||||
Returns: |
|
||||||
(ndarray): The image with annotated boxes and tracks. |
|
||||||
""" |
""" |
||||||
if tracks[0].boxes.id is None: |
self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator |
||||||
return im0 |
self.extract_tracks(im0) # Extract tracks |
||||||
|
|
||||||
boxes = tracks[0].boxes.xyxy.cpu() |
self.annotator.draw_region( |
||||||
clss = tracks[0].boxes.cls.cpu().tolist() |
reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2 |
||||||
t_ids = tracks[0].boxes.id.int().cpu().tolist() |
) # Draw region |
||||||
annotator = Annotator(im0, line_width=self.tf) |
|
||||||
annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 255), thickness=self.tf * 2) |
|
||||||
|
|
||||||
for box, t_id, cls in zip(boxes, t_ids, clss): |
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss): |
||||||
track = self.trk_history[t_id] |
self.store_tracking_history(track_id, box) # Store track history |
||||||
bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)) |
|
||||||
track.append(bbox_center) |
|
||||||
|
|
||||||
if len(track) > 30: |
# Check if track_id is already in self.trk_pp or trk_pt initialize if not |
||||||
track.pop(0) |
if track_id not in self.trk_pt: |
||||||
|
self.trk_pt[track_id] = 0 |
||||||
|
if track_id not in self.trk_pp: |
||||||
|
self.trk_pp[track_id] = self.track_line[-1] |
||||||
|
|
||||||
trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2)) |
speed_label = f"{int(self.spd[track_id])} km/h" if track_id in self.spd else self.names[int(cls)] |
||||||
|
self.annotator.box_label(box, label=speed_label, color=colors(track_id, True)) # Draw bounding box |
||||||
|
|
||||||
if t_id not in self.trk_pt: |
# Draw tracks of objects |
||||||
self.trk_pt[t_id] = 0 |
self.annotator.draw_centroid_and_tracks( |
||||||
|
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width |
||||||
|
) |
||||||
|
|
||||||
speed_label = f"{int(self.spd[t_id])} km/h" if t_id in self.spd else self.names[int(cls)] |
# Calculate object speed and direction based on region intersection |
||||||
bbox_color = colors(int(t_id), True) |
if LineString([self.trk_pp[track_id], self.track_line[-1]]).intersects(self.l_s): |
||||||
|
|
||||||
annotator.box_label(box, speed_label, bbox_color) |
|
||||||
cv2.polylines(im0, [trk_pts], isClosed=False, color=bbox_color, thickness=self.tf) |
|
||||||
cv2.circle(im0, (int(track[-1][0]), int(track[-1][1])), self.tf * 2, bbox_color, -1) |
|
||||||
|
|
||||||
# Calculation of object speed |
|
||||||
if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]: |
|
||||||
return |
|
||||||
if self.reg_pts[1][1] - self.spdl < track[-1][1] < self.reg_pts[1][1] + self.spdl: |
|
||||||
direction = "known" |
|
||||||
elif self.reg_pts[0][1] - self.spdl < track[-1][1] < self.reg_pts[0][1] + self.spdl: |
|
||||||
direction = "known" |
direction = "known" |
||||||
else: |
else: |
||||||
direction = "unknown" |
direction = "unknown" |
||||||
|
|
||||||
if self.trk_pt.get(t_id) != 0 and direction != "unknown" and t_id not in self.trkd_ids: |
# Perform speed calculation and tracking updates if direction is valid |
||||||
self.trkd_ids.append(t_id) |
if direction == "known" and track_id not in self.trkd_ids: |
||||||
|
self.trkd_ids.append(track_id) |
||||||
time_difference = time() - self.trk_pt[t_id] |
time_difference = time() - self.trk_pt[track_id] |
||||||
if time_difference > 0: |
if time_difference > 0: |
||||||
self.spd[t_id] = np.abs(track[-1][1] - self.trk_pp[t_id][1]) / time_difference |
self.spd[track_id] = np.abs(self.track_line[-1][1] - self.trk_pp[track_id][1]) / time_difference |
||||||
|
|
||||||
self.trk_pt[t_id] = time() |
|
||||||
self.trk_pp[t_id] = track[-1] |
|
||||||
|
|
||||||
if self.view_img and self.env_check: |
|
||||||
cv2.imshow("Ultralytics Speed Estimation", im0) |
|
||||||
if cv2.waitKey(1) & 0xFF == ord("q"): |
|
||||||
return |
|
||||||
|
|
||||||
return im0 |
self.trk_pt[track_id] = time() |
||||||
|
self.trk_pp[track_id] = self.track_line[-1] |
||||||
|
|
||||||
|
self.display_output(im0) # display output with base class function |
||||||
|
|
||||||
if __name__ == "__main__": |
return im0 # return output image for more usage |
||||||
names = {0: "person", 1: "car"} # example class names |
|
||||||
speed_estimator = SpeedEstimator(names) |
|
||||||
|
|||||||
Loading…
Reference in new issue