diff --git a/docs/en/guides/distance-calculation.md b/docs/en/guides/distance-calculation.md new file mode 100644 index 000000000..b19d3976f --- /dev/null +++ b/docs/en/guides/distance-calculation.md @@ -0,0 +1,89 @@ +--- +comments: true +description: Distance Calculation Using Ultralytics YOLOv8 +keywords: Ultralytics, YOLOv8, Object Detection, Distance Calculation, Object Tracking, Notebook, IPython Kernel, CLI, Python SDK +--- + +# Distance Calculation using Ultralytics YOLOv8 🚀 + +## What is Distance Calculation? + +Measuring the gap between two objects is known as distance calculation within a specified space. In the case of [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics), the bounding box centroid is employed to calculate the distance for bounding boxes highlighted by the user. + +## Advantages of Distance Calculation? + +- **Localization Precision:** Enhances accurate spatial positioning in computer vision tasks. +- **Size Estimation:** Allows estimation of physical sizes for better contextual understanding. +- **Scene Understanding:** Contributes to a 3D understanding of the environment for improved decision-making. + +???+ tip "Distance Calculation" + + - Click on any two bounding boxes with Left Mouse click for distance calculation + +!!! Example "Distance Calculation using YOLOv8 Example" + + === "Video Stream" + ```python + from ultralytics import YOLO + from ultralytics.solutions import distance_calculation + import cv2 + + model = YOLO("yolov8n.pt") + names = model.model.names + + cap = cv2.VideoCapture("path/to/video/file.mp4") + assert cap.isOpened(), "Error reading video file" + + # Video writer + video_writer = cv2.VideoWriter("distance_calculation.avi", + cv2.VideoWriter_fourcc(*'mp4v'), + int(cap.get(5)), + (int(cap.get(3)), int(cap.get(4)))) + + # Init distance-calculation obj + dist_obj = distance_calculation.DistanceCalculation() + dist_obj.set_args(names=names, view_img=True) + + 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 = dist_obj.start_process(im0, tracks) + video_writer.write(im0) + + cap.release() + video_writer.release() + cv2.destroyAllWindows() + + ``` + +???+ tip "Note" + + - Mouse Right Click will delete all drawn points + - Mouse Left Click can be used to draw points + + +### Optional Arguments `set_args` + +| 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 | + +### Arguments `model.track` + +| Name | Type | Default | Description | +|-----------|---------|----------------|-------------------------------------------------------------| +| `source` | `im0` | `None` | source directory for images or videos | +| `persist` | `bool` | `False` | persisting tracks between frames | +| `tracker` | `str` | `botsort.yaml` | Tracking method 'bytetrack' or 'botsort' | +| `conf` | `float` | `0.3` | Confidence Threshold | +| `iou` | `float` | `0.5` | IOU Threshold | +| `classes` | `list` | `None` | filter results by class, i.e. classes=0, or classes=[0,2,3] | +| `verbose` | `bool` | `True` | Display the object tracking results | diff --git a/docs/en/guides/heatmaps.md b/docs/en/guides/heatmaps.md index a02676eb2..d4c74b7aa 100644 --- a/docs/en/guides/heatmaps.md +++ b/docs/en/guides/heatmaps.md @@ -31,16 +31,13 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult | Transportation | Retail | |:-----------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------:| -| ![Ultralytics YOLOv8 Transportation Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/288d7053-622b-4452-b4e4-1f41aeb764aa) | ![Ultralytics YOLOv8 Retail Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/a9139af0-2cb7-41fe-a0d5-29a300dee768) | +| ![Ultralytics YOLOv8 Transportation Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/288d7053-622b-4452-b4e4-1f41aeb764aa) | ![Ultralytics YOLOv8 Retail Heatmap](https://github.com/RizwanMunawar/ultralytics/assets/62513924/edef75ad-50a7-4c0a-be4a-a66cdfc12802) | | Ultralytics YOLOv8 Transportation Heatmap | Ultralytics YOLOv8 Retail Heatmap | -???+ tip "heatmap_alpha" - heatmap_alpha value should be in range (0.0 - 1.0) - -???+ tip "decay_factor" - - Used for removal of heatmap after object removed from frame, value should be in range (0.0 - 1.0) +???+ tip "Heatmap Configuration" + - `heatmap_alpha`: Ensure this value is within the range (0.0 - 1.0). + - `decay_factor`: Used for removing heatmap after an object is no longer in the frame, its value should also be in the range (0.0 - 1.0). !!! Example "Heatmaps using Ultralytics YOLOv8 Example" diff --git a/docs/en/guides/index.md b/docs/en/guides/index.md index e4efcea38..2f230661e 100644 --- a/docs/en/guides/index.md +++ b/docs/en/guides/index.md @@ -37,6 +37,8 @@ Here's a compilation of in-depth guides to help you master different aspects of * [Heatmaps](heatmaps.md) 🚀 NEW: Elevate your understanding of data with our Detection Heatmaps! These intuitive visual tools use vibrant color gradients to vividly illustrate the intensity of data values across a matrix. Essential in computer vision, heatmaps are skillfully designed to highlight areas of interest, providing an immediate, impactful way to interpret spatial information. * [Instance Segmentation with Object Tracking](instance-segmentation-and-tracking.md) 🚀 NEW: Explore our feature on Object Segmentation in Bounding Boxes Shape, providing a visual representation of precise object boundaries for enhanced understanding and analysis. * [VisionEye View Objects Mapping](vision-eye.md) 🚀 NEW: This feature aim computers to discern and focus on specific objects, much like the way the human eye observes details from a particular viewpoint. +* [Speed Estimation](speed-estimation.md) 🚀 NEW: Speed estimation in computer vision relies on analyzing object motion through techniques like [object tracking](https://docs.ultralytics.com/modes/track/), crucial for applications like autonomous vehicles and traffic monitoring. +* [Distance Calculation](distance-calculation.md) 🚀 NEW: Distance calculation, which involves measuring the separation between two objects within a defined space, is a crucial aspect. In the context of Ultralytics YOLOv8, the method employed for this involves using the bounding box centroid to determine the distance associated with user-highlighted bounding boxes. ## Contribute to Our Guides diff --git a/docs/en/guides/region-counting.md b/docs/en/guides/region-counting.md index 5e8b82889..698989f16 100644 --- a/docs/en/guides/region-counting.md +++ b/docs/en/guides/region-counting.md @@ -8,7 +8,7 @@ keywords: Ultralytics, YOLOv8, Object Detection, Object Counting, Object Trackin ## What is Object Counting in Regions? -Object counting in regions with [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics/) involves precisely determining the number of objects within specified areas using advanced computer vision. This approach is valuable for optimizing processes, enhancing security, and improving efficiency in various applications. +[Object counting](https://docs.ultralytics.com/guides/object-counting/) in regions with [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics/) involves precisely determining the number of objects within specified areas using advanced computer vision. This approach is valuable for optimizing processes, enhancing security, and improving efficiency in various applications.


