Update speed-estimation solution (#16798)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar
GitHub
parent
28f31f14e8
commit
30f64bd35f
4 changed files with 74 additions and 129 deletions
|
|
@ -1,116 +1,76 @@
|
|||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
from collections import defaultdict
|
||||
from time import time
|
||||
|
||||
import cv2
|
||||
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
|
||||
|
||||
|
||||
class SpeedEstimator:
|
||||
class SpeedEstimator(BaseSolution):
|
||||
"""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):
|
||||
"""
|
||||
Initializes the SpeedEstimator with the given parameters.
|
||||
def __init__(self, **kwargs):
|
||||
"""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.initialize_region() # Initialize speed region
|
||||
|
||||
self.names = names # Classes names
|
||||
|
||||
# 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.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_pp = {} # set for tracks previous point
|
||||
|
||||
# Check if the environment supports imshow
|
||||
self.env_check = check_imshow(warn=True)
|
||||
|
||||
def estimate_speed(self, im0, tracks):
|
||||
def estimate_speed(self, im0):
|
||||
"""
|
||||
Estimates the speed of objects based on tracking data.
|
||||
|
||||
Args:
|
||||
im0 (ndarray): Image.
|
||||
tracks (list): List of tracks obtained from the object tracking process.
|
||||
|
||||
Returns:
|
||||
(ndarray): The image with annotated boxes and tracks.
|
||||
im0 (ndarray): The input image that will be used for processing
|
||||
Returns
|
||||
im0 (ndarray): The processed image for more usage
|
||||
"""
|
||||
if tracks[0].boxes.id is None:
|
||||
return im0
|
||||
self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator
|
||||
self.extract_tracks(im0) # Extract tracks
|
||||
|
||||
boxes = tracks[0].boxes.xyxy.cpu()
|
||||
clss = tracks[0].boxes.cls.cpu().tolist()
|
||||
t_ids = tracks[0].boxes.id.int().cpu().tolist()
|
||||
annotator = Annotator(im0, line_width=self.tf)
|
||||
annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 255), thickness=self.tf * 2)
|
||||
self.annotator.draw_region(
|
||||
reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2
|
||||
) # Draw region
|
||||
|
||||
for box, t_id, cls in zip(boxes, t_ids, clss):
|
||||
track = self.trk_history[t_id]
|
||||
bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
|
||||
track.append(bbox_center)
|
||||
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
|
||||
self.store_tracking_history(track_id, box) # Store track history
|
||||
|
||||
if len(track) > 30:
|
||||
track.pop(0)
|
||||
# Check if track_id is already in self.trk_pp or trk_pt initialize if not
|
||||
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:
|
||||
self.trk_pt[t_id] = 0
|
||||
# Draw tracks of objects
|
||||
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)]
|
||||
bbox_color = colors(int(t_id), True)
|
||||
|
||||
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:
|
||||
# Calculate object speed and direction based on region intersection
|
||||
if LineString([self.trk_pp[track_id], self.track_line[-1]]).intersects(self.l_s):
|
||||
direction = "known"
|
||||
else:
|
||||
direction = "unknown"
|
||||
|
||||
if self.trk_pt.get(t_id) != 0 and direction != "unknown" and t_id not in self.trkd_ids:
|
||||
self.trkd_ids.append(t_id)
|
||||
|
||||
time_difference = time() - self.trk_pt[t_id]
|
||||
# Perform speed calculation and tracking updates if direction is valid
|
||||
if direction == "known" and track_id not in self.trkd_ids:
|
||||
self.trkd_ids.append(track_id)
|
||||
time_difference = time() - self.trk_pt[track_id]
|
||||
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]
|
||||
self.trk_pt[track_id] = time()
|
||||
self.trk_pp[track_id] = self.track_line[-1]
|
||||
|
||||
if self.view_img and self.env_check:
|
||||
cv2.imshow("Ultralytics Speed Estimation", im0)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
return
|
||||
self.display_output(im0) # display output with base class function
|
||||
|
||||
return im0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
names = {0: "person", 1: "car"} # example class names
|
||||
speed_estimator = SpeedEstimator(names)
|
||||
return im0 # return output image for more usage
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue