Heatmaps bug fix (ultralytics#17634)

triple-Mu · Nov 19, 2024 · bde3105 · bde3105
1 parent 35b6fed
commit bde3105
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py
@@ -104,12 +104,12 @@ def generate_heatmap(self, im0):
                 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
@@ -63,7 +63,7 @@ def count_objects(self, current_centroid, track_id, prev_position, cls):
             >>> 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