Skip to content

Commit

Permalink
Refactor all Ultralytics Solutions (ultralytics#12790)
Browse files Browse the repository at this point in the history 
Signed-off-by: Glenn Jocher <[email protected]>
Co-authored-by: UltralyticsAssistant <[email protected]>
Co-authored-by: RizwanMunawar <[email protected]>
  • Loading branch information
@glenn-jocher @UltralyticsAssistant @RizwanMunawar
3 people authored May 18, 2024
1 parent a2ecb24 commit 2af71d1
Show file tree
Hide file tree
Showing 134 changed files with 838 additions and 1,013 deletions.
2 changes: 1 addition & 1 deletion docs/build_reference.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def create_markdown(py_filepath: Path, module_path: str, classes: list, function
f"# Reference for `{module_path}.py`\n\n"
f"!!! Note\n\n"
f" This file is available at [{url}]({url}). If you spot a problem please help fix it by [contributing]"
f"(https://docs.ultralytics.com/help/contributing/) a [Pull Request]({edit}) 🛠️. Thank you 🙏!\n\n"
f"(/help/contributing.md) a [Pull Request]({edit}) 🛠️. Thank you 🙏!\n\n"
)
md_content = ["<br><br>\n"] + [f"## ::: {module_name}.{class_name}\n\n<br><br>\n" for class_name in classes]
md_content.extend(f"## ::: {module_name}.{func_name}\n\n<br><br>\n" for func_name in functions)
Expand Down
28 changes: 12 additions & 16 deletions docs/en/guides/distance-calculation.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,7 @@ Measuring the gap between two objects is known as distance calculation within a
=== "Video Stream"

```python
from ultralytics import YOLO
from ultralytics.solutions import distance_calculation
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8n.pt")
Expand All @@ -54,14 +53,10 @@ Measuring the gap between two objects is known as distance calculation within a
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("distance_calculation.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
fps,
(w, h))
video_writer = cv2.VideoWriter("distance_calculation.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

# Init distance-calculation obj
dist_obj = distance_calculation.DistanceCalculation()
dist_obj.set_args(names=names, view_img=True)
dist_obj = solutions.DistanceCalculation(names=names, view_img=True)

while cap.isOpened():
success, im0 = cap.read()
Expand All @@ -84,15 +79,16 @@ Measuring the gap between two objects is known as distance calculation within a
- Mouse Right Click will delete all drawn points
- Mouse Left Click can be used to draw points

### Optional Arguments `set_args`
### Arguments `DistanceCalculation()`

| Name | Type | Default | Description |
|------------------|--------|-----------------|--------------------------------------------------------|
| `names` | `dict` | `None` | Classes names |
| `view_img` | `bool` | `False` | Display frames with counts |
| `line_thickness` | `int` | `2` | Increase bounding boxes thickness |
| `line_color` | `RGB` | `(255, 255, 0)` | Line Color for centroids mapping on two bounding boxes |
| `centroid_color` | `RGB` | `(255, 0, 255)` | Centroid color for each bounding box |
| `Name` | `Type` | `Default` | Description |
|--------------------|---------|-----------------|-----------------------------------------------------------|
| `names` | `dict` | `None` | Dictionary mapping class indices to class names. |
| `pixels_per_meter` | `int` | `10` | Conversion factor from pixels to meters. |
| `view_img` | `bool` | `False` | Flag to indicate if the video stream should be displayed. |
| `line_thickness` | `int` | `2` | Thickness of the lines drawn on the image. |
| `line_color` | `tuple` | `(255, 255, 0)` | Color of the lines drawn on the image (BGR format). |
| `centroid_color` | `tuple` | `(255, 0, 255)` | Color of the centroids drawn (BGR format). |

### Arguments `model.track`

Expand Down
159 changes: 61 additions & 98 deletions docs/en/guides/heatmaps.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
=== "Heatmap"

```python
from ultralytics import YOLO
from ultralytics.solutions import heatmap
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8n.pt")
Expand All @@ -54,19 +53,13 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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))
video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

# Init heatmap
heatmap_obj = heatmap.Heatmap()
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA,
imw=w,
imh=h,
view_img=True,
shape="circle",
classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names)

while cap.isOpened():
success, im0 = cap.read()
Expand All @@ -87,8 +80,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
=== "Line Counting"

