Update the `heatmap` module code + docs (#7045)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/7040/head^2
Muhammad Rizwan Munawar 1 year ago committed by GitHub
parent 34b10b2db3
commit 1d0da53a32
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 30
      docs/en/guides/heatmaps.md
  2. 9
      ultralytics/solutions/heatmap.py

@ -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`

@ -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)

Loading…
Cancel
Save