def__init__(self,names,reg_pts=None,view_img=False,line_thickness=2,region_thickness=5,spdl_dist_thresh=10):""" Initializes the SpeedEstimator with the given parameters. 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. region_thickness (int, optional): Thickness of the region lines. Defaults to 5. spdl_dist_thresh (int, optional): Distance threshold for speed calculation. Defaults to 10. """# Visual & image informationself.im0=Noneself.annotator=Noneself.view_img=view_img# Region informationself.reg_pts=reg_ptsifreg_ptsisnotNoneelse[(20,400),(1260,400)]self.region_thickness=region_thickness# Tracking informationself.clss=Noneself.names=namesself.boxes=Noneself.trk_ids=Noneself.trk_pts=Noneself.line_thickness=line_thicknessself.trk_history=defaultdict(list)# Speed estimation informationself.current_time=0self.dist_data={}self.trk_idslist=[]self.spdl_dist_thresh=spdl_dist_threshself.trk_previous_times={}self.trk_previous_points={}# Check if the environment supports imshowself.env_check=check_imshow(warn=True)
calculate_speed
calculate_speed(trk_id,track)
Calculates the speed of an object.
Parameters:
Name
Type
Description
Default
trk_id
int
Object track id.
required
track
list
Tracking history for drawing tracks path.
required
Source code in ultralytics/solutions/speed_estimation.py
defcalculate_speed(self,trk_id,track):""" Calculates the speed of an object. Args: trk_id (int): Object track id. track (list): Tracking history for drawing tracks path. """ifnotself.reg_pts[0][0]<track[-1][0]<self.reg_pts[1][0]:returnifself.reg_pts[1][1]-self.spdl_dist_thresh<track[-1][1]<self.reg_pts[1][1]+self.spdl_dist_thresh:direction="known"elifself.reg_pts[0][1]-self.spdl_dist_thresh<track[-1][1]<self.reg_pts[0][1]+self.spdl_dist_thresh:direction="known"else:direction="unknown"ifself.trk_previous_times.get(trk_id)!=0anddirection!="unknown"andtrk_idnotinself.trk_idslist:self.trk_idslist.append(trk_id)time_difference=time()-self.trk_previous_times[trk_id]iftime_difference>0:dist_difference=np.abs(track[-1][1]-self.trk_previous_points[trk_id][1])speed=dist_difference/time_differenceself.dist_data[trk_id]=speedself.trk_previous_times[trk_id]=time()self.trk_previous_points[trk_id]=track[-1]
display_frames
display_frames()
Displays the current frame.
Source code in ultralytics/solutions/speed_estimation.py
defestimate_speed(self,im0,tracks,region_color=(255,0,0)):""" Estimates the speed of objects based on tracking data. Args: im0 (ndarray): Image. tracks (list): List of tracks obtained from the object tracking process. region_color (tuple, optional): Color to use when drawing regions. Defaults to (255, 0, 0). Returns: (ndarray): The image with annotated boxes and tracks. """self.im0=im0iftracks[0].boxes.idisNone:ifself.view_imgandself.env_check:self.display_frames()returnim0self.extract_tracks(tracks)self.annotator=Annotator(self.im0,line_width=self.line_thickness)self.annotator.draw_region(reg_pts=self.reg_pts,color=region_color,thickness=self.region_thickness)forbox,trk_id,clsinzip(self.boxes,self.trk_ids,self.clss):track=self.store_track_info(trk_id,box)iftrk_idnotinself.trk_previous_times:self.trk_previous_times[trk_id]=0self.plot_box_and_track(trk_id,box,cls,track)self.calculate_speed(trk_id,track)ifself.view_imgandself.env_check:self.display_frames()returnim0
extract_tracks
extract_tracks(tracks)
Extracts results from the provided tracking data.
Parameters:
Name
Type
Description
Default
tracks
list
List of tracks obtained from the object tracking process.
required
Source code in ultralytics/solutions/speed_estimation.py
defextract_tracks(self,tracks):""" Extracts results from the provided tracking data. Args: tracks (list): List of tracks obtained from the object tracking process. """self.boxes=tracks[0].boxes.xyxy.cpu()self.clss=tracks[0].boxes.cls.cpu().tolist()self.trk_ids=tracks[0].boxes.id.int().cpu().tolist()
plot_box_and_track
plot_box_and_track(track_id,box,cls,track)
Plots track and bounding box.
Parameters:
Name
Type
Description
Default
track_id
int
Object track id.
required
box
list
Object bounding box data.
required
cls
str
Object class name.
required
track
list
Tracking history for drawing tracks path.
required
Source code in ultralytics/solutions/speed_estimation.py