Add HeatMap `guide` in real-world-projects + Code in Solutions Directory (#6796)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>pull/6816/head^2
Muhammad Rizwan Munawar
12 months ago
committed by
GitHub
parent
1e1247ddee
commit
742cbc1b4e
10 changed files with 454 additions and 58 deletions
@ -0,0 +1,142 @@ |
||||
--- |
||||
comments: true |
||||
description: Advanced Data Visualization with Ultralytics YOLOv8 Heatmaps |
||||
keywords: Ultralytics, YOLOv8, Advanced Data Visualization, Heatmap Technology, Object Detection and Tracking, Jupyter Notebook, Python SDK, Command Line Interface |
||||
--- |
||||
|
||||
# Advanced Data Visualization: Heatmaps using Ultralytics YOLOv8 🚀 |
||||
|
||||
## Introduction to Heatmaps |
||||
|
||||
A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics/) transforms complex data into a vibrant, color-coded matrix. This visual tool employs a spectrum of colors to represent varying data values, where warmer hues indicate higher intensities and cooler tones signify lower values. Heatmaps excel in visualizing intricate data patterns, correlations, and anomalies, offering an accessible and engaging approach to data interpretation across diverse domains. |
||||
|
||||
## Why Choose Heatmaps for Data Analysis? |
||||
|
||||
- **Intuitive Data Distribution Visualization:** Heatmaps simplify the comprehension of data concentration and distribution, converting complex datasets into easy-to-understand visual formats. |
||||
- **Efficient Pattern Detection:** By visualizing data in heatmap format, it becomes easier to spot trends, clusters, and outliers, facilitating quicker analysis and insights. |
||||
- **Enhanced Spatial Analysis and Decision Making:** Heatmaps are instrumental in illustrating spatial relationships, aiding in decision-making processes in sectors such as business intelligence, environmental studies, and urban planning. |
||||
|
||||
## Real World Applications |
||||
|
||||
| Transportation | Retail | |
||||
|:-----------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------:| |
||||
|  |  | |
||||
| Ultralytics YOLOv8 Transportation Heatmap | Ultralytics YOLOv8 Retail Heatmap | |
||||
|
||||
???+ tip "heatmap_alpha" |
||||
|
||||
heatmap_alpha value should be in range (0.0 - 1.0) |
||||
|
||||
!!! Example "Heatmap Example" |
||||
|
||||
=== "Heatmap" |
||||
```python |
||||
from ultralytics import YOLO |
||||
from ultralytics.solutions import heatmap |
||||
import cv2 |
||||
|
||||
model = YOLO("yolov8s.pt") |
||||
cap = cv2.VideoCapture("path/to/video/file.mp4") |
||||
if not cap.isOpened(): |
||||
print("Error reading video file") |
||||
exit(0) |
||||
|
||||
heatmap_obj = heatmap.Heatmap() |
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_CIVIDIS, |
||||
imw=cap.get(4), # should same as im0 width |
||||
imh=cap.get(3), # should same as im0 height |
||||
view_img=True) |
||||
|
||||
while cap.isOpened(): |
||||
success, im0 = cap.read() |
||||
if not success: |
||||
exit(0) |
||||
results = model.track(im0, persist=True) |
||||
frame = heatmap_obj.generate_heatmap(im0, tracks=results) |
||||
|
||||
``` |
||||
|
||||
=== "Heatmap with Specific Classes" |
||||
```python |
||||
from ultralytics import YOLO |
||||
from ultralytics.solutions import heatmap |
||||
import cv2 |
||||
|
||||
model = YOLO("yolov8s.pt") |
||||
cap = cv2.VideoCapture("path/to/video/file.mp4") |
||||
if not cap.isOpened(): |
||||
print("Error reading video file") |
||||
exit(0) |
||||
|
||||
classes_for_heatmap = [0, 2] |
||||
|
||||
heatmap_obj = heatmap.Heatmap() |
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_CIVIDIS, |
||||
imw=cap.get(4), # should same as im0 width |
||||
imh=cap.get(3), # should same as im0 height |
||||
view_img=True) |
||||
|
||||
while cap.isOpened(): |
||||
success, im0 = cap.read() |
||||
if not success: |
||||
exit(0) |
||||
results = model.track(im0, persist=True, |
||||
classes=classes_for_heatmap) |
||||
frame = heatmap_obj.generate_heatmap(im0, tracks=results) |
||||
|
||||
``` |
||||
|
||||
=== "Heatmap with Save Output" |
||||
```python |
||||
from ultralytics import YOLO |
||||
import heatmap |
||||
import cv2 |
||||
|
||||
model = YOLO("yolov8n.pt") |
||||
cap = cv2.VideoCapture("path/to/video/file.mp4") |
||||
if not cap.isOpened(): |
||||
print("Error reading video file") |
||||
exit(0) |
||||
|
||||
video_writer = cv2.VideoWriter("heatmap_output.avi", |
||||
cv2.VideoWriter_fourcc(*'mp4v'), |
||||
int(cap.get(5)), |
||||
(int(cap.get(3)), int(cap.get(4)))) |
||||
|
||||
heatmap_obj = heatmap.Heatmap() |
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_CIVIDIS, |
||||
imw=cap.get(4), # should same as im0 width |
||||
imh=cap.get(3), # should same as im0 height |
||||
view_img=True) |
||||
|
||||
while cap.isOpened(): |
||||
success, im0 = cap.read() |
||||
if not success: |
||||
exit(0) |
||||
results = model.track(im0, persist=True) |
||||
frame = heatmap_obj.generate_heatmap(im0, tracks=results) |
||||
video_writer.write(im0) |
||||
|
||||
video_writer.release() |
||||
``` |
||||
|
||||
### 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 | |
||||
|
||||
### Arguments `model.track` |
||||
|
||||
| Name | Type | Default | Description | |
||||
|-----------|---------|----------------|-------------------------------------------------------------| |
||||
| `source` | `im0` | `None` | source directory for images or videos | |
||||
| `persist` | `bool` | `False` | persisting tracks between frames | |
||||
| `tracker` | `str` | `botsort.yaml` | Tracking method 'bytetrack' or 'botsort' | |
||||
| `conf` | `float` | `0.3` | Confidence Threshold | |
||||
| `iou` | `float` | `0.5` | IOU Threshold | |
||||
| `classes` | `list` | `None` | filter results by class, i.e. classes=0, or classes=[0,2,3] | |
||||
@ -0,0 +1,102 @@ |
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license |
||||
|
||||
import cv2 |
||||
import numpy as np |
||||
|
||||
|
||||
class Heatmap: |
||||
"""A class to draw heatmaps in real-time video stream based on their tracks.""" |
||||
|
||||
def __init__(self): |
||||
"""Initializes the heatmap class with default values for Visual, Image, track and heatmap parameters.""" |
||||
|
||||
# Visual Information |
||||
self.annotator = None |
||||
self.view_img = False |
||||
|
||||
# Image Information |
||||
self.imw = None |
||||
self.imh = None |
||||
self.im0 = None |
||||
|
||||
# Heatmap Colormap and heatmap np array |
||||
self.colormap = None |
||||
self.heatmap = None |
||||
self.heatmap_alpha = 0.5 |
||||
|
||||
# Predict/track information |
||||
self.boxes = None |
||||
self.track_ids = None |
||||
self.clss = None |
||||
|
||||
def set_args(self, imw, imh, colormap=cv2.COLORMAP_JET, heatmap_alpha=0.5, view_img=False): |
||||
""" |
||||
Configures the heatmap colormap, width, height and display parameters. |
||||
|
||||
Args: |
||||
colormap (cv2.COLORMAP): The colormap to be set. |
||||
imw (int): The width of the frame. |
||||
imh (int): The height of the frame. |
||||
heatmap_alpha (float): alpha value for heatmap display |
||||
view_img (bool): Flag indicating frame display |
||||
""" |
||||
self.imw = imw |
||||
self.imh = imh |
||||
self.colormap = colormap |
||||
self.heatmap_alpha = heatmap_alpha |
||||
self.view_img = view_img |
||||
|
||||
# Heatmap new frame |
||||
self.heatmap = np.zeros((int(self.imw), int(self.imh)), dtype=np.float32) |
||||
|
||||
def extract_results(self, tracks): |
||||
""" |
||||
Extracts results from the provided data. |
||||
|
||||
Args: |
||||
tracks (list): List of tracks obtained from the object tracking process. |
||||
""" |
||||
if tracks[0].boxes.id is None: |
||||
return |
||||
self.boxes = tracks[0].boxes.xyxy.cpu() |
||||
self.clss = tracks[0].boxes.cls.cpu().tolist() |
||||
self.track_ids = tracks[0].boxes.id.int().cpu().tolist() |
||||
|
||||
def generate_heatmap(self, im0, tracks): |
||||
""" |
||||
Generate heatmap based on tracking data. |
||||
|
||||
Args: |
||||
im0 (nd array): Image |
||||
tracks (list): List of tracks obtained from the object tracking process. |
||||
""" |
||||
self.extract_results(tracks) |
||||
self.im0 = im0 |
||||
|
||||
for box, cls in zip(self.boxes, self.clss): |
||||
self.heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] += 1 |
||||
|
||||
# Normalize, apply colormap to heatmap and combine with original image |
||||
heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX) |
||||
heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap) |
||||
im0_with_heatmap = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0) |
||||
|
||||
if self.view_img: |
||||
self.display_frames(im0_with_heatmap) |
||||
|
||||
return im0_with_heatmap |
||||
|
||||
def display_frames(self, im0_with_heatmap): |
||||
""" |
||||
Display heatmap. |
||||
|
||||
Args: |
||||
im0_with_heatmap (nd array): Original Image with heatmap |
||||
""" |
||||
cv2.imshow('Ultralytics Heatmap', im0_with_heatmap) |
||||
if cv2.waitKey(1) & 0xFF == ord('q'): |
||||
return |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
Heatmap() |
||||
Loading…
Reference in new issue