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

ultralytics 8.2.57 new Solutions Tests and Docs (#14408)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-07-14 23:00:14 +05:00 • committed by GitHub
parent 100a73b6e9
commit c67a3039c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 143 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

2
docs/en/guides/distance-calculation.md
View file

@ -83,7 +83,7 @@ Measuring the gap between two objects is known as distance calculation within a
| `Name` | `Type` | `Default` | Description |
| ------------------ | ------- | --------------- | --------------------------------------------------------- |
| `names` | `dict` | `None` | Dictionary mapping class indices to class names. |
| `names` | `dict` | `None` | Dictionary of classes 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. |

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

@ -61,7 +61,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names,
names=model.names,
)
while cap.isOpened():
@ -102,7 +102,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
view_img=True,
shape="circle",
count_reg_pts=line_points,
classes_names=model.names,
names=model.names,
)
while cap.isOpened():
@ -144,7 +144,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
view_img=True,
shape="circle",
count_reg_pts=region_points,
classes_names=model.names,
names=model.names,
)
while cap.isOpened():
@ -186,7 +186,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
view_img=True,
shape="circle",
count_reg_pts=region_points,
classes_names=model.names,
names=model.names,
)
while cap.isOpened():
@ -221,7 +221,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names,
names=model.names,
)
results = model.track(im0, persist=True)
@ -251,7 +251,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
colormap=cv2.COLORMAP_PARULA,
view_img=True,
shape="circle",
classes_names=model.names,
names=model.names,
)
while cap.isOpened():
@ -273,7 +273,7 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult
| Name | Type | Default | Description |
| ------------------ | ---------------- | ------------------ | ----------------------------------------------------------------- |
| `classes_names` | `dict` | `None` | Dictionary of class names. |
| `names` | `list` | `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. |
@ -348,7 +348,7 @@ from ultralytics import YOLO, solutions
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, view_img=True, shape="circle", classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, view_img=True, shape="circle", names=model.names)
while cap.isOpened():
success, im0 = cap.read()
@ -381,7 +381,7 @@ from ultralytics import YOLO, solutions
model = YOLO("yolov8n.pt")
cap = cv2.VideoCapture("path/to/video/file.mp4")
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, view_img=True, shape="circle", classes_names=model.names)
heatmap_obj = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, view_img=True, shape="circle", names=model.names)
classes_for_heatmap = [0, 2] # Classes to visualize
while cap.isOpened():

20
docs/en/guides/instance-segmentation-and-tracking.md
View file

@ -64,7 +64,9 @@ There are two types of instance segmentation tracking available in the Ultralyti
clss = results[0].boxes.cls.cpu().tolist()
masks = results[0].masks.xy
for mask, cls in zip(masks, clss):
annotator.seg_bbox(mask=mask, mask_color=colors(int(cls), True), det_label=names[int(cls)])
color = colors(int(cls), True)
txt_color = annotator.get_txt_color(color)
annotator.seg_bbox(mask=mask, mask_color=color, label=names[int(cls)], txt_color=txt_color)
out.write(im0)
cv2.imshow("instance-segmentation", im0)
@ -110,7 +112,9 @@ There are two types of instance segmentation tracking available in the Ultralyti
track_ids = results[0].boxes.id.int().cpu().tolist()
for mask, track_id in zip(masks, track_ids):
annotator.seg_bbox(mask=mask, mask_color=colors(track_id, True), track_label=str(track_id))
color = colors(int(track_id), True)
txt_color = annotator.get_txt_color(color)
annotator.seg_bbox(mask=mask, mask_color=color, label=str(track_id), txt_color=txt_color)
out.write(im0)
cv2.imshow("instance-segmentation-object-tracking", im0)
@ -125,12 +129,12 @@ There are two types of instance segmentation tracking available in the Ultralyti
### `seg_bbox` Arguments
| Name | Type | Default | Description |
| ------------- | ------- | --------------- | -------------------------------------- |
| `mask` | `array` | `None` | Segmentation mask coordinates |
| `mask_color` | `tuple` | `(255, 0, 255)` | Mask color for every segmented box |
| `det_label` | `str` | `None` | Label for segmented object |
| `track_label` | `str` | `None` | Label for segmented and tracked object |
| Name | Type | Default | Description |
| ------------ | ------- | --------------- | -------------------------------------------- |
| `mask` | `array` | `None` | Segmentation mask coordinates |
| `mask_color` | `RGB` | `(255, 0, 255)` | Mask color for every segmented box |
| `label` | `str` | `None` | Label for segmented object |
| `txt_color` | `RGB` | `None` | Label color for segmented and tracked object |
## Note

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

@ -70,7 +70,7 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=region_points,
classes_names=model.names,
names=model.names,
draw_tracks=True,
line_thickness=2,
)
@ -112,7 +112,7 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=region_points,
classes_names=model.names,
names=model.names,
draw_tracks=True,
line_thickness=2,
)
@ -154,7 +154,7 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=line_points,
classes_names=model.names,
names=model.names,
draw_tracks=True,
line_thickness=2,
)
@ -196,7 +196,7 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter = solutions.ObjectCounter(
view_img=True,
reg_pts=line_points,
classes_names=model.names,
names=model.names,
draw_tracks=True,
line_thickness=2,
)
@ -226,7 +226,7 @@ Here's a table with the `ObjectCounter` arguments:
| Name | Type | Default | Description |
| -------------------- | ------- | -------------------------- | ---------------------------------------------------------------------- |
| `classes_names` | `dict` | `None` | Dictionary of class names. |
| `names` | `dict` | `None` | Dictionary of classes 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. |
@ -283,7 +283,7 @@ def count_objects_in_region(video_path, output_video_path, model_path):
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
counter = solutions.ObjectCounter(
view_img=True, reg_pts=region_points, classes_names=model.names, draw_tracks=True, line_thickness=2
view_img=True, reg_pts=region_points, names=model.names, draw_tracks=True, line_thickness=2
)
while cap.isOpened():
@ -334,7 +334,7 @@ def count_specific_classes(video_path, output_video_path, model_path, classes_to
line_points = [(20, 400), (1080, 400)]
video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
counter = solutions.ObjectCounter(
view_img=True, reg_pts=line_points, classes_names=model.names, draw_tracks=True, line_thickness=2
view_img=True, reg_pts=line_points, names=model.names, draw_tracks=True, line_thickness=2
)
while cap.isOpened():

10
docs/en/guides/queue-management.md
View file

@ -53,7 +53,7 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager(
classes_names=model.names,
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
@ -97,7 +97,7 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager(
classes_names=model.names,
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
@ -127,7 +127,7 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
| Name | Type | Default | Description |
| ------------------- | ---------------- | -------------------------- | ----------------------------------------------------------------------------------- |
| `classes_names` | `dict` | `model.names` | A dictionary mapping class IDs to class names. |
| `names` | `dict` | `model.names` | A dictionary mapping class IDs to class names. |
| `reg_pts` | `list of tuples` | `[(20, 400), (1260, 400)]` | Points defining the counting region polygon. Defaults to a predefined rectangle. |
| `line_thickness` | `int` | `2` | Thickness of the annotation lines. |
| `track_thickness` | `int` | `2` | Thickness of the track lines. |
@ -175,7 +175,7 @@ cap = cv2.VideoCapture("path/to/video.mp4")
queue_region = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
queue = solutions.QueueManager(
classes_names=model.names,
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
@ -228,7 +228,7 @@ Example for airports:
```python
queue_region_airport = [(50, 600), (1200, 600), (1200, 550), (50, 550)]
queue_airport = solutions.QueueManager(
classes_names=model.names,
names=model.names,
reg_pts=queue_region_airport,
line_thickness=3,
fontsize=1.0,

90
tests/test_solutions.py Normal file
View file

@ -0,0 +1,90 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import cv2
import pytest
from ultralytics import YOLO, solutions
from ultralytics.utils.downloads import safe_download
MAJOR_SOLUTIONS_DEMO = "https://github.com/ultralytics/assets/releases/download/v0.0.0/solutions_ci_demo.mp4"
WORKOUTS_SOLUTION_DEMO = "https://github.com/ultralytics/assets/releases/download/v0.0.0/solution_ci_pose_demo.mp4"
@pytest.mark.slow
def test_major_solutions():
"""Test the object counting, heatmap, speed estimation and queue management solution."""
safe_download(url=MAJOR_SOLUTIONS_DEMO)
model = YOLO("yolov8n.pt")
names = model.names
cap = cv2.VideoCapture("solutions_ci_demo.mp4")
assert cap.isOpened(), "Error reading video file"
region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
counter = solutions.ObjectCounter(reg_pts=region_points, names=names, view_img=False)
heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, names=names, view_img=False)
speed = solutions.SpeedEstimator(reg_pts=region_points, names=names, view_img=False)
queue = solutions.QueueManager(names=names, reg_pts=region_points, view_img=False)
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
original_im0 = im0.copy()
tracks = model.track(im0, persist=True, show=False)
_ = counter.start_counting(original_im0.copy(), tracks)
_ = heatmap.generate_heatmap(original_im0.copy(), tracks)
_ = speed.estimate_speed(original_im0.copy(), tracks)
_ = queue.process_queue(original_im0.copy(), tracks)
cap.release()
cv2.destroyAllWindows()
@pytest.mark.slow
def test_aigym():
"""Test the workouts monitoring solution."""
safe_download(url=WORKOUTS_SOLUTION_DEMO)
model = YOLO("yolov8n-pose.pt")
cap = cv2.VideoCapture("solution_ci_pose_demo.mp4")
assert cap.isOpened(), "Error reading video file"
gym_object = solutions.AIGym(line_thickness=2, pose_type="squat", kpts_to_check=[5, 11, 13])
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
results = model.track(im0, verbose=False)
_ = gym_object.start_counting(im0, results)
cap.release()
cv2.destroyAllWindows()
@pytest.mark.slow
def test_instance_segmentation():
"""Test the instance segmentation solution."""
from ultralytics.utils.plotting import Annotator, colors
model = YOLO("yolov8n-seg.pt")
names = model.names
cap = cv2.VideoCapture("solutions_ci_demo.mp4")
assert cap.isOpened(), "Error reading video file"
while cap.isOpened():
success, im0 = cap.read()
if not success:
break
results = model.predict(im0)
annotator = Annotator(im0, line_width=2)
if results[0].masks is not None:
clss = results[0].boxes.cls.cpu().tolist()
masks = results[0].masks.xy
for mask, cls in zip(masks, clss):
color = colors(int(cls), True)
annotator.seg_bbox(mask=mask, mask_color=color, label=names[int(cls)])
cap.release()
cv2.destroyAllWindows()
@pytest.mark.slow
def test_streamlit_predict():
"""Test streamlit predict live inference solution."""
solutions.inference()

