xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 21:44:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add speed_estimation and distance_calculation in ultralytics solutions (#7325)

Browse Source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-01-05 14:38:13 +05:00 committed by GitHub
parent 2f9ec8c0b4
commit 61fa12460d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
12 changed files with 642 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
docs/en/guides/distance-calculation.md Normal file
View File

@ -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 |

11
docs/en/guides/heatmaps.md
View File

@ -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"

2
docs/en/guides/index.md
View File

@ -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

2
docs/en/guides/region-counting.md
View File

@ -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.
<p align="center">
<br>

98
docs/en/guides/speed-estimation.md Normal file
View File

@ -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 |

16
docs/en/reference/solutions/distance_calculation.md Normal file
View File

@ -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 🙏!
<br><br>
## ::: ultralytics.solutions.distance_calculation.DistanceCalculation
<br><br>

16
docs/en/reference/solutions/speed_estimation.md Normal file
View File

@ -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 🙏!
<br><br>
## ::: ultralytics.solutions.speed_estimation.SpeedEstimator
<br><br>

4
docs/mkdocs.yml
View File

@ -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

187
ultralytics/solutions/distance_calculation.py Normal file
View File

@ -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()

24
ultralytics/solutions/heatmap.py
View File

@ -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)
self.display_frames()
return im0_with_heatmap
return self.im0
@staticmethod
def display_frames(im0_with_heatmap):
"""
Display heatmap.
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

11
ultralytics/solutions/object_counter.py
View File

@ -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:
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

203
ultralytics/solutions/speed_estimation.py Normal file
View File

@ -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()