diff --git a/docs/en/guides/speed-estimation.md b/docs/en/guides/speed-estimation.md new file mode 100644 index 000000000..a30f89f4f --- /dev/null +++ b/docs/en/guides/speed-estimation.md @@ -0,0 +1,98 @@ +--- +comments: true +description: Speed Estimation Using Ultralytics YOLOv8 +keywords: Ultralytics, YOLOv8, Object Detection, Speed Estimation, Object Tracking, Notebook, IPython Kernel, CLI, Python SDK +--- + +# Speed Estimation using Ultralytics YOLOv8 🚀 + +## What is Speed Estimation? + +Speed estimation is the process of calculating the rate of movement of an object within a given context, often employed in computer vision applications. Using [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics/) you can now calculate the speed of object using [object tracking](https://docs.ultralytics.com/modes/track/) alongside distance and time data, crucial for tasks like traffic and surveillance. The accuracy of speed estimation directly influences the efficiency and reliability of various applications, making it a key component in the advancement of intelligent systems and real-time decision-making processes. + +## Advantages of Speed Estimation? + +- **Efficient Traffic Control:** Accurate speed estimation aids in managing traffic flow, enhancing safety, and reducing congestion on roadways. +- **Precise Autonomous Navigation:** In autonomous systems like self-driving cars, reliable speed estimation ensures safe and accurate vehicle navigation. +- **Enhanced Surveillance Security:** Speed estimation in surveillance analytics helps identify unusual behaviors or potential threats, improving the effectiveness of security measures. + +## Real World Applications + +| Transportation | Transportation | +|:-------------------------------------------------------------------------------------------------------------------------------------------------------:|:---------------------------------------------------------------------------------------------------------------------------------------------------------:| +| ![Speed Estimation on Road using Ultralytics YOLOv8](https://github.com/RizwanMunawar/ultralytics/assets/62513924/c8a0fd4a-d394-436d-8de3-d5b754755fc7) | ![Speed Estimation on Bridge using Ultralytics YOLOv8](https://github.com/RizwanMunawar/ultralytics/assets/62513924/cee10e02-b268-4304-b73a-5b9cb42da669) | +| Speed Estimation on Road using Ultralytics YOLOv8 | Speed Estimation on Bridge using Ultralytics YOLOv8 | + +!!! Example "Speed Estimation using YOLOv8 Example" + + === "Speed Estimation" + ```python + from ultralytics import YOLO + from ultralytics.solutions import speed_estimation + import cv2 + + model = YOLO("yolov8n.pt") + names = model.model.names + + cap = cv2.VideoCapture("path/to/video/file.mp4") + assert cap.isOpened(), "Error reading video file" + + # Video writer + video_writer = cv2.VideoWriter("speed_estimation.avi", + cv2.VideoWriter_fourcc(*'mp4v'), + int(cap.get(5)), + (int(cap.get(3)), int(cap.get(4)))) + + line_pts = [(0, 360), (1280, 360)] + + # Init speed-estimation obj + speed_obj = speed_estimation.SpeedEstimator() + speed_obj.set_args(reg_pts=line_pts, + names=names, + view_img=True) + + 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 = speed_obj.estimate_speed(im0, tracks) + video_writer.write(im0) + + cap.release() + video_writer.release() + cv2.destroyAllWindows() + + ``` + +???+ warning "Speed is Estimate" + + Speed will be an estimate and may not be completely accurate. Additionally, the estimation can vary depending on GPU speed. + + +### Optional Arguments `set_args` + +| Name | Type | Default | Description | +|---------------------|-------------|----------------------------|---------------------------------------------------| +| reg_pts | `list` | `[(20, 400), (1260, 400)]` | Points defining the Region Area | +| names | `dict` | `None` | Classes names | +| view_img | `bool` | `False` | Display frames with counts | +| line_thickness | `int` | `2` | Increase bounding boxes thickness | +| region_thickness | `int` | `5` | Thickness for object counter region or line | +| spdl_dist_thresh | `int` | `10` | Euclidean Distance threshold for speed check line | + +### Arguments `model.track` + +| Name | Type | Default | Description | +|-----------|---------|----------------|-------------------------------------------------------------| +| `source` | `im0` | `None` | source directory for images or videos | +| `persist` | `bool` | `False` | persisting tracks between frames | +| `tracker` | `str` | `botsort.yaml` | Tracking method 'bytetrack' or 'botsort' | +| `conf` | `float` | `0.3` | Confidence Threshold | +| `iou` | `float` | `0.5` | IOU Threshold | +| `classes` | `list` | `None` | filter results by class, i.e. classes=0, or classes=[0,2,3] | +| `verbose` | `bool` | `True` | Display the object tracking results | diff --git a/docs/en/reference/solutions/distance_calculation.md b/docs/en/reference/solutions/distance_calculation.md new file mode 100644 index 000000000..5376a2d1b --- /dev/null +++ b/docs/en/reference/solutions/distance_calculation.md @@ -0,0 +1,16 @@ +--- +description: Explore Ultralytics YOLO's distance calculation feature designed for advance analytics, providing an immediate, impactful way to interpret computer vision data. +keywords: Ultralytics, YOLO, distance calculation, object tracking, data visualization, real-time tracking, machine learning, object counting, computer vision, vehicle analytics, YOLOv8, artificial intelligence +--- + +# Reference for `ultralytics/solutions/distance_calculation.py` + +!!! Note + + This file is available at [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/distance_calculation.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/distance_calculation.py). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request](https://github.com/ultralytics/ultralytics/edit/main/ultralytics/solutions/heatmap.py) 🛠️. Thank you 🙏! + +

+ +## ::: ultralytics.solutions.distance_calculation.DistanceCalculation + +

diff --git a/docs/en/reference/solutions/speed_estimation.md b/docs/en/reference/solutions/speed_estimation.md new file mode 100644 index 000000000..86d29b1d0 --- /dev/null +++ b/docs/en/reference/solutions/speed_estimation.md @@ -0,0 +1,16 @@ +--- +description: Transform speed estimation with Ultralytics YOLO speed estimation featuring cutting-edge technology for precise real-time counting in video streams. +keywords: Ultralytics YOLO, speed estimation software, real-time vehicle tracking solutions, video stream analysis, YOLOv8 object detection, smart counting technology, computer vision, AI-powered tracking, video analytics tools, automated monitoring. +--- + +# Reference for `ultralytics/solutions/speed_estimation.py` + +!!! Note + + This file is available at [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/speed_estimation.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/speed_estimation.py). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request](https://github.com/ultralytics/ultralytics/edit/main/ultralytics/solutions/object_counter.py) 🛠️. Thank you 🙏! + +

+ +## ::: ultralytics.solutions.speed_estimation.SpeedEstimator + +

diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 9401924ca..b65ed0825 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -281,6 +281,8 @@ nav: - Heatmaps: guides/heatmaps.md - Instance Segmentation with Object Tracking: guides/instance-segmentation-and-tracking.md - VisionEye Mapping: guides/vision-eye.md + - Speed Estimation: guides/speed-estimation.md + - Distance Calculation: guides/distance-calculation.md - Integrations: - integrations/index.md - Comet ML: integrations/comet.md @@ -429,6 +431,8 @@ nav: - ai_gym: reference/solutions/ai_gym.md - heatmap: reference/solutions/heatmap.md - object_counter: reference/solutions/object_counter.md + - speed_estimation: reference/solutions/speed_estimation.md + - distance_calculation: reference/solutions/distance_calculation.md - trackers: - basetrack: reference/trackers/basetrack.md - bot_sort: reference/trackers/bot_sort.md diff --git a/ultralytics/solutions/distance_calculation.py b/ultralytics/solutions/distance_calculation.py new file mode 100644 index 000000000..684b504f5 --- /dev/null +++ b/ultralytics/solutions/distance_calculation.py @@ -0,0 +1,187 @@ +# Ultralytics YOLO 🚀, AGPL-3.0 license + +import math + +import cv2 + +from ultralytics.utils.plotting import Annotator, colors + + +class DistanceCalculation: + """A class to calculate distance between two objects in real-time video stream based on their tracks.""" + + def __init__(self): + """Initializes the distance calculation class with default values for Visual, Image, track and distance + parameters. + """ + + # Visual & im0 information + self.im0 = None + self.annotator = None + self.view_img = False + self.line_color = (255, 255, 0) + self.centroid_color = (255, 0, 255) + + # Predict/track information + self.clss = None + self.names = None + self.boxes = None + self.line_thickness = 2 + self.trk_ids = None + + # Distance calculation information + self.centroids = [] + self.pixel_per_meter = 10 + + # Mouse event + self.left_mouse_count = 0 + self.selected_boxes = {} + + def set_args(self, + names, + pixels_per_meter=10, + view_img=False, + line_thickness=2, + line_color=(255, 255, 0), + centroid_color=(255, 0, 255)): + """ + Configures the distance calculation and display parameters. + + Args: + names (dict): object detection classes names + pixels_per_meter (int): Number of pixels in meter + view_img (bool): Flag indicating frame display + line_thickness (int): Line thickness for bounding boxes. + line_color (RGB): color of centroids line + centroid_color (RGB): colors of bbox centroids + """ + self.names = names + self.pixel_per_meter = pixels_per_meter + self.view_img = view_img + self.line_thickness = line_thickness + self.line_color = line_color + self.centroid_color = centroid_color + + def mouse_event_for_distance(self, event, x, y, flags, param): + """ + This function is designed to move region with mouse events in a real-time video stream. + + Args: + event (int): The type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.). + x (int): The x-coordinate of the mouse pointer. + y (int): The y-coordinate of the mouse pointer. + flags (int): Any flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, + cv2.EVENT_FLAG_SHIFTKEY, etc.). + param (dict): Additional parameters you may want to pass to the function. + """ + global selected_boxes + global left_mouse_count + if event == cv2.EVENT_LBUTTONDOWN: + self.left_mouse_count += 1 + if self.left_mouse_count <= 2: + for box, track_id in zip(self.boxes, self.trk_ids): + if box[0] < x < box[2] and box[1] < y < box[3]: + if track_id not in self.selected_boxes: + self.selected_boxes[track_id] = [] + self.selected_boxes[track_id] = box + + if event == cv2.EVENT_RBUTTONDOWN: + self.selected_boxes = {} + self.left_mouse_count = 0 + + def extract_tracks(self, tracks): + """ + Extracts results from the provided data. + + Args: + tracks (list): List of tracks obtained from the object tracking process. + """ + self.boxes = tracks[0].boxes.xyxy.cpu() + self.clss = tracks[0].boxes.cls.cpu().tolist() + self.trk_ids = tracks[0].boxes.id.int().cpu().tolist() + + def calculate_centroid(self, box): + """ + Calculate the centroid of bounding box + Args: + box (list): Bounding box data + """ + return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2) + + def calculate_distance(self, centroid1, centroid2): + """ + Calculate distance between two centroids + Args: + centroid1 (point): First bounding box data + centroid2 (point): Second bounding box data + """ + pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2) + return pixel_distance / self.pixel_per_meter + + def plot_distance_and_line(self, distance): + """ + Plot the distance and line on frame + Args: + distance (float): Distance between two centroids + """ + cv2.rectangle(self.im0, (15, 25), (280, 70), (255, 255, 255), -1) + cv2.putText(self.im0, f'Distance : {distance:.2f}m', (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2, + cv2.LINE_AA) + cv2.line(self.im0, self.centroids[0], self.centroids[1], self.line_color, 3) + cv2.circle(self.im0, self.centroids[0], 6, self.centroid_color, -1) + cv2.circle(self.im0, self.centroids[1], 6, self.centroid_color, -1) + + def start_process(self, im0, tracks): + """ + Calculate distance between two bounding boxes based on tracking data + Args: + im0 (nd array): Image + tracks (list): List of tracks obtained from the object tracking process. + """ + self.im0 = im0 + if tracks[0].boxes.id is None: + if self.view_img: + self.display_frames() + return + else: + return + + self.extract_tracks(tracks) + + self.annotator = Annotator(self.im0, line_width=2) + + for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids): + self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)]) + + if len(self.selected_boxes) == 2: + for trk_id, _ in self.selected_boxes.items(): + if trk_id == track_id: + self.selected_boxes[track_id] = box + + if len(self.selected_boxes) == 2: + for trk_id, box in self.selected_boxes.items(): + centroid = self.calculate_centroid(self.selected_boxes[trk_id]) + self.centroids.append(centroid) + + distance = self.calculate_distance(self.centroids[0], self.centroids[1]) + self.plot_distance_and_line(distance) + + self.centroids = [] + + if self.view_img: + self.display_frames() + + return im0 + + def display_frames(self): + """Display frame.""" + cv2.namedWindow('Ultralytics Distance Estimation') + cv2.setMouseCallback('Ultralytics Distance Estimation', self.mouse_event_for_distance) + cv2.imshow('Ultralytics Distance Estimation', self.im0) + + if cv2.waitKey(1) & 0xFF == ord('q'): + return + + +if __name__ == '__main__': + DistanceCalculation() diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py index 1131a5b18..1f3d66952 100644 --- a/ultralytics/solutions/heatmap.py +++ b/ultralytics/solutions/heatmap.py @@ -158,7 +158,11 @@ class Heatmap: """ self.im0 = im0 if tracks[0].boxes.id is None: - return self.im0 + if self.view_img and self.env_check: + self.display_frames() + return + else: + return self.heatmap *= self.decay_factor # decay factor self.extract_results(tracks) @@ -240,22 +244,16 @@ class Heatmap: txt_color=self.count_txt_color, color=self.count_color) - im0_with_heatmap = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0) + self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0) if self.env_check and self.view_img: - self.display_frames(im0_with_heatmap) - - return im0_with_heatmap + self.display_frames() - @staticmethod - def display_frames(im0_with_heatmap): - """ - Display heatmap. + return self.im0 - Args: - im0_with_heatmap (nd array): Original Image with heatmap - """ - cv2.imshow('Ultralytics Heatmap', im0_with_heatmap) + def display_frames(self): + """Display frame.""" + cv2.imshow('Ultralytics Heatmap', self.im0) if cv2.waitKey(1) & 0xFF == ord('q'): return diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index 5b6b45922..f9b86239e 100644 --- a/ultralytics/solutions/object_counter.py +++ b/ultralytics/solutions/object_counter.py @@ -198,7 +198,9 @@ class ObjectCounter: txt_color=self.count_txt_color, color=self.count_color) - if self.env_check and self.view_img: + def display_frames(self): + """Display frame.""" + if self.env_check: cv2.namedWindow('Ultralytics YOLOv8 Object Counter') if len(self.reg_pts) == 4: # only add mouse event If user drawn region cv2.setMouseCallback('Ultralytics YOLOv8 Object Counter', self.mouse_event_for_region, @@ -219,8 +221,15 @@ class ObjectCounter: self.im0 = im0 # store image if tracks[0].boxes.id is None: - return + if self.view_img: + self.display_frames() + return + else: + return self.extract_and_process_tracks(tracks) + + if self.view_img: + self.display_frames() return self.im0 diff --git a/ultralytics/solutions/speed_estimation.py b/ultralytics/solutions/speed_estimation.py new file mode 100644 index 000000000..7260141fe --- /dev/null +++ b/ultralytics/solutions/speed_estimation.py @@ -0,0 +1,203 @@ +# Ultralytics YOLO 🚀, AGPL-3.0 license + +from collections import defaultdict +from time import time + +import cv2 +import numpy as np + +from ultralytics.utils.checks import check_imshow +from ultralytics.utils.plotting import Annotator, colors + + +class SpeedEstimator: + """A class to estimation speed of objects in real-time video stream based on their tracks.""" + + def __init__(self): + """Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters.""" + + # Visual & im0 information + self.im0 = None + self.annotator = None + self.view_img = False + + # Region information + self.reg_pts = [(20, 400), (1260, 400)] + self.region_thickness = 3 + + # Predict/track information + self.clss = None + self.names = None + self.boxes = None + self.trk_ids = None + self.trk_pts = None + self.line_thickness = 2 + self.trk_history = defaultdict(list) + + # Speed estimator information + self.current_time = 0 + self.dist_data = {} + self.trk_idslist = [] + self.spdl_dist_thresh = 10 + self.trk_previous_times = {} + self.trk_previous_points = {} + + # Check if environment support imshow + self.env_check = check_imshow(warn=True) + + def set_args( + self, + reg_pts, + names, + view_img=False, + line_thickness=2, + region_thickness=5, + spdl_dist_thresh=10, + ): + """ + Configures the speed estimation and display parameters. + + Args: + reg_pts (list): Initial list of points defining the speed calculation region. + names (dict): object detection classes names + view_img (bool): Flag indicating frame display + line_thickness (int): Line thickness for bounding boxes. + region_thickness (int): Speed estimation region thickness + spdl_dist_thresh (int): Euclidean distance threshold for speed line + """ + if reg_pts is None: + print('Region points not provided, using default values') + else: + self.reg_pts = reg_pts + self.names = names + self.view_img = view_img + self.line_thickness = line_thickness + self.region_thickness = region_thickness + self.spdl_dist_thresh = spdl_dist_thresh + + def extract_tracks(self, tracks): + """ + Extracts results from the provided data. + + Args: + tracks (list): List of tracks obtained from the object tracking process. + """ + self.boxes = tracks[0].boxes.xyxy.cpu() + self.clss = tracks[0].boxes.cls.cpu().tolist() + self.trk_ids = tracks[0].boxes.id.int().cpu().tolist() + + def store_track_info(self, track_id, box): + """ + Store track data. + + Args: + track_id (int): object track id. + box (list): object bounding box data + """ + track = self.trk_history[track_id] + bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)) + track.append(bbox_center) + + if len(track) > 30: + track.pop(0) + + self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2)) + return track + + def plot_box_and_track(self, track_id, box, cls, track): + """ + Plot track and bounding box. + + Args: + track_id (int): object track id. + box (list): object bounding box data + cls (str): object class name + track (list): tracking history for tracks path drawing + """ + speed_label = str(int( + self.dist_data[track_id])) + 'km/ph' if track_id in self.dist_data else self.names[int(cls)] + bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255) + + self.annotator.box_label(box, speed_label, bbox_color) + + cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1) + cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1) + + def calculate_speed(self, trk_id, track): + """ + Calculation of object speed + Args: + trk_id (int): object track id. + track (list): tracking history for tracks path drawing + """ + + if self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]: + + if (self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh): + direction = 'known' + + elif (self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < + self.reg_pts[0][1] + self.spdl_dist_thresh): + direction = 'known' + + else: + direction = 'unknown' + + if self.trk_previous_times[trk_id] != 0 and direction != 'unknown': + + if trk_id not in self.trk_idslist: + self.trk_idslist.append(trk_id) + + time_difference = time() - self.trk_previous_times[trk_id] + if time_difference > 0: + dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1]) + speed = dist_difference / time_difference + self.dist_data[trk_id] = speed + + self.trk_previous_times[trk_id] = time() + self.trk_previous_points[trk_id] = track[-1] + + def estimate_speed(self, im0, tracks): + """ + Calculate object based on tracking data + Args: + im0 (nd array): Image + tracks (list): List of tracks obtained from the object tracking process. + """ + self.im0 = im0 + if tracks[0].boxes.id is None: + if self.view_img and self.env_check: + self.display_frames() + return + else: + return + + self.extract_tracks(tracks) + + self.annotator = Annotator(self.im0, line_width=2) + self.annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 0), thickness=self.region_thickness) + + for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss): + + track = self.store_track_info(trk_id, box) + + if trk_id not in self.trk_previous_times: + self.trk_previous_times[trk_id] = 0 + + self.plot_box_and_track(trk_id, box, cls, track) + self.calculate_speed(trk_id, track) + + if self.view_img and self.env_check: + self.display_frames() + + return im0 + + def display_frames(self): + """Display frame.""" + cv2.imshow('Ultralytics Speed Estimation', self.im0) + if cv2.waitKey(1) & 0xFF == ord('q'): + return + + +if __name__ == '__main__': + SpeedEstimator()