2
ultralytics/__init__.py
View file

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
__version__ = "8.2.56"
__version__ = "8.2.57"
import os

2
ultralytics/solutions/distance_calculation.py
View file

@ -24,7 +24,7 @@ class DistanceCalculation:
Initializes the DistanceCalculation class with the given parameters.
Args:
names (dict): Dictionary mapping class indices to class names.
names (dict): Dictionary of classes names.
pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.

4
ultralytics/solutions/heatmap.py
View file

@ -18,7 +18,7 @@ class Heatmap:
def __init__(
self,
classes_names,
names,
imw=0,
imh=0,
colormap=cv2.COLORMAP_JET,
@ -44,7 +44,7 @@ class Heatmap:
self.shape = shape
self.initialized = False
self.names = classes_names # Classes names
self.names = names # Classes names
# Image information
self.imw = imw

6
ultralytics/solutions/object_counter.py
View file

@ -17,7 +17,7 @@ class ObjectCounter:
def __init__(
self,
classes_names,
names,
reg_pts=None,
count_reg_color=(255, 0, 255),
count_txt_color=(0, 0, 0),
@ -37,7 +37,7 @@ class ObjectCounter:
Initializes the ObjectCounter with various tracking and counting parameters.
Args:
classes_names (dict): Dictionary of class names.
names (dict): Dictionary of class names.
reg_pts (list): List of points defining the counting region.
count_reg_color (tuple): RGB color of the counting region.
count_txt_color (tuple): RGB color of the count text.
@ -72,7 +72,7 @@ class ObjectCounter:
self.view_in_counts = view_in_counts
self.view_out_counts = view_out_counts
self.names = classes_names # Classes names
self.names = names # Classes names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Object Counter"

8
ultralytics/solutions/queue_management.py
View file

@ -17,7 +17,7 @@ class QueueManager:
def __init__(
self,
classes_names,
names,
reg_pts=None,
line_thickness=2,
track_thickness=2,
@ -34,7 +34,7 @@ class QueueManager:
Initializes the QueueManager with specified parameters for tracking and counting objects.
Args:
classes_names (dict): A dictionary mapping class IDs to class names.
names (dict): A dictionary mapping class IDs to class names.
reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
rectangle.
line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
@ -69,7 +69,7 @@ class QueueManager:
self.view_queue_counts = view_queue_counts
self.fontsize = fontsize
self.names = classes_names # Class names
self.names = names # Class names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Queue Manager"
@ -139,7 +139,7 @@ class QueueManager:
def display_frames(self):
"""Displays the current frame with annotations."""
if self.env_check:
if self.env_check and self.view_img:
self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color)
cv2.namedWindow(self.window_name)
cv2.imshow(self.window_name, self.im0)

17
ultralytics/utils/plotting.py
View file

@ -726,20 +726,18 @@ class Annotator:
)
cv2.putText(self.im, stage_text, stage_text_position, 0, self.sf, txt_color, self.tf)
def seg_bbox(self, mask, mask_color=(255, 0, 255), det_label=None, track_label=None):
def seg_bbox(self, mask, mask_color=(255, 0, 255), label=None, txt_color=(255, 255, 255)):
"""
Function for drawing segmented object in bounding box shape.
Args:
mask (list): masks data list for instance segmentation area plotting
mask_color (tuple): mask foreground color
det_label (str): Detection label text
track_label (str): Tracking label text
mask_color (RGB): mask foreground color
label (str): Detection label text
txt_color (RGB): text color
"""
cv2.polylines(self.im, [np.int32([mask])], isClosed=True, color=mask_color, thickness=2)
label = f"Track ID: {track_label}" if track_label else det_label
text_size, _ = cv2.getTextSize(label, 0, self.sf, self.tf)
cv2.rectangle(
@ -750,9 +748,10 @@ class Annotator:
-1,
)
cv2.putText(
self.im, label, (int(mask[0][0]) - text_size[0] // 2, int(mask[0][1])), 0, self.sf, (255, 255, 255), self.tf
)
if label:
cv2.putText(
self.im, label, (int(mask[0][0]) - text_size[0] // 2, int(mask[0][1])), 0, self.sf, txt_color, self.tf
)
def plot_distance_and_line(self, distance_m, distance_mm, centroids, line_color, centroid_color):
"""