From bde310581347dfa0e45eb1ea51a444ce694aa321 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Tue, 19 Nov 2024 16:40:42 +0500 Subject: [PATCH] Heatmaps bug fix (#17634) --- ultralytics/solutions/heatmap.py | 4 ++-- ultralytics/solutions/object_counter.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py index 39352a9bd7..c9dd808798 100644 --- a/ultralytics/solutions/heatmap.py +++ b/ultralytics/solutions/heatmap.py @@ -104,12 +104,12 @@ class Heatmap(ObjectCounter): self.annotator.draw_region(reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2) self.store_tracking_history(track_id, box) # Store track history self.store_classwise_counts(cls) # store classwise counts in dict - + current_centroid = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2) # Store tracking previous position and perform object counting prev_position = None if len(self.track_history[track_id]) > 1: prev_position = self.track_history[track_id][-2] - self.count_objects(self.track_line, box, track_id, prev_position, cls) # Perform object counting + self.count_objects(current_centroid, track_id, prev_position, cls) # Perform object counting if self.region is not None: self.display_counts(im0) # Display the counts on the frame diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index ac372e5f6a..94e7d87caa 100644 --- a/ultralytics/solutions/object_counter.py +++ b/ultralytics/solutions/object_counter.py @@ -63,7 +63,7 @@ class ObjectCounter(BaseSolution): >>> track_id = 1 >>> prev_position = (120, 220) >>> cls = 0 - >>> counter.count_objects(track_line, box, track_id, prev_position, cls) + >>> counter.count_objects(current_centroid, track_id, prev_position, cls) """ if prev_position is None or track_id in self.counted_ids: return