Merge branch 'main' into custom-formatter

custom-formatter
Ultralytics Assistant 4 weeks ago committed by GitHub
commit 7d524ebeeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 4
      .github/workflows/ci.yml
  2. 15
      docs/en/datasets/detect/medical-pills.md
  3. 12
      docs/en/guides/kfold-cross-validation.md
  4. 2
      docs/en/guides/raspberry-pi.md
  5. 2
      docs/en/guides/triton-inference-server.md
  6. 8
      docs/en/hub/inference-api.md
  7. 6
      docs/en/integrations/ibm-watsonx.md
  8. 2
      docs/en/integrations/index.md
  9. 68
      docs/en/integrations/rockchip-rknn.md
  10. 110
      docs/en/integrations/seeedstudio-recamera.md
  11. 2
      docs/en/integrations/tensorrt.md
  12. 36
      docs/en/macros/export-table.md
  13. 2
      docs/en/models/mobile-sam.md
  14. 4
      docs/en/reference/engine/exporter.md
  15. 2
      examples/YOLOv8-ONNXRuntime-CPP/inference.cpp
  16. 3
      examples/tutorial.ipynb
  17. 1
      mkdocs.yml
  18. 64
      tests/test_exports.py
  19. 2
      ultralytics/__init__.py
  20. 7
      ultralytics/cfg/__init__.py
  21. 9
      ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml
  22. 2
      ultralytics/data/augment.py
  23. 6
      ultralytics/data/split_dota.py
  24. 2
      ultralytics/data/utils.py
  25. 154
      ultralytics/engine/exporter.py
  26. 5
      ultralytics/engine/results.py
  27. 3
      ultralytics/engine/trainer.py
  28. 8
      ultralytics/models/nas/val.py
  29. 48
      ultralytics/models/yolo/detect/predict.py
  30. 4
      ultralytics/models/yolo/detect/val.py
  31. 41
      ultralytics/models/yolo/obb/predict.py
  32. 14
      ultralytics/models/yolo/obb/val.py
  33. 43
      ultralytics/models/yolo/pose/predict.py
  34. 13
      ultralytics/models/yolo/pose/val.py
  35. 71
      ultralytics/models/yolo/segment/predict.py
  36. 11
      ultralytics/models/yolo/segment/val.py
  37. 17
      ultralytics/nn/autobackend.py
  38. 4
      ultralytics/nn/modules/block.py
  39. 2
      ultralytics/nn/modules/conv.py
  40. 6
      ultralytics/nn/tasks.py
  41. 3
      ultralytics/trackers/track.py
  42. 11
      ultralytics/utils/__init__.py
  43. 8
      ultralytics/utils/benchmarks.py
  44. 28
      ultralytics/utils/ops.py

@ -98,7 +98,7 @@ jobs:
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-15, windows-latest]
os: [ubuntu-latest, windows-latest, macos-15, ubuntu-24.04-arm]
python-version: ["3.11"]
model: [yolo11n]
steps:
@ -160,7 +160,7 @@ jobs:
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-15, windows-latest]
os: [ubuntu-latest, macos-15, windows-latest, ubuntu-24.04-arm]
python-version: ["3.11"]
torch: [latest]
include:

