Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

New 🌟 per-class object counting feature and updates #9443

Merged
merged 5 commits into from
Mar 31, 2024
Merged

New 🌟 per-class object counting feature and updates #9443

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
79567cd
Add classwise counts display support + fix object counts bug + docs u…
RizwanMunawar Mar 31, 2024
23b149c
revert mkdocs_github_authors.yaml
RizwanMunawar Mar 31, 2024
c3408eb
Auto-format by https://ultralytics.com/actions
UltralyticsAssistant Mar 31, 2024
9a6ea58
set view_in_counts=True by default
RizwanMunawar Mar 31, 2024
b863689
Merge branch 'main' into solutions_update
glenn-jocher Mar 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 76 additions & 21 deletions docs/en/guides/heatmaps.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand All @@ -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"

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand All @@ -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`

Expand Down
6 changes: 4 additions & 2 deletions docs/en/guides/object-counting.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`

Expand Down
Loading
Loading