Add line counting and circular heatmaps in Ultralytics Solutions (#7113)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Muhammad Rizwan Munawar
GitHub
parent
a5735724c5
commit
38eaf5e29f
5 changed files with 526 additions and 247 deletions
|
|
@ -20,14 +20,19 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
|
||||
| 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"
|
||||
???+ tip "decay_factor"
|
||||
|
||||
Used for removal of heatmap after object removed from frame, value should be in range (0.0 - 1.0)
|
||||
|
||||
|
||||
!!! Example "Heatmaps using Ultralytics YOLOv8 Example"
|
||||
|
||||
=== "Heatmap"
|
||||
```python
|
||||
|
|
@ -35,31 +40,126 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8s.pt") # YOLOv8 custom/pretrained model
|
||||
model = YOLO("yolov8n.pt")
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
# Heatmap Init
|
||||
# Video writer
|
||||
video_writer = cv2.VideoWriter("heatmap_output.avi",
|
||||
cv2.VideoWriter_fourcc(*'mp4v'),
|
||||
int(cap.get(5)),
|
||||
(int(cap.get(3)), int(cap.get(4))))
|
||||
|
||||
# Init heatmap
|
||||
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
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA ,
|
||||
imw=cap.get(4), # should same as cap height
|
||||
imh=cap.get(3), # should same as cap width
|
||||
view_img=True,
|
||||
decay_factor=0.99)
|
||||
shape="circle")
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
tracks = model.track(im0, persist=True, show=False)
|
||||
|
||||
results = model.track(im0, persist=True)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks)
|
||||
video_writer.write(im0)
|
||||
|
||||
cap.release()
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
```
|
||||
|
||||
=== "Line Counting"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
# Video writer
|
||||
video_writer = cv2.VideoWriter("heatmap_output.avi",
|
||||
cv2.VideoWriter_fourcc(*'mp4v'),
|
||||
int(cap.get(5)),
|
||||
(int(cap.get(3)), int(cap.get(4))))
|
||||
|
||||
line_points = [(256, 409), (694, 532)] # line for object counting
|
||||
|
||||
# Init heatmap
|
||||
heatmap_obj = heatmap.Heatmap()
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA ,
|
||||
imw=cap.get(4), # should same as cap height
|
||||
imh=cap.get(3), # should same as cap width
|
||||
view_img=True,
|
||||
shape="circle",
|
||||
count_reg_pts=line_points)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
tracks = model.track(im0, persist=True, show=False)
|
||||
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks)
|
||||
video_writer.write(im0)
|
||||
|
||||
cap.release()
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Heatmap with im0"
|
||||
=== "Region Counting"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
# Video writer
|
||||
video_writer = cv2.VideoWriter("heatmap_output.avi",
|
||||
cv2.VideoWriter_fourcc(*'mp4v'),
|
||||
int(cap.get(5)),
|
||||
(int(cap.get(3)), int(cap.get(4))))
|
||||
|
||||
# Define region points
|
||||
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
|
||||
|
||||
# Init heatmap
|
||||
heatmap_obj = heatmap.Heatmap()
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA ,
|
||||
imw=cap.get(4), # should same as cap height
|
||||
imh=cap.get(3), # should same as cap width
|
||||
view_img=True,
|
||||
shape="circle",
|
||||
count_reg_pts=region_points)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
tracks = model.track(im0, persist=True, show=False)
|
||||
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks)
|
||||
video_writer.write(im0)
|
||||
|
||||
cap.release()
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Im0"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
|
|
@ -71,10 +171,11 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
|
||||
# Heatmap Init
|
||||
heatmap_obj = heatmap.Heatmap()
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_JET,
|
||||
imw=im0.shape[0], # should same as im0 width
|
||||
imh=im0.shape[1], # should same as im0 height
|
||||
view_img=True)
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA ,
|
||||
imw=cap.get(4), # should same as cap height
|
||||
imh=cap.get(3), # should same as cap width
|
||||
view_img=True,
|
||||
shape="circle")
|
||||
|
||||
|
||||
results = model.track(im0, persist=True)
|
||||
|
|
@ -82,43 +183,13 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
cv2.imwrite("ultralytics_output.png", im0)
|
||||
```
|
||||
|
||||
=== "Heatmap with Specific Classes"
|
||||
=== "Specific Classes"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8s.pt") # YOLOv8 custom/pretrained model
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
classes_for_heatmap = [0, 2]
|
||||
|
||||
# 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)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
results = model.track(im0, persist=True, classes=classes_for_heatmap)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Heatmap with Save Output"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8s.pt") # YOLOv8 custom/pretrained model
|
||||
model = YOLO("yolov8n.pt")
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
|
|
@ -128,74 +199,50 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
int(cap.get(5)),
|
||||
(int(cap.get(3)), int(cap.get(4))))
|
||||
|
||||
# Heatmap init
|
||||
classes_for_heatmap = [0, 2] # classes for heatmap
|
||||
|
||||
# Init heatmap
|
||||
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)
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA ,
|
||||
imw=cap.get(4), # should same as cap height
|
||||
imh=cap.get(3), # should same as cap width
|
||||
view_img=True,
|
||||
shape="circle")
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
results = model.track(im0, persist=True)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
tracks = model.track(im0, persist=True, show=False,
|
||||
classes=classes_for_heatmap)
|
||||
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks)
|
||||
video_writer.write(im0)
|
||||
|
||||
cap.release()
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Heatmap with Object Counting"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.solutions import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8s.pt") # YOLOv8 custom/pretrained model
|
||||
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4") # Video file Path, webcam 0
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
|
||||
# Region for object counting
|
||||
count_reg_pts = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
|
||||
|
||||
# Heatmap Init
|
||||
heatmap_obj = heatmap.Heatmap()
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_JET,
|
||||
imw=cap.get(4), # should same as cap width
|
||||
imh=cap.get(3), # should same as cap height
|
||||
view_img=True,
|
||||
count_reg_pts=count_reg_pts)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
if not success:
|
||||
print("Video frame is empty or video processing has been successfully completed.")
|
||||
break
|
||||
results = model.track(im0, persist=True)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
```
|
||||
|
||||
### 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 |
|
||||
| decay_factor | `float` | `0.99` | Decay factor for heatmap area removal after specific time |
|
||||
| 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_txt_color | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text |
|
||||
| count_color | `RGB Color` | `(255, 255, 255)` | Background color for Object counts text |
|
||||
| count_reg_color | `RGB Color` | `(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 |
|
||||
| shape | `str` | `circle` | Heatmap shape for display "rect" or "circle" supported |
|
||||
| line_dist_thresh | `int` | `15` | Euclidean Distance threshold for line counter |
|
||||
|
||||
### Arguments `model.track`
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue