diff --git a/docs/en/guides/heatmaps.md b/docs/en/guides/heatmaps.md index e11ee4d8..7db8c3da 100644 --- a/docs/en/guides/heatmaps.md +++ b/docs/en/guides/heatmaps.md @@ -42,9 +42,10 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult # Heatmap Init heatmap_obj = heatmap.Heatmap() heatmap_obj.set_args(colormap=cv2.COLORMAP_CIVIDIS, - imw=cap.get(4), # should same as cap width - imh=cap.get(3), # should same as cap height - view_img=True) + imw=cap.get(4), # should same as cap width + imh=cap.get(3), # should same as cap height + view_img=True, + decay_factor=0.99) while cap.isOpened(): success, im0 = cap.read() @@ -183,17 +184,18 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult ### Arguments `set_args` -| Name | Type | Default | Description | -|---------------------|----------------|-----------------|---------------------------------| -| view_img | `bool` | `False` | Display the frame with heatmap | -| colormap | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap | -| imw | `int` | `None` | Width of Heatmap | -| imh | `int` | `None` | Height of Heatmap | -| heatmap_alpha | `float` | `0.5` | Heatmap alpha value | -| count_reg_pts | `list` | `None` | Object counting region points | -| count_txt_thickness | `int` | `2` | Count values text size | -| count_reg_color | `tuple` | `(255, 0, 255)` | Counting region color | -| region_thickness | `int` | `5` | Counting region thickness value | +| Name | Type | Default | Description | +|---------------------|----------------|-----------------|-----------------------------------------------------------| +| view_img | `bool` | `False` | Display the frame with heatmap | +| colormap | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap | +| imw | `int` | `None` | Width of Heatmap | +| imh | `int` | `None` | Height of Heatmap | +| heatmap_alpha | `float` | `0.5` | Heatmap alpha value | +| count_reg_pts | `list` | `None` | Object counting region points | +| count_txt_thickness | `int` | `2` | Count values text size | +| count_reg_color | `tuple` | `(255, 0, 255)` | Counting region color | +| region_thickness | `int` | `5` | Counting region thickness value | +| decay_factor | `float` | `0.99` | Decay factor for heatmap area removal after specific time | ### Arguments `model.track` diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py index c12c9f42..abcef69c 100644 --- a/ultralytics/solutions/heatmap.py +++ b/ultralytics/solutions/heatmap.py @@ -50,6 +50,9 @@ class Heatmap: self.count_reg_color = (0, 255, 0) self.region_thickness = 5 + # Decay factor + self.decay_factor = 0.99 + # Check if environment support imshow self.env_check = check_imshow(warn=True) @@ -62,7 +65,8 @@ class Heatmap: count_reg_pts=None, count_txt_thickness=2, count_reg_color=(255, 0, 255), - region_thickness=5): + region_thickness=5, + decay_factor=0.99): """ Configures the heatmap colormap, width, height and display parameters. @@ -76,6 +80,7 @@ class Heatmap: count_txt_thickness (int): Text thickness for object counting display count_reg_color (RGB color): Color of object counting region region_thickness (int): Object counting Region thickness + decay_factor (float): value for removing heatmap area after object passed """ self.imw = imw self.imh = imh @@ -93,6 +98,7 @@ class Heatmap: self.count_txt_thickness = count_txt_thickness # Counting text thickness self.count_reg_color = count_reg_color self.region_thickness = region_thickness + self.decay_factor = decay_factor def extract_results(self, tracks): """ @@ -117,6 +123,7 @@ class Heatmap: if tracks[0].boxes.id is None: return self.im0 + self.heatmap *= self.decay_factor # decay factor self.extract_results(tracks) self.annotator = Annotator(self.im0, self.count_txt_thickness, None)