diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3a6bd37b79..9324a32ca3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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: diff --git a/docs/en/datasets/detect/medical-pills.md b/docs/en/datasets/detect/medical-pills.md index 77585b2c3e..c32aabf2f7 100644 --- a/docs/en/datasets/detect/medical-pills.md +++ b/docs/en/datasets/detect/medical-pills.md @@ -8,7 +8,20 @@ keywords: medical-pills dataset, pill detection, pharmaceutical imaging, AI in h Open Medical Pills Dataset In Colab -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. + +

+
+ +
+ Watch: How to train Ultralytics YOLO11 Model on Medical Pills Detection Dataset in Google Colab +

+ +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 diff --git a/docs/en/guides/kfold-cross-validation.md b/docs/en/guides/kfold-cross-validation.md index 44ba8d82e8..bb8efb7d86 100644 --- a/docs/en/guides/kfold-cross-validation.md +++ b/docs/en/guides/kfold-cross-validation.md @@ -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`. diff --git a/docs/en/guides/raspberry-pi.md b/docs/en/guides/raspberry-pi.md index 4268287ff7..00b8d31572 100644 --- a/docs/en/guides/raspberry-pi.md +++ b/docs/en/guides/raspberry-pi.md @@ -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 diff --git a/docs/en/guides/triton-inference-server.md b/docs/en/guides/triton-inference-server.md index 67d419bf52..68aa3cd87b 100644 --- a/docs/en/guides/triton-inference-server.md +++ b/docs/en/guides/triton-inference-server.md @@ -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 = [] diff --git a/docs/en/hub/inference-api.md b/docs/en/hub/inference-api.md index b532e8150c..fce59c8b21 100644 --- a/docs/en/hub/inference-api.md +++ b/docs/en/hub/inference-api.md @@ -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 diff --git a/docs/en/integrations/ibm-watsonx.md b/docs/en/integrations/ibm-watsonx.md index 16ebaa2a62..0e77bc5e1b 100644 --- a/docs/en/integrations/ibm-watsonx.md +++ b/docs/en/integrations/ibm-watsonx.md @@ -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.

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", diff --git a/docs/en/integrations/index.md b/docs/en/integrations/index.md index f94b295593..4b91b18f2e 100644 --- a/docs/en/integrations/index.md +++ b/docs/en/integrations/index.md @@ -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: diff --git a/docs/en/integrations/rockchip-rknn.md b/docs/en/integrations/rockchip-rknn.md index 269b301fda..ea5c4df7d6 100644 --- a/docs/en/integrations/rockchip-rknn.md +++ b/docs/en/integrations/rockchip-rknn.md @@ -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. +

+ RKNN +

+ !!! 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. -

- RKNN -

- ## 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. diff --git a/docs/en/integrations/seeedstudio-recamera.md b/docs/en/integrations/seeedstudio-recamera.md new file mode 100644 index 0000000000..dcad49351d --- /dev/null +++ b/docs/en/integrations/seeedstudio-recamera.md @@ -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 + +

+ reCamera Toolchain +

+ +### 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) diff --git a/docs/en/integrations/tensorrt.md b/docs/en/integrations/tensorrt.md index cac4ac325c..59dbb280b6 100644 --- a/docs/en/integrations/tensorrt.md +++ b/docs/en/integrations/tensorrt.md @@ -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 diff --git a/docs/en/macros/export-table.md b/docs/en/macros/export-table.md index a8e8e26041..8bf018e53b 100644 --- a/docs/en/macros/export-table.md +++ b/docs/en/macros/export-table.md @@ -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` | diff --git a/docs/en/models/mobile-sam.md b/docs/en/models/mobile-sam.md index a65587de87..34740c6c07 100644 --- a/docs/en/models/mobile-sam.md +++ b/docs/en/models/mobile-sam.md @@ -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 diff --git a/docs/en/reference/engine/exporter.md b/docs/en/reference/engine/exporter.md index a0d1822dce..a650b314e8 100644 --- a/docs/en/reference/engine/exporter.md +++ b/docs/en/reference/engine/exporter.md @@ -19,6 +19,10 @@ keywords: YOLOv8, export formats, ONNX, TensorRT, CoreML, machine learning model


+## ::: ultralytics.engine.exporter.NMSModel + +


+ ## ::: ultralytics.engine.exporter.export_formats