```python
from ultralytics import YOLO
from ultralytics.solutions import heatmap
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8n.pt")
Expand All @@ -97,30 +89,24 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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))
video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

line_points = [(20, 400), (1080, 404)] # line for object counting

# 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=line_points,
classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
count_reg_pts=line_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)

tracks = model.track(im0, persist=True, show=False)
im0 = heatmap_obj.generate_heatmap(im0, tracks)
video_writer.write(im0)

Expand All @@ -131,8 +117,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult

=== "Polygon Counting"
```python
from ultralytics import YOLO
import heatmap
from ultralytics import YOLO, solutions
import cv2
model = YOLO("yolov8n.pt")
Expand All @@ -150,22 +135,19 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
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)
Expand All @@ -177,8 +159,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
=== "Region Counting"

```python
from ultralytics import YOLO
from ultralytics.solutions import heatmap
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8n.pt")
Expand All @@ -187,31 +168,25 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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))
video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

# 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=w,
imh=h,
view_img=True,
shape="circle",
count_reg_pts=region_points,
classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
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)

Expand All @@ -223,8 +198,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
=== "Im0"

```python
from ultralytics import YOLO
from ultralytics.solutions import heatmap
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8s.pt") # YOLOv8 custom/pretrained model
Expand All @@ -233,13 +207,10 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
h, w = im0.shape[:2] # image height and width
# Heatmap Init
heatmap_obj = heatmap.Heatmap()
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA,
imw=w,
imh=h,
view_img=True,
shape="circle",
classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names)

results = model.track(im0, persist=True)
im0 = heatmap_obj.generate_heatmap(im0, tracks=results)
Expand All @@ -249,8 +220,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
=== "Specific Classes"

```python
from ultralytics import YOLO
from ultralytics.solutions import heatmap
from ultralytics import YOLO, solutions
import cv2

model = YOLO("yolov8n.pt")
Expand All @@ -259,21 +229,15 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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))
video_writer = cv2.VideoWriter("heatmap_output.avi", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

classes_for_heatmap = [0, 2] # classes for heatmap

# Init heatmap
heatmap_obj = heatmap.Heatmap()
heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA,
imw=w,
imh=h,
view_img=True,
shape="circle",
classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names)

while cap.isOpened():
success, im0 = cap.read()
Expand All @@ -291,28 +255,27 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
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 |
| `line_thickness` | `int` | `2` | Increase bounding boxes and count text thickness |
| `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_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 |
| `count_bg_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |
### Arguments `Heatmap()`

| Name | Type | Default | Description |
|--------------------|------------------|--------------------|-------------------------------------------------------------------|
| `classes_names` | `dict` | `None` | Dictionary of class names. |
| `imw` | `int` | `0` | Image width. |
| `imh` | `int` | `0` | Image height. |
| `colormap` | `int` | `cv2.COLORMAP_JET` | Colormap to use for the heatmap. |
| `heatmap_alpha` | `float` | `0.5` | Alpha blending value for heatmap overlay. |
| `view_img` | `bool` | `False` | Whether to display the image with the heatmap overlay. |
| `view_in_counts` | `bool` | `True` | Whether to display the count of objects entering the region. |
| `view_out_counts` | `bool` | `True` | Whether to display the count of objects exiting the region. |
| `count_reg_pts` | `list` or `None` | `None` | Points defining the counting region (either a line or a polygon). |
| `count_txt_color` | `tuple` | `(0, 0, 0)` | Text color for displaying counts. |
| `count_bg_color` | `tuple` | `(255, 255, 255)` | Background color for displaying counts. |
| `count_reg_color` | `tuple` | `(255, 0, 255)` | Color for the counting region. |
| `region_thickness` | `int` | `5` | Thickness of the region line. |
| `line_dist_thresh` | `int` | `15` | Distance threshold for line-based counting. |
| `line_thickness` | `int` | `2` | Thickness of the lines used in drawing. |
| `decay_factor` | `float` | `0.99` | Decay factor for the heatmap to reduce intensity over time. |
| `shape` | `str` | `"circle"` | Shape of the heatmap blobs ('circle' or 'rect'). |

### Arguments `model.track`

Expand Down
Loading

0 comments on commit 2af71d1

Please sign in to comment.