Refactor all Ultralytics Solutions (#12790)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: RizwanMunawar <chr043416@gmail.com>
This commit is contained in:
Glenn Jocher 2024-05-18 18:14:42 +02:00 committed by GitHub
parent a2ecb24176
commit 2af71d15a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
134 changed files with 845 additions and 1020 deletions

View file

@ -51,42 +51,39 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
=== "Count in Region"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
from ultralytics import YOLO, solutions
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))
# 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'),
fps,
(w, h))
video_writer = cv2.VideoWriter("object_counting_output.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2)
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2,
)
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()
@ -95,9 +92,8 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
=== "Count in Polygon"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
from ultralytics import YOLO, solutions
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
@ -108,18 +104,16 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360), (20, 400)]
# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
fps,
(w, h))
video_writer = cv2.VideoWriter("object_counting_output.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2)
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2,
)
while cap.isOpened():
success, im0 = cap.read()
@ -139,42 +133,39 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
=== "Count in Line"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
from ultralytics import YOLO, solutions
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))
# Define line points
line_points = [(20, 400), (1080, 400)]
# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
cv2.VideoWriter_fourcc(*'mp4v'),
fps,
(w, h))
video_writer = cv2.VideoWriter("object_counting_output.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2)
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2,
)
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()
@ -183,43 +174,39 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
=== "Specific Classes"
```python
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
from ultralytics import YOLO, solutions
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))
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'),
fps,
(w, h))
video_writer = cv2.VideoWriter("object_counting_output.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2)
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True,
line_thickness=2,
)
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)
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()
@ -229,24 +216,27 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
You can move the region anywhere in the frame by clicking on its edges
### Optional Arguments `set_args`
### Argument `ObjectCounter`
| Name | Type | Default | Description |
|--------------------|-------------|----------------------------|--------------------------------------------------|
| `view_img` | `bool` | `False` | Display frames with counts |
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
| `line_thickness` | `int` | `2` | Increase bounding boxes and count text thickness |
| `reg_pts` | `list` | `[(20, 400), (1260, 400)]` | Points defining the Region Area |
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
| `count_reg_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_color` | `RGB Color` | `(255, 255, 255)` | Foreground color for Object counts text |
| `region_thickness` | `int` | `5` | Thickness for object counter region or line |
| `count_bg_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
Here's a table with the `ObjectCounter` arguments:
| Name | Type | Default | Description |
|----------------------|---------|----------------------------|------------------------------------------------------------------------|
| `classes_names` | `dict` | `None` | Dictionary of class names. |
| `reg_pts` | `list` | `[(20, 400), (1260, 400)]` | List of points defining the counting region. |
| `count_reg_color` | `tuple` | `(255, 0, 255)` | RGB color of the counting region. |
| `count_txt_color` | `tuple` | `(0, 0, 0)` | RGB color of the count text. |
| `count_bg_color` | `tuple` | `(255, 255, 255)` | RGB color of the count text background. |
| `line_thickness` | `int` | `2` | Line thickness for bounding boxes. |
| `track_thickness` | `int` | `2` | Thickness of the track lines. |
| `view_img` | `bool` | `False` | Flag to control whether to display the video stream. |
| `view_in_counts` | `bool` | `True` | Flag to control whether to display the in counts on the video stream. |
| `view_out_counts` | `bool` | `True` | Flag to control whether to display the out counts on the video stream. |
| `draw_tracks` | `bool` | `False` | Flag to control whether to draw the object tracks. |
| `track_color` | `tuple` | `None` | RGB color of the tracks. |
| `region_thickness` | `int` | `5` | Thickness of the object counting region. |
| `line_dist_thresh` | `int` | `15` | Euclidean distance threshold for line counter. |
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count. |
### Arguments `model.track`