diff --git a/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp b/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp index a65391f5d7..168df490c2 100644 --- a/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp +++ b/examples/YOLOv8-ONNXRuntime-CPP/inference.cpp @@ -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); diff --git a/examples/tutorial.ipynb b/examples/tutorial.ipynb index 9ed5dc32b4..20c8e677d0 100644 --- a/examples/tutorial.ipynb +++ b/examples/tutorial.ipynb @@ -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" diff --git a/mkdocs.yml b/mkdocs.yml index f8cbbf1c91..76cdc28831 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -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: diff --git a/tests/test_exports.py b/tests/test_exports.py index c4c9eb4419..adfd3917c5 100644 --- a/tests/test_exports.py +++ b/tests/test_exports.py @@ -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.""" diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 7eaddb8e8f..e7be7b7625 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license -__version__ = "8.3.65" +__version__ = "8.3.67" import os diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index a98de9497a..efad56d4e3 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -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}." diff --git a/ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml b/ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml index baedcb5dc5..e2fbcfac10 100644 --- a/ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +++ b/ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml @@ -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: diff --git a/ultralytics/data/augment.py b/ultralytics/data/augment.py index 85b90148c2..1ab14a647e 100644 --- a/ultralytics/data/augment.py +++ b/ultralytics/data/augment.py @@ -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 diff --git a/ultralytics/data/split_dota.py b/ultralytics/data/split_dota.py index a8f17d4115..1d1ca3105a 100644 --- a/ultralytics/data/split_dota.py +++ b/ultralytics/data/split_dota.py @@ -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) diff --git a/ultralytics/data/utils.py b/ultralytics/data/utils.py index 50b597d861..3a75e5a13e 100644 --- a/ultralytics/data/utils.py +++ b/ultralytics/data/utils.py @@ -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}" ) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 867e319d33..88e6e53300 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -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 diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index 9fc9e6e1f7..b7f7dd72ca 100644 --- a/ultralytics/engine/results.py +++ b/ultralytics/engine/results.py @@ -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): """ diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index 2c083d1eec..8c3c9d7797 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -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 diff --git a/ultralytics/models/nas/val.py b/ultralytics/models/nas/val.py index c3d0f37e37..ca01e94e00 100644 --- a/ultralytics/models/nas/val.py +++ b/ultralytics/models/nas/val.py @@ -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, ) diff --git a/ultralytics/models/yolo/detect/predict.py b/ultralytics/models/yolo/detect/predict.py index 4d9da8966e..172e54d3f9 100644 --- a/ultralytics/models/yolo/detect/predict.py +++ b/ultralytics/models/yolo/detect/predict.py @@ -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]) diff --git a/ultralytics/models/yolo/detect/val.py b/ultralytics/models/yolo/detect/val.py index d5fcbfe5bc..ec809d55fc 100644 --- a/ultralytics/models/yolo/detect/val.py +++ b/ultralytics/models/yolo/detect/val.py @@ -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): diff --git a/ultralytics/models/yolo/obb/predict.py b/ultralytics/models/yolo/obb/predict.py index ebbd7530c7..ef6214d421 100644 --- a/ultralytics/models/yolo/obb/predict.py +++ b/ultralytics/models/yolo/obb/predict.py @@ -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) diff --git a/ultralytics/models/yolo/obb/val.py b/ultralytics/models/yolo/obb/val.py index a839285897..b5cb89f145 100644 --- a/ultralytics/models/yolo/obb/val.py +++ b/ultralytics/models/yolo/obb/val.py @@ -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. diff --git a/ultralytics/models/yolo/pose/predict.py b/ultralytics/models/yolo/pose/predict.py index 4d9315b825..75334b758c 100644 --- a/ultralytics/models/yolo/pose/predict.py +++ b/ultralytics/models/yolo/pose/predict.py @@ -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 diff --git a/ultralytics/models/yolo/pose/val.py b/ultralytics/models/yolo/pose/val.py index 2acdaa3ee8..9fc872f9c1 100644 --- a/ultralytics/models/yolo/pose/val.py +++ b/ultralytics/models/yolo/pose/val.py @@ -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) diff --git a/ultralytics/models/yolo/segment/predict.py b/ultralytics/models/yolo/segment/predict.py index 782816258c..4e0adc7c44 100644 --- a/ultralytics/models/yolo/segment/predict.py +++ b/ultralytics/models/yolo/segment/predict.py @@ -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) diff --git a/ultralytics/models/yolo/segment/val.py b/ultralytics/models/yolo/segment/val.py index 81f3d01b64..8be870e3c5 100644 --- a/ultralytics/models/yolo/segment/val.py +++ b/ultralytics/models/yolo/segment/val.py @@ -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 diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index c979bc60eb..f536a6025c 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -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 diff --git a/ultralytics/nn/modules/block.py b/ultralytics/nn/modules/block.py index 1edb9c801e..beb03c20c4 100644 --- a/ultralytics/nn/modules/block.py +++ b/ultralytics/nn/modules/block.py @@ -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' diff --git a/ultralytics/nn/modules/conv.py b/ultralytics/nn/modules/conv.py index 5184707883..6c15e1d66c 100644 --- a/ultralytics/nn/modules/conv.py +++ b/ultralytics/nn/modules/conv.py @@ -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 diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 91a0ec57b8..a754f5e79b 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -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] diff --git a/ultralytics/trackers/track.py b/ultralytics/trackers/track.py index e55db6d43d..6e422f0db8 100644 --- a/ultralytics/trackers/track.py +++ b/ultralytics/trackers/track.py @@ -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 diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py index 4ff46044fc..3afa07a973 100644 --- a/ultralytics/utils/__init__.py +++ b/ultralytics/utils/__init__.py @@ -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) diff --git a/ultralytics/utils/benchmarks.py b/ultralytics/utils/benchmarks.py index 5ba02195a4..c183203e1a 100644 --- a/ultralytics/utils/benchmarks.py +++ b/ultralytics/utils/benchmarks.py @@ -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 diff --git a/ultralytics/utils/ops.py b/ultralytics/utils/ops.py index 52b5155217..af41ffee3d 100644 --- a/ultralytics/utils/ops.py +++ b/ultralytics/utils/ops.py @@ -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]