`ultralytics 8.0.41` TF SavedModel and EdgeTPU export (#1034)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Noobtoss <96134731+Noobtoss@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
single_channel
Glenn Jocher 2 years ago committed by GitHub
parent 4b866c9718
commit f6e393c1d2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -17,7 +17,7 @@ jobs:
strategy: strategy:
fail-fast: false fail-fast: false
matrix: matrix:
os: [ubuntu-latest] os: [ubuntu-latest, windows-latest, macos-latest]
python-version: ['3.10'] # requires python<=3.9 python-version: ['3.10'] # requires python<=3.9
model: [yolov8n] model: [yolov8n]
steps: steps:
@ -34,7 +34,8 @@ jobs:
- name: Install requirements - name: Install requirements
run: | run: |
python -m pip install --upgrade pip wheel python -m pip install --upgrade pip wheel
pip install -e '.[export]' --extra-index-url https://download.pytorch.org/whl/cpu pip install -e . coremltools openvino-dev tensorflow-cpu paddlepaddle x2paddle --extra-index-url https://download.pytorch.org/whl/cpu
yolo export format=tflite
- name: Check environment - name: Check environment
run: | run: |
echo "RUNNER_OS is ${{ runner.os }}" echo "RUNNER_OS is ${{ runner.os }}"
@ -46,29 +47,21 @@ jobs:
python --version python --version
pip --version pip --version
pip list pip list
- name: TF Lite export
run: |
yolo export model=${{ matrix.model }}.pt format=tflite
yolo task=detect mode=predict model=yolov8n_saved_model/yolov8n_float16.tflite imgsz=640
- name: TF *.pb export
run: |
yolo export model=${{ matrix.model }}.pt format=pb
yolo task=detect mode=predict model=yolov8n.pb imgsz=640
- name: TF Lite Edge TPU export
run: |
yolo export model=${{ matrix.model }}.pt format=edgetpu
- name: TF.js export
run: |
yolo export model=${{ matrix.model }}.pt format=tfjs
- name: Benchmark DetectionModel - name: Benchmark DetectionModel
shell: python
run: | run: |
# yolo benchmark model=${{ matrix.model }}.pt imgsz=320 min_metric=0.29 from ultralytics.yolo.utils.benchmarks import run_benchmarks
run_benchmarks(model='yolov8n.pt', imgsz=160, half=False, hard_fail=False)
- name: Benchmark SegmentationModel - name: Benchmark SegmentationModel
shell: python
run: | run: |
# yolo benchmark model=${{ matrix.model }}-seg.pt imgsz=320 min_metric=0.29 from ultralytics.yolo.utils.benchmarks import run_benchmarks
run_benchmarks(model='yolov8n-seg.pt', imgsz=160, half=False, hard_fail=False)
- name: Benchmark ClassificationModel - name: Benchmark ClassificationModel
shell: python
run: | run: |
# yolo benchmark model=${{ matrix.model }}-cls.pt imgsz=224 min_metric=0.29 from ultralytics.yolo.utils.benchmarks import run_benchmarks
run_benchmarks(model='yolov8n-cls.pt', imgsz=160, half=False, hard_fail=False)
Tests: Tests:
timeout-minutes: 60 timeout-minutes: 60
@ -76,7 +69,7 @@ jobs:
strategy: strategy:
fail-fast: false fail-fast: false
matrix: matrix:
os: [ubuntu-latest, windows-latest, macos-latest] os: [ubuntu-latest]
python-version: ['3.10'] python-version: ['3.10']
model: [yolov8n] model: [yolov8n]
torch: [latest] torch: [latest]
@ -155,9 +148,9 @@ jobs:
- name: Test classification - name: Test classification
shell: bash # for Windows compatibility shell: bash # for Windows compatibility
run: | run: |
yolo task=classify mode=train data=mnist160 model=yolov8n-cls.yaml epochs=1 imgsz=32 yolo task=classify mode=train data=imagenet10 model=yolov8n-cls.yaml epochs=1 imgsz=32
yolo task=classify mode=train data=mnist160 model=yolov8n-cls.pt epochs=1 imgsz=32 yolo task=classify mode=train data=imagenet10 model=yolov8n-cls.pt epochs=1 imgsz=32
yolo task=classify mode=val data=mnist160 model=runs/classify/train/weights/last.pt imgsz=32 yolo task=classify mode=val data=imagenet10 model=runs/classify/train/weights/last.pt imgsz=32
yolo task=classify mode=predict model=runs/classify/train/weights/last.pt imgsz=32 source=ultralytics/assets/bus.jpg yolo task=classify mode=predict model=runs/classify/train/weights/last.pt imgsz=32 source=ultralytics/assets/bus.jpg
yolo mode=export model=runs/classify/train/weights/last.pt imgsz=32 format=torchscript yolo mode=export model=runs/classify/train/weights/last.pt imgsz=32 format=torchscript
- name: Pytest tests - name: Pytest tests

@ -59,7 +59,9 @@ To allow your work to be integrated as seamlessly as possible, we advise you to:
### Docstrings ### Docstrings
Not all functions or classes require docstrings but when they do, we follow [google-style docstrings format](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings). Here is an example: Not all functions or classes require docstrings but when they do, we
follow [google-style docstrings format](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings).
Here is an example:
```python ```python
""" """

@ -222,7 +222,10 @@ See [Classification Docs](https://docs.ultralytics.com/tasks/classification/) fo
## <div align="center">Ultralytics HUB</div> ## <div align="center">Ultralytics HUB</div>
Experience seamless AI with [Ultralytics HUB](https://bit.ly/ultralytics_hub) ⭐, the all-in-one solution for data visualization, YOLOv5 and YOLOv8 (coming soon) 🚀 model training and deployment, without any coding. Transform images into actionable insights and bring your AI visions to life with ease using our cutting-edge platform and user-friendly [Ultralytics App](https://ultralytics.com/app_install). Start your journey for **Free** now! Experience seamless AI with [Ultralytics HUB](https://bit.ly/ultralytics_hub) ⭐, the all-in-one solution for data
visualization, YOLOv5 and YOLOv8 (coming soon) 🚀 model training and deployment, without any coding. Transform images
into actionable insights and bring your AI visions to life with ease using our cutting-edge platform and
user-friendly [Ultralytics App](https://ultralytics.com/app_install). Start your journey for **Free** now!
<a href="https://bit.ly/ultralytics_hub" target="_blank"> <a href="https://bit.ly/ultralytics_hub" target="_blank">
<img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a> <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a>
@ -250,7 +253,9 @@ YOLOv8 is available under two different licenses:
## <div align="center">Contact</div> ## <div align="center">Contact</div>
For YOLOv8 bug reports and feature requests please visit [GitHub Issues](https://github.com/ultralytics/ultralytics/issues) or the [Ultralytics Community Forum](https://community.ultralytics.com/). For YOLOv8 bug reports and feature requests please
visit [GitHub Issues](https://github.com/ultralytics/ultralytics/issues) or
the [Ultralytics Community Forum](https://community.ultralytics.com/).
<br> <br>
<div align="center"> <div align="center">

@ -18,7 +18,9 @@
</div> </div>
<br> <br>
[Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics) 是由 [Ultralytics](https://ultralytics.com) 开发的一个前沿的 SOTA 模型。它在以前成功的 YOLO 版本基础上引入了新的功能和改进进一步提升了其性能和灵活性。YOLOv8 基于快速、准确和易于使用的设计理念,使其成为广泛的目标检测、图像分割和图像分类任务的绝佳选择。 [Ultralytics YOLOv8](https://github.com/ultralytics/ultralytics) 是由 [Ultralytics](https://ultralytics.com) 开发的一个前沿的
SOTA 模型。它在以前成功的 YOLO 版本基础上引入了新的功能和改进进一步提升了其性能和灵活性。YOLOv8
基于快速、准确和易于使用的设计理念,使其成为广泛的目标检测、图像分割和图像分类任务的绝佳选择。
如果要申请企业许可证,请填写 [Ultralytics 许可](https://ultralytics.com/license)。 如果要申请企业许可证,请填写 [Ultralytics 许可](https://ultralytics.com/license)。
@ -53,7 +55,9 @@
<details open> <details open>
<summary>安装</summary> <summary>安装</summary>
Pip 安装包含所有 [requirements.txt](https://github.com/ultralytics/ultralytics/blob/main/requirements.txt) 的 ultralytics 包,环境要求 [**Python>=3.7**](https://www.python.org/),且 [**PyTorch>=1.7**](https://pytorch.org/get-started/locally/)。 Pip 安装包含所有 [requirements.txt](https://github.com/ultralytics/ultralytics/blob/main/requirements.txt) 的
ultralytics 包,环境要求 [**Python>=3.7**](https://www.python.org/),且 [\*\*PyTorch>=1.7
\*\*](https://pytorch.org/get-started/locally/)。
```bash ```bash
pip install ultralytics pip install ultralytics
@ -70,7 +74,8 @@ YOLOv8 可以直接在命令行界面CLI中使用 `yolo` 命令运行:
yolo predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" yolo predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg"
``` ```
`yolo`可以用于各种任务和模式,并接受额外的参数,例如 `imgsz=640`。参见 YOLOv8 [文档](https://docs.ultralytics.com)中可用`yolo`[参数](https://docs.ultralytics.com/cfg/)的完整列表。 `yolo`可以用于各种任务和模式,并接受额外的参数,例如 `imgsz=640`。参见 YOLOv8 [文档](https://docs.ultralytics.com)
中可用`yolo`[参数](https://docs.ultralytics.com/cfg/)的完整列表。
```bash ```bash
yolo task=detect mode=train model=yolov8n.pt args... yolo task=detect mode=train model=yolov8n.pt args...
@ -95,11 +100,13 @@ results = model("https://ultralytics.com/images/bus.jpg") # 预测图像
success = model.export(format="onnx") # 将模型导出为 ONNX 格式 success = model.export(format="onnx") # 将模型导出为 ONNX 格式
``` ```
[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。 [模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从
Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。
### 已知问题 / 待办事项 ### 已知问题 / 待办事项
我们仍在努力完善 YOLOv8 的几个部分!我们的目标是尽快完成这些工作,使 YOLOv8 的功能设置达到YOLOv5 的水平,包括对所有相同格式的导出和推理。我们还在写一篇 YOLOv8 的论文,一旦完成,我们将提交给 [arxiv.org](https://arxiv.org)。 我们仍在努力完善 YOLOv8 的几个部分!我们的目标是尽快完成这些工作,使 YOLOv8 的功能设置达到YOLOv5
的水平,包括对所有相同格式的导出和推理。我们还在写一篇 YOLOv8 的论文,一旦完成,我们将提交给 [arxiv.org](https://arxiv.org)。
- [x] TensorFlow 导出 - [x] TensorFlow 导出
- [x] DDP 恢复训练 - [x] DDP 恢复训练
@ -111,7 +118,8 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
所有 YOLOv8 的预训练模型都可以在这里找到。目标检测和分割模型是在 COCO 数据集上预训练的,而分类模型是在 ImageNet 数据集上预训练的。 所有 YOLOv8 的预训练模型都可以在这里找到。目标检测和分割模型是在 COCO 数据集上预训练的,而分类模型是在 ImageNet 数据集上预训练的。
第一次使用时,[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。 第一次使用时,[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从
Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。
<details open><summary>目标检测</summary> <details open><summary>目标检测</summary>
@ -125,7 +133,8 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。 - **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
<br>复现命令 `yolo val detect data=coco.yaml device=0` <br>复现命令 `yolo val detect data=coco.yaml device=0`
- **推理速度**使用 COCO 验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。 - **推理速度**使用 COCO
验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。
<br>复现命令 `yolo val detect data=coco128.yaml batch=1 device=0/cpu` <br>复现命令 `yolo val detect data=coco128.yaml batch=1 device=0/cpu`
</details> </details>
@ -142,7 +151,8 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
- **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。 - **mAP<sup>val</sup>** 结果都在 [COCO val2017](http://cocodataset.org) 数据集上,使用单模型单尺度测试得到。
<br>复现命令 `yolo val segment data=coco.yaml device=0` <br>复现命令 `yolo val segment data=coco.yaml device=0`
- **推理速度**使用 COCO 验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。 - **推理速度**使用 COCO
验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。
<br>复现命令 `yolo val segment data=coco128-seg.yaml batch=1 device=0/cpu` <br>复现命令 `yolo val segment data=coco128-seg.yaml batch=1 device=0/cpu`
</details> </details>
@ -159,7 +169,8 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
- **acc** 都在 [ImageNet](https://www.image-net.org/) 数据集上,使用单模型单尺度测试得到。 - **acc** 都在 [ImageNet](https://www.image-net.org/) 数据集上,使用单模型单尺度测试得到。
<br>复现命令 `yolo val classify data=path/to/ImageNet device=0` <br>复现命令 `yolo val classify data=path/to/ImageNet device=0`
- **推理速度**使用 ImageNet 验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。 - **推理速度**使用 ImageNet
验证集图片推理时间进行平均得到,测试环境使用 [Amazon EC2 P4d](https://aws.amazon.com/ec2/instance-types/p4/) 实例。
<br>复现命令 `yolo val classify data=path/to/ImageNet batch=1 device=0/cpu` <br>复现命令 `yolo val classify data=path/to/ImageNet batch=1 device=0/cpu`
</details> </details>
@ -192,14 +203,18 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
## <div align="center">Ultralytics HUB</div> ## <div align="center">Ultralytics HUB</div>
[Ultralytics HUB](https://bit.ly/ultralytics_hub) 是我们⭐ **新**的无代码解决方案,用于可视化数据集,训练 YOLOv8🚀 模型,并以无缝体验方式部署到现实世界。现在开始**免费**! 还可以通过下载 [Ultralytics App](https://ultralytics.com/app_install) 在你的 iOS 或 Android 设备上运行 YOLOv8 模型! [Ultralytics HUB](https://bit.ly/ultralytics_hub) 是我们⭐ **新**的无代码解决方案,用于可视化数据集,训练 YOLOv8🚀
模型,并以无缝体验方式部署到现实世界。现在开始**免费**!
还可以通过下载 [Ultralytics App](https://ultralytics.com/app_install) 在你的 iOS 或 Android 设备上运行 YOLOv8 模型!
<a href="https://bit.ly/ultralytics_hub" target="_blank"> <a href="https://bit.ly/ultralytics_hub" target="_blank">
<img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a> <img width="100%" src="https://github.com/ultralytics/assets/raw/main/im/ultralytics-hub.png"></a>
## <div align="center">贡献</div> ## <div align="center">贡献</div>
我们喜欢您的意见或建议!我们希望尽可能简单和透明地为 YOLOv8 做出贡献。请看我们的 [贡献指南](CONTRIBUTING.md) ,并填写 [调查问卷](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey) 向我们发送您的体验反馈。感谢我们所有的贡献者! 我们喜欢您的意见或建议!我们希望尽可能简单和透明地为 YOLOv8 做出贡献。请看我们的 [贡献指南](CONTRIBUTING.md)
,并填写 [调查问卷](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey)
向我们发送您的体验反馈。感谢我们所有的贡献者!
<!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 --> <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
@ -210,11 +225,13 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
- YOLOv8 在两种不同的 License 下可用: - YOLOv8 在两种不同的 License 下可用:
- **GPL-3.0 License** 查看 [License](https://github.com/ultralytics/ultralytics/blob/main/LICENSE) 文件的详细信息。 - **GPL-3.0 License** 查看 [License](https://github.com/ultralytics/ultralytics/blob/main/LICENSE) 文件的详细信息。
- **企业License**:在没有 GPL-3.0 开源要求的情况下为商业产品开发提供更大的灵活性。典型用例是将 Ultralytics 软件和 AI 模型嵌入到商业产品和应用程序中。在以下位置申请企业许可证 [Ultralytics 许可](https://ultralytics.com/license) 。 - **企业License**:在没有 GPL-3.0 开源要求的情况下为商业产品开发提供更大的灵活性。典型用例是将 Ultralytics 软件和 AI
模型嵌入到商业产品和应用程序中。在以下位置申请企业许可证 [Ultralytics 许可](https://ultralytics.com/license) 。
## <div align="center">联系我们</div> ## <div align="center">联系我们</div>
请访问 [GitHub Issues](https://github.com/ultralytics/ultralytics/issues) 或 [Ultralytics Community Forum](https://community.ultralytis.com) 以报告 YOLOv8 错误和请求功能。 请访问 [GitHub Issues](https://github.com/ultralytics/ultralytics/issues)
或 [Ultralytics Community Forum](https://community.ultralytis.com) 以报告 YOLOv8 错误和请求功能。
<br> <br>
<div align="center"> <div align="center">

@ -13,7 +13,7 @@ ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Aria
ENV DEBIAN_FRONTEND noninteractive ENV DEBIAN_FRONTEND noninteractive
RUN apt update RUN apt update
RUN TZ=Etc/UTC apt install -y tzdata RUN TZ=Etc/UTC apt install -y tzdata
RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++
# RUN alias python=python3 # RUN alias python=python3
# Security updates # Security updates

@ -12,7 +12,7 @@ ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Aria
ENV DEBIAN_FRONTEND noninteractive ENV DEBIAN_FRONTEND noninteractive
RUN apt update RUN apt update
RUN TZ=Etc/UTC apt install -y tzdata RUN TZ=Etc/UTC apt install -y tzdata
RUN apt install --no-install-recommends -y python3-pip git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg RUN apt install --no-install-recommends -y python3-pip git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg g++
# RUN alias python=python3 # RUN alias python=python3
# Create working directory # Create working directory

@ -1,17 +1,26 @@
At [Ultralytics](https://ultralytics.com), the security of our users' data and systems is of utmost importance. To ensure the safety and security of our [open-source projects](https://github.com/ultralytics), we have implemented several measures to detect and prevent security vulnerabilities. At [Ultralytics](https://ultralytics.com), the security of our users' data and systems is of utmost importance. To
ensure the safety and security of our [open-source projects](https://github.com/ultralytics), we have implemented
several measures to detect and prevent security vulnerabilities.
[![ultralytics](https://snyk.io/advisor/python/ultralytics/badge.svg)](https://snyk.io/advisor/python/ultralytics) [![ultralytics](https://snyk.io/advisor/python/ultralytics/badge.svg)](https://snyk.io/advisor/python/ultralytics)
## Snyk Scanning ## Snyk Scanning
We use [Snyk](https://snyk.io/advisor/python/ultralytics) to regularly scan the YOLOv8 repository for vulnerabilities and security issues. Our goal is to identify and remediate any potential threats as soon as possible, to minimize any risks to our users. We use [Snyk](https://snyk.io/advisor/python/ultralytics) to regularly scan the YOLOv8 repository for vulnerabilities
and security issues. Our goal is to identify and remediate any potential threats as soon as possible, to minimize any
risks to our users.
## GitHub CodeQL Scanning ## GitHub CodeQL Scanning
In addition to our Snyk scans, we also use GitHub's [CodeQL](https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/about-code-scanning-with-codeql) scans to proactively identify and address security vulnerabilities. In addition to our Snyk scans, we also use
GitHub's [CodeQL](https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/about-code-scanning-with-codeql)
scans to proactively identify and address security vulnerabilities.
## Reporting Security Issues ## Reporting Security Issues
If you suspect or discover a security vulnerability in the YOLOv8 repository, please let us know immediately. You can reach out to us directly via our [contact form](https://ultralytics.com/contact) or via [security@ultralytics.com](mailto:security@ultralytics.com). Our security team will investigate and respond as soon as possible. If you suspect or discover a security vulnerability in the YOLOv8 repository, please let us know immediately. You can
reach out to us directly via our [contact form](https://ultralytics.com/contact) or
via [security@ultralytics.com](mailto:security@ultralytics.com). Our security team will investigate and respond as soon
as possible.
We appreciate your help in keeping the YOLOv8 repository secure and safe for everyone. We appreciate your help in keeping the YOLOv8 repository secure and safe for everyone.

@ -1,10 +1,15 @@
## Callbacks ## Callbacks
Ultralytics framework supports callbacks as entry points in strategic stages of train, val, export, and predict modes. Each callback accepts a `Trainer`, `Validator`, or `Predictor` object depending on the operation type. All properties of these objects can be found in Reference section of the docs.
Ultralytics framework supports callbacks as entry points in strategic stages of train, val, export, and predict modes.
Each callback accepts a `Trainer`, `Validator`, or `Predictor` object depending on the operation type. All properties of
these objects can be found in Reference section of the docs.
## Examples ## Examples
### Returning additional information with Prediction ### Returning additional information with Prediction
In this example, we want to return the original frame with each result object. Here's how we can do that In this example, we want to return the original frame with each result object. Here's how we can do that
```python ```python
def on_predict_batch_end(predictor): def on_predict_batch_end(predictor):
# results -> List[batch_size] # results -> List[batch_size]
@ -19,8 +24,11 @@ for (result, frame) in model.track/predict():
``` ```
## All callbacks ## All callbacks
Here are all supported callbacks. Here are all supported callbacks.
### Trainer ### Trainer
`on_pretrain_routine_start` `on_pretrain_routine_start`
`on_pretrain_routine_end` `on_pretrain_routine_end`
@ -50,6 +58,7 @@ Here are all supported callbacks.
`teardown` `teardown`
### Validator ### Validator
`on_val_start` `on_val_start`
`on_val_batch_start` `on_val_batch_start`
@ -59,6 +68,7 @@ Here are all supported callbacks.
`on_val_end` `on_val_end`
### Predictor ### Predictor
`on_predict_start` `on_predict_start`
`on_predict_batch_start` `on_predict_batch_start`
@ -70,6 +80,7 @@ Here are all supported callbacks.
`on_predict_end` `on_predict_end`
### Exporter ### Exporter
`on_export_start` `on_export_start`
`on_export_end` `on_export_end`

@ -39,6 +39,7 @@ the [Configuration](cfg.md) page.
yolo detect train data=coco128.yaml model=yolov8n.pt epochs=100 imgsz=640 yolo detect train data=coco128.yaml model=yolov8n.pt epochs=100 imgsz=640
yolo detect train resume model=last.pt # resume training yolo detect train resume model=last.pt # resume training
``` ```
## Val ## Val
Validate trained YOLOv8n model accuracy on the COCO128 dataset. No argument need to passed as the `model` retains it's Validate trained YOLOv8n model accuracy on the COCO128 dataset. No argument need to passed as the `model` retains it's
@ -50,6 +51,7 @@ training `data` and arguments as model attributes.
yolo detect val model=yolov8n.pt # val official model yolo detect val model=yolov8n.pt # val official model
yolo detect val model=path/to/best.pt # val custom model yolo detect val model=path/to/best.pt # val custom model
``` ```
## Predict ## Predict
Use a trained YOLOv8n model to run predictions on images. Use a trained YOLOv8n model to run predictions on images.
@ -60,6 +62,7 @@ Use a trained YOLOv8n model to run predictions on images.
yolo detect predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model yolo detect predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model
yolo detect predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model yolo detect predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model
``` ```
## Export ## Export
Export a YOLOv8n model to a different format like ONNX, CoreML, etc. Export a YOLOv8n model to a different format like ONNX, CoreML, etc.

@ -96,12 +96,15 @@ Class reference documentation for `Results` module and its components can be fou
## Visualizing results ## Visualizing results
You can use `visualize()` function of `Result` object to get a visualization. It plots all components(boxes, masks, classification logits, etc) found in the results object You can use `visualize()` function of `Result` object to get a visualization. It plots all components(boxes, masks,
classification logits, etc) found in the results object
```python ```python
res = model(img) res = model(img)
res_plotted = res[0].visualize() res_plotted = res[0].visualize()
cv2.imshow("result", res_plotted) cv2.imshow("result", res_plotted)
``` ```
!!! example "`visualize()` arguments" !!! example "`visualize()` arguments"
`show_conf (bool)`: Show confidence `show_conf (bool)`: Show confidence

@ -4,7 +4,7 @@ This is a list of real-world applications and walkthroughs. These can be folders
| Title | Format | Contributor | | Title | Format | Contributor |
| --------------------------------------------------------------- | ------------------ | --------------------------------------------------- | | --------------------------------------------------------------- | ------------------ | --------------------------------------------------- |
| [Yolov8/yolov5 ONNX Inference with C++](./Yolov8_CPP_Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) | | [Yolov8/yolov5 ONNX Inference with C++](./YOLOv8-CPP-Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) |
| [YOLOv8-OpenCV-ONNX-Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) | | [YOLOv8-OpenCV-ONNX-Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) |
## How can you contribute ? ## How can you contribute ?

@ -48,6 +48,10 @@ yolov5s.onnx:
![image](https://user-images.githubusercontent.com/40023722/217357005-07464492-d1da-42e3-98a7-fc753f87d5e6.png) ![image](https://user-images.githubusercontent.com/40023722/217357005-07464492-d1da-42e3-98a7-fc753f87d5e6.png)
This repository is based on OpenCVs dnn API to run an ONNX exported model of either yolov5/yolov8 (In theory should work for yolov6 and yolov7 but not tested). Note that for this example the networks are exported as rectangular (640x480) resolutions, but it would work for any resolution that you export as although you might want to use the letterBox approach for square images depending on your use-case. This repository is based on OpenCVs dnn API to run an ONNX exported model of either yolov5/yolov8 (In theory should work
for yolov6 and yolov7 but not tested). Note that for this example the networks are exported as rectangular (640x480)
resolutions, but it would work for any resolution that you export as although you might want to use the letterBox
approach for square images depending on your use-case.
The **main** branch version is based on using Qt as a GUI wrapper the main interest here is the **Inference** class file which shows how to transpose yolov8 models to work as yolov5 models. The **main** branch version is based on using Qt as a GUI wrapper the main interest here is the **Inference** class file
which shows how to transpose yolov8 models to work as yolov5 models.

@ -49,7 +49,7 @@
"source": [ "source": [
"# Setup\n", "# Setup\n",
"\n", "\n",
"Pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/requirements.txt) and check PyTorch and GPU." "Pip install `ultralytics` and [dependencies](https://github.com/ultralytics/ultralytics/blob/main/requirements.txt) and check software and hardware."
] ]
}, },
{ {
@ -62,7 +62,6 @@
"outputId": "9bda69d4-e57f-404b-b6fe-117234e24677" "outputId": "9bda69d4-e57f-404b-b6fe-117234e24677"
}, },
"source": [ "source": [
"# Pip install method (recommended)\n",
"%pip install ultralytics\n", "%pip install ultralytics\n",
"import ultralytics\n", "import ultralytics\n",
"ultralytics.checks()" "ultralytics.checks()"
@ -80,17 +79,11 @@
] ]
}, },
{ {
"cell_type": "code", "cell_type": "markdown",
"source": [ "source": [],
"# Git clone method (for development)\n",
"!git clone https://github.com/ultralytics/ultralytics\n",
"%pip install -qe ultralytics"
],
"metadata": { "metadata": {
"id": "TUFPge7f_1ms" "id": "ZOwTlorPd8-D"
}, }
"execution_count": null,
"outputs": []
}, },
{ {
"cell_type": "markdown", "cell_type": "markdown",
@ -655,6 +648,19 @@
"Additional content below." "Additional content below."
] ]
}, },
{
"cell_type": "code",
"source": [
"# Git clone install (for development)\n",
"!git clone https://github.com/ultralytics/ultralytics\n",
"%pip install -qe ultralytics"
],
"metadata": {
"id": "uRKlwxSJdhd1"
},
"execution_count": null,
"outputs": []
},
{ {
"cell_type": "code", "cell_type": "code",
"metadata": { "metadata": {

@ -30,7 +30,7 @@ def test_train_seg():
def test_train_cls(): def test_train_cls():
run(f'yolo train classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1') run(f'yolo train classify model={CFG}-cls.yaml data=imagenet10 imgsz=32 epochs=1')
# Val checks ----------------------------------------------------------------------------------------------------------- # Val checks -----------------------------------------------------------------------------------------------------------
@ -43,7 +43,7 @@ def test_val_segment():
def test_val_classify(): def test_val_classify():
run(f'yolo val classify model={MODEL}-cls.pt data=mnist160 imgsz=32') run(f'yolo val classify model={MODEL}-cls.pt data=imagenet10 imgsz=32')
# Predict checks ------------------------------------------------------------------------------------------------------- # Predict checks -------------------------------------------------------------------------------------------------------

@ -74,8 +74,14 @@ def test_segment():
def test_classify(): def test_classify():
overrides = {'data': 'mnist160', 'model': 'yolov8n-cls.yaml', 'imgsz': 32, 'epochs': 1, 'batch': 64, 'save': False} overrides = {
CFG.data = 'mnist160' 'data': 'imagenet10',
'model': 'yolov8n-cls.yaml',
'imgsz': 32,
'epochs': 1,
'batch': 64,
'save': False}
CFG.data = 'imagenet10'
CFG.imgsz = 32 CFG.imgsz = 32
CFG.batch = 64 CFG.batch = 64
# YOLO(CFG_SEG).train(**overrides) # works # YOLO(CFG_SEG).train(**overrides) # works

@ -154,7 +154,7 @@ def test_predict_callback_and_setup():
bs = [predictor.dataset.bs for _ in range(len(path))] bs = [predictor.dataset.bs for _ in range(len(path))]
predictor.results = zip(predictor.results, im0s, bs) predictor.results = zip(predictor.results, im0s, bs)
model = YOLO('yolov8n.pt') model = YOLO(MODEL)
model.add_callback('on_predict_batch_end', on_predict_batch_end) model.add_callback('on_predict_batch_end', on_predict_batch_end)
dataset = load_inference_source(source=SOURCE, transforms=model.transforms) dataset = load_inference_source(source=SOURCE, transforms=model.transforms)
@ -169,9 +169,16 @@ def test_predict_callback_and_setup():
def test_result(): def test_result():
model = YOLO('yolov8n-seg.pt') model = YOLO('yolov8n-seg.pt')
img = str(ROOT / 'assets/bus.jpg') res = model([SOURCE, SOURCE])
res = model([img, img])
res[0].numpy() res[0].numpy()
res[0].cpu().numpy() res[0].cpu().numpy()
resimg = res[0].visualize(show_conf=False) resimg = res[0].visualize(show_conf=False)
print(resimg) print(resimg)
model = YOLO('yolov8n.pt')
res = model(SOURCE)
res[0].visualize()
model = YOLO('yolov8n-cls.pt')
res = model(SOURCE)
res[0].visualize()

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
__version__ = '8.0.40' __version__ = '8.0.41'
from ultralytics.yolo.engine.model import YOLO from ultralytics.yolo.engine.model import YOLO
from ultralytics.yolo.utils.checks import check_yolo as checks from ultralytics.yolo.utils.checks import check_yolo as checks

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ by Argo AI
# Example usage: python train.py --data Argoverse.yaml # Example usage: yolo train data=Argoverse.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# Global Wheat 2020 dataset http://www.global-wheat.com/ by University of Saskatchewan # Global Wheat 2020 dataset http://www.global-wheat.com/ by University of Saskatchewan
# Example usage: python train.py --data GlobalWheat2020.yaml # Example usage: yolo train data=GlobalWheat2020.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,7 +1,7 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University # ImageNet-1k dataset https://www.image-net.org/index.php by Stanford University
# Simplified class names from https://github.com/anishathalye/imagenet-simple-labels # Simplified class names from https://github.com/anishathalye/imagenet-simple-labels
# Example usage: python classify/train.py --data imagenet # Example usage: yolo train task=classify data=imagenet
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets
@ -2022,4 +2022,4 @@ map:
# Download script/URL (optional) # Download script/URL (optional)
download: data/scripts/get_imagenet.sh download: yolo/data/scripts/get_imagenet.sh

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# Objects365 dataset https://www.objects365.org/ by Megvii # Objects365 dataset https://www.objects365.org/ by Megvii
# Example usage: python train.py --data Objects365.yaml # Example usage: yolo train data=Objects365.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 by Trax Retail # SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 by Trax Retail
# Example usage: python train.py --data SKU-110K.yaml # Example usage: yolo train data=SKU-110K.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
# Example usage: python train.py --data VOC.yaml # Example usage: yolo train data=VOC.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University # VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset by Tianjin University
# Example usage: python train.py --data VisDrone.yaml # Example usage: yolo train data=VisDrone.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft # COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml # Example usage: yolo train data=coco.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics # COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml # Example usage: yolo train data=coco128.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml # Example usage: yolo train data=coco128.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# COCO8-seg dataset (first 8 images from COCO train2017) by Ultralytics # COCO8-seg dataset (first 8 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco8-seg.yaml # Example usage: yolo train data=coco8-seg.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# COCO8 dataset (first 8 images from COCO train2017) by Ultralytics # COCO8 dataset (first 8 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco8.yaml # Example usage: yolo train data=coco8.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -1,7 +1,7 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
# DIUx xView 2018 Challenge https://challenge.xviewdataset.org by U.S. National Geospatial-Intelligence Agency (NGA) # DIUx xView 2018 Challenge https://challenge.xviewdataset.org by U.S. National Geospatial-Intelligence Agency (NGA)
# -------- DOWNLOAD DATA MANUALLY and jar xf val_images.zip to 'datasets/xView' before running train command! -------- # -------- DOWNLOAD DATA MANUALLY and jar xf val_images.zip to 'datasets/xView' before running train command! --------
# Example usage: python train.py --data xView.yaml # Example usage: yolo train data=xView.yaml
# parent # parent
# ├── yolov5 # ├── yolov5
# └── datasets # └── datasets

@ -11,9 +11,9 @@ from random import random
import requests import requests
from ultralytics.yolo.utils import (DEFAULT_CFG_DICT, ENVIRONMENT, LOGGER, RANK, SETTINGS, TryExcept, __version__, from ultralytics.yolo.utils import (DEFAULT_CFG_DICT, ENVIRONMENT, LOGGER, RANK, SETTINGS, TESTS_RUNNING, TryExcept,
colorstr, emojis, get_git_origin_url, is_colab, is_git_dir, is_github_actions_ci, __version__, colorstr, emojis, get_git_origin_url, is_colab, is_git_dir,
is_pip_package, is_pytest_running) is_pip_package)
from ultralytics.yolo.utils.checks import check_online from ultralytics.yolo.utils.checks import check_online
PREFIX = colorstr('Ultralytics: ') PREFIX = colorstr('Ultralytics: ')
@ -157,8 +157,7 @@ class Traces:
self.enabled = SETTINGS['sync'] and \ self.enabled = SETTINGS['sync'] and \
RANK in {-1, 0} and \ RANK in {-1, 0} and \
check_online() and \ check_online() and \
not is_pytest_running() and \ not TESTS_RUNNING and \
not is_github_actions_ci() and \
(is_pip_package() or get_git_origin_url() == 'https://github.com/ultralytics/ultralytics.git') (is_pip_package() or get_git_origin_url() == 'https://github.com/ultralytics/ultralytics.git')
def __call__(self, cfg, all_keys=False, traces_sample_rate=1.0): def __call__(self, cfg, all_keys=False, traces_sample_rate=1.0):

@ -28,13 +28,18 @@ def check_class_names(names):
if not all(isinstance(k, int) for k in names.keys()): # convert string keys to int, i.e. '0' to 0 if not all(isinstance(k, int) for k in names.keys()): # convert string keys to int, i.e. '0' to 0
names = {int(k): v for k, v in names.items()} names = {int(k): v for k, v in names.items()}
if isinstance(names[0], str) and names[0].startswith('n0'): # imagenet class codes, i.e. 'n01440764' if isinstance(names[0], str) and names[0].startswith('n0'): # imagenet class codes, i.e. 'n01440764'
map = yaml_load(ROOT / 'yolo/data/datasets/ImageNet.yaml')['map'] # human-readable names map = yaml_load(ROOT / 'datasets/ImageNet.yaml')['map'] # human-readable names
names = {k: map[v] for k, v in names.items()} names = {k: map[v] for k, v in names.items()}
return names return names
class AutoBackend(nn.Module): class AutoBackend(nn.Module):
def _apply_default_class_names(self, data):
with contextlib.suppress(Exception):
return yaml_load(check_yaml(data))['names']
return {i: f'class{i}' for i in range(999)} # return default if above errors
def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True): def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
""" """
MultiBackend class for python inference on various platforms using Ultralytics YOLO. MultiBackend class for python inference on various platforms using Ultralytics YOLO.
@ -53,7 +58,7 @@ class AutoBackend(nn.Module):
| PyTorch | *.pt | | PyTorch | *.pt |
| TorchScript | *.torchscript | | TorchScript | *.torchscript |
| ONNX Runtime | *.onnx | | ONNX Runtime | *.onnx |
| ONNX OpenCV DNN | *.onnx --dnn | | ONNX OpenCV DNN | *.onnx dnn=True |
| OpenVINO | *.xml | | OpenVINO | *.xml |
| CoreML | *.mlmodel | | CoreML | *.mlmodel |
| TensorRT | *.engine | | TensorRT | *.engine |
@ -142,13 +147,9 @@ class AutoBackend(nn.Module):
logger = trt.Logger(trt.Logger.INFO) logger = trt.Logger(trt.Logger.INFO)
# Read file # Read file
with open(w, 'rb') as f, trt.Runtime(logger) as runtime: with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
# Read metadata length meta_len = int.from_bytes(f.read(4), byteorder='little') # read metadata length
meta_len = int.from_bytes(f.read(4), byteorder='little') meta = json.loads(f.read(meta_len).decode('utf-8')) # read metadata
# Read metadata model = runtime.deserialize_cuda_engine(f.read()) # read engine
meta = json.loads(f.read(meta_len).decode('utf-8'))
stride, names = int(meta['stride']), meta['names']
# Read engine
model = runtime.deserialize_cuda_engine(f.read())
context = model.create_execution_context() context = model.create_execution_context()
bindings = OrderedDict() bindings = OrderedDict()
output_names = [] output_names = []
@ -170,6 +171,7 @@ class AutoBackend(nn.Module):
bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr())) bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items()) binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size
stride, names = int(meta['stride']), meta['names']
elif coreml: # CoreML elif coreml: # CoreML
LOGGER.info(f'Loading {w} for CoreML inference...') LOGGER.info(f'Loading {w} for CoreML inference...')
import coremltools as ct import coremltools as ct
@ -179,6 +181,7 @@ class AutoBackend(nn.Module):
import tensorflow as tf import tensorflow as tf
keras = False # assume TF1 saved_model keras = False # assume TF1 saved_model
model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w) model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
w = Path(w) / 'metadata.yaml'
elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...') LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
import tensorflow as tf import tensorflow as tf
@ -265,7 +268,7 @@ class AutoBackend(nn.Module):
# Check names # Check names
if 'names' not in locals(): # names missing if 'names' not in locals(): # names missing
names = yaml_load(check_yaml(data))['names'] if data else {i: f'class{i}' for i in range(999)} # assign names = self._apply_default_class_names(data)
names = check_class_names(names) names = check_class_names(names)
self.__dict__.update(locals()) # assign all variables to self self.__dict__.update(locals()) # assign all variables to self
@ -324,7 +327,9 @@ class AutoBackend(nn.Module):
box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float) conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1) y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
else: elif len(y) == 1: # classification model
y = list(y.values())
elif len(y) == 2: # segmentation model
y = list(reversed(y.values())) # reversed for segmentation models (pred, proto) y = list(reversed(y.values())) # reversed for segmentation models (pred, proto)
elif self.paddle: # PaddlePaddle elif self.paddle: # PaddlePaddle
im = im.cpu().numpy().astype(np.float32) im = im.cpu().numpy().astype(np.float32)
@ -337,8 +342,14 @@ class AutoBackend(nn.Module):
im = im.cpu().numpy() im = im.cpu().numpy()
if self.saved_model: # SavedModel if self.saved_model: # SavedModel
y = self.model(im, training=False) if self.keras else self.model(im) y = self.model(im, training=False) if self.keras else self.model(im)
if not isinstance(y, list):
y = [y]
elif self.pb: # GraphDef elif self.pb: # GraphDef
y = self.frozen_func(x=self.tf.constant(im)) y = self.frozen_func(x=self.tf.constant(im))
if len(y) == 2 and len(self.names) == 999: # segments and names not defined
ip, ib = (0, 1) if len(y[0].shape) == 4 else (1, 0) # index of protos, boxes
nc = y[ib].shape[1] - y[ip].shape[3] - 4 # y = (1, 160, 160, 32), (1, 116, 8400)
self.names = {i: f'class{i}' for i in range(nc)}
else: # Lite or Edge TPU else: # Lite or Edge TPU
input = self.input_details[0] input = self.input_details[0]
int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
@ -355,11 +366,15 @@ class AutoBackend(nn.Module):
x = (x.astype(np.float32) - zero_point) * scale # re-scale x = (x.astype(np.float32) - zero_point) * scale # re-scale
y.append(x) y.append(x)
# TF segment fixes: export is reversed vs ONNX export and protos are transposed # TF segment fixes: export is reversed vs ONNX export and protos are transposed
if len(self.output_details) == 2: # segment if len(y) == 2: # segment with (det, proto) output order reversed
y = [y[1], np.transpose(y[0], (0, 3, 1, 2))] if len(y[1].shape) != 4:
y = list(reversed(y)) # should be y = (1, 116, 8400), (1, 160, 160, 32)
y[1] = np.transpose(y[1], (0, 3, 1, 2)) # should be y = (1, 116, 8400), (1, 32, 160, 160)
y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y] y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
# y[0][..., :4] *= [w, h, w, h] # xywh normalized to pixels # y[0][..., :4] *= [w, h, w, h] # xywh normalized to pixels
# for x in y:
# print(type(x), len(x)) if isinstance(x, (list, tuple)) else print(type(x), x.shape) # debug shapes
if isinstance(y, (list, tuple)): if isinstance(y, (list, tuple)):
return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y] return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
else: else:
@ -375,7 +390,7 @@ class AutoBackend(nn.Module):
Returns: Returns:
(torch.Tensor): The converted tensor (torch.Tensor): The converted tensor
""" """
return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
def warmup(self, imgsz=(1, 3, 640, 640)): def warmup(self, imgsz=(1, 3, 640, 640)):
""" """

@ -79,7 +79,7 @@ class AutoShape(nn.Module):
with amp.autocast(autocast): with amp.autocast(autocast):
return self.model(ims.to(p.device).type_as(p), augment=augment) # inference return self.model(ims.to(p.device).type_as(p), augment=augment) # inference
# Pre-process # Preprocess
n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims]) # number, list of images n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims]) # number, list of images
shape0, shape1, files = [], [], [] # image and inference shapes, filenames shape0, shape1, files = [], [], [] # image and inference shapes, filenames
for i, im in enumerate(ims): for i, im in enumerate(ims):
@ -108,7 +108,7 @@ class AutoShape(nn.Module):
with dt[1]: with dt[1]:
y = self.model(x, augment=augment) # forward y = self.model(x, augment=augment) # forward
# Post-process # Postprocess
with dt[2]: with dt[2]:
y = non_max_suppression(y if self.dmb else y[0], y = non_max_suppression(y if self.dmb else y[0],
self.conf, self.conf,
@ -181,7 +181,7 @@ class Detections:
self.ims[i] = np.asarray(im) self.ims[i] = np.asarray(im)
if pprint: if pprint:
s = s.lstrip('\n') s = s.lstrip('\n')
return f'{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t return f'{s}\nSpeed: %.1fms preprocess, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t
if crop: if crop:
if save: if save:
LOGGER.info(f'Saved results to {save_dir}\n') LOGGER.info(f'Saved results to {save_dir}\n')

@ -174,7 +174,7 @@ class C2(nn.Module):
self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))) self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))
def forward(self, x): def forward(self, x):
a, b = self.cv1(x).split((self.c, self.c), 1) a, b = self.cv1(x).chunk(2, 1)
return self.cv2(torch.cat((self.m(a), b), 1)) return self.cv2(torch.cat((self.m(a), b), 1))
@ -188,6 +188,11 @@ class C2f(nn.Module):
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)) self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x): def forward(self, x):
y = list(self.cv1(x).chunk(2, 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
def forward_split(self, x):
y = list(self.cv1(x).split((self.c, self.c), 1)) y = list(self.cv1(x).split((self.c, self.c), 1))
y.extend(m(y[-1]) for m in self.m) y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1)) return self.cv2(torch.cat(y, 1))
@ -405,6 +410,11 @@ class Detect(nn.Module):
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5)) self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape self.shape = shape
if self.export and self.format == 'edgetpu': # FlexSplitV ops issue
x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
box = x_cat[:, :self.reg_max * 4]
cls = x_cat[:, self.reg_max * 4:]
else:
box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.nc), 1) box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.nc), 1)
dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
y = torch.cat((dbox, cls.sigmoid()), 1) y = torch.cat((dbox, cls.sigmoid()), 1)

@ -1,6 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
import ast
import contextlib import contextlib
from copy import deepcopy from copy import deepcopy
from pathlib import Path from pathlib import Path
@ -338,7 +337,7 @@ def torch_safe_load(weight):
file = attempt_download_asset(weight) # search online if missing locally file = attempt_download_asset(weight) # search online if missing locally
try: try:
return torch.load(file, map_location='cpu') # load return torch.load(file, map_location='cpu'), file # load
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
if e.name == 'omegaconf': # e.name is missing module name if e.name == 'omegaconf': # e.name is missing module name
LOGGER.warning(f'WARNING ⚠️ {weight} requires {e.name}, which is not in ultralytics requirements.' LOGGER.warning(f'WARNING ⚠️ {weight} requires {e.name}, which is not in ultralytics requirements.'
@ -347,7 +346,7 @@ def torch_safe_load(weight):
f'download updated models from https://github.com/ultralytics/assets/releases/tag/v0.0.0') f'download updated models from https://github.com/ultralytics/assets/releases/tag/v0.0.0')
if e.name != 'models': if e.name != 'models':
check_requirements(e.name) # install missing module check_requirements(e.name) # install missing module
return torch.load(file, map_location='cpu') # load return torch.load(file, map_location='cpu'), file # load
def attempt_load_weights(weights, device=None, inplace=True, fuse=False): def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
@ -355,13 +354,13 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
ensemble = Ensemble() ensemble = Ensemble()
for w in weights if isinstance(weights, list) else [weights]: for w in weights if isinstance(weights, list) else [weights]:
ckpt = torch_safe_load(w) # load ckpt ckpt, w = torch_safe_load(w) # load ckpt
args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} # combine model and default args, preferring model args args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} # combine model and default args, preferring model args
model = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model model = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model
# Model compatibility updates # Model compatibility updates
model.args = args # attach args to model model.args = args # attach args to model
model.pt_path = weights # attach *.pt file path to model model.pt_path = w # attach *.pt file path to model
model.task = guess_model_task(model) model.task = guess_model_task(model)
if not hasattr(model, 'stride'): if not hasattr(model, 'stride'):
model.stride = torch.tensor([32.]) model.stride = torch.tensor([32.])
@ -392,7 +391,7 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False): def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
# Loads a single model weights # Loads a single model weights
ckpt = torch_safe_load(weight) # load ckpt ckpt, weight = torch_safe_load(weight) # load ckpt
args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} # combine model and default args, preferring model args args = {**DEFAULT_CFG_DICT, **ckpt['train_args']} # combine model and default args, preferring model args
model = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model model = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model

@ -29,4 +29,5 @@ yolo segment track source=... tracker=...
``` ```
By default, trackers will use the configuration in `ultralytics/tracker/cfg`. By default, trackers will use the configuration in `ultralytics/tracker/cfg`.
We also support using a modified tracker config file. Please refer to the tracker config files in `ultralytics/tracker/cfg`. We also support using a modified tracker config file. Please refer to the tracker config files
in `ultralytics/tracker/cfg`.

@ -256,7 +256,7 @@ def entrypoint(debug=''):
# Defaults # Defaults
task2model = dict(detect='yolov8n.pt', segment='yolov8n-seg.pt', classify='yolov8n-cls.pt') task2model = dict(detect='yolov8n.pt', segment='yolov8n-seg.pt', classify='yolov8n-cls.pt')
task2data = dict(detect='coco128.yaml', segment='coco128-seg.yaml', classify='mnist160') task2data = dict(detect='coco128.yaml', segment='coco128-seg.yaml', classify='imagenet100')
# Mode # Mode
mode = overrides.get('mode', None) mode = overrides.get('mode', None)

@ -583,6 +583,7 @@ class Albumentations:
# TODO: add supports of segments and keypoints # TODO: add supports of segments and keypoints
if self.transform and random.random() < self.p: if self.transform and random.random() < self.p:
new = self.transform(image=im, bboxes=bboxes, class_labels=cls) # transformed new = self.transform(image=im, bboxes=bboxes, class_labels=cls) # transformed
if len(new['class_labels']) > 0: # skip update if no bbox in new im
labels['img'] = new['image'] labels['img'] = new['image']
labels['cls'] = np.array(new['class_labels']) labels['cls'] = np.array(new['class_labels'])
bboxes = np.array(new['bboxes']) bboxes = np.array(new['bboxes'])

@ -18,19 +18,18 @@ TensorFlow.js | `tfjs` | yolov8n_web_model/
PaddlePaddle | `paddle` | yolov8n_paddle_model/ PaddlePaddle | `paddle` | yolov8n_paddle_model/
Requirements: Requirements:
$ pip install -r requirements.txt coremltools onnx onnxsim onnxruntime openvino-dev tensorflow-cpu # CPU $ pip install ultralytics[export]
$ pip install -r requirements.txt coremltools onnx onnxsim onnxruntime-gpu openvino-dev tensorflow # GPU
Python: Python:
from ultralytics import YOLO from ultralytics import YOLO
model = YOLO('yolov8n.yaml') model = YOLO('yolov8n.pt')
results = model.export(format='onnx') results = model.export(format='onnx')
CLI: CLI:
$ yolo mode=export model=yolov8n.pt format=onnx $ yolo mode=export model=yolov8n.pt format=onnx
Inference: Inference:
$ python detect.py --weights yolov8n.pt # PyTorch $ yolo predict model=yolov8n.pt # PyTorch
yolov8n.torchscript # TorchScript yolov8n.torchscript # TorchScript
yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn
yolov8n_openvino_model # OpenVINO yolov8n_openvino_model # OpenVINO
@ -64,12 +63,12 @@ import pandas as pd
import torch import torch
from ultralytics.nn.autobackend import check_class_names from ultralytics.nn.autobackend import check_class_names
from ultralytics.nn.modules import Detect, Segment from ultralytics.nn.modules import C2f, Detect, Segment
from ultralytics.nn.tasks import DetectionModel, SegmentationModel from ultralytics.nn.tasks import DetectionModel, SegmentationModel
from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages
from ultralytics.yolo.data.utils import IMAGENET_MEAN, IMAGENET_STD, check_det_dataset from ultralytics.yolo.data.utils import IMAGENET_MEAN, IMAGENET_STD, check_det_dataset
from ultralytics.yolo.utils import (DEFAULT_CFG, LINUX, LOGGER, MACOS, WINDOWS, __version__, callbacks, colorstr, from ultralytics.yolo.utils import (DEFAULT_CFG, LINUX, LOGGER, MACOS, __version__, callbacks, colorstr,
get_default_args, yaml_save) get_default_args, yaml_save)
from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version, check_yaml from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version, check_yaml
from ultralytics.yolo.utils.files import file_size from ultralytics.yolo.utils.files import file_size
@ -77,6 +76,7 @@ from ultralytics.yolo.utils.ops import Profile
from ultralytics.yolo.utils.torch_utils import get_latest_opset, select_device, smart_inference_mode from ultralytics.yolo.utils.torch_utils import get_latest_opset, select_device, smart_inference_mode
CUDA = torch.cuda.is_available() CUDA = torch.cuda.is_available()
ARM64 = platform.machine() in ('arm64', 'aarch64')
def export_formats(): def export_formats():
@ -157,11 +157,10 @@ class Exporter:
# Load PyTorch model # Load PyTorch model
self.device = select_device('cpu' if self.args.device is None else self.args.device) self.device = select_device('cpu' if self.args.device is None else self.args.device)
if self.args.half: if self.args.half and onnx and self.device.type == 'cpu':
if self.device.type == 'cpu' and not coreml and not xml: LOGGER.warning('WARNING ⚠️ half=True only compatible with GPU export, i.e. use device=0')
LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
self.args.half = False self.args.half = False
assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic' assert not self.args.dynamic, 'half=True not compatible with dynamic=True, i.e. use only one.'
# Checks # Checks
model.names = check_class_names(model.names) model.names = check_class_names(model.names)
@ -188,11 +187,15 @@ class Exporter:
if isinstance(m, (Detect, Segment)): if isinstance(m, (Detect, Segment)):
m.dynamic = self.args.dynamic m.dynamic = self.args.dynamic
m.export = True m.export = True
m.format = self.args.format
elif isinstance(m, C2f) and not edgetpu:
# EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
m.forward = m.forward_split
y = None y = None
for _ in range(2): for _ in range(2):
y = model(im) # dry runs y = model(im) # dry runs
if self.args.half and not coreml and not xml: if self.args.half and (engine or onnx) and self.device.type != 'cpu':
im, model = im.half(), model.half() # to FP16 im, model = im.half(), model.half() # to FP16
# Warnings # Warnings
@ -207,7 +210,7 @@ class Exporter:
self.output_shape = tuple(y.shape) if isinstance(y, torch.Tensor) else tuple(tuple(x.shape) for x in y) self.output_shape = tuple(y.shape) if isinstance(y, torch.Tensor) else tuple(tuple(x.shape) for x in y)
self.pretty_name = self.file.stem.replace('yolo', 'YOLO') self.pretty_name = self.file.stem.replace('yolo', 'YOLO')
self.metadata = { self.metadata = {
'description': f'Ultralytics {self.pretty_name} model trained on {self.args.data}', 'description': f'Ultralytics {self.pretty_name} model trained on {Path(self.args.data).name}',
'author': 'Ultralytics', 'author': 'Ultralytics',
'license': 'GPL-3.0 https://ultralytics.com/license', 'license': 'GPL-3.0 https://ultralytics.com/license',
'version': __version__, 'version': __version__,
@ -233,19 +236,16 @@ class Exporter:
LOGGER.warning('WARNING ⚠️ YOLOv8 TensorFlow export is still under development. ' LOGGER.warning('WARNING ⚠️ YOLOv8 TensorFlow export is still under development. '
'Please consider contributing to the effort if you have TF expertise. Thank you!') 'Please consider contributing to the effort if you have TF expertise. Thank you!')
nms = False nms = False
self.args.int8 |= edgetpu
f[5], s_model = self._export_saved_model(nms=nms or self.args.agnostic_nms or tfjs, f[5], s_model = self._export_saved_model(nms=nms or self.args.agnostic_nms or tfjs,
agnostic_nms=self.args.agnostic_nms or tfjs) agnostic_nms=self.args.agnostic_nms or tfjs)
if pb or tfjs: # pb prerequisite to tfjs if pb or tfjs: # pb prerequisite to tfjs
f[6], _ = self._export_pb(s_model) f[6], _ = self._export_pb(s_model)
if tflite or edgetpu: if tflite:
f[7] = str(Path(f[5]) / (self.file.stem + '_float16.tflite')) f[7], _ = self._export_tflite(s_model, nms=nms, agnostic_nms=self.args.agnostic_nms)
# f[7], _ = self._export_tflite(s_model,
# int8=self.args.int8 or edgetpu,
# data=self.args.data,
# nms=nms,
# agnostic_nms=self.args.agnostic_nms)
if edgetpu: if edgetpu:
f[8], _ = self._export_edgetpu(tflite_model=f[7]) f[8], _ = self._export_edgetpu(tflite_model=str(
Path(f[5]) / (self.file.stem + '_full_integer_quant.tflite'))) # int8 in/out
if tfjs: if tfjs:
f[9], _ = self._export_tfjs() f[9], _ = self._export_tfjs()
if paddle: # PaddlePaddle if paddle: # PaddlePaddle
@ -263,8 +263,8 @@ class Exporter:
LOGGER.info( LOGGER.info(
f'\nExport complete ({time.time() - t:.1f}s)' f'\nExport complete ({time.time() - t:.1f}s)'
f"\nResults saved to {colorstr('bold', file.parent.resolve())}" f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
f'\nPredict: yolo task={model.task} mode=predict model={f} imgsz={imgsz} {data}' f'\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {data}'
f'\nValidate: yolo task={model.task} mode=val model={f} imgsz={imgsz} data={self.args.data} {s}' f'\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={self.args.data} {s}'
f'\nVisualize: https://netron.app') f'\nVisualize: https://netron.app')
self.run_callbacks('on_export_end') self.run_callbacks('on_export_end')
@ -319,25 +319,27 @@ class Exporter:
# Checks # Checks
model_onnx = onnx.load(f) # load onnx model model_onnx = onnx.load(f) # load onnx model
onnx.checker.check_model(model_onnx) # check onnx model # onnx.checker.check_model(model_onnx) # check onnx model
# Metadata
d = {'stride': int(max(self.model.stride)), 'names': self.model.names}
for k, v in d.items():
meta = model_onnx.metadata_props.add()
meta.key, meta.value = k, str(v)
onnx.save(model_onnx, f)
# Simplify # Simplify
if self.args.simplify: if self.args.simplify:
try: try:
check_requirements('onnxsim') check_requirements(('onnxsim', 'onnxruntime-gpu' if CUDA else 'onnxruntime'))
import onnxsim import onnxsim
LOGGER.info(f'{prefix} simplifying with onnxsim {onnxsim.__version__}...') LOGGER.info(f'{prefix} simplifying with onnxsim {onnxsim.__version__}...')
subprocess.run(f'onnxsim {f} {f}', shell=True) # subprocess.run(f'onnxsim {f} {f}', shell=True)
model_onnx, check = onnxsim.simplify(model_onnx)
assert check, 'Simplified ONNX model could not be validated'
except Exception as e: except Exception as e:
LOGGER.info(f'{prefix} simplifier failure: {e}') LOGGER.info(f'{prefix} simplifier failure: {e}')
# Metadata
for k, v in self.metadata.items():
meta = model_onnx.metadata_props.add()
meta.key, meta.value = k, str(v)
onnx.save(model_onnx, f)
return f, model_onnx return f, model_onnx
@try_export @try_export
@ -402,7 +404,7 @@ class Exporter:
if self.model.task == 'classify': if self.model.task == 'classify':
bias = [-x for x in IMAGENET_MEAN] bias = [-x for x in IMAGENET_MEAN]
scale = 1 / 255 / (sum(IMAGENET_STD) / 3) scale = 1 / 255 / (sum(IMAGENET_STD) / 3)
classifier_config = ct.ClassifierConfig(list(self.model.names.values())) classifier_config = ct.ClassifierConfig(list(self.model.names.values())) if self.args.nms else None
else: else:
bias = [0.0, 0.0, 0.0] bias = [0.0, 0.0, 0.0]
scale = 1 / 255 scale = 1 / 255
@ -414,10 +416,7 @@ class Exporter:
classifier_config=classifier_config) classifier_config=classifier_config)
bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None) bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None)
if bits < 32: if bits < 32:
if MACOS: # quantization only supported on macOS
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
else:
LOGGER.info(f'{prefix} quantization only supported on macOS, skipping...')
if self.args.nms: if self.args.nms:
ct_model = self._pipeline_coreml(ct_model) ct_model = self._pipeline_coreml(ct_model)
@ -440,11 +439,11 @@ class Exporter:
import tensorrt as trt # noqa import tensorrt as trt # noqa
check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=8.0.0 check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=8.0.0
self._export_onnx() self.args.simplify = True
onnx = self.file.with_suffix('.onnx') f_onnx, _ = self._export_onnx()
LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
assert onnx.exists(), f'failed to export ONNX file: {onnx}' assert Path(f_onnx).exists(), f'failed to export ONNX file: {f_onnx}'
f = self.file.with_suffix('.engine') # TensorRT engine file f = self.file.with_suffix('.engine') # TensorRT engine file
logger = trt.Logger(trt.Logger.INFO) logger = trt.Logger(trt.Logger.INFO)
if verbose: if verbose:
@ -458,8 +457,8 @@ class Exporter:
flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag) network = builder.create_network(flag)
parser = trt.OnnxParser(network, logger) parser = trt.OnnxParser(network, logger)
if not parser.parse_from_file(str(onnx)): if not parser.parse_from_file(f_onnx):
raise RuntimeError(f'failed to load ONNX file: {onnx}') raise RuntimeError(f'failed to load ONNX file: {f_onnx}')
inputs = [network.get_input(i) for i in range(network.num_inputs)] inputs = [network.get_input(i) for i in range(network.num_inputs)]
outputs = [network.get_output(i) for i in range(network.num_outputs)] outputs = [network.get_output(i) for i in range(network.num_outputs)]
@ -507,77 +506,37 @@ class Exporter:
try: try:
import tensorflow as tf # noqa import tensorflow as tf # noqa
except ImportError: except ImportError:
check_requirements(f"tensorflow{'' if CUDA else '-macos' if MACOS else '-cpu' if LINUX else ''}") check_requirements(f"tensorflow{'-macos' if MACOS else '-aarch64' if ARM64 else '' if CUDA else '-cpu'}")
import tensorflow as tf # noqa import tensorflow as tf # noqa
check_requirements(('onnx', 'onnx2tf', 'sng4onnx', 'onnxsim', 'onnx_graphsurgeon', 'tflite_support'), check_requirements(('onnx', 'onnx2tf', 'sng4onnx', 'onnxsim', 'onnx_graphsurgeon', 'tflite_support',
'onnxruntime-gpu' if CUDA else 'onnxruntime'),
cmds='--extra-index-url https://pypi.ngc.nvidia.com') cmds='--extra-index-url https://pypi.ngc.nvidia.com')
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
f = str(self.file).replace(self.file.suffix, '_saved_model') f = Path(str(self.file).replace(self.file.suffix, '_saved_model'))
if f.is_dir():
import shutil
shutil.rmtree(f) # delete output folder
# Export to ONNX # Export to ONNX
self._export_onnx() self.args.simplify = True
onnx = self.file.with_suffix('.onnx') f_onnx, _ = self._export_onnx()
# Export to TF SavedModel # Export to TF
subprocess.run(f'onnx2tf -i {onnx} -o {f} --non_verbose', shell=True) int8 = '-oiqt -qt per-tensor' if self.args.int8 else ''
yaml_save(Path(f) / 'metadata.yaml', self.metadata) # add metadata.yaml cmd = f'onnx2tf -i {f_onnx} -o {f} --non_verbose {int8}'
LOGGER.info(f'\n{prefix} running {cmd}')
subprocess.run(cmd, shell=True)
yaml_save(f / 'metadata.yaml', self.metadata) # add metadata.yaml
# Add TFLite metadata # Add TFLite metadata
for file in Path(f).rglob('*.tflite'): for file in f.rglob('*.tflite'):
self._add_tflite_metadata(file) self._add_tflite_metadata(file)
# Load saved_model # Load saved_model
keras_model = tf.saved_model.load(f, tags=None, options=None) keras_model = tf.saved_model.load(f, tags=None, options=None)
return f, keras_model return str(f), keras_model
@try_export
def _export_saved_model_OLD(self,
nms=False,
agnostic_nms=False,
topk_per_class=100,
topk_all=100,
iou_thres=0.45,
conf_thres=0.25,
prefix=colorstr('TensorFlow SavedModel:')):
# YOLOv8 TensorFlow SavedModel export
try:
import tensorflow as tf # noqa
except ImportError:
check_requirements(f"tensorflow{'' if CUDA else '-macos' if MACOS else '-cpu' if LINUX else ''}")
import tensorflow as tf # noqa
# from models.tf import TFModel
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
f = str(self.file).replace(self.file.suffix, '_saved_model')
batch_size, ch, *imgsz = list(self.im.shape) # BCHW
tf_models = None # TODO: no TF modules available
tf_model = tf_models.TFModel(cfg=self.model.yaml, model=self.model.cpu(), nc=self.model.nc, imgsz=imgsz)
im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow
_ = tf_model.predict(im, nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if self.args.dynamic else batch_size)
outputs = tf_model.predict(inputs, nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
keras_model.trainable = False
keras_model.summary()
if self.args.keras:
keras_model.save(f, save_format='tf')
else:
spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
m = tf.function(lambda x: keras_model(x)) # full model
m = m.get_concrete_function(spec)
frozen_func = convert_variables_to_constants_v2(m)
tfm = tf.Module()
tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if nms else frozen_func(x), [spec])
tfm.__call__(im)
tf.saved_model.save(tfm,
f,
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
if check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions())
return f, keras_model
@try_export @try_export
def _export_pb(self, keras_model, prefix=colorstr('TensorFlow GraphDef:')): def _export_pb(self, keras_model, prefix=colorstr('TensorFlow GraphDef:')):
@ -596,8 +555,18 @@ class Exporter:
return f, None return f, None
@try_export @try_export
def _export_tflite(self, keras_model, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')): def _export_tflite(self, keras_model, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
# YOLOv8 TensorFlow Lite export # YOLOv8 TensorFlow Lite export
saved_model = Path(str(self.file).replace(self.file.suffix, '_saved_model'))
if self.args.int8:
f = saved_model / (self.file.stem + 'yolov8n_integer_quant.tflite') # fp32 in/out
elif self.args.half:
f = saved_model / (self.file.stem + '_float16.tflite')
else:
f = saved_model / (self.file.stem + '_float32.tflite')
return str(f), None # noqa
# OLD VERSION BELOW ---------------------------------------------------------------
import tensorflow as tf # noqa import tensorflow as tf # noqa
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
@ -608,7 +577,7 @@ class Exporter:
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.target_spec.supported_types = [tf.float16] converter.target_spec.supported_types = [tf.float16]
converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.optimizations = [tf.lite.Optimize.DEFAULT]
if int8: if self.args.int8:
def representative_dataset_gen(dataset, n_images=100): def representative_dataset_gen(dataset, n_images=100):
# Dataset generator for use with converter.representative_dataset, returns a generator of np arrays # Dataset generator for use with converter.representative_dataset, returns a generator of np arrays
@ -620,7 +589,7 @@ class Exporter:
if n >= n_images: if n >= n_images:
break break
dataset = LoadImages(check_det_dataset(check_yaml(data))['train'], imgsz=imgsz, auto=False) dataset = LoadImages(check_det_dataset(check_yaml(self.args.data))['train'], imgsz=imgsz, auto=False)
converter.representative_dataset = lambda: representative_dataset_gen(dataset, n_images=100) converter.representative_dataset = lambda: representative_dataset_gen(dataset, n_images=100)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.target_spec.supported_types = [] converter.target_spec.supported_types = []
@ -641,7 +610,7 @@ class Exporter:
cmd = 'edgetpu_compiler --version' cmd = 'edgetpu_compiler --version'
help_url = 'https://coral.ai/docs/edgetpu/compiler/' help_url = 'https://coral.ai/docs/edgetpu/compiler/'
assert LINUX, f'export only supported on Linux. See {help_url}' assert LINUX, f'export only supported on Linux. See {help_url}'
if subprocess.run(f'{cmd} >/dev/null', shell=True).returncode != 0: if subprocess.run(f'{cmd} > /dev/null', shell=True).returncode != 0:
LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system
for c in ( for c in (
@ -656,7 +625,7 @@ class Exporter:
LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...') LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
f = str(tflite_model).replace('.tflite', '_edgetpu.tflite') # Edge TPU model f = str(tflite_model).replace('.tflite', '_edgetpu.tflite') # Edge TPU model
cmd = f'edgetpu_compiler -s -d -k 10 --out_dir {self.file.parent} {tflite_model}' cmd = f'edgetpu_compiler -s -d -k 10 --out_dir {Path(f).parent} {tflite_model}'
subprocess.run(cmd.split(), check=True) subprocess.run(cmd.split(), check=True)
self._add_tflite_metadata(f) self._add_tflite_metadata(f)
return f, None return f, None
@ -674,7 +643,7 @@ class Exporter:
cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \ cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \
f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}' f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}'
subprocess.run(cmd.split()) subprocess.run(cmd.split(), check=True)
with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
subst = re.sub( subst = re.sub(
@ -698,14 +667,23 @@ class Exporter:
from tflite_support import metadata as _metadata # noqa from tflite_support import metadata as _metadata # noqa
from tflite_support import metadata_schema_py_generated as _metadata_fb # noqa from tflite_support import metadata_schema_py_generated as _metadata_fb # noqa
# Creates model info. # Create model info
model_meta = _metadata_fb.ModelMetadataT() model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = self.metadata['description'] model_meta.name = self.metadata['description']
model_meta.version = self.metadata['version'] model_meta.version = self.metadata['version']
model_meta.author = self.metadata['author'] model_meta.author = self.metadata['author']
model_meta.license = self.metadata['license'] model_meta.license = self.metadata['license']
# Creates input info. # Label file
tmp_file = file.parent / 'temp_meta.txt'
with open(tmp_file, 'w') as f:
f.write(str(self.metadata))
label_file = _metadata_fb.AssociatedFileT()
label_file.name = tmp_file.name
label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
# Create input info
input_meta = _metadata_fb.TensorMetadataT() input_meta = _metadata_fb.TensorMetadataT()
input_meta.name = 'image' input_meta.name = 'image'
input_meta.description = 'Input image to be detected.' input_meta.description = 'Input image to be detected.'
@ -714,25 +692,21 @@ class Exporter:
input_meta.content.contentProperties.colorSpace = _metadata_fb.ColorSpaceType.RGB input_meta.content.contentProperties.colorSpace = _metadata_fb.ColorSpaceType.RGB
input_meta.content.contentPropertiesType = _metadata_fb.ContentProperties.ImageProperties input_meta.content.contentPropertiesType = _metadata_fb.ContentProperties.ImageProperties
# Creates output info. # Create output info
output_meta = _metadata_fb.TensorMetadataT() output1 = _metadata_fb.TensorMetadataT()
output_meta.name = 'output' output1.name = 'output'
output_meta.description = 'Coordinates of detected objects, class labels, and confidence score.' output1.description = 'Coordinates of detected objects, class labels, and confidence score'
output1.associatedFiles = [label_file]
# Label file if self.model.task == 'segment':
tmp_file = Path('/tmp/meta.txt') output2 = _metadata_fb.TensorMetadataT()
with open(tmp_file, 'w') as meta_f: output2.name = 'output'
meta_f.write(str(self.metadata)) output2.description = 'Mask protos'
output2.associatedFiles = [label_file]
label_file = _metadata_fb.AssociatedFileT()
label_file.name = tmp_file.name # Create subgraph info
label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
output_meta.associatedFiles = [label_file]
# Creates subgraph info.
subgraph = _metadata_fb.SubGraphMetadataT() subgraph = _metadata_fb.SubGraphMetadataT()
subgraph.inputTensorMetadata = [input_meta] subgraph.inputTensorMetadata = [input_meta]
subgraph.outputTensorMetadata = [output_meta] subgraph.outputTensorMetadata = [output1, output2] if self.model.task == 'segment' else [output1]
model_meta.subgraphMetadata = [subgraph] model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0) b = flatbuffers.Builder(0)

@ -29,14 +29,45 @@ MODEL_MAP = {
class YOLO: class YOLO:
""" """
YOLO YOLO (You Only Look Once) object detection model.
A python interface which emulates a model-like behaviour by wrapping trainers. Args:
model (str or Path): Path to the model file to load or create.
type (str): Type/version of models to use. Defaults to "v8".
Attributes:
type (str): Type/version of models being used.
ModelClass (Any): Model class.
TrainerClass (Any): Trainer class.
ValidatorClass (Any): Validator class.
PredictorClass (Any): Predictor class.
predictor (Any): Predictor object.
model (Any): Model object.
trainer (Any): Trainer object.
task (str): Type of model task.
ckpt (Any): Checkpoint object if model loaded from *.pt file.
cfg (str): Model configuration if loaded from *.yaml file.
ckpt_path (str): Checkpoint file path.
overrides (dict): Overrides for trainer object.
metrics_data (Any): Data for metrics.
Methods:
__call__(): Alias for predict method.
_new(cfg, verbose=True): Initializes a new model and infers the task type from the model definitions.
_load(weights): Initializes a new model and infers the task type from the model head.
_check_is_pytorch_model(): Raises TypeError if model is not a PyTorch model.
reset(): Resets the model modules.
info(verbose=False): Logs model info.
fuse(): Fuse model for faster inference.
predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model.
Returns:
List[ultralytics.yolo.engine.results.Results]: The prediction results.
""" """
def __init__(self, model='yolov8n.pt', type='v8') -> None: def __init__(self, model='yolov8n.pt', type='v8') -> None:
""" """
Initializes the YOLO object. Initializes the YOLO model.
Args: Args:
model (str, Path): model to load or create model (str, Path): model to load or create
@ -97,8 +128,9 @@ class YOLO:
self.task = self.model.args['task'] self.task = self.model.args['task']
self.overrides = self.model.args self.overrides = self.model.args
self._reset_ckpt_args(self.overrides) self._reset_ckpt_args(self.overrides)
self.ckpt_path = self.model.pt_path
else: else:
check_file(weights) weights = check_file(weights)
self.model, self.ckpt = weights, None self.model, self.ckpt = weights, None
self.task = guess_model_task(weights) self.task = guess_model_task(weights)
self.ckpt_path = weights self.ckpt_path = weights
@ -204,7 +236,6 @@ class YOLO:
return validator.metrics return validator.metrics
@smart_inference_mode()
def export(self, **kwargs): def export(self, **kwargs):
""" """
Export model. Export model.
@ -279,6 +310,13 @@ class YOLO:
""" """
return self.model.names if hasattr(self.model, 'names') else None return self.model.names if hasattr(self.model, 'names') else None
@property
def device(self):
"""
Returns device if PyTorch model
"""
return next(self.model.parameters()).device if isinstance(self.model, nn.Module) else None
@property @property
def transforms(self): def transforms(self):
""" """
@ -293,7 +331,6 @@ class YOLO:
""" """
if not self.metrics_data: if not self.metrics_data:
LOGGER.info('No metrics data found! Run training or validation operation first.') LOGGER.info('No metrics data found! Run training or validation operation first.')
return self.metrics_data return self.metrics_data
@staticmethod @staticmethod
@ -306,7 +343,7 @@ class YOLO:
@staticmethod @staticmethod
def _reset_ckpt_args(args): def _reset_ckpt_args(args):
for arg in 'augment', 'verbose', 'project', 'name', 'exist_ok', 'resume', 'batch', 'epochs', 'cache', \ for arg in 'augment', 'verbose', 'project', 'name', 'exist_ok', 'resume', 'batch', 'epochs', 'cache', \
'save_json', 'half', 'v5loader', 'device', 'cfg', 'save', 'rect', 'plots', 'opset': 'save_json', 'half', 'v5loader', 'device', 'cfg', 'save', 'rect', 'plots', 'opset', 'simplify':
args.pop(arg, None) args.pop(arg, None)
@staticmethod @staticmethod

@ -1,8 +1,9 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
""" """
Run prediction on images, videos, directories, globs, YouTube, webcam, streams, etc. Run prediction on images, videos, directories, globs, YouTube, webcam, streams, etc.
Usage - sources: Usage - sources:
$ yolo task=... mode=predict model=s.pt --source 0 # webcam $ yolo mode=predict model=yolov8n.pt --source 0 # webcam
img.jpg # image img.jpg # image
vid.mp4 # video vid.mp4 # video
screen # screenshot screen # screenshot
@ -12,10 +13,11 @@ Usage - sources:
'path/*.jpg' # glob 'path/*.jpg' # glob
'https://youtu.be/Zgi9g1ksQHc' # YouTube 'https://youtu.be/Zgi9g1ksQHc' # YouTube
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
Usage - formats: Usage - formats:
$ yolo task=... mode=predict --weights yolov8n.pt # PyTorch $ yolo mode=predict model=yolov8n.pt # PyTorch
yolov8n.torchscript # TorchScript yolov8n.torchscript # TorchScript
yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov8n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
yolov8n_openvino_model # OpenVINO yolov8n_openvino_model # OpenVINO
yolov8n.engine # TensorRT yolov8n.engine # TensorRT
yolov8n.mlmodel # CoreML (macOS-only) yolov8n.mlmodel # CoreML (macOS-only)
@ -24,7 +26,7 @@ Usage - formats:
yolov8n.tflite # TensorFlow Lite yolov8n.tflite # TensorFlow Lite
yolov8n_edgetpu.tflite # TensorFlow Edge TPU yolov8n_edgetpu.tflite # TensorFlow Edge TPU
yolov8n_paddle_model # PaddlePaddle yolov8n_paddle_model # PaddlePaddle
""" """
import platform import platform
from collections import defaultdict from collections import defaultdict
from pathlib import Path from pathlib import Path
@ -200,9 +202,9 @@ class BasePredictor:
# Print results # Print results
if self.args.verbose and self.seen: if self.args.verbose and self.seen:
t = tuple(x.t / self.seen * 1E3 for x in self.dt) # speeds per image t = tuple(x.t / self.seen * 1E3 for x in self.dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms postprocess per image at shape ' LOGGER.info(f'Speed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape '
f'{(1, 3, *self.imgsz)}' % t) f'{(1, 3, *self.imgsz)}' % t)
if self.args.save_txt or self.args.save: if self.args.save or self.args.save_txt or self.args.save_crop:
nl = len(list(self.save_dir.glob('labels/*.txt'))) # number of labels nl = len(list(self.save_dir.glob('labels/*.txt'))) # number of labels
s = f"\n{nl} label{'s' * (nl > 1)} saved to {self.save_dir / 'labels'}" if self.args.save_txt else '' s = f"\n{nl} label{'s' * (nl > 1)} saved to {self.save_dir / 'labels'}" if self.args.save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}") LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}")

@ -4,7 +4,6 @@ from functools import lru_cache
import numpy as np import numpy as np
import torch import torch
import torchvision.transforms.functional as F import torchvision.transforms.functional as F
from PIL import Image
from ultralytics.yolo.utils import LOGGER, ops from ultralytics.yolo.utils import LOGGER, ops
from ultralytics.yolo.utils.plotting import Annotator, colors from ultralytics.yolo.utils.plotting import Annotator, colors
@ -136,7 +135,7 @@ class Results:
img = deepcopy(self.orig_img) img = deepcopy(self.orig_img)
annotator = Annotator(img, line_width, font_size, font, pil, example) annotator = Annotator(img, line_width, font_size, font, pil, example)
boxes = self.boxes boxes = self.boxes
masks = self.masks.data masks = self.masks
logits = self.probs logits = self.probs
names = self.names names = self.names
if boxes is not None: if boxes is not None:

@ -1,8 +1,10 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
""" """
Simple training loop; Boilerplate that could apply to any arbitrary neural network, Train a model on a dataset
"""
Usage:
$ yolo mode=train model=yolov8n.pt data=coco128.yaml imgsz=640 epochs=100 batch=16
"""
import os import os
import subprocess import subprocess
import time import time

@ -1,5 +1,23 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
"""
Check a model's accuracy on a test or val split of a dataset
Usage:
$ yolo mode=val model=yolov8n.pt data=coco128.yaml imgsz=640
Usage - formats:
$ yolo mode=val model=yolov8n.pt # PyTorch
yolov8n.torchscript # TorchScript
yolov8n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
yolov8n_openvino_model # OpenVINO
yolov8n.engine # TensorRT
yolov8n.mlmodel # CoreML (macOS-only)
yolov8n_saved_model # TensorFlow SavedModel
yolov8n.pb # TensorFlow GraphDef
yolov8n.tflite # TensorFlow Lite
yolov8n_edgetpu.tflite # TensorFlow Edge TPU
yolov8n_paddle_model # PaddlePaddle
"""
import json import json
from collections import defaultdict from collections import defaultdict
from pathlib import Path from pathlib import Path
@ -105,8 +123,7 @@ class BaseValidator:
self.device = model.device self.device = model.device
if not pt and not jit: if not pt and not jit:
self.args.batch = 1 # export.py models default to batch-size 1 self.args.batch = 1 # export.py models default to batch-size 1
self.logger.info( self.logger.info(f'Forcing batch=1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
if isinstance(self.args.data, str) and self.args.data.endswith('.yaml'): if isinstance(self.args.data, str) and self.args.data.endswith('.yaml'):
self.data = check_det_dataset(self.args.data) self.data = check_det_dataset(self.args.data)
@ -136,7 +153,7 @@ class BaseValidator:
for batch_i, batch in enumerate(bar): for batch_i, batch in enumerate(bar):
self.run_callbacks('on_val_batch_start') self.run_callbacks('on_val_batch_start')
self.batch_i = batch_i self.batch_i = batch_i
# pre-process # preprocess
with dt[0]: with dt[0]:
batch = self.preprocess(batch) batch = self.preprocess(batch)
@ -149,7 +166,7 @@ class BaseValidator:
if self.training: if self.training:
self.loss += trainer.criterion(preds, batch)[1] self.loss += trainer.criterion(preds, batch)[1]
# pre-process predictions # postprocess
with dt[3]: with dt[3]:
preds = self.postprocess(preds) preds = self.postprocess(preds)
@ -163,13 +180,14 @@ class BaseValidator:
self.check_stats(stats) self.check_stats(stats)
self.print_results() self.print_results()
self.speed = tuple(x.t / len(self.dataloader.dataset) * 1E3 for x in dt) # speeds per image self.speed = tuple(x.t / len(self.dataloader.dataset) * 1E3 for x in dt) # speeds per image
self.finalize_metrics()
self.run_callbacks('on_val_end') self.run_callbacks('on_val_end')
if self.training: if self.training:
model.float() model.float()
results = {**stats, **trainer.label_loss_items(self.loss.cpu() / len(self.dataloader), prefix='val')} results = {**stats, **trainer.label_loss_items(self.loss.cpu() / len(self.dataloader), prefix='val')}
return {k: round(float(v), 5) for k, v in results.items()} # return results as 5 decimal place floats return {k: round(float(v), 5) for k, v in results.items()} # return results as 5 decimal place floats
else: else:
self.logger.info('Speed: %.1fms pre-process, %.1fms inference, %.1fms loss, %.1fms post-process per image' % self.logger.info('Speed: %.1fms preprocess, %.1fms inference, %.1fms loss, %.1fms postprocess per image' %
self.speed) self.speed)
if self.args.save_json and self.jdict: if self.args.save_json and self.jdict:
with open(str(self.save_dir / 'predictions.json'), 'w') as f: with open(str(self.save_dir / 'predictions.json'), 'w') as f:
@ -197,6 +215,9 @@ class BaseValidator:
def update_metrics(self, preds, batch): def update_metrics(self, preds, batch):
pass pass
def finalize_metrics(self, *args, **kwargs):
pass
def get_stats(self): def get_stats(self):
return {} return {}

@ -97,6 +97,7 @@ HELP_MSG = \
torch.set_printoptions(linewidth=320, precision=5, profile='long') torch.set_printoptions(linewidth=320, precision=5, profile='long')
np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
pd.options.display.max_columns = 10 pd.options.display.max_columns = 10
pd.options.display.width = 120
cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' # for deterministic training os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' # for deterministic training
@ -287,9 +288,7 @@ def is_pytest_running():
Returns: Returns:
(bool): True if pytest is running, False otherwise. (bool): True if pytest is running, False otherwise.
""" """
with contextlib.suppress(Exception): return ('PYTEST_CURRENT_TEST' in os.environ) or ('pytest' in sys.modules) or ('pytest' in Path(sys.argv[0]).stem)
return 'pytest' in sys.modules
return False
def is_github_actions_ci() -> bool: def is_github_actions_ci() -> bool:
@ -530,8 +529,7 @@ def set_sentry():
if SETTINGS['sync'] and \ if SETTINGS['sync'] and \
RANK in {-1, 0} and \ RANK in {-1, 0} and \
Path(sys.argv[0]).name == 'yolo' and \ Path(sys.argv[0]).name == 'yolo' and \
not is_pytest_running() and \ not TESTS_RUNNING and \
not is_github_actions_ci() and \
((is_pip_package() and not is_git_dir()) or ((is_pip_package() and not is_git_dir()) or
(get_git_origin_url() == 'https://github.com/ultralytics/ultralytics.git' and get_git_branch() == 'main')): (get_git_origin_url() == 'https://github.com/ultralytics/ultralytics.git' and get_git_branch() == 'main')):
@ -625,4 +623,5 @@ SETTINGS = get_settings()
DATASETS_DIR = Path(SETTINGS['datasets_dir']) # global datasets directory DATASETS_DIR = Path(SETTINGS['datasets_dir']) # global datasets directory
ENVIRONMENT = 'Colab' if is_colab() else 'Kaggle' if is_kaggle() else 'Jupyter' if is_jupyter() else \ ENVIRONMENT = 'Colab' if is_colab() else 'Kaggle' if is_kaggle() else 'Jupyter' if is_jupyter() else \
'Docker' if is_docker() else platform.system() 'Docker' if is_docker() else platform.system()
TESTS_RUNNING = is_pytest_running() or is_github_actions_ci()
set_sentry() set_sentry()

@ -0,0 +1,101 @@
# Ultralytics YOLO 🚀, GPL-3.0 license
"""
Benchmark a YOLO model formats for speed and accuracy
Usage:
from ultralytics.yolo.utils.benchmarks import run_benchmarks
run_benchmarks(model='yolov8n.pt', imgsz=160)
Format | `format=argument` | Model
--- | --- | ---
PyTorch | - | yolov8n.pt
TorchScript | `torchscript` | yolov8n.torchscript
ONNX | `onnx` | yolov8n.onnx
OpenVINO | `openvino` | yolov8n_openvino_model/
TensorRT | `engine` | yolov8n.engine
CoreML | `coreml` | yolov8n.mlmodel
TensorFlow SavedModel | `saved_model` | yolov8n_saved_model/
TensorFlow GraphDef | `pb` | yolov8n.pb
TensorFlow Lite | `tflite` | yolov8n.tflite
TensorFlow Edge TPU | `edgetpu` | yolov8n_edgetpu.tflite
TensorFlow.js | `tfjs` | yolov8n_web_model/
PaddlePaddle | `paddle` | yolov8n_paddle_model/
"""
import platform
import time
from pathlib import Path
import pandas as pd
import torch
from ultralytics import YOLO
from ultralytics.yolo.engine.exporter import export_formats
from ultralytics.yolo.utils import LOGGER, SETTINGS
from ultralytics.yolo.utils.checks import check_yolo
from ultralytics.yolo.utils.files import file_size
def run_benchmarks(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt',
imgsz=640,
half=False,
device='cpu',
hard_fail=False):
device = torch.device(int(device) if device.isnumeric() else device)
model = YOLO(model)
y = []
t0 = time.time()
for i, (name, format, suffix, cpu, gpu) in export_formats().iterrows(): # index, (name, format, suffix, CPU, GPU)
try:
assert i not in (9, 10), 'inference not supported' # Edge TPU and TF.js are unsupported
assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13' # CoreML
if 'cpu' in device.type:
assert cpu, 'inference not supported on CPU'
if 'cuda' in device.type:
assert gpu, 'inference not supported on GPU'
# Export
if format == '-':
filename = model.ckpt_path
export = model # PyTorch format
else:
filename = model.export(imgsz=imgsz, format=format, half=half, device=device) # all others
export = YOLO(filename)
assert suffix in str(filename), 'export failed'
# Validate
if model.task == 'detect':
data, key = 'coco128.yaml', 'metrics/mAP50-95(B)'
elif model.task == 'segment':
data, key = 'coco128-seg.yaml', 'metrics/mAP50-95(M)'
elif model.task == 'classify':
data, key = 'imagenet100', 'metrics/accuracy_top5'
results = export.val(data=data, batch=1, imgsz=imgsz, plots=False, device=device, half=half, verbose=False)
metric, speed = results.results_dict[key], results.speed['inference']
y.append([name, '', round(file_size(filename), 1), round(metric, 4), round(speed, 2)])
except Exception as e:
if hard_fail:
assert type(e) is AssertionError, f'Benchmark --hard-fail for {name}: {e}'
LOGGER.warning(f'ERROR ❌️ Benchmark failure for {name}: {e}')
y.append([name, '', None, None, None]) # mAP, t_inference
# Print results
LOGGER.info('\n')
check_yolo(device=device) # print system info
c = ['Format', 'Status❔', 'Size (MB)', key, 'Inference time (ms/im)'] if map else ['Format', 'Export', '', '']
df = pd.DataFrame(y, columns=c)
LOGGER.info(f'\nBenchmarks complete for {Path(model.ckpt_path).name} on {data} at imgsz={imgsz} '
f'({time.time() - t0:.2f}s)')
LOGGER.info(str(df if map else df.iloc[:, :2]))
if hard_fail and isinstance(hard_fail, str):
metrics = df[key].array # values to compare to floor
floor = eval(hard_fail) # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: metric < floor {floor}'
if __name__ == '__main__':
run_benchmarks()

@ -1,5 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
try: try:
@ -7,6 +7,7 @@ try:
from clearml import Task from clearml import Task
assert clearml.__version__ # verify package is not directory assert clearml.__version__ # verify package is not directory
assert not TESTS_RUNNING # do not log pytest
except (ImportError, AssertionError): except (ImportError, AssertionError):
clearml = None clearml = None
@ -19,7 +20,7 @@ def _log_images(imgs_dict, group='', step=0):
def on_pretrain_routine_start(trainer): def on_pretrain_routine_start(trainer):
# TODO: reuse existing task try:
task = Task.init(project_name=trainer.args.project or 'YOLOv8', task = Task.init(project_name=trainer.args.project or 'YOLOv8',
task_name=trainer.args.name, task_name=trainer.args.name,
tags=['YOLOv8'], tags=['YOLOv8'],
@ -27,6 +28,8 @@ def on_pretrain_routine_start(trainer):
reuse_last_task_id=False, reuse_last_task_id=False,
auto_connect_frameworks={'pytorch': False}) auto_connect_frameworks={'pytorch': False})
task.connect(vars(trainer.args), name='General') task.connect(vars(trainer.args), name='General')
except Exception as e:
LOGGER.warning(f'WARNING ⚠️ ClearML not initialized correctly, not logging this run. {e}')
def on_train_epoch_end(trainer): def on_train_epoch_end(trainer):
@ -35,18 +38,19 @@ def on_train_epoch_end(trainer):
def on_fit_epoch_end(trainer): def on_fit_epoch_end(trainer):
if trainer.epoch == 0: task = Task.current_task()
if task and trainer.epoch == 0:
model_info = { model_info = {
'Parameters': get_num_params(trainer.model), 'Parameters': get_num_params(trainer.model),
'GFLOPs': round(get_flops(trainer.model), 3), 'GFLOPs': round(get_flops(trainer.model), 3),
'Inference speed (ms/img)': round(trainer.validator.speed[1], 3)} 'Inference speed (ms/img)': round(trainer.validator.speed[1], 3)}
Task.current_task().connect(model_info, name='Model') task.connect(model_info, name='Model')
def on_train_end(trainer): def on_train_end(trainer):
Task.current_task().update_output_model(model_path=str(trainer.best), task = Task.current_task()
model_name=trainer.args.name, if task:
auto_delete_file=False) task.update_output_model(model_path=str(trainer.best), model_name=trainer.args.name, auto_delete_file=False)
callbacks = { callbacks = {

@ -1,21 +1,27 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
try: try:
import comet_ml import comet_ml
except ImportError: assert not TESTS_RUNNING # do not log pytest
assert comet_ml.__version__ # verify package is not directory
except (ImportError, AssertionError):
comet_ml = None comet_ml = None
def on_pretrain_routine_start(trainer): def on_pretrain_routine_start(trainer):
try:
experiment = comet_ml.Experiment(project_name=trainer.args.project or 'YOLOv8') experiment = comet_ml.Experiment(project_name=trainer.args.project or 'YOLOv8')
experiment.log_parameters(vars(trainer.args)) experiment.log_parameters(vars(trainer.args))
except Exception as e:
LOGGER.warning(f'WARNING ⚠️ Comet not initialized correctly, not logging this run. {e}')
def on_train_epoch_end(trainer): def on_train_epoch_end(trainer):
experiment = comet_ml.get_global_experiment() experiment = comet_ml.get_global_experiment()
if experiment:
experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix='train'), step=trainer.epoch + 1) experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix='train'), step=trainer.epoch + 1)
if trainer.epoch == 1: if trainer.epoch == 1:
for f in trainer.save_dir.glob('train_batch*.jpg'): for f in trainer.save_dir.glob('train_batch*.jpg'):
@ -24,6 +30,7 @@ def on_train_epoch_end(trainer):
def on_fit_epoch_end(trainer): def on_fit_epoch_end(trainer):
experiment = comet_ml.get_global_experiment() experiment = comet_ml.get_global_experiment()
if experiment:
experiment.log_metrics(trainer.metrics, step=trainer.epoch + 1) experiment.log_metrics(trainer.metrics, step=trainer.epoch + 1)
if trainer.epoch == 0: if trainer.epoch == 0:
model_info = { model_info = {
@ -35,6 +42,7 @@ def on_fit_epoch_end(trainer):
def on_train_end(trainer): def on_train_end(trainer):
experiment = comet_ml.get_global_experiment() experiment = comet_ml.get_global_experiment()
if experiment:
experiment.log_model('YOLOv8', file_or_folder=str(trainer.best), file_name='best.pt', overwrite=True) experiment.log_model('YOLOv8', file_or_folder=str(trainer.best), file_name='best.pt', overwrite=True)

@ -4,11 +4,11 @@ import json
from time import time from time import time
from ultralytics.hub.utils import PREFIX, traces from ultralytics.hub.utils import PREFIX, traces
from ultralytics.yolo.utils import LOGGER from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
def on_pretrain_routine_end(trainer): def on_pretrain_routine_end(trainer):
session = getattr(trainer, 'hub_session', None) session = not TESTS_RUNNING and getattr(trainer, 'hub_session', None)
if session: if session:
# Start timer for upload rate limit # Start timer for upload rate limit
LOGGER.info(f'{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀') LOGGER.info(f'{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀')

@ -194,6 +194,10 @@ def check_requirements(requirements=ROOT.parent / 'requirements.txt', exclude=()
try: try:
pkg.require(r) pkg.require(r)
except (pkg.VersionConflict, pkg.DistributionNotFound): # exception if requirements not met except (pkg.VersionConflict, pkg.DistributionNotFound): # exception if requirements not met
try: # attempt to import (slower but more accurate)
import importlib
importlib.import_module(next(pkg.parse_requirements(r)).name)
except ImportError:
s += f'"{r}" ' s += f'"{r}" '
n += 1 n += 1
@ -250,7 +254,7 @@ def check_file(file, suffix='', download=True):
return file return file
else: # search else: # search
files = [] files = []
for d in 'models', 'yolo/data', 'tracker/cfg': # search directories for d in 'models', 'datasets', 'tracker/cfg': # search directories
files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file
if not files: if not files:
raise FileNotFoundError(f"'{file}' does not exist") raise FileNotFoundError(f"'{file}' does not exist")
@ -280,7 +284,7 @@ def check_imshow(warn=False):
return False return False
def check_yolo(verbose=True): def check_yolo(verbose=True, device=''):
from ultralytics.yolo.utils.torch_utils import select_device from ultralytics.yolo.utils.torch_utils import select_device
if is_colab(): if is_colab():
@ -298,7 +302,7 @@ def check_yolo(verbose=True):
else: else:
s = '' s = ''
select_device(newline=False) select_device(device=device, newline=False)
LOGGER.info(f'Setup complete ✅ {s}') LOGGER.info(f'Setup complete ✅ {s}')

@ -512,6 +512,7 @@ class DetMetrics:
self.plot = plot self.plot = plot
self.names = names self.names = names
self.box = Metric() self.box = Metric()
self.speed = {'preprocess': 0.0, 'inference': 0.0, 'loss': 0.0, 'postprocess': 0.0}
def process(self, tp, conf, pred_cls, target_cls): def process(self, tp, conf, pred_cls, target_cls):
results = ap_per_class(tp, conf, pred_cls, target_cls, plot=self.plot, save_dir=self.save_dir, results = ap_per_class(tp, conf, pred_cls, target_cls, plot=self.plot, save_dir=self.save_dir,
@ -554,6 +555,7 @@ class SegmentMetrics:
self.names = names self.names = names
self.box = Metric() self.box = Metric()
self.seg = Metric() self.seg = Metric()
self.speed = {'preprocess': 0.0, 'inference': 0.0, 'loss': 0.0, 'postprocess': 0.0}
def process(self, tp_m, tp_b, conf, pred_cls, target_cls): def process(self, tp_m, tp_b, conf, pred_cls, target_cls):
results_mask = ap_per_class(tp_m, results_mask = ap_per_class(tp_m,
@ -612,6 +614,7 @@ class ClassifyMetrics:
def __init__(self) -> None: def __init__(self) -> None:
self.top1 = 0 self.top1 = 0
self.top5 = 0 self.top5 = 0
self.speed = {'preprocess': 0.0, 'inference': 0.0, 'loss': 0.0, 'postprocess': 0.0}
def process(self, targets, pred): def process(self, targets, pred):
# target classes and predicted classes # target classes and predicted classes

@ -154,7 +154,7 @@ class Annotator:
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True): def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
xyxy = torch.tensor(xyxy).view(-1, 4) xyxy = torch.Tensor(xyxy).view(-1, 4)
b = xyxy2xywh(xyxy) # boxes b = xyxy2xywh(xyxy) # boxes
if square: if square:
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square

@ -223,7 +223,7 @@ def make_anchors(feats, strides, grid_cell_offset=0.5):
def dist2bbox(distance, anchor_points, xywh=True, dim=-1): def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
"""Transform distance(ltrb) to box(xywh or xyxy).""" """Transform distance(ltrb) to box(xywh or xyxy)."""
lt, rb = torch.split(distance, 2, dim) lt, rb = distance.chunk(2, dim)
x1y1 = anchor_points - lt x1y1 = anchor_points - lt
x2y2 = anchor_points + rb x2y2 = anchor_points + rb
if xywh: if xywh:
@ -235,5 +235,5 @@ def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
def bbox2dist(anchor_points, bbox, reg_max): def bbox2dist(anchor_points, bbox, reg_max):
"""Transform bbox(xyxy) to dist(ltrb).""" """Transform bbox(xyxy) to dist(ltrb)."""
x1y1, x2y2 = torch.split(bbox, 2, -1) x1y1, x2y2 = bbox.chunk(2, -1)
return torch.cat((anchor_points - x1y1, x2y2 - anchor_points), -1).clamp(0, reg_max - 0.01) # dist (lt, rb) return torch.cat((anchor_points - x1y1, x2y2 - anchor_points), -1).clamp(0, reg_max - 0.01) # dist (lt, rb)

@ -22,7 +22,7 @@ class ClassificationPredictor(BasePredictor):
results = [] results = []
for i, pred in enumerate(preds): for i, pred in enumerate(preds):
orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img
results.append(Results(probs=pred.softmax(0), orig_img=orig_img, names=self.model.names)) results.append(Results(probs=pred, orig_img=orig_img, names=self.model.names))
return results return results

@ -30,6 +30,9 @@ class ClassificationValidator(BaseValidator):
self.pred.append(preds.argsort(1, descending=True)[:, :5]) self.pred.append(preds.argsort(1, descending=True)[:, :5])
self.targets.append(batch['cls']) self.targets.append(batch['cls'])
def finalize_metrics(self, *args, **kwargs):
self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
def get_stats(self): def get_stats(self):
self.metrics.process(self.targets, self.pred) self.metrics.process(self.targets, self.pred)
return self.metrics.results_dict return self.metrics.results_dict

@ -111,6 +111,9 @@ class DetectionValidator(BaseValidator):
# if self.args.save_txt: # if self.args.save_txt:
# save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt') # save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
def finalize_metrics(self, *args, **kwargs):
self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
def get_stats(self): def get_stats(self):
stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)] # to numpy stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)] # to numpy
if len(stats) and stats[0].any(): if len(stats) and stats[0].any():

@ -1,6 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license # Ultralytics YOLO 🚀, GPL-3.0 license
import os
from multiprocessing.pool import ThreadPool from multiprocessing.pool import ThreadPool
from pathlib import Path from pathlib import Path
@ -10,7 +9,7 @@ import torch.nn.functional as F
from ultralytics.yolo.utils import DEFAULT_CFG, NUM_THREADS, ops from ultralytics.yolo.utils import DEFAULT_CFG, NUM_THREADS, ops
from ultralytics.yolo.utils.checks import check_requirements from ultralytics.yolo.utils.checks import check_requirements
from ultralytics.yolo.utils.metrics import ConfusionMatrix, SegmentMetrics, box_iou, mask_iou from ultralytics.yolo.utils.metrics import SegmentMetrics, box_iou, mask_iou
from ultralytics.yolo.utils.plotting import output_to_target, plot_images from ultralytics.yolo.utils.plotting import output_to_target, plot_images
from ultralytics.yolo.v8.detect import DetectionValidator from ultralytics.yolo.v8.detect import DetectionValidator
@ -120,6 +119,9 @@ class SegmentationValidator(DetectionValidator):
# if self.args.save_txt: # if self.args.save_txt:
# save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt') # save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
def finalize_metrics(self, *args, **kwargs):
self.metrics.speed = dict(zip(self.metrics.speed.keys(), self.speed))
def _process_batch(self, detections, labels, pred_masks=None, gt_masks=None, overlap=False, masks=False): def _process_batch(self, detections, labels, pred_masks=None, gt_masks=None, overlap=False, masks=False):
""" """
Return correct prediction matrix Return correct prediction matrix

Loading…
Cancel
Save