From a5605b5e4b2e33d840f37dbff463028f66a164c4 Mon Sep 17 00:00:00 2001 From: Chi Date: Tue, 20 Feb 2024 23:11:26 +0530 Subject: [PATCH] Add `estimate_speed(region_color)` parameter (#8285) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/solutions/speed_estimation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ultralytics/solutions/speed_estimation.py b/ultralytics/solutions/speed_estimation.py index b91ccfc6..f3f17959 100644 --- a/ultralytics/solutions/speed_estimation.py +++ b/ultralytics/solutions/speed_estimation.py @@ -154,13 +154,14 @@ class SpeedEstimator: self.trk_previous_times[trk_id] = time() 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. Args: im0 (nd array): Image tracks (list): List of tracks obtained from the object tracking process. + region_color (tuple): Color to use when drawing regions. """ self.im0 = im0 if tracks[0].boxes.id is None: @@ -170,7 +171,7 @@ class SpeedEstimator: self.extract_tracks(tracks) 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): track = self.store_track_info(trk_id, box)