@ -8,7 +8,20 @@ keywords: medical-pills dataset, pill detection, pharmaceutical imaging, AI in h
<a href="https://colab.research.google.com/github/ultralytics/notebooks/blob/main/notebooks/how-to-train-ultralytics-yolo-on-medical-pills-dataset.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open Medical Pills Dataset In Colab"></a>
The medical-pills detection dataset is a proof-of-concept (POC) dataset, carefully curated to demonstrate the potential of AI in pharmaceutical applications. It contains labeled images specifically designed to train [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) [models](https://docs.ultralytics.com/models/) for identifying medical-pills. This dataset serves as a foundational resource for automating essential [tasks](https://docs.ultralytics.com/tasks/) such as quality control, packaging automation, and efficient sorting in pharmaceutical workflows. By integrating this dataset into projects, researchers and developers can explore innovative [solutions](https://docs.ultralytics.com/solutions/) that enhance [accuracy](https://www.ultralytics.com/glossary/accuracy), streamline operations, and ultimately contribute to improved healthcare outcomes.
The medical-pills detection dataset is a proof-of-concept (POC) dataset, carefully curated to demonstrate the potential of AI in pharmaceutical applications. It contains labeled images specifically designed to train [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) [models](https://docs.ultralytics.com/models/) for identifying medical-pills.
<p align="center">
<br>
<iframe loading="lazy" width="720" height="405" src="https://www.youtube.com/embed/8gePl_Zcs5c"
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> How to train Ultralytics YOLO11 Model on Medical Pills Detection Dataset in <a href="https://colab.research.google.com/github/ultralytics/notebooks/blob/main/notebooks/how-to-train-ultralytics-yolo-on-medical-pills-dataset.ipynb">Google Colab</a>
</p>
This dataset serves as a foundational resource for automating essential [tasks](https://docs.ultralytics.com/tasks/) such as quality control, packaging automation, and efficient sorting in pharmaceutical workflows. By integrating this dataset into projects, researchers and developers can explore innovative [solutions](https://docs.ultralytics.com/solutions/) that enhance [accuracy](https://www.ultralytics.com/glossary/accuracy), streamline operations, and ultimately contribute to improved healthcare outcomes.
## Dataset Structure

@ -82,8 +82,8 @@ Without further ado, let's dive in!
```python
import pandas as pd
indx = [label.stem for label in labels] # uses base filename as ID (no extension)
labels_df = pd.DataFrame([], columns=cls_idx, index=indx)
index = [label.stem for label in labels] # uses base filename as ID (no extension)
labels_df = pd.DataFrame([], columns=cls_idx, index=index)
```
5. Count the instances of each class-label present in the annotation files.
@ -146,11 +146,11 @@ The rows index the label files, each corresponding to an image in your dataset,
```python
folds = [f"split_{n}" for n in range(1, ksplit + 1)]
folds_df = pd.DataFrame(index=indx, columns=folds)
folds_df = pd.DataFrame(index=index, columns=folds)
for idx, (train, val) in enumerate(kfolds, start=1):
folds_df[f"split_{idx}"].loc[labels_df.iloc[train].index] = "train"
folds_df[f"split_{idx}"].loc[labels_df.iloc[val].index] = "val"
for i, (train, val) in enumerate(kfolds, start=1):
folds_df[f"split_{i}"].loc[labels_df.iloc[train].index] = "train"
folds_df[f"split_{i}"].loc[labels_df.iloc[val].index] = "val"
```
3. Now we will calculate the distribution of class labels for each fold as a ratio of the classes present in `val` to those present in `train`.

@ -95,7 +95,7 @@ Here we will install Ultralytics package on the Raspberry Pi with optional depen
## Use NCNN on Raspberry Pi
Out of all the model export formats supported by Ultralytics, [NCNN](https://docs.ultralytics.com/integrations/ncnn/) delivers the best inference performance when working with Raspberry Pi devices because NCNN is highly optimized for mobile/ embedded platforms (such as ARM architecture). Therefor our recommendation is to use NCNN with Raspberry Pi.
Out of all the model export formats supported by Ultralytics, [NCNN](https://docs.ultralytics.com/integrations/ncnn/) delivers the best inference performance when working with Raspberry Pi devices because NCNN is highly optimized for mobile/ embedded platforms (such as ARM architecture).
## Convert Model to NCNN and Run Inference

@ -48,7 +48,7 @@ from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt") # load an official model
# Retreive metadata during export
# Retrieve metadata during export
metadata = []

@ -49,15 +49,9 @@ To shut down the dedicated endpoint, click on the **Stop Endpoint** button.
To use the [Ultralytics HUB](https://www.ultralytics.com/hub) Shared Inference API, follow the guides below.
Free users have the following usage limits:
The [Ultralytics HUB](https://www.ultralytics.com/hub) Shared Inference API has the following usage limits:
- 100 calls / hour
- 1000 calls / month
[Pro](./pro.md) users have the following usage limits:
- 1000 calls / hour
- 10000 calls / month
## Python

@ -133,7 +133,7 @@ After loading the dataset, we printed and saved our working directory. We have a
If you see "trash_ICRA19" among the directory's contents, then it has loaded successfully. You should see three files/folders: a `config.yaml` file, a `videos_for_testing` directory, and a `dataset` directory. We will ignore the `videos_for_testing` directory, so feel free to delete it.
We will use the config.yaml file and the contents of the dataset directory to train our [object detection](https://www.ultralytics.com/glossary/object-detection) model. Here is a sample image from our marine litter data set.
We will use the `config.yaml` file and the contents of the dataset directory to train our [object detection](https://www.ultralytics.com/glossary/object-detection) model. Here is a sample image from our marine litter data set.
<p align="center">
<img width="400" src="https://github.com/ultralytics/docs/releases/download/0/marine-litter-bounding-box.avif" alt="Marine Litter with Bounding Box">
@ -205,14 +205,14 @@ names:
2: rov
```
Run the following script to delete the current contents of config.yaml and replace it with the above contents that reflect our new data set directory structure. Be certain to replace the work_dir portion of the root directory path in line 4 with your own working directory path we retrieved earlier. Leave the train, val, and test subdirectory definitions. Also, do not change {work_dir} in line 23 of the code.
Run the following script to delete the current contents of `config.yaml` and replace it with the above contents that reflect our new data set directory structure. Be certain to replace the work_dir portion of the root directory path in line 4 with your own working directory path we retrieved earlier. Leave the train, val, and test subdirectory definitions. Also, do not change {work_dir} in line 23 of the code.
!!! example "Edit the .yaml File"
=== "Python"
```python
# Contents of new confg.yaml file
# Contents of new config.yaml file
def update_yaml_file(file_path):
data = {
"path": "work_dir/trash_ICRA19/dataset",

@ -97,6 +97,8 @@ Welcome to the Ultralytics Integrations page! This page provides an overview of
- [Rockchip RKNN](rockchip-rknn.md): Developed by [Rockchip](https://www.rock-chips.com/), RKNN is a specialized neural network inference framework optimized for Rockchip's hardware platforms, particularly their NPUs. It facilitates efficient deployment of AI models on edge devices, enabling high-performance inference in real-time applications.
- [Seeed Studio reCamera](seeedstudio-recamera.md): Developed by [Seeed Studio](https://www.seeedstudio.com/), the reCamera is a cutting-edge edge AI device designed for real-time computer vision applications. Powered by the RISC-V-based SG200X processor, it delivers high-performance AI inference with energy efficiency. Its modular design, advanced video processing capabilities, and support for flexible deployment make it an ideal choice for various use cases, including safety monitoring, environmental applications, and manufacturing.
### Export Formats
We also support a variety of model export formats for deployment in different environments. Here are the available formats:

@ -4,18 +4,18 @@ description: Learn how to export YOLO11 models to RKNN format for efficient depl
keywords: YOLO11, RKNN, model export, Ultralytics, Rockchip, machine learning, model deployment, computer vision, deep learning
---
# RKNN Export for Ultralytics YOLO11 Models
# Rockchip RKNN Export for Ultralytics YOLO11 Models
When deploying computer vision models on embedded devices, especially those powered by Rockchip processors, having a compatible model format is essential. Exporting [Ultralytics YOLO11](https://github.com/ultralytics/ultralytics) models to RKNN format ensures optimized performance and compatibility with Rockchip's hardware. This guide will walk you through converting your YOLO11 models to RKNN format, enabling efficient deployment on Rockchip platforms.
<p align="center">
<img width="50%" src="https://github.com/ultralytics/assets/releases/download/v0.0.0/rockchip-rknn.avif" alt="RKNN">
</p>
!!! note
This guide has been tested with [Radxa Rock 5B](https://radxa.com/products/rock5/5b) which is based on Rockchip RK3588 and [Radxa Zero 3W](https://radxa.com/products/zeros/zero3w) which is based on Rockchip RK3566. It is expected to work across other Rockchip-based devices which supports [rknn-toolkit2](https://github.com/airockchip/rknn-toolkit2) such as RK3576, RK3568, RK3562, RV1103, RV1106, RV1103B, RV1106B and RK2118.
<p align="center">
<img width="100%" src="https://www.rock-chips.com/Images/web/solution/AI/chip_s.png" alt="RKNN">
</p>
## What is Rockchip?
Renowned for delivering versatile and power-efficient solutions, Rockchip designs advanced System-on-Chips (SoCs) that power a wide range of consumer electronics, industrial applications, and AI technologies. With ARM-based architecture, built-in Neural Processing Units (NPUs), and high-resolution multimedia support, Rockchip SoCs enable cutting-edge performance for devices like tablets, smart TVs, IoT systems, and edge AI applications. Companies like Radxa, ASUS, Pine64, Orange Pi, Odroid, Khadas, and Banana Pi offer a variety of products based on Rockchip SoCs, further extending their reach and impact across diverse markets.
@ -79,15 +79,15 @@ For detailed instructions and best practices related to the installation process
model = YOLO("yolo11n.pt")
# Export the model to RKNN format
# Here name can be one of rk3588, rk3576, rk3566, rk3568, rk3562, rv1103, rv1106, rv1103b, rv1106b, rk2118
model.export(format="rknn", args={"name": "rk3588"}) # creates '/yolo11n_rknn_model'
# 'name' can be one of rk3588, rk3576, rk3566, rk3568, rk3562, rv1103, rv1106, rv1103b, rv1106b, rk2118
model.export(format="rknn", name="rk3588") # creates '/yolo11n_rknn_model'
```
=== "CLI"
```bash
# Export a YOLO11n PyTorch model to RKNN format
# Here name can be one of rk3588, rk3576, rk3566, rk3568, rk3562, rv1103, rv1106, rv1103b, rv1106b, rk2118
# 'name' can be one of rk3588, rk3576, rk3566, rk3568, rk3562, rv1103, rv1106, rv1103b, rv1106b, rk2118
yolo export model=yolo11n.pt format=rknn name=rk3588 # creates '/yolo11n_rknn_model'
```
@ -139,11 +139,11 @@ YOLO11 benchmarks below were run by the Ultralytics team on Radxa Rock 5B based
| Model | Format | Status | Size (MB) | mAP50-95(B) | Inference time (ms/im) |
| ------- | ------ | ------ | --------- | ----------- | ---------------------- |
| YOLO11n | rknn | ✅ | 7.4 | 0.61 | 99.5 |
| YOLO11s | rknn | ✅ | 20.7 | 0.741 | 122.3 |
| YOLO11m | rknn | ✅ | 41.9 | 0.764 | 298.0 |
| YOLO11l | rknn | ✅ | 53.3 | 0.72 | 319.6 |
| YOLO11x | rknn | ✅ | 114.6 | 0.828 | 632.1 |
| YOLO11n | `rknn` | ✅ | 7.4 | 0.61 | 99.5 |
| YOLO11s | `rknn` | ✅ | 20.7 | 0.741 | 122.3 |
| YOLO11m | `rknn` | ✅ | 41.9 | 0.764 | 298.0 |
| YOLO11l | `rknn` | ✅ | 53.3 | 0.72 | 319.6 |
| YOLO11x | `rknn` | ✅ | 114.6 | 0.828 | 632.1 |
!!! note
@ -156,3 +156,45 @@ In this guide, you've learned how to export Ultralytics YOLO11 models to RKNN fo
For further details on usage, visit the [RKNN official documentation](https://github.com/airockchip/rknn-toolkit2).
Also, if you'd like to know more about other Ultralytics YOLO11 integrations, visit our [integration guide page](../integrations/index.md). You'll find plenty of useful resources and insights there.
## FAQ
### How do I export my Ultralytics YOLO model to RKNN format?
You can easily export your Ultralytics YOLO model to RKNN format using the `export()` method in the Ultralytics Python package or via the command-line interface (CLI). Ensure you are using an x86-based Linux PC for the export process, as ARM64 devices like Rockchip are not supported for this operation. You can specify the target Rockchip platform using the `name` argument, such as `rk3588`, `rk3566`, or others. This process generates an optimized RKNN model ready for deployment on your Rockchip device, taking advantage of its Neural Processing Unit (NPU) for accelerated inference.
!!! Example
=== "Python"
```python
from ultralytics import YOLO
# Load your YOLO model
model = YOLO("yolo11n.pt")
# Export to RKNN format for a specific Rockchip platform
model.export(format="rknn", name="rk3588")
```
=== "CLI"
```bash
yolo export model=yolo11n.pt format=rknn name=rk3588
```
### What are the benefits of using RKNN models on Rockchip devices?
RKNN models are specifically designed to leverage the hardware acceleration capabilities of Rockchip's Neural Processing Units (NPUs). This optimization results in significantly faster inference speeds and reduced latency compared to running generic model formats like ONNX or TensorFlow Lite on the same hardware. Using RKNN models allows for more efficient use of the device's resources, leading to lower power consumption and better overall performance, especially critical for real-time applications on edge devices. By converting your Ultralytics YOLO models to RKNN, you can achieve optimal performance on devices powered by Rockchip SoCs like the RK3588, RK3566, and others.
### Can I deploy RKNN models on devices from other manufacturers like NVIDIA or Google?
RKNN models are specifically optimized for Rockchip platforms and their integrated NPUs. While you can technically run an RKNN model on other platforms using software emulation, you will not benefit from the hardware acceleration provided by Rockchip devices. For optimal performance on other platforms, it's recommended to export your Ultralytics YOLO models to formats specifically designed for those platforms, such as TensorRT for NVIDIA GPUs or [TensorFlow Lite](https://docs.ultralytics.com/integrations/tflite/) for Google's Edge TPU. Ultralytics supports exporting to a wide range of formats, ensuring compatibility with various hardware accelerators.
### What Rockchip platforms are supported for RKNN model deployment?
The Ultralytics YOLO export to RKNN format supports a wide range of Rockchip platforms, including the popular RK3588, RK3576, RK3566, RK3568, RK3562, RV1103, RV1106, RV1103B, RV1106B, and RK2118. These platforms are commonly found in devices from manufacturers like Radxa, ASUS, Pine64, Orange Pi, Odroid, Khadas, and Banana Pi. This broad support ensures that you can deploy your optimized RKNN models on various Rockchip-powered devices, from single-board computers to industrial systems, taking full advantage of their AI acceleration capabilities for enhanced performance in your computer vision applications.
### How does the performance of RKNN models compare to other formats on Rockchip devices?
RKNN models generally outperform other formats like ONNX or TensorFlow Lite on Rockchip devices due to their optimization for Rockchip's NPUs. For instance, benchmarks on the Radxa Rock 5B (RK3588) show that [YOLO11n](https://www.ultralytics.com/blog/all-you-need-to-know-about-ultralytics-yolo11-and-its-applications) in RKNN format achieves an inference time of 99.5 ms/image, significantly faster than other formats. This performance advantage is consistent across various YOLO11 model sizes, as demonstrated in the [benchmarks section](#benchmarks). By leveraging the dedicated NPU hardware, RKNN models minimize latency and maximize throughput, making them ideal for real-time applications on Rockchip-based edge devices.

@ -0,0 +1,110 @@
---
comments: true
description: Discover how to get started with Seeed Studio reCamera for edge AI applications using Ultralytics YOLO11. Learn about its powerful features, real-world applications, and how to export YOLO11 models to ONNX format for seamless integration.
keywords: Seeed Studio reCamera, YOLO11, ONNX export, edge AI, computer vision, real-time detection, personal protective equipment detection, fire detection, waste detection, fall detection, modular AI devices, Ultralytics
---
# Quick Start Guide: Seeed Studio reCamera with Ultralytics YOLO11
[reCamera](https://www.seeedstudio.com/recamera) was introduced for the AI community at [YOLO Vision 2024 (YV24)](https://www.youtube.com/watch?v=rfI5vOo3-_A), [Ultralytics](https://ultralytics.com/) annual hybrid event. It is mainly designed for edge AI applications, offering powerful processing capabilities and effortless deployment.
With support for diverse hardware configurations and open-source resources, it serves as an ideal platform for prototyping and deploying innovative [computer vision](https://www.ultralytics.com/glossary/computer-vision-cv) [solutions](https://docs.ultralytics.com/solutions/#solutions) at the edge.
![Seeed Studio reCamera](https://github.com/ultralytics/docs/releases/download/0/saeed-studio-recamera.avif)
## Why Choose reCamera?
reCamera series is purpose-built for edge AI applications, tailored to meet the needs of developers and innovators. Here's why it stands out:
- **RISC-V Powered Performance**: At its core is the SG200X processor, built on the RISC-V architecture, delivering exceptional performance for edge AI tasks while maintaining energy efficiency. With the ability to execute 1 trillion operations per second (1 TOPS), it handles demanding tasks like real-time object detection easily.
- **Optimized Video Technologies**: Supports advanced video compression standards, including H.264 and H.265, to reduce storage and bandwidth requirements without sacrificing quality. Features like HDR imaging, 3D noise reduction, and lens correction ensure professional visuals, even in challenging environments.
- **Energy-Efficient Dual Processing**: While the SG200X handles complex AI tasks, a smaller 8-bit microcontroller manages simpler operations to conserve power, making the reCamera ideal for battery-operated or low-power setups.
- **Modular and Upgradable Design**: The reCamera is built with a modular structure, consisting of three main components: the core board, sensor board, and baseboard. This design allows developers to easily swap or upgrade components, ensuring flexibility and future-proofing for evolving projects.
## Quick Hardware Setup of reCamera
Please follow [reCamera Quick Start Guide](https://wiki.seeedstudio.com/recamera_getting_started) for initial onboarding of the device such as connecting the device to a WiFi network and access the [Node-RED](https://nodered.org) web UI for quick previewing of detection redsults with the pre-installed Ultralytics YOLO models.
## Export to cvimodel: Converting Your YOLO11 Model
Here we will first convert `PyTorch` model to `ONNX` and then convert it to `MLIR` model format. Finally `MLIR` will be converted to `cvimodel` in order to inference on-device
<p align="center">
<img width="80%" src="https://github.com/ultralytics/assets/releases/download/v0.0.0/recamera-toolchain-workflow.avif" alt="reCamera Toolchain">
</p>
### Export to ONNX
Export an Ultralytics YOLO11 model to ONNX model format.
#### Installation
To install the required packages, run:
!!! Tip "Installation"
=== "CLI"
```bash
pip install ultralytics
```
For detailed instructions and best practices related to the installation process, check our [Ultralytics Installation guide](../quickstart.md). While installing the required packages for YOLO11, if you encounter any difficulties, consult our [Common Issues guide](../guides/yolo-common-issues.md) for solutions and tips.
#### Usage
!!! Example "Usage"
=== "Python"
```python
from ultralytics import YOLO
# Load the YOLO11 model
model = YOLO("yolo11n.pt")
# Export the model to ONNX format
model.export(format="onnx") # creates 'yolo11n.onnx'
```
=== "CLI"
```bash
# Export a YOLO11n PyTorch model to ONNX format
yolo export model=yolo11n.pt format=onnx # creates 'yolo11n.onnx'
```
For more details about the export process, visit the [Ultralytics documentation page on exporting](../modes/export.md).
### Export ONNX to MLIR and cvimodel
After obtaining an ONNX model, refer to [Convert and Quantize AI Models](https://wiki.seeedstudio.com/recamera_model_conversion) page to convert the ONNX model to MLIR and then to cvimodel.
!!! note
We're actively working on adding reCamera support directly into the Ultralytics package, and it will be available soon. In the meantime, check out our blog on [Integrating Ultralytics YOLO Models with Seeed Studio's reCamera](https://www.ultralytics.com/blog/integrating-ultralytics-yolo-models-on-seeed-studios-recamera) for more insights.
## Benchmarks
Coming soon.
## Real-World Applications of reCamera
reCamera advanced computer vision capabilities and modular design make it suitable for a wide range of real-world scenarios, helping developers and businesses tackle unique challenges with ease.
- **Fall Detection**: Designed for safety and healthcare applications, the reCamera can detect falls in real-time, making it ideal for elderly care, hospitals, and industrial settings where rapid response is critical.
- **Personal Protective Equipment Detection**: The reCamera can be used to ensure workplace safety by detecting PPE compliance in real-time. It helps identify whether workers are wearing helmets, gloves, or other safety gear, reducing risks in industrial environments.
![Personal protective equipment detection](https://github.com/ultralytics/docs/releases/download/0/personal-protective-equipment-detection.avif)
- **Fire Detection**: The reCamera's real-time processing capabilities make it an excellent choice for fire detection in industrial and residential areas, providing early warnings to prevent potential disasters.
- **Waste Detection**: It can also be utilized for waste detection applications, making it an excellent tool for environmental monitoring and waste management.
- **Car Parts Detection**: In manufacturing and automotive industries, it aids in detecting and analyzing car parts for quality control, assembly line monitoring, and inventory management.
![Car parts detection](https://github.com/ultralytics/docs/releases/download/0/carparts-detection.avif)

@ -185,7 +185,7 @@ Experimentation by NVIDIA led them to recommend using at least 500 calibration i
???+ warning "Calibration Cache"
TensorRT will generate a calibration `.cache` which can be re-used to speed up export of future model weights using the same data, but this may result in poor calibration when the data is vastly different or if the `batch` value is changed drastically. In these circumstances, the existing `.cache` should be renamed and moved to a different directory or deleted entirely.
TensorRT will generate a calibration `.cache` which can be reused to speed up export of future model weights using the same data, but this may result in poor calibration when the data is vastly different or if the `batch` value is changed drastically. In these circumstances, the existing `.cache` should be renamed and moved to a different directory or deleted entirely.
#### Advantages of using YOLO with TensorRT INT8

@ -1,18 +1,18 @@
| Format | `format` Argument | Model | Metadata | Arguments |
| ------------------------------------------------- | ----------------- | ----------------------------------------------- | -------- | -------------------------------------------------------------------- |
| [PyTorch](https://pytorch.org/) | - | `{{ model_name or "yolo11n" }}.pt` | ✅ | - |
| [TorchScript](../integrations/torchscript.md) | `torchscript` | `{{ model_name or "yolo11n" }}.torchscript` | ✅ | `imgsz`, `optimize`, `batch` |
| [ONNX](../integrations/onnx.md) | `onnx` | `{{ model_name or "yolo11n" }}.onnx` | ✅ | `imgsz`, `half`, `dynamic`, `simplify`, `opset`, `batch` |
| [OpenVINO](../integrations/openvino.md) | `openvino` | `{{ model_name or "yolo11n" }}_openvino_model/` | ✅ | `imgsz`, `half`, `dynamic`, `int8`, `batch` |
| [TensorRT](../integrations/tensorrt.md) | `engine` | `{{ model_name or "yolo11n" }}.engine` | ✅ | `imgsz`, `half`, `dynamic`, `simplify`, `workspace`, `int8`, `batch` |
| [CoreML](../integrations/coreml.md) | `coreml` | `{{ model_name or "yolo11n" }}.mlpackage` | ✅ | `imgsz`, `half`, `int8`, `nms`, `batch` |
| [TF SavedModel](../integrations/tf-savedmodel.md) | `saved_model` | `{{ model_name or "yolo11n" }}_saved_model/` | ✅ | `imgsz`, `keras`, `int8`, `batch` |
| [TF GraphDef](../integrations/tf-graphdef.md) | `pb` | `{{ model_name or "yolo11n" }}.pb` | ❌ | `imgsz`, `batch` |
| [TF Lite](../integrations/tflite.md) | `tflite` | `{{ model_name or "yolo11n" }}.tflite` | ✅ | `imgsz`, `half`, `int8`, `batch` |
| [TF Edge TPU](../integrations/edge-tpu.md) | `edgetpu` | `{{ model_name or "yolo11n" }}_edgetpu.tflite` | ✅ | `imgsz` |
| [TF.js](../integrations/tfjs.md) | `tfjs` | `{{ model_name or "yolo11n" }}_web_model/` | ✅ | `imgsz`, `half`, `int8`, `batch` |
| [PaddlePaddle](../integrations/paddlepaddle.md) | `paddle` | `{{ model_name or "yolo11n" }}_paddle_model/` | ✅ | `imgsz`, `batch` |
| [MNN](../integrations/mnn.md) | `mnn` | `{{ model_name or "yolo11n" }}.mnn` | ✅ | `imgsz`, `batch`, `int8`, `half` |
| [NCNN](../integrations/ncnn.md) | `ncnn` | `{{ model_name or "yolo11n" }}_ncnn_model/` | ✅ | `imgsz`, `half`, `batch` |
| [IMX500](../integrations/sony-imx500.md) | `imx` | `{{ model_name or "yolov8n" }}_imx_model/` | ✅ | `imgsz`, `int8` |
| [RKNN](../integrations/rockchip-rknn.md) | `rknn` | `{{ model_name or "yolo11n" }}_rknn_model/` | ✅ | `imgsz`, `batch`, `name` |
| Format | `format` Argument | Model | Metadata | Arguments |
| ------------------------------------------------- | ----------------- | ----------------------------------------------- | -------- | --------------------------------------------------------------------------- |
| [PyTorch](https://pytorch.org/) | - | `{{ model_name or "yolo11n" }}.pt` | ✅ | - |
| [TorchScript](../integrations/torchscript.md) | `torchscript` | `{{ model_name or "yolo11n" }}.torchscript` | ✅ | `imgsz`, `optimize`, `nms`, `batch` |
| [ONNX](../integrations/onnx.md) | `onnx` | `{{ model_name or "yolo11n" }}.onnx` | ✅ | `imgsz`, `half`, `dynamic`, `simplify`, `opset`, `nms`, `batch` |
| [OpenVINO](../integrations/openvino.md) | `openvino` | `{{ model_name or "yolo11n" }}_openvino_model/` | ✅ | `imgsz`, `half`, `dynamic`, `int8`, `nms`, `batch` |
| [TensorRT](../integrations/tensorrt.md) | `engine` | `{{ model_name or "yolo11n" }}.engine` | ✅ | `imgsz`, `half`, `dynamic`, `simplify`, `workspace`, `int8`, `nms`, `batch` |
| [CoreML](../integrations/coreml.md) | `coreml` | `{{ model_name or "yolo11n" }}.mlpackage` | ✅ | `imgsz`, `half`, `int8`, `nms`, `batch` |
| [TF SavedModel](../integrations/tf-savedmodel.md) | `saved_model` | `{{ model_name or "yolo11n" }}_saved_model/` | ✅ | `imgsz`, `keras`, `int8`, `nms`, `batch` |
| [TF GraphDef](../integrations/tf-graphdef.md) | `pb` | `{{ model_name or "yolo11n" }}.pb` | ❌ | `imgsz`, `batch` |
| [TF Lite](../integrations/tflite.md) | `tflite` | `{{ model_name or "yolo11n" }}.tflite` | ✅ | `imgsz`, `half`, `int8`, `nms`, `batch` |
| [TF Edge TPU](../integrations/edge-tpu.md) | `edgetpu` | `{{ model_name or "yolo11n" }}_edgetpu.tflite` | ✅ | `imgsz` |
| [TF.js](../integrations/tfjs.md) | `tfjs` | `{{ model_name or "yolo11n" }}_web_model/` | ✅ | `imgsz`, `half`, `int8`, `nms`, `batch` |
| [PaddlePaddle](../integrations/paddlepaddle.md) | `paddle` | `{{ model_name or "yolo11n" }}_paddle_model/` | ✅ | `imgsz`, `batch` |
| [MNN](../integrations/mnn.md) | `mnn` | `{{ model_name or "yolo11n" }}.mnn` | ✅ | `imgsz`, `batch`, `int8`, `half` |
| [NCNN](../integrations/ncnn.md) | `ncnn` | `{{ model_name or "yolo11n" }}_ncnn_model/` | ✅ | `imgsz`, `half`, `batch` |
| [IMX500](../integrations/sony-imx500.md) | `imx` | `{{ model_name or "yolov8n" }}_imx_model/` | ✅ | `imgsz`, `int8` |
| [RKNN](../integrations/rockchip-rknn.md) | `rknn` | `{{ model_name or "yolo11n" }}_rknn_model/` | ✅ | `imgsz`, `batch`, `name` |

@ -118,7 +118,7 @@ You can download the model [here](https://github.com/ChaoningZhang/MobileSAM/blo
# Predict a segment based on a single point prompt
model.predict("ultralytics/assets/zidane.jpg", points=[900, 370], labels=[1])
# Predict mutiple segments based on multiple points prompt
# Predict multiple segments based on multiple points prompt
model.predict("ultralytics/assets/zidane.jpg", points=[[400, 370], [900, 370]], labels=[1, 1])
# Predict a segment based on multiple points prompt per object

@ -19,6 +19,10 @@ keywords: YOLOv8, export formats, ONNX, TensorRT, CoreML, machine learning model
<br><br><hr><br>
## ::: ultralytics.engine.exporter.NMSModel
<br><br><hr><br>
## ::: ultralytics.engine.exporter.export_formats
<br><br><hr><br>

@ -107,11 +107,11 @@ char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) {
iouThreshold = iParams.iouThreshold;
imgSize = iParams.imgSize;
modelType = iParams.modelType;
cudaEnable = iParams.cudaEnable;
env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo");
Ort::SessionOptions sessionOption;
if (iParams.cudaEnable)
{
cudaEnable = iParams.cudaEnable;
OrtCUDAProviderOptions cudaOption;
cudaOption.device_id = 0;
sessionOption.AppendExecutionProvider_CUDA(cudaOption);

@ -378,7 +378,8 @@
"| [PaddlePaddle](https://docs.ultralytics.com/integrations/paddlepaddle) | `paddle` | `yolo11n_paddle_model/` | ✅ | `imgsz`, `batch` |\n",
"| [MNN](https://docs.ultralytics.com/integrations/mnn) | `mnn` | `yolo11n.mnn` | ✅ | `imgsz`, `batch`, `int8`, `half` |\n",
"| [NCNN](https://docs.ultralytics.com/integrations/ncnn) | `ncnn` | `yolo11n_ncnn_model/` | ✅ | `imgsz`, `half`, `batch` |\n",
"| [IMX500](https://docs.ultralytics.com/integrations/sony-imx500) | `imx` | `yolov8n_imx_model/` | ✅ | `imgsz`, `int8` |"
"| [IMX500](https://docs.ultralytics.com/integrations/sony-imx500) | `imx` | `yolov8n_imx_model/` | ✅ | `imgsz`, `int8` |\n",
"| [RKNN](https://docs.ultralytics.com/integrations/rockchip-rknn) | `rknn` | `yolo11n_rknn_model/` | ✅ | `imgsz`, `batch`, `name` |"
],
"metadata": {
"id": "nPZZeNrLCQG6"

@ -425,6 +425,7 @@ nav:
- Albumentations: integrations/albumentations.md
- SONY IMX500: integrations/sony-imx500.md
- Rockchip RKNN: integrations/rockchip-rknn.md
- Seeed Studio reCamera: integrations/seeedstudio-recamera.md
- HUB:
- hub/index.md
- Web:

@ -11,6 +11,7 @@ from tests import MODEL, SOURCE
from ultralytics import YOLO
from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
from ultralytics.utils import (
ARM64,
IS_RASPBERRYPI,
LINUX,
MACOS,
@ -42,23 +43,19 @@ def test_export_openvino():
@pytest.mark.slow
@pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
"task, dynamic, int8, half, batch, nms",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
(task, dynamic, int8, half, batch, nms)
for task, dynamic, int8, half, batch, nms in product(
TASKS, [True, False], [True, False], [True, False], [1, 2], [True, False]
)
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_openvino_matrix(task, dynamic, int8, half, batch):
def test_export_openvino_matrix(task, dynamic, int8, half, batch, nms):
"""Test YOLO model exports to OpenVINO under various configuration matrix conditions."""
file = YOLO(TASK2MODEL[task]).export(
format="openvino",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
data=TASK2DATA[task],
format="openvino", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, data=TASK2DATA[task], nms=nms
)
if WINDOWS:
# Use unique filenames due to Windows file permissions bug possibly due to latent threaded use
@ -71,34 +68,26 @@ def test_export_openvino_matrix(task, dynamic, int8, half, batch):
@pytest.mark.slow
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch, simplify", product(TASKS, [True, False], [False], [False], [1, 2], [True, False])
"task, dynamic, int8, half, batch, simplify, nms",
product(TASKS, [True, False], [False], [False], [1, 2], [True, False], [True, False]),
)
def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify):
def test_export_onnx_matrix(task, dynamic, int8, half, batch, simplify, nms):
"""Test YOLO exports to ONNX format with various configurations and parameters."""
file = YOLO(TASK2MODEL[task]).export(
format="onnx",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
simplify=simplify,
format="onnx", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, simplify=simplify, nms=nms
)
YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32) # exported model inference
Path(file).unlink() # cleanup
@pytest.mark.slow
@pytest.mark.parametrize("task, dynamic, int8, half, batch", product(TASKS, [False], [False], [False], [1, 2]))
def test_export_torchscript_matrix(task, dynamic, int8, half, batch):
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch, nms", product(TASKS, [False], [False], [False], [1, 2], [True, False])
)
def test_export_torchscript_matrix(task, dynamic, int8, half, batch, nms):
"""Tests YOLO model exports to TorchScript format under varied configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="torchscript",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
format="torchscript", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, nms=nms
)
YOLO(file)([SOURCE] * 3, imgsz=64 if dynamic else 32) # exported model inference at batch=3
Path(file).unlink() # cleanup
@ -134,22 +123,19 @@ def test_export_coreml_matrix(task, dynamic, int8, half, batch):
@pytest.mark.skipif(not checks.IS_PYTHON_MINIMUM_3_10, reason="TFLite export requires Python>=3.10")
@pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
@pytest.mark.parametrize(
"task, dynamic, int8, half, batch",
"task, dynamic, int8, half, batch, nms",
[ # generate all combinations but exclude those where both int8 and half are True
(task, dynamic, int8, half, batch)
for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
(task, dynamic, int8, half, batch, nms)
for task, dynamic, int8, half, batch, nms in product(
TASKS, [False], [True, False], [True, False], [1], [True, False]
)
if not (int8 and half) # exclude cases where both int8 and half are True
],
)
def test_export_tflite_matrix(task, dynamic, int8, half, batch):
def test_export_tflite_matrix(task, dynamic, int8, half, batch, nms):
"""Test YOLO exports to TFLite format considering various export configurations."""
file = YOLO(TASK2MODEL[task]).export(
format="tflite",
imgsz=32,
dynamic=dynamic,
int8=int8,
half=half,
batch=batch,
format="tflite", imgsz=32, dynamic=dynamic, int8=int8, half=half, batch=batch, nms=nms
)
YOLO(file)([SOURCE] * batch, imgsz=32) # exported model inference at batch=3
Path(file).unlink() # cleanup
@ -157,7 +143,7 @@ def test_export_tflite_matrix(task, dynamic, int8, half, batch):
@pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
@pytest.mark.skipif(WINDOWS, reason="CoreML not supported on Windows") # RuntimeError: BlobWriter not loaded
@pytest.mark.skipif(IS_RASPBERRYPI, reason="CoreML not supported on Raspberry Pi")
@pytest.mark.skipif(LINUX and ARM64, reason="CoreML not supported on aarch64 Linux")
@pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="CoreML not supported in Python 3.12")
def test_export_coreml():
"""Test YOLO exports to CoreML format, optimized for macOS only."""

@ -1,6 +1,6 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
__version__ = "8.3.65"
__version__ = "8.3.67"
import os

@ -921,12 +921,7 @@ def entrypoint(debug=""):
# Task
task = overrides.pop("task", None)
if task:
if task == "classify" and mode == "track":
raise ValueError(
f"❌ Classification doesn't support 'mode=track'. Valid modes for classification are"
f" {MODES - {'track'}}.\n{CLI_HELP_MSG}"
)
elif task not in TASKS:
if task not in TASKS:
if task == "track":
LOGGER.warning(
"WARNING ⚠ invalid 'task=track', setting 'task=detect' and 'mode=track'. Valid tasks are {TASKS}.\n{CLI_HELP_MSG}."

@ -6,18 +6,11 @@
# Parameters
nc: 10 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolo11n-cls.yaml' will call yolo11-cls.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024]
s: [0.33, 0.50, 1024]
m: [0.67, 0.75, 1024]
l: [1.00, 1.00, 1024]
x: [1.00, 1.25, 1024]
# ResNet18 backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, TorchVision, [512, "resnet18", "DEFAULT", True, 2]] # truncate two layers from the end
- [-1, 1, TorchVision, [512, resnet18, DEFAULT, True, 2]] # truncate two layers from the end
# YOLO11n head
head:

@ -1850,7 +1850,7 @@ class Albumentations:
A.CLAHE(p=0.01),
A.RandomBrightnessContrast(p=0.0),
A.RandomGamma(p=0.0),
A.ImageCompression(quality_lower=75, p=0.0),
A.ImageCompression(quality_range=(75, 100), p=0.0),
]
# Compose transforms

@ -8,9 +8,9 @@ from pathlib import Path
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
from ultralytics.data.utils import exif_size, img2label_paths
from ultralytics.utils import TQDM
from ultralytics.utils.checks import check_requirements
@ -221,7 +221,7 @@ def split_images_and_labels(data_root, save_dir, split="train", crop_sizes=(1024
lb_dir.mkdir(parents=True, exist_ok=True)
annos = load_yolo_dota(data_root, split=split)
for anno in tqdm(annos, total=len(annos), desc=split):
for anno in TQDM(annos, total=len(annos), desc=split):
windows = get_windows(anno["ori_size"], crop_sizes, gaps)
window_objs = get_window_obj(anno, windows)
crop_and_save(anno, windows, window_objs, str(im_dir), str(lb_dir))
@ -281,7 +281,7 @@ def split_test(data_root, save_dir, crop_size=1024, gap=200, rates=(1.0,)):
im_dir = Path(data_root) / "images" / "test"
assert im_dir.exists(), f"Can't find {im_dir}, please check your data root."
im_files = glob(str(im_dir / "*"))
for im_file in tqdm(im_files, total=len(im_files), desc="test"):
for im_file in TQDM(im_files, total=len(im_files), desc="test"):
w, h = exif_size(Image.open(im_file))
windows = get_windows((h, w), crop_sizes=crop_sizes, gaps=gaps)
im = cv2.imread(im_file)

@ -136,7 +136,7 @@ def verify_image_label(args):
# All labels
max_cls = lb[:, 0].max() # max label count
assert max_cls <= num_cls, (
assert max_cls < num_cls, (
f"Label class {int(max_cls)} exceeds dataset class count {num_cls}. "
f"Possible class labels are 0-{num_cls - 1}"
)

@ -103,7 +103,7 @@ from ultralytics.utils.checks import (
)
from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
from ultralytics.utils.files import file_size, spaces_in_path
from ultralytics.utils.ops import Profile
from ultralytics.utils.ops import Profile, nms_rotated, xywh2xyxy
from ultralytics.utils.torch_utils import TORCH_1_13, get_latest_opset, select_device
@ -111,16 +111,16 @@ def export_formats():
"""Ultralytics YOLO export formats."""
x = [
["PyTorch", "-", ".pt", True, True, []],
["TorchScript", "torchscript", ".torchscript", True, True, ["batch", "optimize"]],
["ONNX", "onnx", ".onnx", True, True, ["batch", "dynamic", "half", "opset", "simplify"]],
["OpenVINO", "openvino", "_openvino_model", True, False, ["batch", "dynamic", "half", "int8"]],
["TensorRT", "engine", ".engine", False, True, ["batch", "dynamic", "half", "int8", "simplify"]],
["TorchScript", "torchscript", ".torchscript", True, True, ["batch", "optimize", "nms"]],
["ONNX", "onnx", ".onnx", True, True, ["batch", "dynamic", "half", "opset", "simplify", "nms"]],
["OpenVINO", "openvino", "_openvino_model", True, False, ["batch", "dynamic", "half", "int8", "nms"]],
["TensorRT", "engine", ".engine", False, True, ["batch", "dynamic", "half", "int8", "simplify", "nms"]],
["CoreML", "coreml", ".mlpackage", True, False, ["batch", "half", "int8", "nms"]],
["TensorFlow SavedModel", "saved_model", "_saved_model", True, True, ["batch", "int8", "keras"]],
["TensorFlow SavedModel", "saved_model", "_saved_model", True, True, ["batch", "int8", "keras", "nms"]],
["TensorFlow GraphDef", "pb", ".pb", True, True, ["batch"]],
["TensorFlow Lite", "tflite", ".tflite", True, False, ["batch", "half", "int8"]],
["TensorFlow Lite", "tflite", ".tflite", True, False, ["batch", "half", "int8", "nms"]],
["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", True, False, []],
["TensorFlow.js", "tfjs", "_web_model", True, False, ["batch", "half", "int8"]],
["TensorFlow.js", "tfjs", "_web_model", True, False, ["batch", "half", "int8", "nms"]],
["PaddlePaddle", "paddle", "_paddle_model", True, True, ["batch"]],
["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]],
["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
@ -281,6 +281,11 @@ class Exporter:
)
if self.args.int8 and tflite:
assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
if self.args.nms:
if getattr(model, "end2end", False):
LOGGER.warning("WARNING ⚠ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
self.args.nms = False
self.args.conf = self.args.conf or 0.25 # set conf default value for nms export
if edgetpu:
if not LINUX:
raise SystemError("Edge TPU export only supported on Linux. See https://coral.ai/docs/edgetpu/compiler")
@ -344,8 +349,8 @@ class Exporter:
)
y = None
for _ in range(2):
y = model(im) # dry runs
for _ in range(2): # dry runs
y = NMSModel(model, self.args)(im) if self.args.nms and not coreml else model(im)
if self.args.half and onnx and self.device.type != "cpu":
im, model = im.half(), model.half() # to FP16
@ -476,7 +481,7 @@ class Exporter:
LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
f = self.file.with_suffix(".torchscript")
ts = torch.jit.trace(self.model, self.im, strict=False)
ts = torch.jit.trace(NMSModel(self.model, self.args) if self.args.nms else self.model, self.im, strict=False)
extra_files = {"config.txt": json.dumps(self.metadata)} # torch._C.ExtraFilesMap()
if self.args.optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
LOGGER.info(f"{prefix} optimizing for mobile...")
@ -499,7 +504,6 @@ class Exporter:
opset_version = self.args.opset or get_latest_opset()
LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset_version}...")
f = str(self.file.with_suffix(".onnx"))
output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
dynamic = self.args.dynamic
if dynamic:
@ -509,9 +513,18 @@ class Exporter:
dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic["output0"] = {0: "batch", 2: "anchors"} # shape(1, 84, 8400)
if self.args.nms: # only batch size is dynamic with NMS
dynamic["output0"].pop(2)
if self.args.nms and self.model.task == "obb":
self.args.opset = opset_version # for NMSModel
# OBB error https://github.com/pytorch/pytorch/issues/110859#issuecomment-1757841865
torch.onnx.register_custom_op_symbolic("aten::lift_fresh", lambda g, x: x, opset_version)
check_requirements("onnxslim>=0.1.46") # Older versions has bug with OBB
torch.onnx.export(
self.model.cpu() if dynamic else self.model, # dynamic=True only compatible with cpu
NMSModel(self.model.cpu() if dynamic else self.model, self.args)
if self.args.nms
else self.model, # dynamic=True only compatible with cpu
self.im.cpu() if dynamic else self.im,
f,
verbose=False,
@ -553,7 +566,7 @@ class Exporter:
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
assert TORCH_1_13, f"OpenVINO export requires torch>=1.13.0 but torch=={torch.__version__} is installed"
ov_model = ov.convert_model(
self.model,
NMSModel(self.model, self.args) if self.args.nms else self.model,
input=None if self.args.dynamic else [self.im.shape],
example_input=self.im,
)
@ -736,9 +749,6 @@ class Exporter:
f = self.file.with_suffix(".mlmodel" if mlmodel else ".mlpackage")
if f.is_dir():
shutil.rmtree(f)
if self.args.nms and getattr(self.model, "end2end", False):
LOGGER.warning(f"{prefix} WARNING ⚠ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
self.args.nms = False
bias = [0.0, 0.0, 0.0]
scale = 1 / 255
@ -1159,21 +1169,19 @@ class Exporter:
from rknn.api import RKNN
f, _ = self.export_onnx()
platform = self.args.name
export_path = Path(f"{Path(f).stem}_rknn_model")
export_path.mkdir(exist_ok=True)
rknn = RKNN(verbose=False)
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
_ = rknn.load_onnx(model=f)
_ = rknn.build(do_quantization=False) # TODO: Add quantization support
f = f.replace(".onnx", f"-{platform}.rknn")
_ = rknn.export_rknn(f"{export_path / f}")
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=self.args.name)
rknn.load_onnx(model=f)
rknn.build(do_quantization=False) # TODO: Add quantization support
f = f.replace(".onnx", f"-{self.args.name}.rknn")
rknn.export_rknn(f"{export_path / f}")
yaml_save(export_path / "metadata.yaml", self.metadata)
return export_path, None
@try_export
def export_imx(self, prefix=colorstr("IMX:")):
"""YOLO IMX export."""
gptq = False
@ -1191,6 +1199,8 @@ class Exporter:
import onnx
from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")
try:
out = subprocess.run(
["java", "--version"], check=True, capture_output=True
@ -1286,7 +1296,7 @@ class Exporter:
f = Path(str(self.file).replace(self.file.suffix, "_imx_model"))
f.mkdir(exist_ok=True)
onnx_model = f / Path(str(self.file).replace(self.file.suffix, "_imx.onnx")) # js dir
onnx_model = f / Path(str(self.file.name).replace(self.file.suffix, "_imx.onnx")) # js dir
mct.exporter.pytorch_export_model(
model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
)
@ -1438,8 +1448,8 @@ class Exporter:
nms.coordinatesOutputFeatureName = "coordinates"
nms.iouThresholdInputFeatureName = "iouThreshold"
nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
nms.iouThreshold = 0.45
nms.confidenceThreshold = 0.25
nms.iouThreshold = self.args.iou
nms.confidenceThreshold = self.args.conf
nms.pickTop.perClass = True
nms.stringClassLabels.vector.extend(names.values())
nms_model = ct.models.MLModel(nms_spec)
@ -1507,3 +1517,91 @@ class IOSDetectModel(torch.nn.Module):
"""Normalize predictions of object detection model with input size-dependent factors."""
xywh, cls = self.model(x)[0].transpose(0, 1).split((4, self.nc), 1)
return cls, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4)
class NMSModel(torch.nn.Module):
"""Model wrapper with embedded NMS for Detect, Segment, Pose and OBB."""
def __init__(self, model, args):
"""
Initialize the NMSModel.
Args:
model (torch.nn.module): The model to wrap with NMS postprocessing.
args (Namespace): The export arguments.
"""
super().__init__()
self.model = model
self.args = args
self.obb = model.task == "obb"
self.is_tf = self.args.format in frozenset({"saved_model", "tflite", "tfjs"})
def forward(self, x):
"""
Performs inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
Args:
x (torch.tensor): The preprocessed tensor with shape (N, 3, H, W).
Returns:
out (torch.tensor): The post-processed results with shape (N, max_det, 4 + 2 + extra_shape).
"""
from functools import partial
from torchvision.ops import nms
preds = self.model(x)
pred = preds[0] if isinstance(preds, tuple) else preds
pred = pred.transpose(-1, -2) # shape(1,84,6300) to shape(1,6300,84)
extra_shape = pred.shape[-1] - (4 + self.model.nc) # extras from Segment, OBB, Pose
boxes, scores, extras = pred.split([4, self.model.nc, extra_shape], dim=2)
scores, classes = scores.max(dim=-1)
# (N, max_det, 4 coords + 1 class score + 1 class label + extra_shape).
out = torch.zeros(
boxes.shape[0],
self.args.max_det,
boxes.shape[-1] + 2 + extra_shape,
device=boxes.device,
dtype=boxes.dtype,
)
for i, (box, cls, score, extra) in enumerate(zip(boxes, classes, scores, extras)):
mask = score > self.args.conf
if self.is_tf:
# TFLite GatherND error if mask is empty
score *= mask
# Explicit length otherwise reshape error, hardcoded to `self.args.max_det * 5`
mask = score.topk(self.args.max_det * 5).indices
box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
if not self.obb:
box = xywh2xyxy(box)
if self.is_tf:
# TFlite bug returns less boxes
box = torch.nn.functional.pad(box, (0, 0, 0, mask.shape[0] - box.shape[0]))
nmsbox = box.clone()
# `8` is the minimum value experimented to get correct NMS results for obb
multiplier = 8 if self.obb else 1
# Normalize boxes for NMS since large values for class offset causes issue with int8 quantization
if self.args.format == "tflite": # TFLite is already normalized
nmsbox *= multiplier
else:
nmsbox = multiplier * nmsbox / torch.tensor(x.shape[2:], device=box.device, dtype=box.dtype).max()
if not self.args.agnostic_nms: # class-specific NMS
end = 2 if self.obb else 4
# fully explicit expansion otherwise reshape error
# large max_wh causes issues when quantizing
cls_offset = cls.reshape(-1, 1).expand(nmsbox.shape[0], end)
offbox = nmsbox[:, :end] + cls_offset * multiplier
nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
nms_fn = (
partial(nms_rotated, use_triu=not (self.is_tf or (self.args.opset or 14) < 14)) if self.obb else nms
)
keep = nms_fn(
torch.cat([nmsbox, extra], dim=-1) if self.obb else nmsbox,
score,
self.args.iou,
)[: self.args.max_det]
dets = torch.cat([box[keep], score[keep].view(-1, 1), cls[keep].view(-1, 1), extra[keep]], dim=-1)
# Zero-pad to max_det size to avoid reshape error
pad = (0, 0, 0, self.args.max_det - dets.shape[0])
out[i] = torch.nn.functional.pad(dets, pad)
return (out, preds[1]) if self.model.task == "segment" else out

@ -305,7 +305,7 @@ class Results(SimpleClass):
if v is not None:
return len(v)
def update(self, boxes=None, masks=None, probs=None, obb=None):
def update(self, boxes=None, masks=None, probs=None, obb=None, keypoints=None):
"""
Updates the Results object with new detection data.
@ -318,6 +318,7 @@ class Results(SimpleClass):
masks (torch.Tensor | None): A tensor of shape (N, H, W) containing segmentation masks.
probs (torch.Tensor | None): A tensor of shape (num_classes,) containing class probabilities.
obb (torch.Tensor | None): A tensor of shape (N, 5) containing oriented bounding box coordinates.
keypoints (torch.Tensor | None): A tensor of shape (N, 17, 3) containing keypoints.
Examples:
>>> results = model("image.jpg")
@ -332,6 +333,8 @@ class Results(SimpleClass):
self.probs = probs
if obb is not None:
self.obb = OBB(obb, self.orig_shape)
if keypoints is not None:
self.keypoints = Keypoints(keypoints, self.orig_shape)
def _apply(self, fn, *args, **kwargs):
"""

@ -271,7 +271,6 @@ class BaseTrainer:
)
if world_size > 1:
self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[RANK], find_unused_parameters=True)
self.set_model_attributes() # set again after DDP wrapper
# Check imgsz
gs = max(int(self.model.stride.max() if hasattr(self.model, "stride") else 32), 32) # grid size (max stride)
@ -782,7 +781,7 @@ class BaseTrainer:
f"ignoring 'lr0={self.args.lr0}' and 'momentum={self.args.momentum}' and "
f"determining best 'optimizer', 'lr0' and 'momentum' automatically... "
)
nc = getattr(model, "nc", 10) # number of classes
nc = self.data.get("nc", 10) # number of classes
lr_fit = round(0.002 * 5 / (4 + nc), 6) # lr0 fit equation to 6 decimal places
name, lr, momentum = ("SGD", 0.01, 0.9) if iterations > 10000 else ("AdamW", lr_fit, 0.9)
self.args.warmup_bias_lr = 0.0 # no higher than 0.01 for Adam

@ -38,13 +38,7 @@ class NASValidator(DetectionValidator):
"""Apply Non-maximum suppression to prediction outputs."""
boxes = ops.xyxy2xywh(preds_in[0][0])
preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
return ops.non_max_suppression(
return super().postprocess(
preds,
self.args.conf,
self.args.iou,
labels=self.lb,
multi_label=False,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
max_time_img=0.5,
)

@ -20,22 +20,54 @@ class DetectionPredictor(BasePredictor):
```
"""
def postprocess(self, preds, img, orig_imgs):
def postprocess(self, preds, img, orig_imgs, **kwargs):
"""Post-processes predictions and returns a list of Results objects."""
preds = ops.non_max_suppression(
preds,
self.args.conf,
self.args.iou,
agnostic=self.args.agnostic_nms,
self.args.classes,
self.args.agnostic_nms,
max_det=self.args.max_det,
classes=self.args.classes,
nc=len(self.model.names),
end2end=getattr(self.model, "end2end", False),
rotated=self.args.task == "obb",
)
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
return results
return self.construct_results(preds, img, orig_imgs, **kwargs)
def construct_results(self, preds, img, orig_imgs):
"""
Constructs a list of result objects from the predictions.
Args:
preds (List[torch.Tensor]): List of predicted bounding boxes and scores.
img (torch.Tensor): The image after preprocessing.
orig_imgs (List[np.ndarray]): List of original images before preprocessing.
Returns:
(list): List of result objects containing the original images, image paths, class names, and bounding boxes.
"""
return [
self.construct_result(pred, img, orig_img, img_path)
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0])
]
def construct_result(self, pred, img, orig_img, img_path):
"""
Constructs the result object from the prediction.
Args:
pred (torch.Tensor): The predicted bounding boxes and scores.
img (torch.Tensor): The image after preprocessing.
orig_img (np.ndarray): The original image before preprocessing.
img_path (str): The path to the original image.
Returns:
(Results): The result object containing the original image, image path, class names, and bounding boxes.
"""
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6])

@ -78,6 +78,7 @@ class DetectionValidator(BaseValidator):
self.args.save_json |= self.args.val and (self.is_coco or self.is_lvis) and not self.training # run final val
self.names = model.names
self.nc = len(model.names)
self.end2end = getattr(model, "end2end", False)
self.metrics.names = self.names
self.metrics.plot = self.args.plots
self.confusion_matrix = ConfusionMatrix(nc=self.nc, conf=self.args.conf)
@ -96,9 +97,12 @@ class DetectionValidator(BaseValidator):
self.args.conf,
self.args.iou,
labels=self.lb,
nc=self.nc,
multi_label=True,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
end2end=self.end2end,
rotated=self.args.task == "obb",
)
def _prepare_batch(self, si, batch):

@ -27,27 +27,20 @@ class OBBPredictor(DetectionPredictor):
super().__init__(cfg, overrides, _callbacks)
self.args.task = "obb"
def postprocess(self, preds, img, orig_imgs):
"""Post-processes predictions and returns a list of Results objects."""
preds = ops.non_max_suppression(
preds,
self.args.conf,
self.args.iou,
agnostic=self.args.agnostic_nms,
max_det=self.args.max_det,
nc=len(self.model.names),
classes=self.args.classes,
rotated=True,
)
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
# xywh, r, conf, cls
obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
results.append(Results(orig_img, path=img_path, names=self.model.names, obb=obb))
return results
def construct_result(self, pred, img, orig_img, img_path):
"""
Constructs the result object from the prediction.
Args:
pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles.
img (torch.Tensor): The image after preprocessing.
orig_img (np.ndarray): The original image before preprocessing.
img_path (str): The path to the original image.
Returns:
(Results): The result object containing the original image, image path, class names, and oriented bounding boxes.
"""
rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
return Results(orig_img, path=img_path, names=self.model.names, obb=obb)

@ -36,20 +36,6 @@ class OBBValidator(DetectionValidator):
val = self.data.get(self.args.split, "") # validation path
self.is_dota = isinstance(val, str) and "DOTA" in val # is COCO
def postprocess(self, preds):
"""Apply Non-maximum suppression to prediction outputs."""
return ops.non_max_suppression(
preds,
self.args.conf,
self.args.iou,
labels=self.lb,
nc=self.nc,
multi_label=True,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
rotated=True,
)
def _process_batch(self, detections, gt_bboxes, gt_cls):
"""
Perform computation of the correct prediction matrix for a batch of detections and ground truth bounding boxes.

@ -1,6 +1,5 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from ultralytics.engine.results import Results
from ultralytics.models.yolo.detect.predict import DetectionPredictor
from ultralytics.utils import DEFAULT_CFG, LOGGER, ops
@ -30,27 +29,21 @@ class PosePredictor(DetectionPredictor):
"See https://github.com/ultralytics/ultralytics/issues/4031."
)
def postprocess(self, preds, img, orig_imgs):
"""Return detection results for a given input image or list of images."""
preds = ops.non_max_suppression(
preds,
self.args.conf,
self.args.iou,
agnostic=self.args.agnostic_nms,
max_det=self.args.max_det,
classes=self.args.classes,
nc=len(self.model.names),
)
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round()
pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
results.append(
Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts)
)
return results
def construct_result(self, pred, img, orig_img, img_path):
"""
Constructs the result object from the prediction.
Args:
pred (torch.Tensor): The predicted bounding boxes, scores, and keypoints.
img (torch.Tensor): The image after preprocessing.
orig_img (np.ndarray): The original image before preprocessing.
img_path (str): The path to the original image.
Returns:
(Results): The result object containing the original image, image path, class names, bounding boxes, and keypoints.
"""
result = super().construct_result(pred, img, orig_img, img_path)
pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
result.update(keypoints=pred_kpts)
return result

@ -61,19 +61,6 @@ class PoseValidator(DetectionValidator):
"mAP50-95)",
)
def postprocess(self, preds):
"""Apply non-maximum suppression and return detections with high confidence scores."""
return ops.non_max_suppression(
preds,
self.args.conf,
self.args.iou,
labels=self.lb,
multi_label=True,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
nc=self.nc,
)
def init_metrics(self, model):
"""Initiate pose estimation metrics for YOLO model."""
super().init_metrics(model)

@ -27,29 +27,48 @@ class SegmentationPredictor(DetectionPredictor):
def postprocess(self, preds, img, orig_imgs):
"""Applies non-max suppression and processes detections for each image in an input batch."""
p = ops.non_max_suppression(
preds[0],
self.args.conf,
self.args.iou,
agnostic=self.args.agnostic_nms,
max_det=self.args.max_det,
nc=len(self.model.names),
classes=self.args.classes,
)
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
proto = preds[1][-1] if isinstance(preds[1], tuple) else preds[1] # tuple if PyTorch model or array if exported
for i, (pred, orig_img, img_path) in enumerate(zip(p, orig_imgs, self.batch[0])):
if not len(pred): # save empty boxes
masks = None
elif self.args.retina_masks:
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2]) # HWC
else:
masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True) # HWC
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
return results
# tuple if PyTorch model or array if exported
protos = preds[1][-1] if isinstance(preds[1], tuple) else preds[1]
return super().postprocess(preds[0], img, orig_imgs, protos=protos)
def construct_results(self, preds, img, orig_imgs, protos):
"""
Constructs a list of result objects from the predictions.
Args:
preds (List[torch.Tensor]): List of predicted bounding boxes, scores, and masks.
img (torch.Tensor): The image after preprocessing.
orig_imgs (List[np.ndarray]): List of original images before preprocessing.
protos (List[torch.Tensor]): List of prototype masks.
Returns:
(list): List of result objects containing the original images, image paths, class names, bounding boxes, and masks.
"""
return [
self.construct_result(pred, img, orig_img, img_path, proto)
for pred, orig_img, img_path, proto in zip(preds, orig_imgs, self.batch[0], protos)
]
def construct_result(self, pred, img, orig_img, img_path, proto):
"""
Constructs the result object from the prediction.
Args:
pred (np.ndarray): The predicted bounding boxes, scores, and masks.
img (torch.Tensor): The image after preprocessing.
orig_img (np.ndarray): The original image before preprocessing.
img_path (str): The path to the original image.
proto (torch.Tensor): The prototype masks.
Returns:
(Results): The result object containing the original image, image path, class names, bounding boxes, and masks.
"""
if not len(pred): # save empty boxes
masks = None
elif self.args.retina_masks:
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
masks = ops.process_mask_native(proto, pred[:, 6:], pred[:, :4], orig_img.shape[:2]) # HWC
else:
masks = ops.process_mask(proto, pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True) # HWC
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)

@ -70,16 +70,7 @@ class SegmentationValidator(DetectionValidator):
def postprocess(self, preds):
"""Post-processes YOLO predictions and returns output detections with proto."""
p = ops.non_max_suppression(
preds[0],
self.args.conf,
self.args.iou,
labels=self.lb,
multi_label=True,
agnostic=self.args.single_cls or self.args.agnostic_nms,
max_det=self.args.max_det,
nc=self.nc,
)
p = super().postprocess(preds[0])
proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported
return p, proto

@ -132,6 +132,7 @@ class AutoBackend(nn.Module):
fp16 &= pt or jit or onnx or xml or engine or nn_module or triton # FP16
nhwc = coreml or saved_model or pb or tflite or edgetpu or rknn # BHWC formats (vs torch BCWH)
stride = 32 # default stride
end2end = False # default end2end
model, metadata, task = None, None, None
# Set device
@ -222,16 +223,18 @@ class AutoBackend(nn.Module):
output_names = [x.name for x in session.get_outputs()]
metadata = session.get_modelmeta().custom_metadata_map
dynamic = isinstance(session.get_outputs()[0].shape[0], str)
fp16 = True if "float16" in session.get_inputs()[0].type else False
if not dynamic:
io = session.io_binding()
bindings = []
for output in session.get_outputs():
y_tensor = torch.empty(output.shape, dtype=torch.float16 if fp16 else torch.float32).to(device)
out_fp16 = "float16" in output.type
y_tensor = torch.empty(output.shape, dtype=torch.float16 if out_fp16 else torch.float32).to(device)
io.bind_output(
name=output.name,
device_type=device.type,
device_id=device.index if cuda else 0,
element_type=np.float16 if fp16 else np.float32,
element_type=np.float16 if out_fp16 else np.float32,
shape=tuple(y_tensor.shape),
buffer_ptr=y_tensor.data_ptr(),
)
@ -482,7 +485,7 @@ class AutoBackend(nn.Module):
w = next(w.rglob("*.rknn")) # get *.rknn file from *_rknn_model dir
rknn_model = RKNNLite()
rknn_model.load_rknn(w)
ret = rknn_model.init_runtime()
rknn_model.init_runtime()
metadata = Path(w).parent / "metadata.yaml"
# Any other format (unsupported)
@ -501,7 +504,7 @@ class AutoBackend(nn.Module):
for k, v in metadata.items():
if k in {"stride", "batch"}:
metadata[k] = int(v)
elif k in {"imgsz", "names", "kpt_shape"} and isinstance(v, str):
elif k in {"imgsz", "names", "kpt_shape", "args"} and isinstance(v, str):
metadata[k] = eval(v)
stride = metadata["stride"]
task = metadata["task"]
@ -509,6 +512,7 @@ class AutoBackend(nn.Module):
imgsz = metadata["imgsz"]
names = metadata["names"]
kpt_shape = metadata.get("kpt_shape")
end2end = metadata.get("args", {}).get("nms", False)
elif not (pt or triton or nn_module):
LOGGER.warning(f"WARNING ⚠ Metadata not found for 'model={weights}'")
@ -703,9 +707,12 @@ class AutoBackend(nn.Module):
if x.ndim == 3: # if task is not classification, excluding masks (ndim=4) as well
# Denormalize xywh by image size. See https://github.com/ultralytics/ultralytics/pull/1695
# xywh are normalized in TFLite/EdgeTPU to mitigate quantization error of integer models
if x.shape[-1] == 6: # end-to-end model
if x.shape[-1] == 6 or self.end2end: # end-to-end model
x[:, :, [0, 2]] *= w
x[:, :, [1, 3]] *= h
if self.task == "pose":
x[:, :, 6::3] *= w
x[:, :, 7::3] *= h
else:
x[:, [0, 2]] *= w
x[:, [1, 3]] *= h

@ -1120,8 +1120,6 @@ class TorchVision(nn.Module):
m (nn.Module): The loaded torchvision model, possibly truncated and unwrapped.
Args:
c1 (int): Input channels.
c2 (): Output channels.
model (str): Name of the torchvision model to load.
weights (str, optional): Pre-trained weights to load. Default is "DEFAULT".
unwrap (bool, optional): If True, unwraps the model to a sequential containing all but the last `truncate` layers. Default is True.
@ -1129,7 +1127,7 @@ class TorchVision(nn.Module):
split (bool, optional): Returns output from intermediate child modules as list. Default is False.
"""
def __init__(self, c1, c2, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
def __init__(self, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
"""Load the model and weights from torchvision."""
import torchvision # scope for faster 'import ultralytics'

@ -336,7 +336,7 @@ class Concat(nn.Module):
class Index(nn.Module):
"""Returns a particular index of the input."""
def __init__(self, c1, c2, index=0):
def __init__(self, index=0):
"""Returns a particular index of the input."""
super().__init__()
self.index = index

@ -1060,12 +1060,16 @@ def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
m.legacy = legacy
elif m is RTDETRDecoder: # special case, channels arg must be passed in index 1
args.insert(1, [ch[x] for x in f])
elif m in frozenset({CBLinear, TorchVision, Index}):
elif m is CBLinear:
c2 = args[0]
c1 = ch[f]
args = [c1, c2, *args[1:]]
elif m is CBFuse:
c2 = ch[f[-1]]
elif m in frozenset({TorchVision, Index}):
c2 = args[0]
c1 = ch[f]
args = [*args[1:]]
else:
c2 = ch[f]

@ -31,6 +31,9 @@ def on_predict_start(predictor: object, persist: bool = False) -> None:
>>> predictor = SomePredictorClass()
>>> on_predict_start(predictor, persist=True)
"""
if predictor.args.task == "classify":
raise ValueError("❌ Classification doesn't support 'mode=track'")
if hasattr(predictor, "trackers") and persist:
return

@ -13,6 +13,7 @@ import sys
import threading
import time
import uuid
import warnings
from pathlib import Path
from threading import Lock
from types import SimpleNamespace
@ -23,8 +24,8 @@ import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import tqdm
import yaml
from tqdm import tqdm as tqdm_original
from ultralytics import __version__
@ -132,8 +133,11 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" # suppress verbose TF compiler warning
os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR" # suppress "NNPACK.cpp could not initialize NNPACK" warnings
os.environ["KINETO_LOG_LEVEL"] = "5" # suppress verbose PyTorch profiler output when computing FLOPs
if TQDM_RICH := str(os.getenv("YOLO_TQDM_RICH", False)).lower() == "true":
from tqdm import rich
class TQDM(tqdm_original):
class TQDM(rich.tqdm if TQDM_RICH else tqdm.tqdm):
"""
A custom TQDM progress bar class that extends the original tqdm functionality.
@ -176,7 +180,8 @@ class TQDM(tqdm_original):
... # Your code here
... pass
"""
kwargs["disable"] = not VERBOSE or kwargs.get("disable", False) # logical 'and' with default value if passed
warnings.filterwarnings("ignore", category=tqdm.TqdmExperimentalWarning) # suppress tqdm.rich warning
kwargs["disable"] = not VERBOSE or kwargs.get("disable", False)
kwargs.setdefault("bar_format", TQDM_BAR_FORMAT) # override default value if passed
super().__init__(*args, **kwargs)

@ -41,7 +41,7 @@ import yaml
from ultralytics import YOLO, YOLOWorld
from ultralytics.cfg import TASK2DATA, TASK2METRIC
from ultralytics.engine.exporter import export_formats
from ultralytics.utils import ARM64, ASSETS, IS_JETSON, IS_RASPBERRYPI, LINUX, LOGGER, MACOS, TQDM, WEIGHTS_DIR
from ultralytics.utils import ARM64, ASSETS, LINUX, LOGGER, MACOS, TQDM, WEIGHTS_DIR
from ultralytics.utils.checks import IS_PYTHON_3_12, check_requirements, check_yolo, is_rockchip
from ultralytics.utils.downloads import safe_download
from ultralytics.utils.files import file_size
@ -100,9 +100,9 @@ def benchmark(
elif i == 9: # Edge TPU
assert LINUX and not ARM64, "Edge TPU export only supported on non-aarch64 Linux"
elif i in {5, 10}: # CoreML and TF.js
assert MACOS or LINUX, "CoreML and TF.js export only supported on macOS and Linux"
assert not IS_RASPBERRYPI, "CoreML and TF.js export not supported on Raspberry Pi"
assert not IS_JETSON, "CoreML and TF.js export not supported on NVIDIA Jetson"
assert MACOS or (LINUX and not ARM64), (
"CoreML and TF.js export only supported on macOS and non-aarch64 Linux"
)
if i in {5}: # CoreML
assert not IS_PYTHON_3_12, "CoreML not supported on Python 3.12"
if i in {6, 7, 8}: # TF SavedModel, TF GraphDef, and TFLite

@ -143,7 +143,7 @@ def make_divisible(x, divisor):
return math.ceil(x / divisor) * divisor
def nms_rotated(boxes, scores, threshold=0.45):
def nms_rotated(boxes, scores, threshold=0.45, use_triu=True):
"""
NMS for oriented bounding boxes using probiou and fast-nms.
@ -151,16 +151,30 @@ def nms_rotated(boxes, scores, threshold=0.45):
boxes (torch.Tensor): Rotated bounding boxes, shape (N, 5), format xywhr.
scores (torch.Tensor): Confidence scores, shape (N,).
threshold (float, optional): IoU threshold. Defaults to 0.45.
use_triu (bool, optional): Whether to use `torch.triu` operator. It'd be useful for disable it
when exporting obb models to some formats that do not support `torch.triu`.
Returns:
(torch.Tensor): Indices of boxes to keep after NMS.
"""
if len(boxes) == 0:
return np.empty((0,), dtype=np.int8)
sorted_idx = torch.argsort(scores, descending=True)
boxes = boxes[sorted_idx]
ious = batch_probiou(boxes, boxes).triu_(diagonal=1)
pick = torch.nonzero(ious.max(dim=0)[0] < threshold).squeeze_(-1)
ious = batch_probiou(boxes, boxes)
if use_triu:
ious = ious.triu_(diagonal=1)
# pick = torch.nonzero(ious.max(dim=0)[0] < threshold).squeeze_(-1)
# NOTE: handle the case when len(boxes) hence exportable by eliminating if-else condition
pick = torch.nonzero((ious >= threshold).sum(0) <= 0).squeeze_(-1)
else:
n = boxes.shape[0]
row_idx = torch.arange(n, device=boxes.device).view(-1, 1).expand(-1, n)
col_idx = torch.arange(n, device=boxes.device).view(1, -1).expand(n, -1)
upper_mask = row_idx < col_idx
ious = ious * upper_mask
# Zeroing these scores ensures the additional indices would not affect the final results
scores[~((ious >= threshold).sum(0) <= 0)] = 0
# NOTE: return indices with fixed length to avoid TFLite reshape error
pick = torch.topk(scores, scores.shape[0]).indices
return sorted_idx[pick]
@ -179,6 +193,7 @@ def non_max_suppression(
max_wh=7680,
in_place=True,
rotated=False,
end2end=False,
):
"""
Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box.
@ -205,6 +220,7 @@ def non_max_suppression(
max_wh (int): The maximum box width and height in pixels.
in_place (bool): If True, the input prediction tensor will be modified in place.
rotated (bool): If Oriented Bounding Boxes (OBB) are being passed for NMS.
end2end (bool): If the model doesn't require NMS.
Returns:
(List[torch.Tensor]): A list of length batch_size, where each element is a tensor of
@ -221,7 +237,7 @@ def non_max_suppression(
if classes is not None:
classes = torch.tensor(classes, device=prediction.device)
if prediction.shape[-1] == 6: # end-to-end model (BNC, i.e. 1,300,6)
if prediction.shape[-1] == 6 or end2end: # end-to-end model (BNC, i.e. 1,300,6)
output = [pred[pred[:, 4] > conf_thres][:max_det] for pred in prediction]
if classes is not None:
output = [pred[(pred[:, 5:6] == classes).any(1)] for pred in output]

Loading…
Cancel
Save