New 🌟 per-class object counting feature and updates (#9443)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar
GitHub
parent
1703025e8e
commit
18036908d4
5 changed files with 274 additions and 150 deletions
|
|
@ -65,7 +65,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
imw=w,
|
||||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle")
|
||||
shape="circle",
|
||||
classes_names=model.names)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
|
|
@ -110,7 +111,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle",
|
||||
count_reg_pts=line_points)
|
||||
count_reg_pts=line_points,
|
||||
classes_names=model.names)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
|
|
@ -126,6 +128,51 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
=== "Polygon Counting"
|
||||
```python
|
||||
from ultralytics import YOLO
|
||||
import heatmap
|
||||
import cv2
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
cap = cv2.VideoCapture("path/to/video/file.mp4")
|
||||
assert cap.isOpened(), "Error reading video file"
|
||||
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
||||
|
||||
# Video writer
|
||||
video_writer = cv2.VideoWriter("heatmap_output.avi",
|
||||
cv2.VideoWriter_fourcc(*'mp4v'),
|
||||
fps,
|
||||
(w, h))
|
||||
|
||||
# Define polygon points
|
||||
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360), (20, 400)]
|
||||
|
||||
# Init heatmap
|
||||
heatmap_obj = heatmap.Heatmap()
|
||||
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA,
|
||||
imw=w,
|
||||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle",
|
||||
count_reg_pts=region_points,
|
||||
classes_names=model.names)
|
||||
|
||||
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()
|
||||
```
|
||||
|
||||
=== "Region Counting"
|
||||
|
||||
|
|
@ -155,7 +202,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle",
|
||||
count_reg_pts=region_points)
|
||||
count_reg_pts=region_points,
|
||||
classes_names=model.names)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
|
|
@ -190,7 +238,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
imw=w,
|
||||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle")
|
||||
shape="circle",
|
||||
classes_names=model.names)
|
||||
|
||||
results = model.track(im0, persist=True)
|
||||
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
|
||||
|
|
@ -223,7 +272,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
|
|||
imw=w,
|
||||
imh=h,
|
||||
view_img=True,
|
||||
shape="circle")
|
||||
shape="circle",
|
||||
classes_names=model.names)
|
||||
|
||||
while cap.isOpened():
|
||||
success, im0 = cap.read()
|
||||
|
|
@ -243,22 +293,27 @@ 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_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 |
|
||||
| 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 |
|
||||
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
|
||||
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
|
||||
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
|
||||
| `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_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 |
|
||||
| `fontsize` | `float` | `0.6` | Font size of counting text |
|
||||
| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
|
||||
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |
|
||||
|
||||
### Arguments `model.track`
|
||||
|
||||
|
|
|
|||
|
|
@ -229,9 +229,11 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
|
|||
| `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 |
|
||||
| `count_txt_color` | `RGB Color` | `(255, 255, 255)` | Foreground color for Object counts text |
|
||||
| `region_thickness` | `int` | `5` | Thickness for object counter region or line |
|
||||
| `fontsize` | `float` | `0.6` | Font size of counting text |
|
||||
| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
|
||||
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |
|
||||
|
||||
### Arguments `model.track`
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue