Chiebot-Mirror
/
ultralytics
mirror of https://github.com/ultralytics/ultralytics.git
9
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into docs-build

Browse Source
docs-build
Ultralytics Assistant 5 months ago committed by GitHub
parent e4b9c231c5 c2068df9d9
commit 0a17d58021
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 187 additions and 227 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 12
      docs/en/datasets/obb/dota-v2.md
  2. 12
      docs/en/datasets/obb/index.md
  3. 11
      docs/en/guides/model-training-tips.md
  4. 33
      docs/en/guides/queue-management.md
  5. 3
      docs/en/guides/speed-estimation.md
  6. 18
      docs/en/macros/track-args.md
  7. 3
      docs/mkdocs_github_authors.yaml
  8. 17
      mkdocs.yml
  9. 2
      ultralytics/__init__.py
  10. 2
      ultralytics/engine/exporter.py
  11. 2
      ultralytics/models/nas/val.py
  12. 2
      ultralytics/models/yolo/obb/val.py
  13. 2
      ultralytics/models/yolo/pose/val.py
  14. 2
      ultralytics/models/yolo/segment/val.py
  15. 55
      ultralytics/nn/tasks.py
  16. 66
      ultralytics/solutions/queue_management.py
  17. 172
      ultralytics/solutions/speed_estimation.py

12
docs/en/datasets/obb/dota-v2.md
Unescape View File

@ -111,15 +111,15 @@ To train a model on the DOTA v1 dataset, you can utilize the following code snip
# Create a new YOLOv8n-OBB model from scratch
model = YOLO("yolov8n-obb.yaml")
# Train the model on the DOTAv2 dataset
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=640)
# Train the model on the DOTAv1 dataset
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=1024)
```
=== "CLI"
```bash
# Train a new YOLOv8n-OBB model on the DOTAv2 dataset
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=640
# Train a new YOLOv8n-OBB model on the DOTAv1 dataset
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=1024
```
## Sample Data and Annotations
@ -180,14 +180,14 @@ To train a model on the DOTA dataset, you can use the following example with Ult
model = YOLO("yolov8n-obb.yaml")
# Train the model on the DOTAv1 dataset
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=640)
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=1024)
```
=== "CLI"
```bash
# Train a new YOLOv8n-OBB model on the DOTAv1 dataset
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=640
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=1024
```
For more details on how to split and preprocess the DOTA images, refer to the [split DOTA images section](#split-dota-images).

12
docs/en/datasets/obb/index.md
Unescape View File

@ -42,21 +42,23 @@ To train a model using these OBB formats:
# Create a new YOLOv8n-OBB model from scratch
model = YOLO("yolov8n-obb.yaml")
# Train the model on the DOTAv2 dataset
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=640)
# Train the model on the DOTAv1 dataset
results = model.train(data="DOTAv1.yaml", epochs=100, imgsz=1024)
```
=== "CLI"
```bash
# Train a new YOLOv8n-OBB model on the DOTAv2 dataset
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=640
# Train a new YOLOv8n-OBB model on the DOTAv1 dataset
yolo obb train data=DOTAv1.yaml model=yolov8n-obb.pt epochs=100 imgsz=1024
```
## Supported Datasets
Currently, the following datasets with Oriented Bounding Boxes are supported:
- [DOTA-v1](dota-v2.md): The first version of the DOTA dataset, providing a comprehensive set of aerial images with oriented bounding boxes for object detection.
- [DOTA-v1.5](dota-v2.md): An intermediate version of the DOTA dataset, offering additional annotations and improvements over DOTA-v1 for enhanced object detection tasks.
- [DOTA-v2](dota-v2.md): DOTA (A Large-scale Dataset for Object Detection in Aerial Images) version 2, emphasizes detection from aerial perspectives and contains oriented bounding boxes with 1.7 million instances and 11,268 images.
- [DOTA8](dota8.md): A small, 8-image subset of the full DOTA dataset suitable for testing workflows and Continuous Integration (CI) checks of OBB training in the `ultralytics` repository.
@ -133,6 +135,8 @@ This ensures your model leverages the detailed OBB annotations for improved dete
Currently, Ultralytics supports the following datasets for OBB training:
- [DOTA-v1](dota-v2.md): The first version of the DOTA dataset, providing a comprehensive set of aerial images with oriented bounding boxes for object detection.
- [DOTA-v1.5](dota-v2.md): An intermediate version of the DOTA dataset, offering additional annotations and improvements over DOTA-v1 for enhanced object detection tasks.
- [DOTA-v2](dota-v2.md): This dataset includes 1.7 million instances with oriented bounding boxes and 11,268 images, primarily focusing on aerial object detection.
- [DOTA8](dota8.md): A smaller, 8-image subset of the DOTA dataset used for testing and continuous integration (CI) checks.

11
docs/en/guides/model-training-tips.md
Unescape View File

@ -10,6 +10,17 @@ keywords: Model Training Machine Learning, AI Model Training, Number of Epochs,
One of the most important steps when working on a [computer vision project](./steps-of-a-cv-project.md) is model training. Before reaching this step, you need to [define your goals](./defining-project-goals.md) and [collect and annotate your data](./data-collection-and-annotation.md). After [preprocessing the data](./preprocessing_annotated_data.md) to make sure it is clean and consistent, you can move on to training your model.
<p align="center">
<br>
<iframe loading="lazy" width="720" height="405" src="https://www.youtube.com/embed/GIrFEoR5PoU"
title="YouTube video player" frameborder="0"
allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share"
allowfullscreen>
</iframe>
<br>
<strong>Watch:</strong> Model Training Tips | How to Handle Large Datasets | Batch Size, GPU Utilization and Mixed Precision
</p>
So, what is [model training](../modes/train.md)? Model training is the process of teaching your model to recognize visual patterns and make predictions based on your data. It directly impacts the performance and accuracy of your application. In this guide, we'll cover best practices, optimization techniques, and troubleshooting tips to help you train your computer vision models effectively.
## How to Train a Machine Learning Model

33
docs/en/guides/queue-management.md
Unescape View File

@ -56,15 +56,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
)
while cap.isOpened():
success, im0 = cap.read()
if success:
tracks = model.track(im0, show=False, persist=True, verbose=False)
tracks = model.track(im0, persist=True)
out = queue.process_queue(im0, tracks)
video_writer.write(im0)
@ -100,15 +98,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
)
while cap.isOpened():
success, im0 = cap.read()
if success:
tracks = model.track(im0, show=False, persist=True, verbose=False, classes=0) # Only person class
tracks = model.track(im0, persist=True, classes=0) # Only person class
out = queue.process_queue(im0, tracks)
video_writer.write(im0)
@ -125,20 +121,13 @@ Queue management using [Ultralytics YOLOv8](https://github.com/ultralytics/ultra
### Arguments `QueueManager`
| Name | Type | Default | Description |
| ------------------- | ---------------- | -------------------------- | ----------------------------------------------------------------------------------- |
| `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. |
| `view_img` | `bool` | `False` | Whether to display the image frames. |
| `region_color` | `tuple` | `(255, 0, 255)` | Color of the counting region lines (BGR). |
| `view_queue_counts` | `bool` | `True` | Whether to display the queue counts. |
| `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. |
| `count_txt_color` | `tuple` | `(255, 255, 255)` | Color of the count text (BGR). |
| `track_color` | `tuple` | `None` | Color of the tracks. If `None`, different colors will be used for different tracks. |
| `region_thickness` | `int` | `5` | Thickness of the counting region lines. |
| `fontsize` | `float` | `0.7` | Font size for the text annotations. |
| Name | Type | Default | Description |
| ---------------- | ---------------- | -------------------------- | -------------------------------------------------------------------------------- |
| `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. |
| `view_img` | `bool` | `False` | Whether to display the image frames. |
| `draw_tracks` | `bool` | `False` | Whether to draw tracks of the objects. |
### Arguments `model.track`
@ -170,8 +159,6 @@ queue = solutions.QueueManager(
names=model.names,
reg_pts=queue_region,
line_thickness=3,
fontsize=1.0,
region_color=(255, 144, 31),
)
while cap.isOpened():
@ -223,8 +210,6 @@ queue_airport = solutions.QueueManager(
names=model.names,
reg_pts=queue_region_airport,
line_thickness=3,
fontsize=1.0,
region_color=(0, 255, 0),
)
```

3
docs/en/guides/speed-estimation.md
Unescape View File

@ -72,7 +72,7 @@ keywords: Ultralytics YOLOv8, speed estimation, object tracking, computer vision
print("Video frame is empty or video processing has been successfully completed.")
break
tracks = model.track(im0, persist=True, show=False)
tracks = model.track(im0, persist=True)
im0 = speed_obj.estimate_speed(im0, tracks)
video_writer.write(im0)
@ -94,7 +94,6 @@ keywords: Ultralytics YOLOv8, speed estimation, object tracking, computer vision
| `reg_pts` | `list` | `[(20, 400), (1260, 400)]` | List of region points for speed estimation. |
| `view_img` | `bool` | `False` | Whether to display the image with annotations. |
| `line_thickness` | `int` | `2` | Thickness of the lines for drawing boxes and tracks. |
| `region_thickness` | `int` | `5` | Thickness of the region lines. |
| `spdl_dist_thresh` | `int` | `10` | Distance threshold for speed calculation. |
### Arguments `model.track`

18
docs/en/macros/track-args.md
Unescape View File

@ -1,9 +1,9 @@
| 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 |
| Argument | Type | Default | Description |
| --------- | ------- | -------------- | ---------------------------------------------------------------------------------------------------------------------------- |
| `source` | `str` | `None` | Specifies the source directory for images or videos. Supports file paths and URLs. |
| `persist` | `bool` | `False` | Enables persistent tracking of objects between frames, maintaining IDs across video sequences. |
| `tracker` | `str` | `botsort.yaml` | Specifies the tracking algorithm to use, e.g., `bytetrack.yaml` or `botsort.yaml`. |
| `conf` | `float` | `0.3` | Sets the confidence threshold for detections; lower values allow more objects to be tracked but may include false positives. |
| `iou` | `float` | `0.5` | Sets the Intersection over Union (IoU) threshold for filtering overlapping detections. |
| `classes` | `list` | `None` | Filters results by class index. For example, `classes=[0, 2, 3]` only tracks the specified classes. |
| `verbose` | `bool` | `True` | Controls the display of tracking results, providing a visual output of tracked objects. |

3
docs/mkdocs_github_authors.yaml
Unescape View File

@ -19,6 +19,9 @@
17316848+maianumerosky@users.noreply.github.com:
avatar: https://avatars.githubusercontent.com/u/17316848?v=4
username: maianumerosky
32206511+Y-T-G@users.noreply.github.com:
avatar: https://avatars.githubusercontent.com/u/32206511?v=4
username: Y-T-G
34196005+fcakyon@users.noreply.github.com:
avatar: https://avatars.githubusercontent.com/u/34196005?v=4
username: fcakyon

17
mkdocs.yml
Unescape View File

@ -211,8 +211,6 @@ nav:
- solutions/index.md
- Guides:
- guides/index.md
- Explorer:
- datasets/explorer/index.md
- Modes:
- modes/index.md
- Train: modes/train.md
@ -221,6 +219,13 @@ nav:
- Export: modes/export.md
- Track: modes/track.md
- Benchmark: modes/benchmark.md
- Tasks:
- tasks/index.md
- Detect: tasks/detect.md
- Segment: tasks/segment.md
- Classify: tasks/classify.md
- Pose: tasks/pose.md
- OBB: tasks/obb.md
- Tasks:
- tasks/index.md
- Detect: tasks/detect.md
@ -228,6 +233,14 @@ nav:
- Classify: tasks/classify.md
- Pose: tasks/pose.md
- OBB: tasks/obb.md
- Modes:
- modes/index.md
- Train: modes/train.md
- Val: modes/val.md
- Predict: modes/predict.md
- Export: modes/export.md
- Track: modes/track.md
- Benchmark: modes/benchmark.md
- Models:
- models/index.md
- YOLOv3: models/yolov3.md

2
ultralytics/__init__.py
Unescape View File

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

2
ultralytics/engine/exporter.py
Unescape View File

@ -391,7 +391,7 @@ class Exporter:
"""YOLOv8 ONNX export."""
requirements = ["onnx>=1.12.0"]
if self.args.simplify:
requirements += ["onnxslim==0.1.32", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
requirements += ["onnxslim==0.1.34", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
check_requirements(requirements)
import onnx # noqa

2
ultralytics/models/nas/val.py
Unescape View File

@ -44,7 +44,7 @@ class NASValidator(DetectionValidator):
self.args.iou,
labels=self.lb,
multi_label=False,
agnostic=self.args.single_cls,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
max_time_img=0.5,
)

2
ultralytics/models/yolo/obb/val.py
Unescape View File

@ -45,7 +45,7 @@ class OBBValidator(DetectionValidator):
labels=self.lb,
nc=self.nc,
multi_label=True,
agnostic=self.args.single_cls,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
rotated=True,
)

2
ultralytics/models/yolo/pose/val.py
Unescape View File

@ -69,7 +69,7 @@ class PoseValidator(DetectionValidator):
self.args.iou,
labels=self.lb,
multi_label=True,
agnostic=self.args.single_cls,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
nc=self.nc,
)

2
ultralytics/models/yolo/segment/val.py
Unescape View File

@ -76,7 +76,7 @@ class SegmentationValidator(DetectionValidator):
self.args.iou,
labels=self.lb,
multi_label=True,
agnostic=self.args.single_cls,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
nc=self.nc,
)

55
ultralytics/nn/tasks.py
Unescape View File

@ -1,6 +1,8 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import contextlib
import pickle
import types
from copy import deepcopy
from pathlib import Path
@ -750,7 +752,35 @@ def temporary_modules(modules=None, attributes=None):
del sys.modules[old]
def torch_safe_load(weight):
class SafeClass:
"""A placeholder class to replace unknown classes during unpickling."""
def __init__(self, *args, **kwargs):
"""Initialize SafeClass instance, ignoring all arguments."""
pass
class SafeUnpickler(pickle.Unpickler):
"""Custom Unpickler that replaces unknown classes with SafeClass."""
def find_class(self, module, name):
"""Attempt to find a class, returning SafeClass if not among safe modules."""
safe_modules = (
"torch",
"collections",
"collections.abc",
"builtins",
"math",
"numpy",
# Add other modules considered safe
)
if module in safe_modules:
return super().find_class(module, name)
else:
return SafeClass
def torch_safe_load(weight, safe_only=False):
"""
Attempts to load a PyTorch model with the torch.load() function. If a ModuleNotFoundError is raised, it catches the
error, logs a warning message, and attempts to install the missing module via the check_requirements() function.
@ -758,9 +788,18 @@ def torch_safe_load(weight):
Args:
weight (str): The file path of the PyTorch model.
safe_only (bool): If True, replace unknown classes with SafeClass during loading.
Example:
```python
from ultralytics.nn.tasks import torch_safe_load
ckpt, file = torch_safe_load("path/to/best.pt", safe_only=True)
```
Returns:
(dict): The loaded PyTorch model.
ckpt (dict): The loaded model checkpoint.
file (str): The loaded filename
"""
from ultralytics.utils.downloads import attempt_download_asset
@ -779,7 +818,15 @@ def torch_safe_load(weight):
"ultralytics.utils.loss.v10DetectLoss": "ultralytics.utils.loss.E2EDetectLoss", # YOLOv10
},
):
ckpt = torch.load(file, map_location="cpu")
if safe_only:
# Load via custom pickle module
safe_pickle = types.ModuleType("safe_pickle")
safe_pickle.Unpickler = SafeUnpickler
safe_pickle.load = lambda file_obj: SafeUnpickler(file_obj).load()
with open(file, "rb") as f:
ckpt = torch.load(f, pickle_module=safe_pickle)
else:
ckpt = torch.load(file, map_location="cpu")
except ModuleNotFoundError as e: # e.name is missing module name
if e.name == "models":
@ -809,7 +856,7 @@ def torch_safe_load(weight):
)
ckpt = {"model": ckpt.model}
return ckpt, file # load
return ckpt, file
def attempt_load_weights(weights, device=None, inplace=True, fuse=False):

66
ultralytics/solutions/queue_management.py
Unescape View File

@ -20,15 +20,8 @@ class QueueManager:
names,
reg_pts=None,
line_thickness=2,
track_thickness=2,
view_img=False,
region_color=(255, 0, 255),
view_queue_counts=True,
draw_tracks=False,
count_txt_color=(255, 255, 255),
track_color=None,
region_thickness=5,
fontsize=0.7,
):
"""
Initializes the QueueManager with specified parameters for tracking and counting objects.
@ -38,57 +31,35 @@ class QueueManager:
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.
track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
view_img (bool, optional): Whether to display the image frames. Defaults to False.
region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
tracks. Defaults to None.
region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
"""
# Mouse events state
self.is_drawing = False
self.selected_point = None
# Region & Line Information
self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
self.counting_region = (
Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
)
self.region_color = region_color
self.region_thickness = region_thickness
# Image and annotation Information
self.im0 = None
# annotation Information
self.tf = line_thickness
self.view_img = view_img
self.view_queue_counts = view_queue_counts
self.fontsize = fontsize
self.names = names # Class names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Queue Manager"
# Object counting Information
self.counts = 0
self.count_txt_color = count_txt_color
# Tracks info
self.track_history = defaultdict(list)
self.track_thickness = track_thickness
self.draw_tracks = draw_tracks
self.track_color = track_color
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
def extract_and_process_tracks(self, tracks):
def extract_and_process_tracks(self, tracks, im0):
"""Extracts and processes tracks for queue management in a video stream."""
# Initialize annotator and draw the queue region
self.annotator = Annotator(self.im0, self.tf, self.names)
annotator = Annotator(im0, self.tf, self.names)
self.counts = 0 # Reset counts every frame
if tracks[0].boxes.id is not None:
boxes = tracks[0].boxes.xyxy.cpu()
@ -98,7 +69,7 @@ class QueueManager:
# Extract tracks
for box, track_id, cls in zip(boxes, track_ids, clss):
# Draw bounding box
self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
annotator.box_label(box, label=self.names[cls], color=colors(int(track_id), True))
# Update track history
track_line = self.track_history[track_id]
@ -108,10 +79,10 @@ class QueueManager:
# Draw track trails if enabled
if self.draw_tracks:
self.annotator.draw_centroid_and_tracks(
annotator.draw_centroid_and_tracks(
track_line,
color=self.track_color or colors(int(track_id), True),
track_thickness=self.track_thickness,
color=colors(int(track_id), True),
track_thickness=self.line_thickness,
)
prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
@ -125,21 +96,16 @@ class QueueManager:
# Display queue counts
label = f"Queue Counts : {str(self.counts)}"
if label is not None:
self.annotator.queue_counts_display(
annotator.queue_counts_display(
label,
points=self.reg_pts,
region_color=self.region_color,
txt_color=self.count_txt_color,
region_color=(255, 0, 255),
txt_color=(104, 31, 17),
)
self.display_frames()
def display_frames(self):
"""Displays the current frame with annotations."""
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)
annotator.draw_region(reg_pts=self.reg_pts, thickness=self.tf * 2, color=(255, 0, 255))
cv2.imshow("Ultralytics YOLOv8 Queue Manager", im0)
# Close window on 'q' key press
if cv2.waitKey(1) & 0xFF == ord("q"):
return
@ -152,12 +118,8 @@ class QueueManager:
im0 (ndarray): Current frame from the video stream.
tracks (list): List of tracks obtained from the object tracking process.
"""
self.im0 = im0 # Store the current frame
self.extract_and_process_tracks(tracks) # Extract and process tracks
if self.view_img:
self.display_frames() # Display the frame if enabled
return self.im0
self.extract_and_process_tracks(tracks, im0) # Extract and process tracks
return im0
if __name__ == "__main__":

172
ultralytics/solutions/speed_estimation.py
Unescape View File

@ -13,7 +13,7 @@ from ultralytics.utils.plotting import Annotator, colors
class SpeedEstimator:
"""A class to estimate the speed of objects in a real-time video stream based on their tracks."""
def __init__(self, names, reg_pts=None, view_img=False, line_thickness=2, region_thickness=5, spdl_dist_thresh=10):
def __init__(self, names, reg_pts=None, view_img=False, line_thickness=2, spdl_dist_thresh=10):
"""
Initializes the SpeedEstimator with the given parameters.
@ -22,158 +22,94 @@ class SpeedEstimator:
reg_pts (list, optional): List of region points for speed estimation. Defaults to [(20, 400), (1260, 400)].
view_img (bool, optional): Whether to display the image with annotations. Defaults to False.
line_thickness (int, optional): Thickness of the lines for drawing boxes and tracks. Defaults to 2.
region_thickness (int, optional): Thickness of the region lines. Defaults to 5.
spdl_dist_thresh (int, optional): Distance threshold for speed calculation. Defaults to 10.
"""
# Visual & image information
self.im0 = None
self.annotator = None
self.view_img = view_img
# Region information
self.reg_pts = reg_pts if reg_pts is not None else [(20, 400), (1260, 400)]
self.region_thickness = region_thickness
self.names = names # Classes names
# Tracking information
self.clss = None
self.names = names
self.boxes = None
self.trk_ids = None
self.trk_pts = None
self.line_thickness = line_thickness
self.trk_history = defaultdict(list)
# Speed estimation information
self.current_time = 0
self.dist_data = {}
self.trk_idslist = []
self.spdl_dist_thresh = spdl_dist_thresh
self.trk_previous_times = {}
self.trk_previous_points = {}
self.view_img = view_img # bool for displaying inference
self.tf = line_thickness # line thickness for annotator
self.spd = {} # set for speed data
self.trkd_ids = [] # list for already speed_estimated and tracked ID's
self.spdl = spdl_dist_thresh # Speed line distance threshold
self.trk_pt = {} # set for tracks previous time
self.trk_pp = {} # set for tracks previous point
# Check if the environment supports imshow
self.env_check = check_imshow(warn=True)
def extract_tracks(self, tracks):
def estimate_speed(self, im0, tracks):
"""
Extracts results from the provided tracking data.
Estimates the speed of objects based on tracking data.
Args:
im0 (ndarray): Image.
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):
"""
Stores track data.
Args:
track_id (int): Object track id.
box (list): Object bounding box data.
Returns:
(list): Updated tracking history for the given track_id.
(ndarray): The image with annotated boxes and tracks.
"""
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 tracks[0].boxes.id is None:
return im0
if len(track) > 30:
track.pop(0)
boxes = tracks[0].boxes.xyxy.cpu()
clss = tracks[0].boxes.cls.cpu().tolist()
t_ids = tracks[0].boxes.id.int().cpu().tolist()
annotator = Annotator(im0, line_width=self.tf)
annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 255), thickness=self.tf * 2)
self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
return track
for box, t_id, cls in zip(boxes, t_ids, clss):
track = self.trk_history[t_id]
bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
track.append(bbox_center)
def plot_box_and_track(self, track_id, box, cls, track):
"""
Plots track and bounding box.
if len(track) > 30:
track.pop(0)
Args:
track_id (int): Object track id.
box (list): Object bounding box data.
cls (str): Object class name.
track (list): Tracking history for drawing tracks path.
"""
speed_label = f"{int(self.dist_data[track_id])} km/h" 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)
trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
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)
if t_id not in self.trk_pt:
self.trk_pt[t_id] = 0
def calculate_speed(self, trk_id, track):
"""
Calculates the speed of an object.
speed_label = f"{int(self.spd[t_id])} km/h" if t_id in self.spd else self.names[int(cls)]
bbox_color = colors(int(t_id), True)
Args:
trk_id (int): Object track id.
track (list): Tracking history for drawing tracks path.
"""
if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
return
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.get(trk_id) != 0 and direction != "unknown" and 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, region_color=(255, 0, 0)):
"""
Estimates the speed of objects based on tracking data.
annotator.box_label(box, speed_label, bbox_color)
cv2.polylines(im0, [trk_pts], isClosed=False, color=bbox_color, thickness=self.tf)
cv2.circle(im0, (int(track[-1][0]), int(track[-1][1])), self.tf * 2, bbox_color, -1)
Args:
im0 (ndarray): Image.
tracks (list): List of tracks obtained from the object tracking process.
region_color (tuple, optional): Color to use when drawing regions. Defaults to (255, 0, 0).
# Calculation of object speed
if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
return
if self.reg_pts[1][1] - self.spdl < track[-1][1] < self.reg_pts[1][1] + self.spdl:
direction = "known"
elif self.reg_pts[0][1] - self.spdl < track[-1][1] < self.reg_pts[0][1] + self.spdl:
direction = "known"
else:
direction = "unknown"
Returns:
(ndarray): The image with annotated boxes and tracks.
"""
self.im0 = im0
if tracks[0].boxes.id is None:
if self.view_img and self.env_check:
self.display_frames()
return im0
if self.trk_pt.get(t_id) != 0 and direction != "unknown" and t_id not in self.trkd_ids:
self.trkd_ids.append(t_id)
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=self.line_thickness)
self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
time_difference = time() - self.trk_pt[t_id]
if time_difference > 0:
self.spd[t_id] = np.abs(track[-1][1] - self.trk_pp[t_id][1]) / time_difference
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)
self.trk_pt[t_id] = time()
self.trk_pp[t_id] = track[-1]
if self.view_img and self.env_check:
self.display_frames()
cv2.imshow("Ultralytics Speed Estimation", im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
return im0
def display_frames(self):
"""Displays the current frame."""
cv2.imshow("Ultralytics Speed Estimation", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
names = {0: "person", 1: "car"} # example class names

Write Preview
Loading…
Cancel
Save