xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 21:44:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add estimate_speed(region_color) parameter (#8285)

Browse Source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Chi 2024-02-20 23:11:26 +05:30 committed by GitHub
parent 5171d6679d
commit a5605b5e4b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ultralytics/solutions/speed_estimation.py
View File

@ -154,13 +154,14 @@ class SpeedEstimator:
self.trk_previous_times[trk_id] = time() self.trk_previous_times[trk_id] = time()
self.trk_previous_points[trk_id] = track[-1] self.trk_previous_points[trk_id] = track[-1]
def estimate_speed(self, im0, tracks): def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
""" """
Calculate object based on tracking data. Calculate object based on tracking data.
Args: Args:
im0 (nd array): Image im0 (nd array): Image
tracks (list): List of tracks obtained from the object tracking process. tracks (list): List of tracks obtained from the object tracking process.
region_color (tuple): Color to use when drawing regions.
""" """
self.im0 = im0 self.im0 = im0
if tracks[0].boxes.id is None: if tracks[0].boxes.id is None:
@ -170,7 +171,7 @@ class SpeedEstimator:
self.extract_tracks(tracks) self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=2) self.annotator = Annotator(self.im0, line_width=2)
self.annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 0), thickness=self.region_thickness) self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss): for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
track = self.store_track_info(trk_id, box) track = self.store_track_info(trk_id, box)