HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add line counting and circular heatmaps in Ultralytics Solutions (#7113)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Muhammad Rizwan Munawar 2023-12-22 05:56:44 +05:00 • committed by GitHub
parent a5735724c5
commit 38eaf5e29f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 526 additions and 247 deletions
Download patch file Download diff file Expand all files Collapse all files

255
docs/en/guides/heatmaps.md
View file

@ -20,14 +20,19 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
| Transportation | Retail |
|:-----------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------:|
| ![Ultralytics YOLOv8 Transportation Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/50d197b8-c7f6-4ecf-a664-3d4363b073de) | ![Ultralytics YOLOv8 Retail Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/ffd0649f-5ff5-48d2-876d-6bdffeff5c54) |
| ![Ultralytics YOLOv8 Transportation Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/288d7053-622b-4452-b4e4-1f41aeb764aa) | ![Ultralytics YOLOv8 Retail Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/a9139af0-2cb7-41fe-a0d5-29a300dee768) |
| 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`

196
docs/en/guides/object-counting.md
View file

@ -34,9 +34,9 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
| ![Conveyor Belt Packets Counting Using Ultralytics YOLOv8](https://github.com/RizwanMunawar/ultralytics/assets/62513924/70e2d106-510c-4c6c-a57a-d34a765aa757) | ![Fish Counting in Sea using Ultralytics YOLOv8](https://github.com/RizwanMunawar/ultralytics/assets/62513924/c60d047b-3837-435f-8d29-bb9fc95d2191) |
| Conveyor Belt Packets Counting Using Ultralytics YOLOv8 | Fish Counting in Sea using Ultralytics YOLOv8 |
!!! Example "Object Counting Example"
!!! Example "Object Counting using YOLOv8 Example"
=== "Object Counting"
=== "Region"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
@ -46,75 +46,21 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
counter = object_counter.ObjectCounter() # Init Object Counter
# Define region points
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
int(cap.get(5)),
(int(cap.get(3)), int(cap.get(4))))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
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 = counter.start_counting(im0, tracks)
cv2.destroyAllWindows()
```
=== "Object Counting with Specific Classes"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
classes_to_count = [0, 2]
counter = object_counter.ObjectCounter() # Init Object Counter
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
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,
classes=classes_to_count)
im0 = counter.start_counting(im0, tracks)
cv2.destroyAllWindows()
```
=== "Object Counting with Save Output"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
video_writer = cv2.VideoWriter("object_counting.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
int(cap.get(5)),
(int(cap.get(3)), int(cap.get(4))))
counter = object_counter.ObjectCounter() # Init Object Counter
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
while cap.isOpened():
success, im0 = cap.read()
@ -122,9 +68,95 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
print("Video frame is empty or video processing has been successfully completed.")
break
tracks = model.track(im0, persist=True, show=False)
im0 = counter.start_counting(im0, tracks)
video_writer.write(im0)
cap.release()
video_writer.release()
cv2.destroyAllWindows()
```
=== "Line"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
# Define line points
line_points = [(20, 400), (1080, 400)]
# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
int(cap.get(5)),
(int(cap.get(3)), int(cap.get(4))))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True)
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 = counter.start_counting(im0, tracks)
video_writer.write(im0)
cap.release()
video_writer.release()
cv2.destroyAllWindows()
```
=== "Specific Classes"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
line_points = [(20, 400), (1080, 400)] # line or region points
classes_to_count = [0, 2] # person and car classes for count
# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
int(cap.get(5)),
(int(cap.get(3)), int(cap.get(4))))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True)
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,
classes=classes_to_count)
im0 = counter.start_counting(im0, tracks)
video_writer.write(im0)
cap.release()
video_writer.release()
cv2.destroyAllWindows()
```
@ -135,15 +167,22 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
### Optional Arguments `set_args`
| Name | Type | Default | Description |
|-----------------|---------|--------------------------------------------------|---------------------------------------|
| view_img | `bool` | `False` | Display the frame with counts |
| line_thickness | `int` | `2` | Increase the thickness of count value |
| reg_pts | `list` | `(20, 400), (1080, 404), (1080, 360), (20, 360)` | Region Area Points |
| classes_names | `dict` | `model.model.names` | Classes Names Dict |
| region_color | `tuple` | `(0, 255, 0)` | Region Area Color |
| track_thickness | `int` | `2` | Tracking line thickness |
| draw_tracks | `bool` | `False` | Draw Tracks lines |
| Name | Type | Default | Description |
|---------------------|-------------|----------------------------|-----------------------------------------------|
| view_img | `bool` | `False` | Display frames with counts |
| line_thickness | `int` | `2` | Increase bounding boxes thickness |
| reg_pts | `list` | `[(20, 400), (1260, 400)]` | Points defining the Region Area |
| classes_names | `dict` | `model.model.names` | Dictionary of Class Names |
| region_color | `RGB Color` | `(255, 0, 255)` | Color of the Object counting Region or Line |
| track_thickness | `int` | `2` | Thickness of Tracking Lines |
| draw_tracks | `bool` | `False` | Enable drawing Track lines |
| track_color | `RGB Color` | `(0, 255, 0)` | Color for each track line |
| line_dist_thresh | `int` | `15` | Euclidean Distance threshold for line counter |
| count_txt_thickness | `int` | `2` | Thickness of Object counts text |
| 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 |
| region_thickness | `int` | `5` | Thickness for object counter region or line |
### Arguments `model.track`
@ -155,3 +194,4 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
| `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] |
| `verbose` | `bool` | `True` | Display the object tracking results |