From c67a3039c1f7d1b14170910f1b936a61025c085f Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Sun, 14 Jul 2024 23:00:14 +0500 Subject: [PATCH] `ultralytics 8.2.57` new Solutions Tests and Docs (#14408) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- docs/en/guides/distance-calculation.md | 2 +- docs/en/guides/heatmaps.md | 18 ++-- .../instance-segmentation-and-tracking.md | 20 +++-- docs/en/guides/object-counting.md | 14 +-- docs/en/guides/queue-management.md | 10 +-- tests/test_solutions.py | 90 +++++++++++++++++++ ultralytics/__init__.py | 2 +- ultralytics/solutions/distance_calculation.py | 2 +- ultralytics/solutions/heatmap.py | 4 +- ultralytics/solutions/object_counter.py | 6 +- ultralytics/solutions/queue_management.py | 8 +- ultralytics/utils/plotting.py | 17 ++-- 12 files changed, 143 insertions(+), 50 deletions(-) create mode 100644 tests/test_solutions.py diff --git a/docs/en/guides/distance-calculation.md b/docs/en/guides/distance-calculation.md index 8164a8f7..52d6460f 100644 --- a/docs/en/guides/distance-calculation.md +++ b/docs/en/guides/distance-calculation.md @@ -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. | diff --git a/docs/en/guides/heatmaps.md b/docs/en/guides/heatmaps.md index 65adf3e3..70a80bdd 100644 --- a/docs/en/guides/heatmaps.md +++ b/docs/en/guides/heatmaps.md @@ -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(): diff --git a/docs/en/guides/instance-segmentation-and-tracking.md b/docs/en/guides/instance-segmentation-and-tracking.md index 5eb5ba48..a26c2994 100644 --- a/docs/en/guides/instance-segmentation-and-tracking.md +++ b/docs/en/guides/instance-segmentation-and-tracking.md @@ -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 diff --git a/docs/en/guides/object-counting.md b/docs/en/guides/object-counting.md index f0fa5df9..033b845f 100644 --- a/docs/en/guides/object-counting.md +++ b/docs/en/guides/object-counting.md @@ -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(): diff --git a/docs/en/guides/queue-management.md b/docs/en/guides/queue-management.md index 664805a8..84724cb3 100644 --- a/docs/en/guides/queue-management.md +++ b/docs/en/guides/queue-management.md @@ -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, diff --git a/tests/test_solutions.py b/tests/test_solutions.py new file mode 100644 index 00000000..c8f05b04 --- /dev/null +++ b/tests/test_solutions.py @@ -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() diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 0d552a9b..4f65b69e 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license -__version__ = "8.2.56" +__version__ = "8.2.57" import os diff --git a/ultralytics/solutions/distance_calculation.py b/ultralytics/solutions/distance_calculation.py index 2707e6ad..35a791ca 100644 --- a/ultralytics/solutions/distance_calculation.py +++ b/ultralytics/solutions/distance_calculation.py @@ -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. diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py index c40cafb1..e0c6ec9d 100644 --- a/ultralytics/solutions/heatmap.py +++ b/ultralytics/solutions/heatmap.py @@ -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 diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index e7de1cb1..fcafb53f 100644 --- a/ultralytics/solutions/object_counter.py +++ b/ultralytics/solutions/object_counter.py @@ -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" diff --git a/ultralytics/solutions/queue_management.py b/ultralytics/solutions/queue_management.py index 96ac6291..0a2f7320 100644 --- a/ultralytics/solutions/queue_management.py +++ b/ultralytics/solutions/queue_management.py @@ -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) diff --git a/ultralytics/utils/plotting.py b/ultralytics/utils/plotting.py index 4edaf351..f000a8c5 100644 --- a/ultralytics/utils/plotting.py +++ b/ultralytics/utils/plotting.py @@ -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): """