YOLOv5 v6.0 compatibility update (#1857)
* Initial commit * Initial commit * Cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix precommit errors * Remove TF builds from CI * export last.pt * Created using Colaboratory * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
parent
1be31704c9
commit
7eb23e3c1d
@ -8,17 +8,21 @@ coco
|
|||||||
storage.googleapis.com
|
storage.googleapis.com
|
||||||
|
|
||||||
data/samples/*
|
data/samples/*
|
||||||
**/results*.txt
|
**/results*.csv
|
||||||
*.jpg
|
*.jpg
|
||||||
|
|
||||||
# Neural Network weights -----------------------------------------------------------------------------------------------
|
# Neural Network weights -----------------------------------------------------------------------------------------------
|
||||||
**/*.weights
|
|
||||||
**/*.pt
|
**/*.pt
|
||||||
**/*.pth
|
**/*.pth
|
||||||
**/*.onnx
|
**/*.onnx
|
||||||
**/*.mlmodel
|
**/*.mlmodel
|
||||||
**/*.torchscript
|
**/*.torchscript
|
||||||
|
**/*.torchscript.pt
|
||||||
|
**/*.tflite
|
||||||
|
**/*.h5
|
||||||
|
**/*.pb
|
||||||
|
*_saved_model/
|
||||||
|
*_web_model/
|
||||||
|
|
||||||
# Below Copied From .gitignore -----------------------------------------------------------------------------------------
|
# Below Copied From .gitignore -----------------------------------------------------------------------------------------
|
||||||
# Below Copied From .gitignore -----------------------------------------------------------------------------------------
|
# Below Copied From .gitignore -----------------------------------------------------------------------------------------
|
||||||
|
|||||||
55
.github/ISSUE_TEMPLATE/bug-report.md
vendored
55
.github/ISSUE_TEMPLATE/bug-report.md
vendored
@ -1,55 +0,0 @@
|
|||||||
---
|
|
||||||
name: "🐛 Bug report"
|
|
||||||
about: Create a report to help us improve
|
|
||||||
title: ''
|
|
||||||
labels: bug
|
|
||||||
assignees: ''
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
Before submitting a bug report, please be aware that your issue **must be reproducible** with all of the following, otherwise it is non-actionable, and we can not help you:
|
|
||||||
- **Current repo**: run `git fetch && git status -uno` to check and `git pull` to update repo
|
|
||||||
- **Common dataset**: coco.yaml or coco128.yaml
|
|
||||||
- **Common environment**: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov3#environments
|
|
||||||
|
|
||||||
If this is a custom dataset/training question you **must include** your `train*.jpg`, `test*.jpg` and `results.png` figures, or we can not help you. You can generate these with `utils.plot_results()`.
|
|
||||||
|
|
||||||
|
|
||||||
## 🐛 Bug
|
|
||||||
A clear and concise description of what the bug is.
|
|
||||||
|
|
||||||
|
|
||||||
## To Reproduce (REQUIRED)
|
|
||||||
|
|
||||||
Input:
|
|
||||||
```
|
|
||||||
import torch
|
|
||||||
|
|
||||||
a = torch.tensor([5])
|
|
||||||
c = a / 0
|
|
||||||
```
|
|
||||||
|
|
||||||
Output:
|
|
||||||
```
|
|
||||||
Traceback (most recent call last):
|
|
||||||
File "/Users/glennjocher/opt/anaconda3/envs/env1/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
|
|
||||||
exec(code_obj, self.user_global_ns, self.user_ns)
|
|
||||||
File "<ipython-input-5-be04c762b799>", line 5, in <module>
|
|
||||||
c = a / 0
|
|
||||||
RuntimeError: ZeroDivisionError
|
|
||||||
```
|
|
||||||
|
|
||||||
|
|
||||||
## Expected behavior
|
|
||||||
A clear and concise description of what you expected to happen.
|
|
||||||
|
|
||||||
|
|
||||||
## Environment
|
|
||||||
If applicable, add screenshots to help explain your problem.
|
|
||||||
|
|
||||||
- OS: [e.g. Ubuntu]
|
|
||||||
- GPU [e.g. 2080 Ti]
|
|
||||||
|
|
||||||
|
|
||||||
## Additional context
|
|
||||||
Add any other context about the problem here.
|
|
||||||
85
.github/ISSUE_TEMPLATE/bug-report.yml
vendored
Normal file
85
.github/ISSUE_TEMPLATE/bug-report.yml
vendored
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
name: 🐛 Bug Report
|
||||||
|
# title: " "
|
||||||
|
description: Problems with YOLOv3
|
||||||
|
labels: [bug, triage]
|
||||||
|
body:
|
||||||
|
- type: markdown
|
||||||
|
attributes:
|
||||||
|
value: |
|
||||||
|
Thank you for submitting a YOLOv3 🐛 Bug Report!
|
||||||
|
|
||||||
|
- type: checkboxes
|
||||||
|
attributes:
|
||||||
|
label: Search before asking
|
||||||
|
description: >
|
||||||
|
Please search the [issues](https://github.com/ultralytics/yolov3/issues) to see if a similar bug report already exists.
|
||||||
|
options:
|
||||||
|
- label: >
|
||||||
|
I have searched the YOLOv3 [issues](https://github.com/ultralytics/yolov3/issues) and found no similar bug report.
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: dropdown
|
||||||
|
attributes:
|
||||||
|
label: YOLOv3 Component
|
||||||
|
description: |
|
||||||
|
Please select the part of YOLOv3 where you found the bug.
|
||||||
|
multiple: true
|
||||||
|
options:
|
||||||
|
- "Training"
|
||||||
|
- "Validation"
|
||||||
|
- "Detection"
|
||||||
|
- "Export"
|
||||||
|
- "PyTorch Hub"
|
||||||
|
- "Multi-GPU"
|
||||||
|
- "Evolution"
|
||||||
|
- "Integrations"
|
||||||
|
- "Other"
|
||||||
|
validations:
|
||||||
|
required: false
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Bug
|
||||||
|
description: Provide console output with error messages and/or screenshots of the bug.
|
||||||
|
placeholder: |
|
||||||
|
💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
|
||||||
|
validations:
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Environment
|
||||||
|
description: Please specify the software and hardware you used to produce the bug.
|
||||||
|
placeholder: |
|
||||||
|
- YOLO: YOLOv3 🚀 v6.0-67-g60e42e1 torch 1.9.0+cu111 CUDA:0 (A100-SXM4-40GB, 40536MiB)
|
||||||
|
- OS: Ubuntu 20.04
|
||||||
|
- Python: 3.9.0
|
||||||
|
validations:
|
||||||
|
required: false
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Minimal Reproducible Example
|
||||||
|
description: >
|
||||||
|
When asking a question, people will be better able to provide help if you provide code that they can easily understand and use to **reproduce** the problem.
|
||||||
|
This is referred to by community members as creating a [minimal reproducible example](https://stackoverflow.com/help/minimal-reproducible-example).
|
||||||
|
placeholder: |
|
||||||
|
```
|
||||||
|
# Code to reproduce your issue here
|
||||||
|
```
|
||||||
|
validations:
|
||||||
|
required: false
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Additional
|
||||||
|
description: Anything else you would like to share?
|
||||||
|
|
||||||
|
- type: checkboxes
|
||||||
|
attributes:
|
||||||
|
label: Are you willing to submit a PR?
|
||||||
|
description: >
|
||||||
|
(Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov3/pulls) (PR) to help improve YOLOv3 for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||||
|
See the YOLOv3 [Contributing Guide](https://github.com/ultralytics/yolov3/blob/master/CONTRIBUTING.md) to get started.
|
||||||
|
options:
|
||||||
|
- label: Yes I'd like to help by submitting a PR!
|
||||||
8
.github/ISSUE_TEMPLATE/config.yml
vendored
Normal file
8
.github/ISSUE_TEMPLATE/config.yml
vendored
Normal file
@ -0,0 +1,8 @@
|
|||||||
|
blank_issues_enabled: true
|
||||||
|
contact_links:
|
||||||
|
- name: Slack
|
||||||
|
url: https://join.slack.com/t/ultralytics/shared_invite/zt-w29ei8bp-jczz7QYUmDtgo6r6KcMIAg
|
||||||
|
about: Ask on Ultralytics Slack Forum
|
||||||
|
- name: Stack Overflow
|
||||||
|
url: https://stackoverflow.com/search?q=YOLOv3
|
||||||
|
about: Ask on Stack Overflow with 'YOLOv3' tag
|
||||||
27
.github/ISSUE_TEMPLATE/feature-request.md
vendored
27
.github/ISSUE_TEMPLATE/feature-request.md
vendored
@ -1,27 +0,0 @@
|
|||||||
---
|
|
||||||
name: "🚀 Feature request"
|
|
||||||
about: Suggest an idea for this project
|
|
||||||
title: ''
|
|
||||||
labels: enhancement
|
|
||||||
assignees: ''
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 🚀 Feature
|
|
||||||
<!-- A clear and concise description of the feature proposal -->
|
|
||||||
|
|
||||||
## Motivation
|
|
||||||
|
|
||||||
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
|
|
||||||
|
|
||||||
## Pitch
|
|
||||||
|
|
||||||
<!-- A clear and concise description of what you want to happen. -->
|
|
||||||
|
|
||||||
## Alternatives
|
|
||||||
|
|
||||||
<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
|
|
||||||
|
|
||||||
## Additional context
|
|
||||||
|
|
||||||
<!-- Add any other context or screenshots about the feature request here. -->
|
|
||||||
50
.github/ISSUE_TEMPLATE/feature-request.yml
vendored
Normal file
50
.github/ISSUE_TEMPLATE/feature-request.yml
vendored
Normal file
@ -0,0 +1,50 @@
|
|||||||
|
name: 🚀 Feature Request
|
||||||
|
description: Suggest a YOLOv3 idea
|
||||||
|
# title: " "
|
||||||
|
labels: [enhancement]
|
||||||
|
body:
|
||||||
|
- type: markdown
|
||||||
|
attributes:
|
||||||
|
value: |
|
||||||
|
Thank you for submitting a YOLOv3 🚀 Feature Request!
|
||||||
|
|
||||||
|
- type: checkboxes
|
||||||
|
attributes:
|
||||||
|
label: Search before asking
|
||||||
|
description: >
|
||||||
|
Please search the [issues](https://github.com/ultralytics/yolov3/issues) to see if a similar feature request already exists.
|
||||||
|
options:
|
||||||
|
- label: >
|
||||||
|
I have searched the YOLOv3 [issues](https://github.com/ultralytics/yolov3/issues) and found no similar feature requests.
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Description
|
||||||
|
description: A short description of your feature.
|
||||||
|
placeholder: |
|
||||||
|
What new feature would you like to see in YOLOv3?
|
||||||
|
validations:
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Use case
|
||||||
|
description: |
|
||||||
|
Describe the use case of your feature request. It will help us understand and prioritize the feature request.
|
||||||
|
placeholder: |
|
||||||
|
How would this feature be used, and who would use it?
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Additional
|
||||||
|
description: Anything else you would like to share?
|
||||||
|
|
||||||
|
- type: checkboxes
|
||||||
|
attributes:
|
||||||
|
label: Are you willing to submit a PR?
|
||||||
|
description: >
|
||||||
|
(Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov3/pulls) (PR) to help improve YOLOv3 for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||||
|
See the YOLOv3 [Contributing Guide](https://github.com/ultralytics/yolov3/blob/master/CONTRIBUTING.md) to get started.
|
||||||
|
options:
|
||||||
|
- label: Yes I'd like to help by submitting a PR!
|
||||||
13
.github/ISSUE_TEMPLATE/question.md
vendored
13
.github/ISSUE_TEMPLATE/question.md
vendored
@ -1,13 +0,0 @@
|
|||||||
---
|
|
||||||
name: "❓Question"
|
|
||||||
about: Ask a general question
|
|
||||||
title: ''
|
|
||||||
labels: question
|
|
||||||
assignees: ''
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## ❔Question
|
|
||||||
|
|
||||||
|
|
||||||
## Additional context
|
|
||||||
33
.github/ISSUE_TEMPLATE/question.yml
vendored
Normal file
33
.github/ISSUE_TEMPLATE/question.yml
vendored
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
name: ❓ Question
|
||||||
|
description: Ask a YOLOv3 question
|
||||||
|
# title: " "
|
||||||
|
labels: [question]
|
||||||
|
body:
|
||||||
|
- type: markdown
|
||||||
|
attributes:
|
||||||
|
value: |
|
||||||
|
Thank you for asking a YOLOv3 ❓ Question!
|
||||||
|
|
||||||
|
- type: checkboxes
|
||||||
|
attributes:
|
||||||
|
label: Search before asking
|
||||||
|
description: >
|
||||||
|
Please search the [issues](https://github.com/ultralytics/yolov3/issues) and [discussions](https://github.com/ultralytics/yolov3/discussions) to see if a similar question already exists.
|
||||||
|
options:
|
||||||
|
- label: >
|
||||||
|
I have searched the YOLOv3 [issues](https://github.com/ultralytics/yolov3/issues) and [discussions](https://github.com/ultralytics/yolov3/discussions) and found no similar questions.
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Question
|
||||||
|
description: What is your question?
|
||||||
|
placeholder: |
|
||||||
|
💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
|
||||||
|
validations:
|
||||||
|
required: true
|
||||||
|
|
||||||
|
- type: textarea
|
||||||
|
attributes:
|
||||||
|
label: Additional
|
||||||
|
description: Anything else you would like to share?
|
||||||
31
.github/dependabot.yml
vendored
31
.github/dependabot.yml
vendored
@ -1,12 +1,23 @@
|
|||||||
version: 2
|
version: 2
|
||||||
updates:
|
updates:
|
||||||
- package-ecosystem: pip
|
- package-ecosystem: pip
|
||||||
directory: "/"
|
directory: "/"
|
||||||
schedule:
|
schedule:
|
||||||
interval: weekly
|
interval: weekly
|
||||||
time: "04:00"
|
time: "04:00"
|
||||||
open-pull-requests-limit: 10
|
open-pull-requests-limit: 10
|
||||||
reviewers:
|
reviewers:
|
||||||
- glenn-jocher
|
- glenn-jocher
|
||||||
labels:
|
labels:
|
||||||
- dependencies
|
- dependencies
|
||||||
|
|
||||||
|
- package-ecosystem: github-actions
|
||||||
|
directory: "/"
|
||||||
|
schedule:
|
||||||
|
interval: weekly
|
||||||
|
time: "04:00"
|
||||||
|
open-pull-requests-limit: 5
|
||||||
|
reviewers:
|
||||||
|
- glenn-jocher
|
||||||
|
labels:
|
||||||
|
- dependencies
|
||||||
|
|||||||
48
.github/workflows/ci-testing.yml
vendored
48
.github/workflows/ci-testing.yml
vendored
@ -1,6 +1,8 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
name: CI CPU testing
|
name: CI CPU testing
|
||||||
|
|
||||||
on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows
|
on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows
|
||||||
push:
|
push:
|
||||||
branches: [ master ]
|
branches: [ master ]
|
||||||
pull_request:
|
pull_request:
|
||||||
@ -16,9 +18,9 @@ jobs:
|
|||||||
strategy:
|
strategy:
|
||||||
fail-fast: false
|
fail-fast: false
|
||||||
matrix:
|
matrix:
|
||||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
os: [ ubuntu-latest, macos-latest, windows-latest ]
|
||||||
python-version: [3.8]
|
python-version: [ 3.9 ]
|
||||||
model: ['yolov3-tiny'] # models to test
|
model: [ 'yolov3-tiny' ] # models to test
|
||||||
|
|
||||||
# Timeout: https://stackoverflow.com/a/59076067/4521646
|
# Timeout: https://stackoverflow.com/a/59076067/4521646
|
||||||
timeout-minutes: 50
|
timeout-minutes: 50
|
||||||
@ -37,23 +39,27 @@ jobs:
|
|||||||
python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)"
|
python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)"
|
||||||
|
|
||||||
- name: Cache pip
|
- name: Cache pip
|
||||||
uses: actions/cache@v1
|
uses: actions/cache@v2.1.6
|
||||||
with:
|
with:
|
||||||
path: ${{ steps.pip-cache.outputs.dir }}
|
path: ${{ steps.pip-cache.outputs.dir }}
|
||||||
key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }}
|
key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }}
|
||||||
restore-keys: |
|
restore-keys: |
|
||||||
${{ runner.os }}-${{ matrix.python-version }}-pip-
|
${{ runner.os }}-${{ matrix.python-version }}-pip-
|
||||||
|
|
||||||
|
# Known Keras 2.7.0 issue: https://github.com/ultralytics/yolov5/pull/5486
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
run: |
|
run: |
|
||||||
python -m pip install --upgrade pip
|
python -m pip install --upgrade pip
|
||||||
pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html
|
pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html
|
||||||
pip install -q onnx
|
pip install -q onnx tensorflow-cpu keras==2.6.0 # wandb # extras
|
||||||
python --version
|
python --version
|
||||||
pip --version
|
pip --version
|
||||||
pip list
|
pip list
|
||||||
shell: bash
|
shell: bash
|
||||||
|
|
||||||
|
# - name: W&B login
|
||||||
|
# run: wandb login 345011b3fb26dc8337fd9b20e53857c1d403f2aa
|
||||||
|
|
||||||
- name: Download data
|
- name: Download data
|
||||||
run: |
|
run: |
|
||||||
# curl -L -o tmp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
|
# curl -L -o tmp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
|
||||||
@ -63,18 +69,26 @@ jobs:
|
|||||||
- name: Tests workflow
|
- name: Tests workflow
|
||||||
run: |
|
run: |
|
||||||
# export PYTHONPATH="$PWD" # to run '$ python *.py' files in subdirectories
|
# export PYTHONPATH="$PWD" # to run '$ python *.py' files in subdirectories
|
||||||
di=cpu # inference devices # define device
|
di=cpu # device
|
||||||
|
|
||||||
# train
|
# Train
|
||||||
python train.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di
|
python train.py --img 64 --batch 32 --weights ${{ matrix.model }}.pt --cfg ${{ matrix.model }}.yaml --epochs 1 --device $di
|
||||||
# detect
|
# Val
|
||||||
python detect.py --weights weights/${{ matrix.model }}.pt --device $di
|
python val.py --img 64 --batch 32 --weights ${{ matrix.model }}.pt --device $di
|
||||||
|
python val.py --img 64 --batch 32 --weights runs/train/exp/weights/last.pt --device $di
|
||||||
|
# Detect
|
||||||
|
python detect.py --weights ${{ matrix.model }}.pt --device $di
|
||||||
python detect.py --weights runs/train/exp/weights/last.pt --device $di
|
python detect.py --weights runs/train/exp/weights/last.pt --device $di
|
||||||
# test
|
|
||||||
python test.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --device $di
|
|
||||||
python test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
|
|
||||||
|
|
||||||
python hubconf.py # hub
|
python hubconf.py # hub
|
||||||
python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect
|
# Export
|
||||||
python models/export.py --img 128 --batch 1 --weights weights/${{ matrix.model }}.pt # export
|
python models/yolo.py --cfg ${{ matrix.model }}.yaml # build PyTorch model
|
||||||
|
# python models/tf.py --weights ${{ matrix.model }}.pt # build TensorFlow model (YOLOv3 not supported)
|
||||||
|
python export.py --img 64 --batch 1 --weights runs/train/exp/weights/last.pt --include torchscript onnx # export
|
||||||
|
# Python
|
||||||
|
python - <<EOF
|
||||||
|
import torch
|
||||||
|
# Known issue, urllib.error.HTTPError: HTTP Error 403: rate limit exceeded, will be resolved in torch==1.10.0
|
||||||
|
# model = torch.hub.load('ultralytics/yolov3', 'custom', path='runs/train/exp/weights/last.pt')
|
||||||
|
EOF
|
||||||
|
|
||||||
shell: bash
|
shell: bash
|
||||||
|
|||||||
52
.github/workflows/codeql-analysis.yml
vendored
52
.github/workflows/codeql-analysis.yml
vendored
@ -15,40 +15,40 @@ jobs:
|
|||||||
strategy:
|
strategy:
|
||||||
fail-fast: false
|
fail-fast: false
|
||||||
matrix:
|
matrix:
|
||||||
language: [ 'python' ]
|
language: ['python']
|
||||||
# CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ]
|
# CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ]
|
||||||
# Learn more:
|
# Learn more:
|
||||||
# https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed
|
# https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed
|
||||||
|
|
||||||
steps:
|
steps:
|
||||||
- name: Checkout repository
|
- name: Checkout repository
|
||||||
uses: actions/checkout@v2
|
uses: actions/checkout@v2
|
||||||
|
|
||||||
# Initializes the CodeQL tools for scanning.
|
# Initializes the CodeQL tools for scanning.
|
||||||
- name: Initialize CodeQL
|
- name: Initialize CodeQL
|
||||||
uses: github/codeql-action/init@v1
|
uses: github/codeql-action/init@v1
|
||||||
with:
|
with:
|
||||||
languages: ${{ matrix.language }}
|
languages: ${{ matrix.language }}
|
||||||
# If you wish to specify custom queries, you can do so here or in a config file.
|
# If you wish to specify custom queries, you can do so here or in a config file.
|
||||||
# By default, queries listed here will override any specified in a config file.
|
# By default, queries listed here will override any specified in a config file.
|
||||||
# Prefix the list here with "+" to use these queries and those in the config file.
|
# Prefix the list here with "+" to use these queries and those in the config file.
|
||||||
# queries: ./path/to/local/query, your-org/your-repo/queries@main
|
# queries: ./path/to/local/query, your-org/your-repo/queries@main
|
||||||
|
|
||||||
# Autobuild attempts to build any compiled languages (C/C++, C#, or Java).
|
# Autobuild attempts to build any compiled languages (C/C++, C#, or Java).
|
||||||
# If this step fails, then you should remove it and run the build manually (see below)
|
# If this step fails, then you should remove it and run the build manually (see below)
|
||||||
- name: Autobuild
|
- name: Autobuild
|
||||||
uses: github/codeql-action/autobuild@v1
|
uses: github/codeql-action/autobuild@v1
|
||||||
|
|
||||||
# ℹ️ Command-line programs to run using the OS shell.
|
# ℹ️ Command-line programs to run using the OS shell.
|
||||||
# 📚 https://git.io/JvXDl
|
# 📚 https://git.io/JvXDl
|
||||||
|
|
||||||
# ✏️ If the Autobuild fails above, remove it and uncomment the following three lines
|
# ✏️ If the Autobuild fails above, remove it and uncomment the following three lines
|
||||||
# and modify them (or add more) to build your code if your project
|
# and modify them (or add more) to build your code if your project
|
||||||
# uses a compiled language
|
# uses a compiled language
|
||||||
|
|
||||||
#- run: |
|
#- run: |
|
||||||
# make bootstrap
|
# make bootstrap
|
||||||
# make release
|
# make release
|
||||||
|
|
||||||
- name: Perform CodeQL Analysis
|
- name: Perform CodeQL Analysis
|
||||||
uses: github/codeql-action/analyze@v1
|
uses: github/codeql-action/analyze@v1
|
||||||
|
|||||||
17
.github/workflows/greetings.yml
vendored
17
.github/workflows/greetings.yml
vendored
@ -1,3 +1,5 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
name: Greetings
|
name: Greetings
|
||||||
|
|
||||||
on: [pull_request_target, issues]
|
on: [pull_request_target, issues]
|
||||||
@ -11,12 +13,12 @@ jobs:
|
|||||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||||
pr-message: |
|
pr-message: |
|
||||||
👋 Hello @${{ github.actor }}, thank you for submitting a 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:
|
👋 Hello @${{ github.actor }}, thank you for submitting a 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:
|
||||||
- ✅ Verify your PR is **up-to-date with origin/master.** If your PR is behind origin/master update by running the following, replacing 'feature' with the name of your local branch:
|
- ✅ Verify your PR is **up-to-date with upstream/master.** If your PR is behind upstream/master an automatic [GitHub actions](https://github.com/ultralytics/yolov3/blob/master/.github/workflows/rebase.yml) rebase may be attempted by including the /rebase command in a comment body, or by running the following code, replacing 'feature' with the name of your local branch:
|
||||||
```bash
|
```bash
|
||||||
git remote add upstream https://github.com/ultralytics/yolov3.git
|
git remote add upstream https://github.com/ultralytics/yolov3.git
|
||||||
git fetch upstream
|
git fetch upstream
|
||||||
git checkout feature # <----- replace 'feature' with local branch name
|
git checkout feature # <----- replace 'feature' with local branch name
|
||||||
git rebase upstream/master
|
git merge upstream/master
|
||||||
git push -u origin -f
|
git push -u origin -f
|
||||||
```
|
```
|
||||||
- ✅ Verify all Continuous Integration (CI) **checks are passing**.
|
- ✅ Verify all Continuous Integration (CI) **checks are passing**.
|
||||||
@ -29,12 +31,14 @@ jobs:
|
|||||||
|
|
||||||
If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online [W&B logging](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data#visualize) if available.
|
If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online [W&B logging](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data#visualize) if available.
|
||||||
|
|
||||||
For business inquiries or professional support requests please visit https://www.ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.
|
For business inquiries or professional support requests please visit https://ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.
|
||||||
|
|
||||||
## Requirements
|
## Requirements
|
||||||
|
|
||||||
Python 3.8 or later with all [requirements.txt](https://github.com/ultralytics/yolov3/blob/master/requirements.txt) dependencies installed, including `torch>=1.7`. To install run:
|
[**Python>=3.6.0**](https://www.python.org/) with all [requirements.txt](https://github.com/ultralytics/yolov3/blob/master/requirements.txt) installed including [**PyTorch>=1.7**](https://pytorch.org/get-started/locally/). To get started:
|
||||||
```bash
|
```bash
|
||||||
|
$ git clone https://github.com/ultralytics/yolov3
|
||||||
|
$ cd yolov3
|
||||||
$ pip install -r requirements.txt
|
$ pip install -r requirements.txt
|
||||||
```
|
```
|
||||||
|
|
||||||
@ -50,7 +54,6 @@ jobs:
|
|||||||
|
|
||||||
## Status
|
## Status
|
||||||
|
|
||||||
|
<a href="https://github.com/ultralytics/yolov3/actions"><img src="https://github.com/ultralytics/yolov3/workflows/CI%20CPU%20testing/badge.svg" alt="CI CPU testing"></a>
|
||||||
|
|
||||||
If this badge is green, all [YOLOv3 GitHub Actions](https://github.com/ultralytics/yolov3/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv3 training ([train.py](https://github.com/ultralytics/yolov3/blob/master/train.py)), testing ([test.py](https://github.com/ultralytics/yolov3/blob/master/test.py)), inference ([detect.py](https://github.com/ultralytics/yolov3/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov3/blob/master/models/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.
|
|
||||||
|
|
||||||
|
If this badge is green, all [YOLOv3 GitHub Actions](https://github.com/ultralytics/yolov3/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv3 training ([train.py](https://github.com/ultralytics/yolov3/blob/master/train.py)), validation ([val.py](https://github.com/ultralytics/yolov3/blob/master/val.py)), inference ([detect.py](https://github.com/ultralytics/yolov3/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov3/blob/master/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.
|
||||||
|
|||||||
10
.github/workflows/rebase.yml
vendored
10
.github/workflows/rebase.yml
vendored
@ -1,10 +1,9 @@
|
|||||||
name: Automatic Rebase
|
|
||||||
# https://github.com/marketplace/actions/automatic-rebase
|
# https://github.com/marketplace/actions/automatic-rebase
|
||||||
|
|
||||||
|
name: Automatic Rebase
|
||||||
on:
|
on:
|
||||||
issue_comment:
|
issue_comment:
|
||||||
types: [created]
|
types: [created]
|
||||||
|
|
||||||
jobs:
|
jobs:
|
||||||
rebase:
|
rebase:
|
||||||
name: Rebase
|
name: Rebase
|
||||||
@ -14,8 +13,9 @@ jobs:
|
|||||||
- name: Checkout the latest code
|
- name: Checkout the latest code
|
||||||
uses: actions/checkout@v2
|
uses: actions/checkout@v2
|
||||||
with:
|
with:
|
||||||
fetch-depth: 0
|
token: ${{ secrets.ACTIONS_TOKEN }}
|
||||||
|
fetch-depth: 0 # otherwise, you will fail to push refs to dest repo
|
||||||
- name: Automatic Rebase
|
- name: Automatic Rebase
|
||||||
uses: cirrus-actions/rebase@1.3.1
|
uses: cirrus-actions/rebase@1.5
|
||||||
env:
|
env:
|
||||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
GITHUB_TOKEN: ${{ secrets.ACTIONS_TOKEN }}
|
||||||
|
|||||||
10
.github/workflows/stale.yml
vendored
10
.github/workflows/stale.yml
vendored
@ -1,3 +1,5 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
name: Close stale issues
|
name: Close stale issues
|
||||||
on:
|
on:
|
||||||
schedule:
|
schedule:
|
||||||
@ -7,19 +9,19 @@ jobs:
|
|||||||
stale:
|
stale:
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/stale@v3
|
- uses: actions/stale@v4
|
||||||
with:
|
with:
|
||||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||||
stale-issue-message: |
|
stale-issue-message: |
|
||||||
👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.
|
👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.
|
||||||
|
|
||||||
Access additional [YOLOv3](https://ultralytics.com/yolov5) 🚀 resources:
|
Access additional [YOLOv3](https://ultralytics.com/yolov3) 🚀 resources:
|
||||||
- **Wiki** – https://github.com/ultralytics/yolov3/wiki
|
- **Wiki** – https://github.com/ultralytics/yolov3/wiki
|
||||||
- **Tutorials** – https://github.com/ultralytics/yolov3#tutorials
|
- **Tutorials** – https://github.com/ultralytics/yolov3#tutorials
|
||||||
- **Docs** – https://docs.ultralytics.com
|
- **Docs** – https://docs.ultralytics.com
|
||||||
|
|
||||||
Access additional [Ultralytics](https://ultralytics.com) ⚡ resources:
|
Access additional [Ultralytics](https://ultralytics.com) ⚡ resources:
|
||||||
- **Ultralytics HUB** – https://ultralytics.com/pricing
|
- **Ultralytics HUB** – https://ultralytics.com/hub
|
||||||
- **Vision API** – https://ultralytics.com/yolov5
|
- **Vision API** – https://ultralytics.com/yolov5
|
||||||
- **About Us** – https://ultralytics.com/about
|
- **About Us** – https://ultralytics.com/about
|
||||||
- **Join Our Team** – https://ultralytics.com/work
|
- **Join Our Team** – https://ultralytics.com/work
|
||||||
@ -29,7 +31,7 @@ jobs:
|
|||||||
|
|
||||||
Thank you for your contributions to YOLOv3 🚀 and Vision AI ⭐!
|
Thank you for your contributions to YOLOv3 🚀 and Vision AI ⭐!
|
||||||
|
|
||||||
stale-pr-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions YOLOv3 🚀 and Vision AI ⭐.'
|
stale-pr-message: 'This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions YOLOv3 🚀 and Vision AI ⭐.'
|
||||||
days-before-stale: 30
|
days-before-stale: 30
|
||||||
days-before-close: 5
|
days-before-close: 5
|
||||||
exempt-issue-labels: 'documentation,tutorial'
|
exempt-issue-labels: 'documentation,tutorial'
|
||||||
|
|||||||
20
.gitignore
vendored
20
.gitignore
vendored
@ -19,26 +19,19 @@
|
|||||||
*.avi
|
*.avi
|
||||||
*.data
|
*.data
|
||||||
*.json
|
*.json
|
||||||
|
|
||||||
*.cfg
|
*.cfg
|
||||||
|
!setup.cfg
|
||||||
!cfg/yolov3*.cfg
|
!cfg/yolov3*.cfg
|
||||||
|
|
||||||
storage.googleapis.com
|
storage.googleapis.com
|
||||||
runs/*
|
runs/*
|
||||||
data/*
|
data/*
|
||||||
|
!data/hyps/*
|
||||||
!data/images/zidane.jpg
|
!data/images/zidane.jpg
|
||||||
!data/images/bus.jpg
|
!data/images/bus.jpg
|
||||||
!data/coco.names
|
|
||||||
!data/coco_paper.names
|
|
||||||
!data/coco.data
|
|
||||||
!data/coco_*.data
|
|
||||||
!data/coco_*.txt
|
|
||||||
!data/trainvalno5k.shapes
|
|
||||||
!data/*.sh
|
!data/*.sh
|
||||||
|
|
||||||
pycocotools/*
|
results*.csv
|
||||||
results*.txt
|
|
||||||
gcp_test*.sh
|
|
||||||
|
|
||||||
# Datasets -------------------------------------------------------------------------------------------------------------
|
# Datasets -------------------------------------------------------------------------------------------------------------
|
||||||
coco/
|
coco/
|
||||||
@ -53,9 +46,14 @@ VOC/
|
|||||||
# Neural Network weights -----------------------------------------------------------------------------------------------
|
# Neural Network weights -----------------------------------------------------------------------------------------------
|
||||||
*.weights
|
*.weights
|
||||||
*.pt
|
*.pt
|
||||||
|
*.pb
|
||||||
*.onnx
|
*.onnx
|
||||||
*.mlmodel
|
*.mlmodel
|
||||||
*.torchscript
|
*.torchscript
|
||||||
|
*.tflite
|
||||||
|
*.h5
|
||||||
|
*_saved_model/
|
||||||
|
*_web_model/
|
||||||
darknet53.conv.74
|
darknet53.conv.74
|
||||||
yolov3-tiny.conv.15
|
yolov3-tiny.conv.15
|
||||||
|
|
||||||
@ -84,7 +82,7 @@ sdist/
|
|||||||
var/
|
var/
|
||||||
wheels/
|
wheels/
|
||||||
*.egg-info/
|
*.egg-info/
|
||||||
wandb/
|
/wandb/
|
||||||
.installed.cfg
|
.installed.cfg
|
||||||
*.egg
|
*.egg
|
||||||
|
|
||||||
|
|||||||
66
.pre-commit-config.yaml
Normal file
66
.pre-commit-config.yaml
Normal file
@ -0,0 +1,66 @@
|
|||||||
|
# Define hooks for code formations
|
||||||
|
# Will be applied on any updated commit files if a user has installed and linked commit hook
|
||||||
|
|
||||||
|
default_language_version:
|
||||||
|
python: python3.8
|
||||||
|
|
||||||
|
# Define bot property if installed via https://github.com/marketplace/pre-commit-ci
|
||||||
|
ci:
|
||||||
|
autofix_prs: true
|
||||||
|
autoupdate_commit_msg: '[pre-commit.ci] pre-commit suggestions'
|
||||||
|
autoupdate_schedule: quarterly
|
||||||
|
# submodules: true
|
||||||
|
|
||||||
|
repos:
|
||||||
|
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||||
|
rev: v4.0.1
|
||||||
|
hooks:
|
||||||
|
- id: end-of-file-fixer
|
||||||
|
- id: trailing-whitespace
|
||||||
|
- id: check-case-conflict
|
||||||
|
- id: check-yaml
|
||||||
|
- id: check-toml
|
||||||
|
- id: pretty-format-json
|
||||||
|
- id: check-docstring-first
|
||||||
|
|
||||||
|
- repo: https://github.com/asottile/pyupgrade
|
||||||
|
rev: v2.23.1
|
||||||
|
hooks:
|
||||||
|
- id: pyupgrade
|
||||||
|
args: [--py36-plus]
|
||||||
|
name: Upgrade code
|
||||||
|
|
||||||
|
- repo: https://github.com/PyCQA/isort
|
||||||
|
rev: 5.9.3
|
||||||
|
hooks:
|
||||||
|
- id: isort
|
||||||
|
name: Sort imports
|
||||||
|
|
||||||
|
# TODO
|
||||||
|
#- repo: https://github.com/pre-commit/mirrors-yapf
|
||||||
|
# rev: v0.31.0
|
||||||
|
# hooks:
|
||||||
|
# - id: yapf
|
||||||
|
# name: formatting
|
||||||
|
|
||||||
|
# TODO
|
||||||
|
#- repo: https://github.com/executablebooks/mdformat
|
||||||
|
# rev: 0.7.7
|
||||||
|
# hooks:
|
||||||
|
# - id: mdformat
|
||||||
|
# additional_dependencies:
|
||||||
|
# - mdformat-gfm
|
||||||
|
# - mdformat-black
|
||||||
|
# - mdformat_frontmatter
|
||||||
|
|
||||||
|
# TODO
|
||||||
|
#- repo: https://github.com/asottile/yesqa
|
||||||
|
# rev: v1.2.3
|
||||||
|
# hooks:
|
||||||
|
# - id: yesqa
|
||||||
|
|
||||||
|
- repo: https://github.com/PyCQA/flake8
|
||||||
|
rev: 3.9.2
|
||||||
|
hooks:
|
||||||
|
- id: flake8
|
||||||
|
name: PEP8
|
||||||
94
CONTRIBUTING.md
Normal file
94
CONTRIBUTING.md
Normal file
@ -0,0 +1,94 @@
|
|||||||
|
## Contributing to YOLOv3 🚀
|
||||||
|
|
||||||
|
We love your input! We want to make contributing to YOLOv3 as easy and transparent as possible, whether it's:
|
||||||
|
|
||||||
|
- Reporting a bug
|
||||||
|
- Discussing the current state of the code
|
||||||
|
- Submitting a fix
|
||||||
|
- Proposing a new feature
|
||||||
|
- Becoming a maintainer
|
||||||
|
|
||||||
|
YOLOv3 works so well due to our combined community effort, and for every small improvement you contribute you will be
|
||||||
|
helping push the frontiers of what's possible in AI 😃!
|
||||||
|
|
||||||
|
## Submitting a Pull Request (PR) 🛠️
|
||||||
|
|
||||||
|
Submitting a PR is easy! This example shows how to submit a PR for updating `requirements.txt` in 4 steps:
|
||||||
|
|
||||||
|
### 1. Select File to Update
|
||||||
|
|
||||||
|
Select `requirements.txt` to update by clicking on it in GitHub.
|
||||||
|
<p align="center"><img width="800" alt="PR_step1" src="https://user-images.githubusercontent.com/26833433/122260847-08be2600-ced4-11eb-828b-8287ace4136c.png"></p>
|
||||||
|
|
||||||
|
### 2. Click 'Edit this file'
|
||||||
|
|
||||||
|
Button is in top-right corner.
|
||||||
|
<p align="center"><img width="800" alt="PR_step2" src="https://user-images.githubusercontent.com/26833433/122260844-06f46280-ced4-11eb-9eec-b8a24be519ca.png"></p>
|
||||||
|
|
||||||
|
### 3. Make Changes
|
||||||
|
|
||||||
|
Change `matplotlib` version from `3.2.2` to `3.3`.
|
||||||
|
<p align="center"><img width="800" alt="PR_step3" src="https://user-images.githubusercontent.com/26833433/122260853-0a87e980-ced4-11eb-9fd2-3650fb6e0842.png"></p>
|
||||||
|
|
||||||
|
### 4. Preview Changes and Submit PR
|
||||||
|
|
||||||
|
Click on the **Preview changes** tab to verify your updates. At the bottom of the screen select 'Create a **new branch**
|
||||||
|
for this commit', assign your branch a descriptive name such as `fix/matplotlib_version` and click the green **Propose
|
||||||
|
changes** button. All done, your PR is now submitted to YOLOv3 for review and approval 😃!
|
||||||
|
<p align="center"><img width="800" alt="PR_step4" src="https://user-images.githubusercontent.com/26833433/122260856-0b208000-ced4-11eb-8e8e-77b6151cbcc3.png"></p>
|
||||||
|
|
||||||
|
### PR recommendations
|
||||||
|
|
||||||
|
To allow your work to be integrated as seamlessly as possible, we advise you to:
|
||||||
|
|
||||||
|
- ✅ Verify your PR is **up-to-date with upstream/master.** If your PR is behind upstream/master an
|
||||||
|
automatic [GitHub actions](https://github.com/ultralytics/yolov3/blob/master/.github/workflows/rebase.yml) rebase may
|
||||||
|
be attempted by including the /rebase command in a comment body, or by running the following code, replacing 'feature'
|
||||||
|
with the name of your local branch:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git remote add upstream https://github.com/ultralytics/yolov3.git
|
||||||
|
git fetch upstream
|
||||||
|
git checkout feature # <----- replace 'feature' with local branch name
|
||||||
|
git merge upstream/master
|
||||||
|
git push -u origin -f
|
||||||
|
```
|
||||||
|
|
||||||
|
- ✅ Verify all Continuous Integration (CI) **checks are passing**.
|
||||||
|
- ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase
|
||||||
|
but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ — Bruce Lee
|
||||||
|
|
||||||
|
## Submitting a Bug Report 🐛
|
||||||
|
|
||||||
|
If you spot a problem with YOLOv3 please submit a Bug Report!
|
||||||
|
|
||||||
|
For us to start investigating a possible problem we need to be able to reproduce it ourselves first. We've created a few
|
||||||
|
short guidelines below to help users provide what we need in order to get started.
|
||||||
|
|
||||||
|
When asking a question, people will be better able to provide help if you provide **code** that they can easily
|
||||||
|
understand and use to **reproduce** the problem. This is referred to by community members as creating
|
||||||
|
a [minimum reproducible example](https://stackoverflow.com/help/minimal-reproducible-example). Your code that reproduces
|
||||||
|
the problem should be:
|
||||||
|
|
||||||
|
* ✅ **Minimal** – Use as little code as possible that still produces the same problem
|
||||||
|
* ✅ **Complete** – Provide **all** parts someone else needs to reproduce your problem in the question itself
|
||||||
|
* ✅ **Reproducible** – Test the code you're about to provide to make sure it reproduces the problem
|
||||||
|
|
||||||
|
In addition to the above requirements, for [Ultralytics](https://ultralytics.com/) to provide assistance your code
|
||||||
|
should be:
|
||||||
|
|
||||||
|
* ✅ **Current** – Verify that your code is up-to-date with current
|
||||||
|
GitHub [master](https://github.com/ultralytics/yolov3/tree/master), and if necessary `git pull` or `git clone` a new
|
||||||
|
copy to ensure your problem has not already been resolved by previous commits.
|
||||||
|
* ✅ **Unmodified** – Your problem must be reproducible without any modifications to the codebase in this
|
||||||
|
repository. [Ultralytics](https://ultralytics.com/) does not provide support for custom code ⚠️.
|
||||||
|
|
||||||
|
If you believe your problem meets all of the above criteria, please close this issue and raise a new one using the 🐛 **
|
||||||
|
Bug Report** [template](https://github.com/ultralytics/yolov3/issues/new/choose) and providing
|
||||||
|
a [minimum reproducible example](https://stackoverflow.com/help/minimal-reproducible-example) to help us better
|
||||||
|
understand and diagnose your problem.
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
By contributing, you agree that your contributions will be licensed under
|
||||||
|
the [GPL-3.0 license](https://choosealicense.com/licenses/gpl-3.0/)
|
||||||
29
Dockerfile
29
Dockerfile
@ -1,5 +1,7 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
|
# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
|
||||||
FROM nvcr.io/nvidia/pytorch:21.03-py3
|
FROM nvcr.io/nvidia/pytorch:21.10-py3
|
||||||
|
|
||||||
# Install linux packages
|
# Install linux packages
|
||||||
RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
|
RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
|
||||||
@ -8,7 +10,9 @@ RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
|
|||||||
COPY requirements.txt .
|
COPY requirements.txt .
|
||||||
RUN python -m pip install --upgrade pip
|
RUN python -m pip install --upgrade pip
|
||||||
RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof
|
RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof
|
||||||
RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook
|
RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook wandb>=0.12.2
|
||||||
|
RUN pip install --no-cache -U torch torchvision numpy Pillow
|
||||||
|
# RUN pip install --no-cache torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
|
||||||
|
|
||||||
# Create working directory
|
# Create working directory
|
||||||
RUN mkdir -p /usr/src/app
|
RUN mkdir -p /usr/src/app
|
||||||
@ -17,27 +21,29 @@ WORKDIR /usr/src/app
|
|||||||
# Copy contents
|
# Copy contents
|
||||||
COPY . /usr/src/app
|
COPY . /usr/src/app
|
||||||
|
|
||||||
|
# Downloads to user config dir
|
||||||
|
ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/
|
||||||
|
|
||||||
# Set environment variables
|
# Set environment variables
|
||||||
ENV HOME=/usr/src/app
|
# ENV HOME=/usr/src/app
|
||||||
|
|
||||||
|
|
||||||
# --------------------------------------------------- Extras Below ---------------------------------------------------
|
# Usage Examples -------------------------------------------------------------------------------------------------------
|
||||||
|
|
||||||
# Build and Push
|
# Build and Push
|
||||||
# t=ultralytics/yolov3:latest && sudo docker build -t $t . && sudo docker push $t
|
# t=ultralytics/yolov3:latest && sudo docker build -t $t . && sudo docker push $t
|
||||||
# for v in {300..303}; do t=ultralytics/coco:v$v && sudo docker build -t $t . && sudo docker push $t; done
|
|
||||||
|
|
||||||
# Pull and Run
|
# Pull and Run
|
||||||
# t=ultralytics/yolov3:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t
|
# t=ultralytics/yolov3:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t
|
||||||
|
|
||||||
# Pull and Run with local directory access
|
# Pull and Run with local directory access
|
||||||
# t=ultralytics/yolov3:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/coco:/usr/src/coco $t
|
# t=ultralytics/yolov3:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t
|
||||||
|
|
||||||
# Kill all
|
# Kill all
|
||||||
# sudo docker kill $(sudo docker ps -q)
|
# sudo docker kill $(sudo docker ps -q)
|
||||||
|
|
||||||
# Kill all image-based
|
# Kill all image-based
|
||||||
# sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest)
|
# sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov3:latest)
|
||||||
|
|
||||||
# Bash into running container
|
# Bash into running container
|
||||||
# sudo docker exec -it 5a9b5863d93d bash
|
# sudo docker exec -it 5a9b5863d93d bash
|
||||||
@ -45,8 +51,11 @@ ENV HOME=/usr/src/app
|
|||||||
# Bash into stopped container
|
# Bash into stopped container
|
||||||
# id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash
|
# id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash
|
||||||
|
|
||||||
# Send weights to GCP
|
|
||||||
# python -c "from utils.general import *; strip_optimizer('runs/train/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt
|
|
||||||
|
|
||||||
# Clean up
|
# Clean up
|
||||||
# docker system prune -a --volumes
|
# docker system prune -a --volumes
|
||||||
|
|
||||||
|
# Update Ubuntu drivers
|
||||||
|
# https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/
|
||||||
|
|
||||||
|
# DDP test
|
||||||
|
# python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3
|
||||||
|
|||||||
347
README.md
Executable file → Normal file
347
README.md
Executable file → Normal file
@ -1,76 +1,141 @@
|
|||||||
<a align="left" href="https://apps.apple.com/app/id1452689527" target="_blank">
|
<div align="center">
|
||||||
<img width="800" src="https://user-images.githubusercontent.com/26833433/99805965-8f2ca800-2b3d-11eb-8fad-13a96b222a23.jpg"></a>
|
<p>
|
||||||
 
|
<a align="left" href="https://ultralytics.com/yolov3" target="_blank">
|
||||||
|
<img width="850" src="https://user-images.githubusercontent.com/26833433/99805965-8f2ca800-2b3d-11eb-8fad-13a96b222a23.jpg"></a>
|
||||||
|
</p>
|
||||||
|
<br>
|
||||||
|
<div>
|
||||||
|
<a href="https://github.com/ultralytics/yolov3/actions"><img src="https://github.com/ultralytics/yolov3/workflows/CI%20CPU%20testing/badge.svg" alt="CI CPU testing"></a>
|
||||||
|
<a href="https://zenodo.org/badge/latestdoi/264818686"><img src="https://zenodo.org/badge/264818686.svg" alt="YOLOv3 Citation"></a>
|
||||||
|
<a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
|
||||||
|
<br>
|
||||||
|
<a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
|
||||||
|
<a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||||
|
<a href="https://join.slack.com/t/ultralytics/shared_invite/zt-w29ei8bp-jczz7QYUmDtgo6r6KcMIAg"><img src="https://img.shields.io/badge/Slack-Join_Forum-blue.svg?logo=slack" alt="Join Forum"></a>
|
||||||
|
</div>
|
||||||
|
<br>
|
||||||
|
<div align="center">
|
||||||
|
<a href="https://github.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-github.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
<img width="2%" />
|
||||||
|
<a href="https://www.linkedin.com/company/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-linkedin.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
<img width="2%" />
|
||||||
|
<a href="https://twitter.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-twitter.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
<img width="2%" />
|
||||||
|
<a href="https://youtube.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-youtube.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
<img width="2%" />
|
||||||
|
<a href="https://www.facebook.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-facebook.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
<img width="2%" />
|
||||||
|
<a href="https://www.instagram.com/ultralytics/">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-instagram.png" width="2%"/>
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
<a href="https://github.com/ultralytics/yolov3/actions"><img src="https://github.com/ultralytics/yolov3/workflows/CI%20CPU%20testing/badge.svg" alt="CI CPU testing"></a>
|
<br>
|
||||||
|
<p>
|
||||||
|
YOLOv3 🚀 is a family of object detection architectures and models pretrained on the COCO dataset, and represents <a href="https://ultralytics.com">Ultralytics</a>
|
||||||
|
open-source research into future vision AI methods, incorporating lessons learned and best practices evolved over thousands of hours of research and development.
|
||||||
|
</p>
|
||||||
|
|
||||||
This repository represents Ultralytics open-source research into future object detection methods, and incorporates lessons learned and best practices evolved over thousands of hours of training and evolution on anonymized client datasets. **All code and models are under active development, and are subject to modification or deletion without notice.** Use at your own risk.
|
<!--
|
||||||
|
<a align="center" href="https://ultralytics.com/yolov3" target="_blank">
|
||||||
|
<img width="800" src="https://github.com/ultralytics/yolov5/releases/download/v1.0/banner-api.png"></a>
|
||||||
|
-->
|
||||||
|
|
||||||
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/114424655-a0dc1e00-9bb8-11eb-9a2e-cbe21803f05c.png"></p>
|
</div>
|
||||||
<details>
|
|
||||||
<summary>YOLOv5-P5 640 Figure (click to expand)</summary>
|
|
||||||
|
|
||||||
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/114313219-f1d70e00-9af5-11eb-9973-52b1f98d321a.png"></p>
|
## <div align="center">Documentation</div>
|
||||||
</details>
|
|
||||||
<details>
|
|
||||||
<summary>Figure Notes (click to expand)</summary>
|
|
||||||
|
|
||||||
* GPU Speed measures end-to-end time per image averaged over 5000 COCO val2017 images using a V100 GPU with batch size 32, and includes image preprocessing, PyTorch FP16 inference, postprocessing and NMS.
|
See the [YOLOv3 Docs](https://docs.ultralytics.com) for full documentation on training, testing and deployment.
|
||||||
* EfficientDet data from [google/automl](https://github.com/google/automl) at batch size 8.
|
|
||||||
* **Reproduce** by `python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt yolov5l.pt`
|
|
||||||
</details>
|
|
||||||
|
|
||||||
|
## <div align="center">Quick Start Examples</div>
|
||||||
|
|
||||||
## Branch Notice
|
<details open>
|
||||||
|
<summary>Install</summary>
|
||||||
|
|
||||||
|
[**Python>=3.6.0**](https://www.python.org/) is required with all
|
||||||
|
[requirements.txt](https://github.com/ultralytics/yolov3/blob/master/requirements.txt) installed including
|
||||||
|
[**PyTorch>=1.7**](https://pytorch.org/get-started/locally/):
|
||||||
|
<!-- $ sudo apt update && apt install -y libgl1-mesa-glx libsm6 libxext6 libxrender-dev -->
|
||||||
|
|
||||||
The [ultralytics/yolov3](https://github.com/ultralytics/yolov3) repository is now divided into two branches:
|
|
||||||
* [Master branch](https://github.com/ultralytics/yolov3/tree/master): Forward-compatible with all [YOLOv5](https://github.com/ultralytics/yolov5) models and methods (**recommended** ✅).
|
|
||||||
```bash
|
|
||||||
$ git clone https://github.com/ultralytics/yolov3 # master branch (default)
|
|
||||||
```
|
|
||||||
* [Archive branch](https://github.com/ultralytics/yolov3/tree/archive): Backwards-compatible with original [darknet](https://pjreddie.com/darknet/) *.cfg models (**no longer maintained** ⚠️).
|
|
||||||
```bash
|
|
||||||
$ git clone https://github.com/ultralytics/yolov3 -b archive # archive branch
|
|
||||||
```
|
|
||||||
|
|
||||||
## Pretrained Checkpoints
|
|
||||||
|
|
||||||
[assets3]: https://github.com/ultralytics/yolov3/releases
|
|
||||||
[assets5]: https://github.com/ultralytics/yolov5/releases
|
|
||||||
|
|
||||||
Model |size<br><sup>(pixels) |mAP<sup>val<br>0.5:0.95 |mAP<sup>test<br>0.5:0.95 |mAP<sup>val<br>0.5 |Speed<br><sup>V100 (ms) | |params<br><sup>(M) |FLOPS<br><sup>640 (B)
|
|
||||||
--- |--- |--- |--- |--- |--- |---|--- |---
|
|
||||||
[YOLOv3-tiny][assets3] |640 |17.6 |17.6 |34.8 |**1.2** | |8.8 |13.2
|
|
||||||
[YOLOv3][assets3] |640 |43.3 |43.3 |63.0 |4.1 | |61.9 |156.3
|
|
||||||
[YOLOv3-SPP][assets3] |640 |44.3 |44.3 |64.6 |4.1 | |63.0 |157.1
|
|
||||||
| | | | | | || |
|
|
||||||
[YOLOv5l][assets5] |640 |**48.2** |**48.2** |**66.9** |3.7 | |47.0 |115.4
|
|
||||||
|
|
||||||
|
|
||||||
<details>
|
|
||||||
<summary>Table Notes (click to expand)</summary>
|
|
||||||
|
|
||||||
* AP<sup>test</sup> denotes COCO [test-dev2017](http://cocodataset.org/#upload) server results, all other AP results denote val2017 accuracy.
|
|
||||||
* AP values are for single-model single-scale unless otherwise noted. **Reproduce mAP** by `python test.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
|
|
||||||
* Speed<sub>GPU</sub> averaged over 5000 COCO val2017 images using a GCP [n1-standard-16](https://cloud.google.com/compute/docs/machine-types#n1_standard_machine_types) V100 instance, and includes FP16 inference, postprocessing and NMS. **Reproduce speed** by `python test.py --data coco.yaml --img 640 --conf 0.25 --iou 0.45`
|
|
||||||
* All checkpoints are trained to 300 epochs with default settings and hyperparameters (no autoaugmentation).
|
|
||||||
</details>
|
|
||||||
|
|
||||||
|
|
||||||
## Requirements
|
|
||||||
|
|
||||||
Python 3.8 or later with all [requirements.txt](https://github.com/ultralytics/yolov3/blob/master/requirements.txt) dependencies installed, including `torch>=1.7`. To install run:
|
|
||||||
```bash
|
```bash
|
||||||
|
$ git clone https://github.com/ultralytics/yolov3
|
||||||
|
$ cd yolov3
|
||||||
$ pip install -r requirements.txt
|
$ pip install -r requirements.txt
|
||||||
```
|
```
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
## Tutorials
|
<details open>
|
||||||
|
<summary>Inference</summary>
|
||||||
|
|
||||||
|
Inference with YOLOv3 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36). Models automatically download
|
||||||
|
from the [latest YOLOv3 release](https://github.com/ultralytics/yolov3/releases).
|
||||||
|
|
||||||
|
```python
|
||||||
|
import torch
|
||||||
|
|
||||||
|
# Model
|
||||||
|
model = torch.hub.load('ultralytics/yolov3', 'yolov3') # or yolov3-spp, yolov3-tiny, custom
|
||||||
|
|
||||||
|
# Images
|
||||||
|
img = 'https://ultralytics.com/images/zidane.jpg' # or file, Path, PIL, OpenCV, numpy, list
|
||||||
|
|
||||||
|
# Inference
|
||||||
|
results = model(img)
|
||||||
|
|
||||||
|
# Results
|
||||||
|
results.print() # or .show(), .save(), .crop(), .pandas(), etc.
|
||||||
|
```
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary>Inference with detect.py</summary>
|
||||||
|
|
||||||
|
`detect.py` runs inference on a variety of sources, downloading models automatically from
|
||||||
|
the [latest YOLOv3 release](https://github.com/ultralytics/yolov3/releases) and saving results to `runs/detect`.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
$ python detect.py --source 0 # webcam
|
||||||
|
img.jpg # image
|
||||||
|
vid.mp4 # video
|
||||||
|
path/ # directory
|
||||||
|
path/*.jpg # glob
|
||||||
|
'https://youtu.be/Zgi9g1ksQHc' # YouTube
|
||||||
|
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
||||||
|
```
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary>Training</summary>
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
<img width="800" src="https://user-images.githubusercontent.com/26833433/90222759-949d8800-ddc1-11ea-9fa1-1c97eed2b963.png">
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<details open>
|
||||||
|
<summary>Tutorials</summary>
|
||||||
|
|
||||||
* [Train Custom Data](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data) 🚀 RECOMMENDED
|
* [Train Custom Data](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data) 🚀 RECOMMENDED
|
||||||
* [Tips for Best Training Results](https://github.com/ultralytics/yolov5/wiki/Tips-for-Best-Training-Results) ☘️ RECOMMENDED
|
* [Tips for Best Training Results](https://github.com/ultralytics/yolov3/wiki/Tips-for-Best-Training-Results) ☘️
|
||||||
|
RECOMMENDED
|
||||||
* [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289) 🌟 NEW
|
* [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289) 🌟 NEW
|
||||||
* [Supervisely Ecosystem](https://github.com/ultralytics/yolov5/issues/2518) 🌟 NEW
|
* [Roboflow for Datasets, Labeling, and Active Learning](https://github.com/ultralytics/yolov5/issues/4975) 🌟 NEW
|
||||||
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
|
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
|
||||||
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) ⭐ NEW
|
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) ⭐ NEW
|
||||||
* [TorchScript, ONNX, CoreML Export](https://github.com/ultralytics/yolov5/issues/251) 🚀
|
* [TorchScript, ONNX, CoreML Export](https://github.com/ultralytics/yolov5/issues/251) 🚀
|
||||||
@ -81,80 +146,128 @@ $ pip install -r requirements.txt
|
|||||||
* [Transfer Learning with Frozen Layers](https://github.com/ultralytics/yolov5/issues/1314) ⭐ NEW
|
* [Transfer Learning with Frozen Layers](https://github.com/ultralytics/yolov5/issues/1314) ⭐ NEW
|
||||||
* [TensorRT Deployment](https://github.com/wang-xinyu/tensorrtx)
|
* [TensorRT Deployment](https://github.com/wang-xinyu/tensorrtx)
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
## Environments
|
## <div align="center">Environments</div>
|
||||||
|
|
||||||
YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
|
Get started in seconds with our verified environments. Click each icon below for details.
|
||||||
|
|
||||||
- **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
<div align="center">
|
||||||
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
|
<a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb">
|
||||||
- **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-colab-small.png" width="15%"/>
|
||||||
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
|
</a>
|
||||||
|
<a href="https://www.kaggle.com/ultralytics/yolov3">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-kaggle-small.png" width="15%"/>
|
||||||
|
</a>
|
||||||
|
<a href="https://hub.docker.com/r/ultralytics/yolov3">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-docker-small.png" width="15%"/>
|
||||||
|
</a>
|
||||||
|
<a href="https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-aws-small.png" width="15%"/>
|
||||||
|
</a>
|
||||||
|
<a href="https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gcp-small.png" width="15%"/>
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
## <div align="center">Integrations</div>
|
||||||
|
|
||||||
|
<div align="center">
|
||||||
|
<a href="https://wandb.ai/site?utm_campaign=repo_yolo_readme">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-wb-long.png" width="49%"/>
|
||||||
|
</a>
|
||||||
|
<a href="https://roboflow.com/?ref=ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-roboflow-long.png" width="49%"/>
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
|Weights and Biases|Roboflow ⭐ NEW|
|
||||||
|
|:-:|:-:|
|
||||||
|
|Automatically track and visualize all your YOLOv3 training runs in the cloud with [Weights & Biases](https://wandb.ai/site?utm_campaign=repo_yolo_readme)|Label and export your custom datasets directly to YOLOv3 for training with [Roboflow](https://roboflow.com/?ref=ultralytics) |
|
||||||
|
|
||||||
|
|
||||||
## Inference
|
## <div align="center">Why YOLOv5</div>
|
||||||
|
|
||||||
`detect.py` runs inference on a variety of sources, downloading models automatically from the [latest YOLOv3 release](https://github.com/ultralytics/yolov3/releases) and saving results to `runs/detect`.
|
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/136901921-abcfcd9d-f978-4942-9b97-0e3f202907df.png"></p>
|
||||||
```bash
|
<details>
|
||||||
$ python detect.py --source 0 # webcam
|
<summary>YOLOv3-P5 640 Figure (click to expand)</summary>
|
||||||
file.jpg # image
|
|
||||||
file.mp4 # video
|
|
||||||
path/ # directory
|
|
||||||
path/*.jpg # glob
|
|
||||||
'https://youtu.be/NUsoVlDFqZg' # YouTube video
|
|
||||||
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
|
||||||
```
|
|
||||||
|
|
||||||
To run inference on example images in `data/images`:
|
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/136763877-b174052b-c12f-48d2-8bc4-545e3853398e.png"></p>
|
||||||
```bash
|
</details>
|
||||||
$ python detect.py --source data/images --weights yolov3.pt --conf 0.25
|
<details>
|
||||||
```
|
<summary>Figure Notes (click to expand)</summary>
|
||||||
<img width="500" src="https://user-images.githubusercontent.com/26833433/100375993-06b37900-300f-11eb-8d2d-5fc7b22fbfbd.jpg">
|
|
||||||
|
|
||||||
### PyTorch Hub
|
* **COCO AP val** denotes mAP@0.5:0.95 metric measured on the 5000-image [COCO val2017](http://cocodataset.org) dataset over various inference sizes from 256 to 1536.
|
||||||
|
* **GPU Speed** measures average inference time per image on [COCO val2017](http://cocodataset.org) dataset using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) V100 instance at batch-size 32.
|
||||||
|
* **EfficientDet** data from [google/automl](https://github.com/google/automl) at batch size 8.
|
||||||
|
* **Reproduce** by `python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n6.pt yolov5s6.pt yolov5m6.pt yolov5l6.pt yolov5x6.pt`
|
||||||
|
</details>
|
||||||
|
|
||||||
To run **batched inference** with YOLOv3 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36):
|
### Pretrained Checkpoints
|
||||||
```python
|
|
||||||
import torch
|
|
||||||
|
|
||||||
# Model
|
[assets]: https://github.com/ultralytics/yolov5/releases
|
||||||
model = torch.hub.load('ultralytics/yolov3', 'yolov3') # or 'yolov3_spp', 'yolov3_tiny'
|
[TTA]: https://github.com/ultralytics/yolov5/issues/303
|
||||||
|
|
||||||
# Image
|
|Model |size<br><sup>(pixels) |mAP<sup>val<br>0.5:0.95 |mAP<sup>val<br>0.5 |Speed<br><sup>CPU b1<br>(ms) |Speed<br><sup>V100 b1<br>(ms) |Speed<br><sup>V100 b32<br>(ms) |params<br><sup>(M) |FLOPs<br><sup>@640 (B)
|
||||||
img = 'https://ultralytics.com/images/zidane.jpg'
|
|--- |--- |--- |--- |--- |--- |--- |--- |---
|
||||||
|
|[YOLOv5n][assets] |640 |28.4 |46.0 |**45** |**6.3**|**0.6**|**1.9**|**4.5**
|
||||||
|
|[YOLOv5s][assets] |640 |37.2 |56.0 |98 |6.4 |0.9 |7.2 |16.5
|
||||||
|
|[YOLOv5m][assets] |640 |45.2 |63.9 |224 |8.2 |1.7 |21.2 |49.0
|
||||||
|
|[YOLOv5l][assets] |640 |48.8 |67.2 |430 |10.1 |2.7 |46.5 |109.1
|
||||||
|
|[YOLOv5x][assets] |640 |50.7 |68.9 |766 |12.1 |4.8 |86.7 |205.7
|
||||||
|
| | | | | | | | |
|
||||||
|
|[YOLOv5n6][assets] |1280 |34.0 |50.7 |153 |8.1 |2.1 |3.2 |4.6
|
||||||
|
|[YOLOv5s6][assets] |1280 |44.5 |63.0 |385 |8.2 |3.6 |16.8 |12.6
|
||||||
|
|[YOLOv5m6][assets] |1280 |51.0 |69.0 |887 |11.1 |6.8 |35.7 |50.0
|
||||||
|
|[YOLOv5l6][assets] |1280 |53.6 |71.6 |1784 |15.8 |10.5 |76.8 |111.4
|
||||||
|
|[YOLOv5x6][assets]<br>+ [TTA][TTA]|1280<br>1536 |54.7<br>**55.4** |**72.4**<br>72.3 |3136<br>- |26.2<br>- |19.4<br>- |140.7<br>- |209.8<br>-
|
||||||
|
|
||||||
# Inference
|
<details>
|
||||||
results = model(img)
|
<summary>Table Notes (click to expand)</summary>
|
||||||
results.print() # or .show(), .save()
|
|
||||||
```
|
* All checkpoints are trained to 300 epochs with default settings and hyperparameters.
|
||||||
|
* **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.<br>Reproduce by `python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
|
||||||
|
* **Speed** averaged over COCO val images using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) instance. NMS times (~1 ms/img) not included.<br>Reproduce by `python val.py --data coco.yaml --img 640 --conf 0.25 --iou 0.45`
|
||||||
|
* **TTA** [Test Time Augmentation](https://github.com/ultralytics/yolov5/issues/303) includes reflection and scale augmentations.<br>Reproduce by `python val.py --data coco.yaml --img 1536 --iou 0.7 --augment`
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
## <div align="center">Contribute</div>
|
||||||
|
|
||||||
|
We love your input! We want to make contributing to YOLOv3 as easy and transparent as possible. Please see our [Contributing Guide](CONTRIBUTING.md) to get started, and fill out the [YOLOv3 Survey](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey) to send us feedback on your experiences. Thank you to all our contributors!
|
||||||
|
|
||||||
|
<a href="https://github.com/ultralytics/yolov3/graphs/contributors"><img src="https://opencollective.com/ultralytics/contributors.svg?width=990" /></a>
|
||||||
|
|
||||||
|
|
||||||
## Training
|
## <div align="center">Contact</div>
|
||||||
|
|
||||||
Run commands below to reproduce results on [COCO](https://github.com/ultralytics/yolov3/blob/master/data/scripts/get_coco.sh) dataset (dataset auto-downloads on first use). Training times for YOLOv3/YOLOv3-SPP/YOLOv3-tiny are 6/6/2 days on a single V100 (multi-GPU times faster). Use the largest `--batch-size` your GPU allows (batch sizes shown for 16 GB devices).
|
For YOLOv3 bugs and feature requests please visit [GitHub Issues](https://github.com/ultralytics/yolov3/issues). For business inquiries or
|
||||||
```bash
|
professional support requests please visit [https://ultralytics.com/contact](https://ultralytics.com/contact).
|
||||||
$ python train.py --data coco.yaml --cfg yolov3.yaml --weights '' --batch-size 24
|
|
||||||
yolov3-spp.yaml 24
|
|
||||||
yolov3-tiny.yaml 64
|
|
||||||
```
|
|
||||||
<img width="800" src="https://user-images.githubusercontent.com/26833433/100378028-af170c80-3012-11eb-8521-f0d2a8d021bc.png">
|
|
||||||
|
|
||||||
|
<br>
|
||||||
|
|
||||||
## Citation
|
<div align="center">
|
||||||
|
<a href="https://github.com/ultralytics">
|
||||||
[](https://zenodo.org/badge/latestdoi/146165888)
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-github.png" width="3%"/>
|
||||||
|
</a>
|
||||||
|
<img width="3%" />
|
||||||
## About Us
|
<a href="https://www.linkedin.com/company/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-linkedin.png" width="3%"/>
|
||||||
Ultralytics is a U.S.-based particle physics and AI startup with over 6 years of expertise supporting government, academic and business clients. We offer a wide range of vision AI services, spanning from simple expert advice up to delivery of fully customized, end-to-end production solutions, including:
|
</a>
|
||||||
- **Cloud-based AI** systems operating on **hundreds of HD video streams in realtime.**
|
<img width="3%" />
|
||||||
- **Edge AI** integrated into custom iOS and Android apps for realtime **30 FPS video inference.**
|
<a href="https://twitter.com/ultralytics">
|
||||||
- **Custom data training**, hyperparameter evolution, and model exportation to any destination.
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-twitter.png" width="3%"/>
|
||||||
|
</a>
|
||||||
For business inquiries and professional support requests please visit us at https://ultralytics.com.
|
<img width="3%" />
|
||||||
|
<a href="https://youtube.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-youtube.png" width="3%"/>
|
||||||
## Contact
|
</a>
|
||||||
|
<img width="3%" />
|
||||||
**Issues should be raised directly in the repository.** For business inquiries or professional support requests please visit https://ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.
|
<a href="https://www.facebook.com/ultralytics">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-facebook.png" width="3%"/>
|
||||||
|
</a>
|
||||||
|
<img width="3%" />
|
||||||
|
<a href="https://www.instagram.com/ultralytics/">
|
||||||
|
<img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-social-instagram.png" width="3%"/>
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|||||||
67
data/Argoverse.yaml
Normal file
67
data/Argoverse.yaml
Normal file
@ -0,0 +1,67 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
|
||||||
|
# Example usage: python train.py --data Argoverse.yaml
|
||||||
|
# parent
|
||||||
|
# ├── yolov3
|
||||||
|
# └── datasets
|
||||||
|
# └── Argoverse ← downloads here
|
||||||
|
|
||||||
|
|
||||||
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
|
path: ../datasets/Argoverse # dataset root dir
|
||||||
|
train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images
|
||||||
|
val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images
|
||||||
|
test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview
|
||||||
|
|
||||||
|
# Classes
|
||||||
|
nc: 8 # number of classes
|
||||||
|
names: ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign'] # class names
|
||||||
|
|
||||||
|
|
||||||
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
|
download: |
|
||||||
|
import json
|
||||||
|
|
||||||
|
from tqdm import tqdm
|
||||||
|
from utils.general import download, Path
|
||||||
|
|
||||||
|
|
||||||
|
def argoverse2yolo(set):
|
||||||
|
labels = {}
|
||||||
|
a = json.load(open(set, "rb"))
|
||||||
|
for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv3 format..."):
|
||||||
|
img_id = annot['image_id']
|
||||||
|
img_name = a['images'][img_id]['name']
|
||||||
|
img_label_name = img_name[:-3] + "txt"
|
||||||
|
|
||||||
|
cls = annot['category_id'] # instance class id
|
||||||
|
x_center, y_center, width, height = annot['bbox']
|
||||||
|
x_center = (x_center + width / 2) / 1920.0 # offset and scale
|
||||||
|
y_center = (y_center + height / 2) / 1200.0 # offset and scale
|
||||||
|
width /= 1920.0 # scale
|
||||||
|
height /= 1200.0 # scale
|
||||||
|
|
||||||
|
img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']]
|
||||||
|
if not img_dir.exists():
|
||||||
|
img_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
k = str(img_dir / img_label_name)
|
||||||
|
if k not in labels:
|
||||||
|
labels[k] = []
|
||||||
|
labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n")
|
||||||
|
|
||||||
|
for k in labels:
|
||||||
|
with open(k, "w") as f:
|
||||||
|
f.writelines(labels[k])
|
||||||
|
|
||||||
|
|
||||||
|
# Download
|
||||||
|
dir = Path('../datasets/Argoverse') # dataset root dir
|
||||||
|
urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip']
|
||||||
|
download(urls, dir=dir, delete=False)
|
||||||
|
|
||||||
|
# Convert
|
||||||
|
annotations_dir = 'Argoverse-HD/annotations/'
|
||||||
|
(dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images') # rename 'tracking' to 'images'
|
||||||
|
for d in "train.json", "val.json":
|
||||||
|
argoverse2yolo(dir / annotations_dir / d) # convert VisDrone annotations to YOLO labels
|
||||||
@ -1,43 +1,41 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Global Wheat 2020 dataset http://www.global-wheat.com/
|
# Global Wheat 2020 dataset http://www.global-wheat.com/
|
||||||
# Train command: python train.py --data GlobalWheat2020.yaml
|
# Example usage: python train.py --data GlobalWheat2020.yaml
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /datasets/GlobalWheat2020
|
# └── datasets
|
||||||
# /yolov3
|
# └── GlobalWheat2020 ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
train: # 3422 images
|
path: ../datasets/GlobalWheat2020 # dataset root dir
|
||||||
- ../datasets/GlobalWheat2020/images/arvalis_1
|
train: # train images (relative to 'path') 3422 images
|
||||||
- ../datasets/GlobalWheat2020/images/arvalis_2
|
- images/arvalis_1
|
||||||
- ../datasets/GlobalWheat2020/images/arvalis_3
|
- images/arvalis_2
|
||||||
- ../datasets/GlobalWheat2020/images/ethz_1
|
- images/arvalis_3
|
||||||
- ../datasets/GlobalWheat2020/images/rres_1
|
- images/ethz_1
|
||||||
- ../datasets/GlobalWheat2020/images/inrae_1
|
- images/rres_1
|
||||||
- ../datasets/GlobalWheat2020/images/usask_1
|
- images/inrae_1
|
||||||
|
- images/usask_1
|
||||||
|
val: # val images (relative to 'path') 748 images (WARNING: train set contains ethz_1)
|
||||||
|
- images/ethz_1
|
||||||
|
test: # test images (optional) 1276 images
|
||||||
|
- images/utokyo_1
|
||||||
|
- images/utokyo_2
|
||||||
|
- images/nau_1
|
||||||
|
- images/uq_1
|
||||||
|
|
||||||
val: # 748 images (WARNING: train set contains ethz_1)
|
# Classes
|
||||||
- ../datasets/GlobalWheat2020/images/ethz_1
|
nc: 1 # number of classes
|
||||||
|
names: ['wheat_head'] # class names
|
||||||
test: # 1276 images
|
|
||||||
- ../datasets/GlobalWheat2020/images/utokyo_1
|
|
||||||
- ../datasets/GlobalWheat2020/images/utokyo_2
|
|
||||||
- ../datasets/GlobalWheat2020/images/nau_1
|
|
||||||
- ../datasets/GlobalWheat2020/images/uq_1
|
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 1
|
|
||||||
|
|
||||||
# class names
|
|
||||||
names: [ 'wheat_head' ]
|
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
download: |
|
download: |
|
||||||
from utils.general import download, Path
|
from utils.general import download, Path
|
||||||
|
|
||||||
# Download
|
# Download
|
||||||
dir = Path('../datasets/GlobalWheat2020') # dataset directory
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
urls = ['https://zenodo.org/record/4298502/files/global-wheat-codalab-official.zip',
|
urls = ['https://zenodo.org/record/4298502/files/global-wheat-codalab-official.zip',
|
||||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/GlobalWheat2020_labels.zip']
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/GlobalWheat2020_labels.zip']
|
||||||
download(urls, dir=dir)
|
download(urls, dir=dir)
|
||||||
|
|||||||
@ -1,39 +1,39 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19
|
# SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19
|
||||||
# Train command: python train.py --data SKU-110K.yaml
|
# Example usage: python train.py --data SKU-110K.yaml
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /datasets/SKU-110K
|
# └── datasets
|
||||||
# /yolov3
|
# └── SKU-110K ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
train: ../datasets/SKU-110K/train.txt # 8219 images
|
path: ../datasets/SKU-110K # dataset root dir
|
||||||
val: ../datasets/SKU-110K/val.txt # 588 images
|
train: train.txt # train images (relative to 'path') 8219 images
|
||||||
test: ../datasets/SKU-110K/test.txt # 2936 images
|
val: val.txt # val images (relative to 'path') 588 images
|
||||||
|
test: test.txt # test images (optional) 2936 images
|
||||||
|
|
||||||
# number of classes
|
# Classes
|
||||||
nc: 1
|
nc: 1 # number of classes
|
||||||
|
names: ['object'] # class names
|
||||||
# class names
|
|
||||||
names: [ 'object' ]
|
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
download: |
|
download: |
|
||||||
import shutil
|
import shutil
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
from utils.general import np, pd, Path, download, xyxy2xywh
|
from utils.general import np, pd, Path, download, xyxy2xywh
|
||||||
|
|
||||||
# Download
|
# Download
|
||||||
datasets = Path('../datasets') # download directory
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
|
parent = Path(dir.parent) # download dir
|
||||||
urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz']
|
urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz']
|
||||||
download(urls, dir=datasets, delete=False)
|
download(urls, dir=parent, delete=False)
|
||||||
|
|
||||||
# Rename directories
|
# Rename directories
|
||||||
dir = (datasets / 'SKU-110K')
|
|
||||||
if dir.exists():
|
if dir.exists():
|
||||||
shutil.rmtree(dir)
|
shutil.rmtree(dir)
|
||||||
(datasets / 'SKU110K_fixed').rename(dir) # rename dir
|
(parent / 'SKU110K_fixed').rename(dir) # rename dir
|
||||||
(dir / 'labels').mkdir(parents=True, exist_ok=True) # create labels dir
|
(dir / 'labels').mkdir(parents=True, exist_ok=True) # create labels dir
|
||||||
|
|
||||||
# Convert labels
|
# Convert labels
|
||||||
|
|||||||
@ -1,24 +1,24 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset
|
# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset
|
||||||
# Train command: python train.py --data VisDrone.yaml
|
# Example usage: python train.py --data VisDrone.yaml
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /VisDrone
|
# └── datasets
|
||||||
# /yolov3
|
# └── VisDrone ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
train: ../VisDrone/VisDrone2019-DET-train/images # 6471 images
|
path: ../datasets/VisDrone # dataset root dir
|
||||||
val: ../VisDrone/VisDrone2019-DET-val/images # 548 images
|
train: VisDrone2019-DET-train/images # train images (relative to 'path') 6471 images
|
||||||
test: ../VisDrone/VisDrone2019-DET-test-dev/images # 1610 images
|
val: VisDrone2019-DET-val/images # val images (relative to 'path') 548 images
|
||||||
|
test: VisDrone2019-DET-test-dev/images # test images (optional) 1610 images
|
||||||
|
|
||||||
# number of classes
|
# Classes
|
||||||
nc: 10
|
nc: 10 # number of classes
|
||||||
|
names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']
|
||||||
# class names
|
|
||||||
names: [ 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]
|
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
download: |
|
download: |
|
||||||
from utils.general import download, os, Path
|
from utils.general import download, os, Path
|
||||||
|
|
||||||
@ -49,7 +49,7 @@ download: |
|
|||||||
|
|
||||||
|
|
||||||
# Download
|
# Download
|
||||||
dir = Path('../VisDrone') # dataset directory
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip',
|
urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip',
|
||||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
|
||||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip',
|
'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip',
|
||||||
|
|||||||
@ -1,21 +0,0 @@
|
|||||||
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
|
|
||||||
# Train command: python train.py --data argoverse_hd.yaml
|
|
||||||
# Default dataset location is next to YOLOv3:
|
|
||||||
# /parent_folder
|
|
||||||
# /argoverse
|
|
||||||
# /yolov3
|
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional)
|
|
||||||
download: bash data/scripts/get_argoverse_hd.sh
|
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
|
||||||
train: ../argoverse/Argoverse-1.1/images/train/ # 39384 images
|
|
||||||
val: ../argoverse/Argoverse-1.1/images/val/ # 15062 iamges
|
|
||||||
test: ../argoverse/Argoverse-1.1/images/test/ # Submit to: https://eval.ai/web/challenges/challenge-page/800/overview
|
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 8
|
|
||||||
|
|
||||||
# class names
|
|
||||||
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign' ]
|
|
||||||
@ -1,35 +1,44 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# COCO 2017 dataset http://cocodataset.org
|
# COCO 2017 dataset http://cocodataset.org
|
||||||
# Train command: python train.py --data coco.yaml
|
# Example usage: python train.py --data coco.yaml
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /coco
|
# └── datasets
|
||||||
# /yolov3
|
# └── coco ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional)
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
download: bash data/scripts/get_coco.sh
|
path: ../datasets/coco # dataset root dir
|
||||||
|
train: train2017.txt # train images (relative to 'path') 118287 images
|
||||||
|
val: val2017.txt # train images (relative to 'path') 5000 images
|
||||||
|
test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Classes
|
||||||
train: ../coco/train2017.txt # 118287 images
|
nc: 80 # number of classes
|
||||||
val: ../coco/val2017.txt # 5000 images
|
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||||
test: ../coco/test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
|
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||||
|
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||||
|
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||||
|
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||||
|
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||||
|
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||||
|
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||||
|
'hair drier', 'toothbrush'] # class names
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 80
|
|
||||||
|
|
||||||
# class names
|
# Download script/URL (optional)
|
||||||
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
download: |
|
||||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
from utils.general import download, Path
|
||||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
|
||||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
|
||||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
|
||||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
|
||||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
|
||||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
|
||||||
'hair drier', 'toothbrush' ]
|
|
||||||
|
|
||||||
# Print classes
|
# Download labels
|
||||||
# with open('data/coco.yaml') as f:
|
segments = False # segment or box labels
|
||||||
# d = yaml.safe_load(f) # dict
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
# for i, x in enumerate(d['names']):
|
url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
|
||||||
# print(i, x)
|
urls = [url + ('coco2017labels-segments.zip' if segments else 'coco2017labels.zip')] # labels
|
||||||
|
download(urls, dir=dir.parent)
|
||||||
|
|
||||||
|
# Download data
|
||||||
|
urls = ['http://images.cocodataset.org/zips/train2017.zip', # 19G, 118k images
|
||||||
|
'http://images.cocodataset.org/zips/val2017.zip', # 1G, 5k images
|
||||||
|
'http://images.cocodataset.org/zips/test2017.zip'] # 7G, 41k images (optional)
|
||||||
|
download(urls, dir=dir / 'images', threads=3)
|
||||||
|
|||||||
@ -1,28 +1,30 @@
|
|||||||
# COCO 2017 dataset http://cocodataset.org - first 128 training images
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Train command: python train.py --data coco128.yaml
|
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
|
||||||
# Default dataset location is next to YOLOv3:
|
# Example usage: python train.py --data coco128.yaml
|
||||||
# /parent_folder
|
# parent
|
||||||
# /coco128
|
# ├── yolov3
|
||||||
# /yolov3
|
# └── datasets
|
||||||
|
# └── coco128 ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional)
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
|
path: ../datasets/coco128 # dataset root dir
|
||||||
|
train: images/train2017 # train images (relative to 'path') 128 images
|
||||||
|
val: images/train2017 # val images (relative to 'path') 128 images
|
||||||
|
test: # test images (optional)
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Classes
|
||||||
train: ../coco128/images/train2017/ # 128 images
|
nc: 80 # number of classes
|
||||||
val: ../coco128/images/train2017/ # 128 images
|
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||||
|
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||||
|
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||||
|
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||||
|
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||||
|
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||||
|
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||||
|
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||||
|
'hair drier', 'toothbrush'] # class names
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 80
|
|
||||||
|
|
||||||
# class names
|
# Download script/URL (optional)
|
||||||
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
download: https://ultralytics.com/assets/coco128.zip
|
||||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
|
||||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
|
||||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
|
||||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
|
||||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
|
||||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
|
||||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
|
||||||
'hair drier', 'toothbrush' ]
|
|
||||||
|
|||||||
@ -1,38 +0,0 @@
|
|||||||
# Hyperparameters for VOC finetuning
|
|
||||||
# python train.py --batch 64 --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50
|
|
||||||
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
|
||||||
|
|
||||||
|
|
||||||
# Hyperparameter Evolution Results
|
|
||||||
# Generations: 306
|
|
||||||
# P R mAP.5 mAP.5:.95 box obj cls
|
|
||||||
# Metrics: 0.6 0.936 0.896 0.684 0.0115 0.00805 0.00146
|
|
||||||
|
|
||||||
lr0: 0.0032
|
|
||||||
lrf: 0.12
|
|
||||||
momentum: 0.843
|
|
||||||
weight_decay: 0.00036
|
|
||||||
warmup_epochs: 2.0
|
|
||||||
warmup_momentum: 0.5
|
|
||||||
warmup_bias_lr: 0.05
|
|
||||||
box: 0.0296
|
|
||||||
cls: 0.243
|
|
||||||
cls_pw: 0.631
|
|
||||||
obj: 0.301
|
|
||||||
obj_pw: 0.911
|
|
||||||
iou_t: 0.2
|
|
||||||
anchor_t: 2.91
|
|
||||||
# anchors: 3.63
|
|
||||||
fl_gamma: 0.0
|
|
||||||
hsv_h: 0.0138
|
|
||||||
hsv_s: 0.664
|
|
||||||
hsv_v: 0.464
|
|
||||||
degrees: 0.373
|
|
||||||
translate: 0.245
|
|
||||||
scale: 0.898
|
|
||||||
shear: 0.602
|
|
||||||
perspective: 0.0
|
|
||||||
flipud: 0.00856
|
|
||||||
fliplr: 0.5
|
|
||||||
mosaic: 1.0
|
|
||||||
mixup: 0.243
|
|
||||||
@ -1,28 +0,0 @@
|
|||||||
lr0: 0.00258
|
|
||||||
lrf: 0.17
|
|
||||||
momentum: 0.779
|
|
||||||
weight_decay: 0.00058
|
|
||||||
warmup_epochs: 1.33
|
|
||||||
warmup_momentum: 0.86
|
|
||||||
warmup_bias_lr: 0.0711
|
|
||||||
box: 0.0539
|
|
||||||
cls: 0.299
|
|
||||||
cls_pw: 0.825
|
|
||||||
obj: 0.632
|
|
||||||
obj_pw: 1.0
|
|
||||||
iou_t: 0.2
|
|
||||||
anchor_t: 3.44
|
|
||||||
anchors: 3.2
|
|
||||||
fl_gamma: 0.0
|
|
||||||
hsv_h: 0.0188
|
|
||||||
hsv_s: 0.704
|
|
||||||
hsv_v: 0.36
|
|
||||||
degrees: 0.0
|
|
||||||
translate: 0.0902
|
|
||||||
scale: 0.491
|
|
||||||
shear: 0.0
|
|
||||||
perspective: 0.0
|
|
||||||
flipud: 0.0
|
|
||||||
fliplr: 0.5
|
|
||||||
mosaic: 1.0
|
|
||||||
mixup: 0.0
|
|
||||||
34
data/hyps/hyp.scratch-high.yaml
Normal file
34
data/hyps/hyp.scratch-high.yaml
Normal file
@ -0,0 +1,34 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# Hyperparameters for high-augmentation COCO training from scratch
|
||||||
|
# python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
|
||||||
|
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
||||||
|
|
||||||
|
lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||||
|
lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf)
|
||||||
|
momentum: 0.937 # SGD momentum/Adam beta1
|
||||||
|
weight_decay: 0.0005 # optimizer weight decay 5e-4
|
||||||
|
warmup_epochs: 3.0 # warmup epochs (fractions ok)
|
||||||
|
warmup_momentum: 0.8 # warmup initial momentum
|
||||||
|
warmup_bias_lr: 0.1 # warmup initial bias lr
|
||||||
|
box: 0.05 # box loss gain
|
||||||
|
cls: 0.3 # cls loss gain
|
||||||
|
cls_pw: 1.0 # cls BCELoss positive_weight
|
||||||
|
obj: 0.7 # obj loss gain (scale with pixels)
|
||||||
|
obj_pw: 1.0 # obj BCELoss positive_weight
|
||||||
|
iou_t: 0.20 # IoU training threshold
|
||||||
|
anchor_t: 4.0 # anchor-multiple threshold
|
||||||
|
# anchors: 3 # anchors per output layer (0 to ignore)
|
||||||
|
fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
|
||||||
|
hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
|
||||||
|
hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
|
||||||
|
hsv_v: 0.4 # image HSV-Value augmentation (fraction)
|
||||||
|
degrees: 0.0 # image rotation (+/- deg)
|
||||||
|
translate: 0.1 # image translation (+/- fraction)
|
||||||
|
scale: 0.9 # image scale (+/- gain)
|
||||||
|
shear: 0.0 # image shear (+/- deg)
|
||||||
|
perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
|
||||||
|
flipud: 0.0 # image flip up-down (probability)
|
||||||
|
fliplr: 0.5 # image flip left-right (probability)
|
||||||
|
mosaic: 1.0 # image mosaic (probability)
|
||||||
|
mixup: 0.1 # image mixup (probability)
|
||||||
|
copy_paste: 0.1 # segment copy-paste (probability)
|
||||||
34
data/hyps/hyp.scratch-low.yaml
Normal file
34
data/hyps/hyp.scratch-low.yaml
Normal file
@ -0,0 +1,34 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# Hyperparameters for low-augmentation COCO training from scratch
|
||||||
|
# python train.py --batch 64 --cfg yolov5n6.yaml --weights '' --data coco.yaml --img 640 --epochs 300 --linear
|
||||||
|
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
||||||
|
|
||||||
|
lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||||
|
lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf)
|
||||||
|
momentum: 0.937 # SGD momentum/Adam beta1
|
||||||
|
weight_decay: 0.0005 # optimizer weight decay 5e-4
|
||||||
|
warmup_epochs: 3.0 # warmup epochs (fractions ok)
|
||||||
|
warmup_momentum: 0.8 # warmup initial momentum
|
||||||
|
warmup_bias_lr: 0.1 # warmup initial bias lr
|
||||||
|
box: 0.05 # box loss gain
|
||||||
|
cls: 0.5 # cls loss gain
|
||||||
|
cls_pw: 1.0 # cls BCELoss positive_weight
|
||||||
|
obj: 1.0 # obj loss gain (scale with pixels)
|
||||||
|
obj_pw: 1.0 # obj BCELoss positive_weight
|
||||||
|
iou_t: 0.20 # IoU training threshold
|
||||||
|
anchor_t: 4.0 # anchor-multiple threshold
|
||||||
|
# anchors: 3 # anchors per output layer (0 to ignore)
|
||||||
|
fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
|
||||||
|
hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
|
||||||
|
hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
|
||||||
|
hsv_v: 0.4 # image HSV-Value augmentation (fraction)
|
||||||
|
degrees: 0.0 # image rotation (+/- deg)
|
||||||
|
translate: 0.1 # image translation (+/- fraction)
|
||||||
|
scale: 0.5 # image scale (+/- gain)
|
||||||
|
shear: 0.0 # image shear (+/- deg)
|
||||||
|
perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
|
||||||
|
flipud: 0.0 # image flip up-down (probability)
|
||||||
|
fliplr: 0.5 # image flip left-right (probability)
|
||||||
|
mosaic: 1.0 # image mosaic (probability)
|
||||||
|
mixup: 0.0 # image mixup (probability)
|
||||||
|
copy_paste: 0.0 # segment copy-paste (probability)
|
||||||
34
data/hyps/hyp.scratch-med.yaml
Normal file
34
data/hyps/hyp.scratch-med.yaml
Normal file
@ -0,0 +1,34 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# Hyperparameters for medium-augmentation COCO training from scratch
|
||||||
|
# python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
|
||||||
|
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
||||||
|
|
||||||
|
lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||||
|
lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
|
||||||
|
momentum: 0.937 # SGD momentum/Adam beta1
|
||||||
|
weight_decay: 0.0005 # optimizer weight decay 5e-4
|
||||||
|
warmup_epochs: 3.0 # warmup epochs (fractions ok)
|
||||||
|
warmup_momentum: 0.8 # warmup initial momentum
|
||||||
|
warmup_bias_lr: 0.1 # warmup initial bias lr
|
||||||
|
box: 0.05 # box loss gain
|
||||||
|
cls: 0.3 # cls loss gain
|
||||||
|
cls_pw: 1.0 # cls BCELoss positive_weight
|
||||||
|
obj: 0.7 # obj loss gain (scale with pixels)
|
||||||
|
obj_pw: 1.0 # obj BCELoss positive_weight
|
||||||
|
iou_t: 0.20 # IoU training threshold
|
||||||
|
anchor_t: 4.0 # anchor-multiple threshold
|
||||||
|
# anchors: 3 # anchors per output layer (0 to ignore)
|
||||||
|
fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
|
||||||
|
hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
|
||||||
|
hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
|
||||||
|
hsv_v: 0.4 # image HSV-Value augmentation (fraction)
|
||||||
|
degrees: 0.0 # image rotation (+/- deg)
|
||||||
|
translate: 0.1 # image translation (+/- fraction)
|
||||||
|
scale: 0.9 # image scale (+/- gain)
|
||||||
|
shear: 0.0 # image shear (+/- deg)
|
||||||
|
perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
|
||||||
|
flipud: 0.0 # image flip up-down (probability)
|
||||||
|
fliplr: 0.5 # image flip left-right (probability)
|
||||||
|
mosaic: 1.0 # image mosaic (probability)
|
||||||
|
mixup: 0.1 # image mixup (probability)
|
||||||
|
copy_paste: 0.0 # segment copy-paste (probability)
|
||||||
@ -1,10 +1,10 @@
|
|||||||
|
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Hyperparameters for COCO training from scratch
|
# Hyperparameters for COCO training from scratch
|
||||||
# python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300
|
# python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300
|
||||||
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
|
||||||
|
|
||||||
|
|
||||||
lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||||
lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf)
|
lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
|
||||||
momentum: 0.937 # SGD momentum/Adam beta1
|
momentum: 0.937 # SGD momentum/Adam beta1
|
||||||
weight_decay: 0.0005 # optimizer weight decay 5e-4
|
weight_decay: 0.0005 # optimizer weight decay 5e-4
|
||||||
warmup_epochs: 3.0 # warmup epochs (fractions ok)
|
warmup_epochs: 3.0 # warmup epochs (fractions ok)
|
||||||
@ -31,3 +31,4 @@ flipud: 0.0 # image flip up-down (probability)
|
|||||||
fliplr: 0.5 # image flip left-right (probability)
|
fliplr: 0.5 # image flip left-right (probability)
|
||||||
mosaic: 1.0 # image mosaic (probability)
|
mosaic: 1.0 # image mosaic (probability)
|
||||||
mixup: 0.0 # image mixup (probability)
|
mixup: 0.0 # image mixup (probability)
|
||||||
|
copy_paste: 0.0 # segment copy-paste (probability)
|
||||||
@ -1,102 +1,112 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Objects365 dataset https://www.objects365.org/
|
# Objects365 dataset https://www.objects365.org/
|
||||||
# Train command: python train.py --data objects365.yaml
|
# Example usage: python train.py --data Objects365.yaml
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /datasets/objects365
|
# └── datasets
|
||||||
# /yolov3
|
# └── Objects365 ← downloads here
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
|
||||||
train: ../datasets/objects365/images/train # 1742289 images
|
|
||||||
val: ../datasets/objects365/images/val # 5570 images
|
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 365
|
|
||||||
|
|
||||||
# class names
|
|
||||||
names: [ 'Person', 'Sneakers', 'Chair', 'Other Shoes', 'Hat', 'Car', 'Lamp', 'Glasses', 'Bottle', 'Desk', 'Cup',
|
|
||||||
'Street Lights', 'Cabinet/shelf', 'Handbag/Satchel', 'Bracelet', 'Plate', 'Picture/Frame', 'Helmet', 'Book',
|
|
||||||
'Gloves', 'Storage box', 'Boat', 'Leather Shoes', 'Flower', 'Bench', 'Potted Plant', 'Bowl/Basin', 'Flag',
|
|
||||||
'Pillow', 'Boots', 'Vase', 'Microphone', 'Necklace', 'Ring', 'SUV', 'Wine Glass', 'Belt', 'Monitor/TV',
|
|
||||||
'Backpack', 'Umbrella', 'Traffic Light', 'Speaker', 'Watch', 'Tie', 'Trash bin Can', 'Slippers', 'Bicycle',
|
|
||||||
'Stool', 'Barrel/bucket', 'Van', 'Couch', 'Sandals', 'Basket', 'Drum', 'Pen/Pencil', 'Bus', 'Wild Bird',
|
|
||||||
'High Heels', 'Motorcycle', 'Guitar', 'Carpet', 'Cell Phone', 'Bread', 'Camera', 'Canned', 'Truck',
|
|
||||||
'Traffic cone', 'Cymbal', 'Lifesaver', 'Towel', 'Stuffed Toy', 'Candle', 'Sailboat', 'Laptop', 'Awning',
|
|
||||||
'Bed', 'Faucet', 'Tent', 'Horse', 'Mirror', 'Power outlet', 'Sink', 'Apple', 'Air Conditioner', 'Knife',
|
|
||||||
'Hockey Stick', 'Paddle', 'Pickup Truck', 'Fork', 'Traffic Sign', 'Balloon', 'Tripod', 'Dog', 'Spoon', 'Clock',
|
|
||||||
'Pot', 'Cow', 'Cake', 'Dinning Table', 'Sheep', 'Hanger', 'Blackboard/Whiteboard', 'Napkin', 'Other Fish',
|
|
||||||
'Orange/Tangerine', 'Toiletry', 'Keyboard', 'Tomato', 'Lantern', 'Machinery Vehicle', 'Fan',
|
|
||||||
'Green Vegetables', 'Banana', 'Baseball Glove', 'Airplane', 'Mouse', 'Train', 'Pumpkin', 'Soccer', 'Skiboard',
|
|
||||||
'Luggage', 'Nightstand', 'Tea pot', 'Telephone', 'Trolley', 'Head Phone', 'Sports Car', 'Stop Sign',
|
|
||||||
'Dessert', 'Scooter', 'Stroller', 'Crane', 'Remote', 'Refrigerator', 'Oven', 'Lemon', 'Duck', 'Baseball Bat',
|
|
||||||
'Surveillance Camera', 'Cat', 'Jug', 'Broccoli', 'Piano', 'Pizza', 'Elephant', 'Skateboard', 'Surfboard',
|
|
||||||
'Gun', 'Skating and Skiing shoes', 'Gas stove', 'Donut', 'Bow Tie', 'Carrot', 'Toilet', 'Kite', 'Strawberry',
|
|
||||||
'Other Balls', 'Shovel', 'Pepper', 'Computer Box', 'Toilet Paper', 'Cleaning Products', 'Chopsticks',
|
|
||||||
'Microwave', 'Pigeon', 'Baseball', 'Cutting/chopping Board', 'Coffee Table', 'Side Table', 'Scissors',
|
|
||||||
'Marker', 'Pie', 'Ladder', 'Snowboard', 'Cookies', 'Radiator', 'Fire Hydrant', 'Basketball', 'Zebra', 'Grape',
|
|
||||||
'Giraffe', 'Potato', 'Sausage', 'Tricycle', 'Violin', 'Egg', 'Fire Extinguisher', 'Candy', 'Fire Truck',
|
|
||||||
'Billiards', 'Converter', 'Bathtub', 'Wheelchair', 'Golf Club', 'Briefcase', 'Cucumber', 'Cigar/Cigarette',
|
|
||||||
'Paint Brush', 'Pear', 'Heavy Truck', 'Hamburger', 'Extractor', 'Extension Cord', 'Tong', 'Tennis Racket',
|
|
||||||
'Folder', 'American Football', 'earphone', 'Mask', 'Kettle', 'Tennis', 'Ship', 'Swing', 'Coffee Machine',
|
|
||||||
'Slide', 'Carriage', 'Onion', 'Green beans', 'Projector', 'Frisbee', 'Washing Machine/Drying Machine',
|
|
||||||
'Chicken', 'Printer', 'Watermelon', 'Saxophone', 'Tissue', 'Toothbrush', 'Ice cream', 'Hot-air balloon',
|
|
||||||
'Cello', 'French Fries', 'Scale', 'Trophy', 'Cabbage', 'Hot dog', 'Blender', 'Peach', 'Rice', 'Wallet/Purse',
|
|
||||||
'Volleyball', 'Deer', 'Goose', 'Tape', 'Tablet', 'Cosmetics', 'Trumpet', 'Pineapple', 'Golf Ball',
|
|
||||||
'Ambulance', 'Parking meter', 'Mango', 'Key', 'Hurdle', 'Fishing Rod', 'Medal', 'Flute', 'Brush', 'Penguin',
|
|
||||||
'Megaphone', 'Corn', 'Lettuce', 'Garlic', 'Swan', 'Helicopter', 'Green Onion', 'Sandwich', 'Nuts',
|
|
||||||
'Speed Limit Sign', 'Induction Cooker', 'Broom', 'Trombone', 'Plum', 'Rickshaw', 'Goldfish', 'Kiwi fruit',
|
|
||||||
'Router/modem', 'Poker Card', 'Toaster', 'Shrimp', 'Sushi', 'Cheese', 'Notepaper', 'Cherry', 'Pliers', 'CD',
|
|
||||||
'Pasta', 'Hammer', 'Cue', 'Avocado', 'Hamimelon', 'Flask', 'Mushroom', 'Screwdriver', 'Soap', 'Recorder',
|
|
||||||
'Bear', 'Eggplant', 'Board Eraser', 'Coconut', 'Tape Measure/Ruler', 'Pig', 'Showerhead', 'Globe', 'Chips',
|
|
||||||
'Steak', 'Crosswalk Sign', 'Stapler', 'Camel', 'Formula 1', 'Pomegranate', 'Dishwasher', 'Crab',
|
|
||||||
'Hoverboard', 'Meat ball', 'Rice Cooker', 'Tuba', 'Calculator', 'Papaya', 'Antelope', 'Parrot', 'Seal',
|
|
||||||
'Butterfly', 'Dumbbell', 'Donkey', 'Lion', 'Urinal', 'Dolphin', 'Electric Drill', 'Hair Dryer', 'Egg tart',
|
|
||||||
'Jellyfish', 'Treadmill', 'Lighter', 'Grapefruit', 'Game board', 'Mop', 'Radish', 'Baozi', 'Target', 'French',
|
|
||||||
'Spring Rolls', 'Monkey', 'Rabbit', 'Pencil Case', 'Yak', 'Red Cabbage', 'Binoculars', 'Asparagus', 'Barbell',
|
|
||||||
'Scallop', 'Noddles', 'Comb', 'Dumpling', 'Oyster', 'Table Tennis paddle', 'Cosmetics Brush/Eyeliner Pencil',
|
|
||||||
'Chainsaw', 'Eraser', 'Lobster', 'Durian', 'Okra', 'Lipstick', 'Cosmetics Mirror', 'Curling', 'Table Tennis' ]
|
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional) --------------------------------------------------------------------------------------
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
|
path: ../datasets/Objects365 # dataset root dir
|
||||||
|
train: images/train # train images (relative to 'path') 1742289 images
|
||||||
|
val: images/val # val images (relative to 'path') 80000 images
|
||||||
|
test: # test images (optional)
|
||||||
|
|
||||||
|
# Classes
|
||||||
|
nc: 365 # number of classes
|
||||||
|
names: ['Person', 'Sneakers', 'Chair', 'Other Shoes', 'Hat', 'Car', 'Lamp', 'Glasses', 'Bottle', 'Desk', 'Cup',
|
||||||
|
'Street Lights', 'Cabinet/shelf', 'Handbag/Satchel', 'Bracelet', 'Plate', 'Picture/Frame', 'Helmet', 'Book',
|
||||||
|
'Gloves', 'Storage box', 'Boat', 'Leather Shoes', 'Flower', 'Bench', 'Potted Plant', 'Bowl/Basin', 'Flag',
|
||||||
|
'Pillow', 'Boots', 'Vase', 'Microphone', 'Necklace', 'Ring', 'SUV', 'Wine Glass', 'Belt', 'Monitor/TV',
|
||||||
|
'Backpack', 'Umbrella', 'Traffic Light', 'Speaker', 'Watch', 'Tie', 'Trash bin Can', 'Slippers', 'Bicycle',
|
||||||
|
'Stool', 'Barrel/bucket', 'Van', 'Couch', 'Sandals', 'Basket', 'Drum', 'Pen/Pencil', 'Bus', 'Wild Bird',
|
||||||
|
'High Heels', 'Motorcycle', 'Guitar', 'Carpet', 'Cell Phone', 'Bread', 'Camera', 'Canned', 'Truck',
|
||||||
|
'Traffic cone', 'Cymbal', 'Lifesaver', 'Towel', 'Stuffed Toy', 'Candle', 'Sailboat', 'Laptop', 'Awning',
|
||||||
|
'Bed', 'Faucet', 'Tent', 'Horse', 'Mirror', 'Power outlet', 'Sink', 'Apple', 'Air Conditioner', 'Knife',
|
||||||
|
'Hockey Stick', 'Paddle', 'Pickup Truck', 'Fork', 'Traffic Sign', 'Balloon', 'Tripod', 'Dog', 'Spoon', 'Clock',
|
||||||
|
'Pot', 'Cow', 'Cake', 'Dinning Table', 'Sheep', 'Hanger', 'Blackboard/Whiteboard', 'Napkin', 'Other Fish',
|
||||||
|
'Orange/Tangerine', 'Toiletry', 'Keyboard', 'Tomato', 'Lantern', 'Machinery Vehicle', 'Fan',
|
||||||
|
'Green Vegetables', 'Banana', 'Baseball Glove', 'Airplane', 'Mouse', 'Train', 'Pumpkin', 'Soccer', 'Skiboard',
|
||||||
|
'Luggage', 'Nightstand', 'Tea pot', 'Telephone', 'Trolley', 'Head Phone', 'Sports Car', 'Stop Sign',
|
||||||
|
'Dessert', 'Scooter', 'Stroller', 'Crane', 'Remote', 'Refrigerator', 'Oven', 'Lemon', 'Duck', 'Baseball Bat',
|
||||||
|
'Surveillance Camera', 'Cat', 'Jug', 'Broccoli', 'Piano', 'Pizza', 'Elephant', 'Skateboard', 'Surfboard',
|
||||||
|
'Gun', 'Skating and Skiing shoes', 'Gas stove', 'Donut', 'Bow Tie', 'Carrot', 'Toilet', 'Kite', 'Strawberry',
|
||||||
|
'Other Balls', 'Shovel', 'Pepper', 'Computer Box', 'Toilet Paper', 'Cleaning Products', 'Chopsticks',
|
||||||
|
'Microwave', 'Pigeon', 'Baseball', 'Cutting/chopping Board', 'Coffee Table', 'Side Table', 'Scissors',
|
||||||
|
'Marker', 'Pie', 'Ladder', 'Snowboard', 'Cookies', 'Radiator', 'Fire Hydrant', 'Basketball', 'Zebra', 'Grape',
|
||||||
|
'Giraffe', 'Potato', 'Sausage', 'Tricycle', 'Violin', 'Egg', 'Fire Extinguisher', 'Candy', 'Fire Truck',
|
||||||
|
'Billiards', 'Converter', 'Bathtub', 'Wheelchair', 'Golf Club', 'Briefcase', 'Cucumber', 'Cigar/Cigarette',
|
||||||
|
'Paint Brush', 'Pear', 'Heavy Truck', 'Hamburger', 'Extractor', 'Extension Cord', 'Tong', 'Tennis Racket',
|
||||||
|
'Folder', 'American Football', 'earphone', 'Mask', 'Kettle', 'Tennis', 'Ship', 'Swing', 'Coffee Machine',
|
||||||
|
'Slide', 'Carriage', 'Onion', 'Green beans', 'Projector', 'Frisbee', 'Washing Machine/Drying Machine',
|
||||||
|
'Chicken', 'Printer', 'Watermelon', 'Saxophone', 'Tissue', 'Toothbrush', 'Ice cream', 'Hot-air balloon',
|
||||||
|
'Cello', 'French Fries', 'Scale', 'Trophy', 'Cabbage', 'Hot dog', 'Blender', 'Peach', 'Rice', 'Wallet/Purse',
|
||||||
|
'Volleyball', 'Deer', 'Goose', 'Tape', 'Tablet', 'Cosmetics', 'Trumpet', 'Pineapple', 'Golf Ball',
|
||||||
|
'Ambulance', 'Parking meter', 'Mango', 'Key', 'Hurdle', 'Fishing Rod', 'Medal', 'Flute', 'Brush', 'Penguin',
|
||||||
|
'Megaphone', 'Corn', 'Lettuce', 'Garlic', 'Swan', 'Helicopter', 'Green Onion', 'Sandwich', 'Nuts',
|
||||||
|
'Speed Limit Sign', 'Induction Cooker', 'Broom', 'Trombone', 'Plum', 'Rickshaw', 'Goldfish', 'Kiwi fruit',
|
||||||
|
'Router/modem', 'Poker Card', 'Toaster', 'Shrimp', 'Sushi', 'Cheese', 'Notepaper', 'Cherry', 'Pliers', 'CD',
|
||||||
|
'Pasta', 'Hammer', 'Cue', 'Avocado', 'Hamimelon', 'Flask', 'Mushroom', 'Screwdriver', 'Soap', 'Recorder',
|
||||||
|
'Bear', 'Eggplant', 'Board Eraser', 'Coconut', 'Tape Measure/Ruler', 'Pig', 'Showerhead', 'Globe', 'Chips',
|
||||||
|
'Steak', 'Crosswalk Sign', 'Stapler', 'Camel', 'Formula 1', 'Pomegranate', 'Dishwasher', 'Crab',
|
||||||
|
'Hoverboard', 'Meat ball', 'Rice Cooker', 'Tuba', 'Calculator', 'Papaya', 'Antelope', 'Parrot', 'Seal',
|
||||||
|
'Butterfly', 'Dumbbell', 'Donkey', 'Lion', 'Urinal', 'Dolphin', 'Electric Drill', 'Hair Dryer', 'Egg tart',
|
||||||
|
'Jellyfish', 'Treadmill', 'Lighter', 'Grapefruit', 'Game board', 'Mop', 'Radish', 'Baozi', 'Target', 'French',
|
||||||
|
'Spring Rolls', 'Monkey', 'Rabbit', 'Pencil Case', 'Yak', 'Red Cabbage', 'Binoculars', 'Asparagus', 'Barbell',
|
||||||
|
'Scallop', 'Noddles', 'Comb', 'Dumpling', 'Oyster', 'Table Tennis paddle', 'Cosmetics Brush/Eyeliner Pencil',
|
||||||
|
'Chainsaw', 'Eraser', 'Lobster', 'Durian', 'Okra', 'Lipstick', 'Cosmetics Mirror', 'Curling', 'Table Tennis']
|
||||||
|
|
||||||
|
|
||||||
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
download: |
|
download: |
|
||||||
from pycocotools.coco import COCO
|
from pycocotools.coco import COCO
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from utils.general import download, Path
|
from utils.general import Path, download, np, xyxy2xywhn
|
||||||
|
|
||||||
# Make Directories
|
# Make Directories
|
||||||
dir = Path('../datasets/objects365') # dataset directory
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
for p in 'images', 'labels':
|
for p in 'images', 'labels':
|
||||||
(dir / p).mkdir(parents=True, exist_ok=True)
|
(dir / p).mkdir(parents=True, exist_ok=True)
|
||||||
for q in 'train', 'val':
|
for q in 'train', 'val':
|
||||||
(dir / p / q).mkdir(parents=True, exist_ok=True)
|
(dir / p / q).mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
# Download
|
# Train, Val Splits
|
||||||
url = "https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/train/"
|
for split, patches in [('train', 50 + 1), ('val', 43 + 1)]:
|
||||||
download([url + 'zhiyuan_objv2_train.tar.gz'], dir=dir, delete=False) # annotations json
|
print(f"Processing {split} in {patches} patches ...")
|
||||||
download([url + f for f in [f'patch{i}.tar.gz' for i in range(51)]], dir=dir / 'images' / 'train',
|
images, labels = dir / 'images' / split, dir / 'labels' / split
|
||||||
curl=True, delete=False, threads=8)
|
|
||||||
|
|
||||||
# Move
|
# Download
|
||||||
train = dir / 'images' / 'train'
|
url = f"https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/{split}/"
|
||||||
for f in tqdm(train.rglob('*.jpg'), desc=f'Moving images'):
|
if split == 'train':
|
||||||
f.rename(train / f.name) # move to /images/train
|
download([f'{url}zhiyuan_objv2_{split}.tar.gz'], dir=dir, delete=False) # annotations json
|
||||||
|
download([f'{url}patch{i}.tar.gz' for i in range(patches)], dir=images, curl=True, delete=False, threads=8)
|
||||||
|
elif split == 'val':
|
||||||
|
download([f'{url}zhiyuan_objv2_{split}.json'], dir=dir, delete=False) # annotations json
|
||||||
|
download([f'{url}images/v1/patch{i}.tar.gz' for i in range(15 + 1)], dir=images, curl=True, delete=False, threads=8)
|
||||||
|
download([f'{url}images/v2/patch{i}.tar.gz' for i in range(16, patches)], dir=images, curl=True, delete=False, threads=8)
|
||||||
|
|
||||||
# Labels
|
# Move
|
||||||
coco = COCO(dir / 'zhiyuan_objv2_train.json')
|
for f in tqdm(images.rglob('*.jpg'), desc=f'Moving {split} images'):
|
||||||
names = [x["name"] for x in coco.loadCats(coco.getCatIds())]
|
f.rename(images / f.name) # move to /images/{split}
|
||||||
for cid, cat in enumerate(names):
|
|
||||||
catIds = coco.getCatIds(catNms=[cat])
|
|
||||||
imgIds = coco.getImgIds(catIds=catIds)
|
|
||||||
for im in tqdm(coco.loadImgs(imgIds), desc=f'Class {cid + 1}/{len(names)} {cat}'):
|
|
||||||
width, height = im["width"], im["height"]
|
|
||||||
path = Path(im["file_name"]) # image filename
|
|
||||||
try:
|
|
||||||
with open(dir / 'labels' / 'train' / path.with_suffix('.txt').name, 'a') as file:
|
|
||||||
annIds = coco.getAnnIds(imgIds=im["id"], catIds=catIds, iscrowd=None)
|
|
||||||
for a in coco.loadAnns(annIds):
|
|
||||||
x, y, w, h = a['bbox'] # bounding box in xywh (xy top-left corner)
|
|
||||||
x, y = x + w / 2, y + h / 2 # xy to center
|
|
||||||
file.write(f"{cid} {x / width:.5f} {y / height:.5f} {w / width:.5f} {h / height:.5f}\n")
|
|
||||||
|
|
||||||
except Exception as e:
|
# Labels
|
||||||
print(e)
|
coco = COCO(dir / f'zhiyuan_objv2_{split}.json')
|
||||||
|
names = [x["name"] for x in coco.loadCats(coco.getCatIds())]
|
||||||
|
for cid, cat in enumerate(names):
|
||||||
|
catIds = coco.getCatIds(catNms=[cat])
|
||||||
|
imgIds = coco.getImgIds(catIds=catIds)
|
||||||
|
for im in tqdm(coco.loadImgs(imgIds), desc=f'Class {cid + 1}/{len(names)} {cat}'):
|
||||||
|
width, height = im["width"], im["height"]
|
||||||
|
path = Path(im["file_name"]) # image filename
|
||||||
|
try:
|
||||||
|
with open(labels / path.with_suffix('.txt').name, 'a') as file:
|
||||||
|
annIds = coco.getAnnIds(imgIds=im["id"], catIds=catIds, iscrowd=None)
|
||||||
|
for a in coco.loadAnns(annIds):
|
||||||
|
x, y, w, h = a['bbox'] # bounding box in xywh (xy top-left corner)
|
||||||
|
xyxy = np.array([x, y, x + w, y + h])[None] # pixels(1,4)
|
||||||
|
x, y, w, h = xyxy2xywhn(xyxy, w=width, h=height, clip=True)[0] # normalized and clipped
|
||||||
|
file.write(f"{cid} {x:.5f} {y:.5f} {w:.5f} {h:.5f}\n")
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
|||||||
18
data/scripts/download_weights.sh
Executable file
18
data/scripts/download_weights.sh
Executable file
@ -0,0 +1,18 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# Download latest models from https://github.com/ultralytics/yolov3/releases
|
||||||
|
# Example usage: bash path/to/download_weights.sh
|
||||||
|
# parent
|
||||||
|
# └── yolov3
|
||||||
|
# ├── yolov3.pt ← downloads here
|
||||||
|
# ├── yolov3-spp.pt
|
||||||
|
# └── ...
|
||||||
|
|
||||||
|
python - <<EOF
|
||||||
|
from utils.downloads import attempt_download
|
||||||
|
|
||||||
|
models = ['yolov3', 'yolov3-spp', 'yolov3-tiny']
|
||||||
|
for x in models:
|
||||||
|
attempt_download(f'{x}.pt')
|
||||||
|
|
||||||
|
EOF
|
||||||
@ -1,61 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
|
|
||||||
# Download command: bash data/scripts/get_argoverse_hd.sh
|
|
||||||
# Train command: python train.py --data argoverse_hd.yaml
|
|
||||||
# Default dataset location is next to YOLOv3:
|
|
||||||
# /parent_folder
|
|
||||||
# /argoverse
|
|
||||||
# /yolov3
|
|
||||||
|
|
||||||
# Download/unzip images
|
|
||||||
d='../argoverse/' # unzip directory
|
|
||||||
mkdir $d
|
|
||||||
url=https://argoverse-hd.s3.us-east-2.amazonaws.com/
|
|
||||||
f=Argoverse-HD-Full.zip
|
|
||||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f &# download, unzip, remove in background
|
|
||||||
wait # finish background tasks
|
|
||||||
|
|
||||||
cd ../argoverse/Argoverse-1.1/
|
|
||||||
ln -s tracking images
|
|
||||||
|
|
||||||
cd ../Argoverse-HD/annotations/
|
|
||||||
|
|
||||||
python3 - "$@" <<END
|
|
||||||
import json
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
annotation_files = ["train.json", "val.json"]
|
|
||||||
print("Converting annotations to YOLOv3 format...")
|
|
||||||
|
|
||||||
for val in annotation_files:
|
|
||||||
a = json.load(open(val, "rb"))
|
|
||||||
|
|
||||||
label_dict = {}
|
|
||||||
for annot in a['annotations']:
|
|
||||||
img_id = annot['image_id']
|
|
||||||
img_name = a['images'][img_id]['name']
|
|
||||||
img_label_name = img_name[:-3] + "txt"
|
|
||||||
|
|
||||||
cls = annot['category_id'] # instance class id
|
|
||||||
x_center, y_center, width, height = annot['bbox']
|
|
||||||
x_center = (x_center + width / 2) / 1920. # offset and scale
|
|
||||||
y_center = (y_center + height / 2) / 1200. # offset and scale
|
|
||||||
width /= 1920. # scale
|
|
||||||
height /= 1200. # scale
|
|
||||||
|
|
||||||
img_dir = "./labels/" + a['seq_dirs'][a['images'][annot['image_id']]['sid']]
|
|
||||||
|
|
||||||
Path(img_dir).mkdir(parents=True, exist_ok=True)
|
|
||||||
if img_dir + "/" + img_label_name not in label_dict:
|
|
||||||
label_dict[img_dir + "/" + img_label_name] = []
|
|
||||||
|
|
||||||
label_dict[img_dir + "/" + img_label_name].append(f"{cls} {x_center} {y_center} {width} {height}\n")
|
|
||||||
|
|
||||||
for filename in label_dict:
|
|
||||||
with open(filename, "w") as file:
|
|
||||||
for string in label_dict[filename]:
|
|
||||||
file.write(string)
|
|
||||||
|
|
||||||
END
|
|
||||||
|
|
||||||
mv ./labels ../../Argoverse-1.1/
|
|
||||||
@ -1,27 +1,27 @@
|
|||||||
#!/bin/bash
|
#!/bin/bash
|
||||||
# COCO 2017 dataset http://cocodataset.org
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Download command: bash data/scripts/get_coco.sh
|
# Download COCO 2017 dataset http://cocodataset.org
|
||||||
# Train command: python train.py --data coco.yaml
|
# Example usage: bash data/scripts/get_coco.sh
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /coco
|
# └── datasets
|
||||||
# /yolov3
|
# └── coco ← downloads here
|
||||||
|
|
||||||
# Download/unzip labels
|
# Download/unzip labels
|
||||||
d='../' # unzip directory
|
d='../datasets' # unzip directory
|
||||||
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
||||||
f='coco2017labels.zip' # or 'coco2017labels-segments.zip', 68 MB
|
f='coco2017labels.zip' # or 'coco2017labels-segments.zip', 68 MB
|
||||||
echo 'Downloading' $url$f ' ...'
|
echo 'Downloading' $url$f ' ...'
|
||||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f &
|
||||||
|
|
||||||
# Download/unzip images
|
# Download/unzip images
|
||||||
d='../coco/images' # unzip directory
|
d='../datasets/coco/images' # unzip directory
|
||||||
url=http://images.cocodataset.org/zips/
|
url=http://images.cocodataset.org/zips/
|
||||||
f1='train2017.zip' # 19G, 118k images
|
f1='train2017.zip' # 19G, 118k images
|
||||||
f2='val2017.zip' # 1G, 5k images
|
f2='val2017.zip' # 1G, 5k images
|
||||||
f3='test2017.zip' # 7G, 41k images (optional)
|
f3='test2017.zip' # 7G, 41k images (optional)
|
||||||
for f in $f1 $f2; do
|
for f in $f1 $f2; do
|
||||||
echo 'Downloading' $url$f '...'
|
echo 'Downloading' $url$f '...'
|
||||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f &
|
||||||
done
|
done
|
||||||
wait # finish background tasks
|
wait # finish background tasks
|
||||||
|
|||||||
@ -1,17 +1,17 @@
|
|||||||
#!/bin/bash
|
#!/bin/bash
|
||||||
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Download command: bash data/scripts/get_coco128.sh
|
# Download COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
|
||||||
# Train command: python train.py --data coco128.yaml
|
# Example usage: bash data/scripts/get_coco128.sh
|
||||||
# Default dataset location is next to YOLOv3:
|
# parent
|
||||||
# /parent_folder
|
# ├── yolov3
|
||||||
# /coco128
|
# └── datasets
|
||||||
# /yolov3
|
# └── coco128 ← downloads here
|
||||||
|
|
||||||
# Download/unzip images and labels
|
# Download/unzip images and labels
|
||||||
d='../' # unzip directory
|
d='../datasets' # unzip directory
|
||||||
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
||||||
f='coco128.zip' # or 'coco2017labels-segments.zip', 68 MB
|
f='coco128.zip' # or 'coco128-segments.zip', 68 MB
|
||||||
echo 'Downloading' $url$f ' ...'
|
echo 'Downloading' $url$f ' ...'
|
||||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f &
|
||||||
|
|
||||||
wait # finish background tasks
|
wait # finish background tasks
|
||||||
|
|||||||
@ -1,116 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
|
|
||||||
# Download command: bash data/scripts/get_voc.sh
|
|
||||||
# Train command: python train.py --data voc.yaml
|
|
||||||
# Default dataset location is next to YOLOv3:
|
|
||||||
# /parent_folder
|
|
||||||
# /VOC
|
|
||||||
# /yolov3
|
|
||||||
|
|
||||||
start=$(date +%s)
|
|
||||||
mkdir -p ../tmp
|
|
||||||
cd ../tmp/
|
|
||||||
|
|
||||||
# Download/unzip images and labels
|
|
||||||
d='.' # unzip directory
|
|
||||||
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
|
||||||
f1=VOCtrainval_06-Nov-2007.zip # 446MB, 5012 images
|
|
||||||
f2=VOCtest_06-Nov-2007.zip # 438MB, 4953 images
|
|
||||||
f3=VOCtrainval_11-May-2012.zip # 1.95GB, 17126 images
|
|
||||||
for f in $f3 $f2 $f1; do
|
|
||||||
echo 'Downloading' $url$f '...'
|
|
||||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
|
||||||
done
|
|
||||||
wait # finish background tasks
|
|
||||||
|
|
||||||
end=$(date +%s)
|
|
||||||
runtime=$((end - start))
|
|
||||||
echo "Completed in" $runtime "seconds"
|
|
||||||
|
|
||||||
echo "Splitting dataset..."
|
|
||||||
python3 - "$@" <<END
|
|
||||||
import os
|
|
||||||
import xml.etree.ElementTree as ET
|
|
||||||
from os import getcwd
|
|
||||||
|
|
||||||
sets = [('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
|
|
||||||
|
|
||||||
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
|
|
||||||
"horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
|
|
||||||
|
|
||||||
|
|
||||||
def convert_box(size, box):
|
|
||||||
dw = 1. / (size[0])
|
|
||||||
dh = 1. / (size[1])
|
|
||||||
x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
|
|
||||||
return x * dw, y * dh, w * dw, h * dh
|
|
||||||
|
|
||||||
|
|
||||||
def convert_annotation(year, image_id):
|
|
||||||
in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml' % (year, image_id))
|
|
||||||
out_file = open('VOCdevkit/VOC%s/labels/%s.txt' % (year, image_id), 'w')
|
|
||||||
tree = ET.parse(in_file)
|
|
||||||
root = tree.getroot()
|
|
||||||
size = root.find('size')
|
|
||||||
w = int(size.find('width').text)
|
|
||||||
h = int(size.find('height').text)
|
|
||||||
|
|
||||||
for obj in root.iter('object'):
|
|
||||||
difficult = obj.find('difficult').text
|
|
||||||
cls = obj.find('name').text
|
|
||||||
if cls not in classes or int(difficult) == 1:
|
|
||||||
continue
|
|
||||||
cls_id = classes.index(cls)
|
|
||||||
xmlbox = obj.find('bndbox')
|
|
||||||
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
|
|
||||||
float(xmlbox.find('ymax').text))
|
|
||||||
bb = convert_box((w, h), b)
|
|
||||||
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
|
|
||||||
|
|
||||||
|
|
||||||
cwd = getcwd()
|
|
||||||
for year, image_set in sets:
|
|
||||||
if not os.path.exists('VOCdevkit/VOC%s/labels/' % year):
|
|
||||||
os.makedirs('VOCdevkit/VOC%s/labels/' % year)
|
|
||||||
image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt' % (year, image_set)).read().strip().split()
|
|
||||||
list_file = open('%s_%s.txt' % (year, image_set), 'w')
|
|
||||||
for image_id in image_ids:
|
|
||||||
list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n' % (cwd, year, image_id))
|
|
||||||
convert_annotation(year, image_id)
|
|
||||||
list_file.close()
|
|
||||||
END
|
|
||||||
|
|
||||||
cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt >train.txt
|
|
||||||
cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt
|
|
||||||
|
|
||||||
mkdir ../VOC ../VOC/images ../VOC/images/train ../VOC/images/val
|
|
||||||
mkdir ../VOC/labels ../VOC/labels/train ../VOC/labels/val
|
|
||||||
|
|
||||||
python3 - "$@" <<END
|
|
||||||
import os
|
|
||||||
|
|
||||||
print(os.path.exists('../tmp/train.txt'))
|
|
||||||
with open('../tmp/train.txt', 'r') as f:
|
|
||||||
for line in f.readlines():
|
|
||||||
line = "/".join(line.split('/')[-5:]).strip()
|
|
||||||
if os.path.exists("../" + line):
|
|
||||||
os.system("cp ../" + line + " ../VOC/images/train")
|
|
||||||
|
|
||||||
line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
|
|
||||||
if os.path.exists("../" + line):
|
|
||||||
os.system("cp ../" + line + " ../VOC/labels/train")
|
|
||||||
|
|
||||||
print(os.path.exists('../tmp/2007_test.txt'))
|
|
||||||
with open('../tmp/2007_test.txt', 'r') as f:
|
|
||||||
for line in f.readlines():
|
|
||||||
line = "/".join(line.split('/')[-5:]).strip()
|
|
||||||
if os.path.exists("../" + line):
|
|
||||||
os.system("cp ../" + line + " ../VOC/images/val")
|
|
||||||
|
|
||||||
line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
|
|
||||||
if os.path.exists("../" + line):
|
|
||||||
os.system("cp ../" + line + " ../VOC/labels/val")
|
|
||||||
END
|
|
||||||
|
|
||||||
rm -rf ../tmp # remove temporary directory
|
|
||||||
echo "VOC download done."
|
|
||||||
@ -1,21 +1,80 @@
|
|||||||
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
# Train command: python train.py --data voc.yaml
|
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC
|
||||||
# Default dataset location is next to YOLOv3:
|
# Example usage: python train.py --data VOC.yaml
|
||||||
# /parent_folder
|
# parent
|
||||||
# /VOC
|
# ├── yolov3
|
||||||
# /yolov3
|
# └── datasets
|
||||||
|
# └── VOC ← downloads here
|
||||||
|
|
||||||
|
|
||||||
# download command/URL (optional)
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
download: bash data/scripts/get_voc.sh
|
path: ../datasets/VOC
|
||||||
|
train: # train images (relative to 'path') 16551 images
|
||||||
|
- images/train2012
|
||||||
|
- images/train2007
|
||||||
|
- images/val2012
|
||||||
|
- images/val2007
|
||||||
|
val: # val images (relative to 'path') 4952 images
|
||||||
|
- images/test2007
|
||||||
|
test: # test images (optional)
|
||||||
|
- images/test2007
|
||||||
|
|
||||||
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
|
# Classes
|
||||||
train: ../VOC/images/train/ # 16551 images
|
nc: 20 # number of classes
|
||||||
val: ../VOC/images/val/ # 4952 images
|
names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
|
||||||
|
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] # class names
|
||||||
|
|
||||||
# number of classes
|
|
||||||
nc: 20
|
|
||||||
|
|
||||||
# class names
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
names: [ 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
|
download: |
|
||||||
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]
|
import xml.etree.ElementTree as ET
|
||||||
|
|
||||||
|
from tqdm import tqdm
|
||||||
|
from utils.general import download, Path
|
||||||
|
|
||||||
|
|
||||||
|
def convert_label(path, lb_path, year, image_id):
|
||||||
|
def convert_box(size, box):
|
||||||
|
dw, dh = 1. / size[0], 1. / size[1]
|
||||||
|
x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
|
||||||
|
return x * dw, y * dh, w * dw, h * dh
|
||||||
|
|
||||||
|
in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
|
||||||
|
out_file = open(lb_path, 'w')
|
||||||
|
tree = ET.parse(in_file)
|
||||||
|
root = tree.getroot()
|
||||||
|
size = root.find('size')
|
||||||
|
w = int(size.find('width').text)
|
||||||
|
h = int(size.find('height').text)
|
||||||
|
|
||||||
|
for obj in root.iter('object'):
|
||||||
|
cls = obj.find('name').text
|
||||||
|
if cls in yaml['names'] and not int(obj.find('difficult').text) == 1:
|
||||||
|
xmlbox = obj.find('bndbox')
|
||||||
|
bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
|
||||||
|
cls_id = yaml['names'].index(cls) # class id
|
||||||
|
out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')
|
||||||
|
|
||||||
|
|
||||||
|
# Download
|
||||||
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
|
url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/'
|
||||||
|
urls = [url + 'VOCtrainval_06-Nov-2007.zip', # 446MB, 5012 images
|
||||||
|
url + 'VOCtest_06-Nov-2007.zip', # 438MB, 4953 images
|
||||||
|
url + 'VOCtrainval_11-May-2012.zip'] # 1.95GB, 17126 images
|
||||||
|
download(urls, dir=dir / 'images', delete=False)
|
||||||
|
|
||||||
|
# Convert
|
||||||
|
path = dir / f'images/VOCdevkit'
|
||||||
|
for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):
|
||||||
|
imgs_path = dir / 'images' / f'{image_set}{year}'
|
||||||
|
lbs_path = dir / 'labels' / f'{image_set}{year}'
|
||||||
|
imgs_path.mkdir(exist_ok=True, parents=True)
|
||||||
|
lbs_path.mkdir(exist_ok=True, parents=True)
|
||||||
|
|
||||||
|
image_ids = open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt').read().strip().split()
|
||||||
|
for id in tqdm(image_ids, desc=f'{image_set}{year}'):
|
||||||
|
f = path / f'VOC{year}/JPEGImages/{id}.jpg' # old img path
|
||||||
|
lb_path = (lbs_path / f.name).with_suffix('.txt') # new label path
|
||||||
|
f.rename(imgs_path / f.name) # move image
|
||||||
|
convert_label(path, lb_path, year, id) # convert labels to YOLO format
|
||||||
|
|||||||
102
data/xView.yaml
Normal file
102
data/xView.yaml
Normal file
@ -0,0 +1,102 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
# xView 2018 dataset https://challenge.xviewdataset.org
|
||||||
|
# -------- DOWNLOAD DATA MANUALLY from URL above and unzip to 'datasets/xView' before running train command! --------
|
||||||
|
# Example usage: python train.py --data xView.yaml
|
||||||
|
# parent
|
||||||
|
# ├── yolov3
|
||||||
|
# └── datasets
|
||||||
|
# └── xView ← downloads here
|
||||||
|
|
||||||
|
|
||||||
|
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
||||||
|
path: ../datasets/xView # dataset root dir
|
||||||
|
train: images/autosplit_train.txt # train images (relative to 'path') 90% of 847 train images
|
||||||
|
val: images/autosplit_val.txt # train images (relative to 'path') 10% of 847 train images
|
||||||
|
|
||||||
|
# Classes
|
||||||
|
nc: 60 # number of classes
|
||||||
|
names: ['Fixed-wing Aircraft', 'Small Aircraft', 'Cargo Plane', 'Helicopter', 'Passenger Vehicle', 'Small Car', 'Bus',
|
||||||
|
'Pickup Truck', 'Utility Truck', 'Truck', 'Cargo Truck', 'Truck w/Box', 'Truck Tractor', 'Trailer',
|
||||||
|
'Truck w/Flatbed', 'Truck w/Liquid', 'Crane Truck', 'Railway Vehicle', 'Passenger Car', 'Cargo Car',
|
||||||
|
'Flat Car', 'Tank car', 'Locomotive', 'Maritime Vessel', 'Motorboat', 'Sailboat', 'Tugboat', 'Barge',
|
||||||
|
'Fishing Vessel', 'Ferry', 'Yacht', 'Container Ship', 'Oil Tanker', 'Engineering Vehicle', 'Tower crane',
|
||||||
|
'Container Crane', 'Reach Stacker', 'Straddle Carrier', 'Mobile Crane', 'Dump Truck', 'Haul Truck',
|
||||||
|
'Scraper/Tractor', 'Front loader/Bulldozer', 'Excavator', 'Cement Mixer', 'Ground Grader', 'Hut/Tent', 'Shed',
|
||||||
|
'Building', 'Aircraft Hangar', 'Damaged Building', 'Facility', 'Construction Site', 'Vehicle Lot', 'Helipad',
|
||||||
|
'Storage Tank', 'Shipping container lot', 'Shipping Container', 'Pylon', 'Tower'] # class names
|
||||||
|
|
||||||
|
|
||||||
|
# Download script/URL (optional) ---------------------------------------------------------------------------------------
|
||||||
|
download: |
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
from utils.datasets import autosplit
|
||||||
|
from utils.general import download, xyxy2xywhn
|
||||||
|
|
||||||
|
|
||||||
|
def convert_labels(fname=Path('xView/xView_train.geojson')):
|
||||||
|
# Convert xView geoJSON labels to YOLO format
|
||||||
|
path = fname.parent
|
||||||
|
with open(fname) as f:
|
||||||
|
print(f'Loading {fname}...')
|
||||||
|
data = json.load(f)
|
||||||
|
|
||||||
|
# Make dirs
|
||||||
|
labels = Path(path / 'labels' / 'train')
|
||||||
|
os.system(f'rm -rf {labels}')
|
||||||
|
labels.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# xView classes 11-94 to 0-59
|
||||||
|
xview_class2index = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, -1, 9, 10, 11,
|
||||||
|
12, 13, 14, 15, -1, -1, 16, 17, 18, 19, 20, 21, 22, -1, 23, 24, 25, -1, 26, 27, -1, 28, -1,
|
||||||
|
29, 30, 31, 32, 33, 34, 35, 36, 37, -1, 38, 39, 40, 41, 42, 43, 44, 45, -1, -1, -1, -1, 46,
|
||||||
|
47, 48, 49, -1, 50, 51, -1, 52, -1, -1, -1, 53, 54, -1, 55, -1, -1, 56, -1, 57, -1, 58, 59]
|
||||||
|
|
||||||
|
shapes = {}
|
||||||
|
for feature in tqdm(data['features'], desc=f'Converting {fname}'):
|
||||||
|
p = feature['properties']
|
||||||
|
if p['bounds_imcoords']:
|
||||||
|
id = p['image_id']
|
||||||
|
file = path / 'train_images' / id
|
||||||
|
if file.exists(): # 1395.tif missing
|
||||||
|
try:
|
||||||
|
box = np.array([int(num) for num in p['bounds_imcoords'].split(",")])
|
||||||
|
assert box.shape[0] == 4, f'incorrect box shape {box.shape[0]}'
|
||||||
|
cls = p['type_id']
|
||||||
|
cls = xview_class2index[int(cls)] # xView class to 0-60
|
||||||
|
assert 59 >= cls >= 0, f'incorrect class index {cls}'
|
||||||
|
|
||||||
|
# Write YOLO label
|
||||||
|
if id not in shapes:
|
||||||
|
shapes[id] = Image.open(file).size
|
||||||
|
box = xyxy2xywhn(box[None].astype(np.float), w=shapes[id][0], h=shapes[id][1], clip=True)
|
||||||
|
with open((labels / id).with_suffix('.txt'), 'a') as f:
|
||||||
|
f.write(f"{cls} {' '.join(f'{x:.6f}' for x in box[0])}\n") # write label.txt
|
||||||
|
except Exception as e:
|
||||||
|
print(f'WARNING: skipping one label for {file}: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
# Download manually from https://challenge.xviewdataset.org
|
||||||
|
dir = Path(yaml['path']) # dataset root dir
|
||||||
|
# urls = ['https://d307kc0mrhucc3.cloudfront.net/train_labels.zip', # train labels
|
||||||
|
# 'https://d307kc0mrhucc3.cloudfront.net/train_images.zip', # 15G, 847 train images
|
||||||
|
# 'https://d307kc0mrhucc3.cloudfront.net/val_images.zip'] # 5G, 282 val images (no labels)
|
||||||
|
# download(urls, dir=dir, delete=False)
|
||||||
|
|
||||||
|
# Convert labels
|
||||||
|
convert_labels(dir / 'xView_train.geojson')
|
||||||
|
|
||||||
|
# Move images
|
||||||
|
images = Path(dir / 'images')
|
||||||
|
images.mkdir(parents=True, exist_ok=True)
|
||||||
|
Path(dir / 'train_images').rename(dir / 'images' / 'train')
|
||||||
|
Path(dir / 'val_images').rename(dir / 'images' / 'val')
|
||||||
|
|
||||||
|
# Split
|
||||||
|
autosplit(dir / 'images' / 'train')
|
||||||
248
detect.py
248
detect.py
@ -1,98 +1,147 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Run inference on images, videos, directories, streams, etc.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
$ python path/to/detect.py --weights yolov3.pt --source 0 # webcam
|
||||||
|
img.jpg # image
|
||||||
|
vid.mp4 # video
|
||||||
|
path/ # directory
|
||||||
|
path/*.jpg # glob
|
||||||
|
'https://youtu.be/Zgi9g1ksQHc' # YouTube
|
||||||
|
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
||||||
|
"""
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import time
|
import os
|
||||||
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import torch
|
import torch
|
||||||
import torch.backends.cudnn as cudnn
|
import torch.backends.cudnn as cudnn
|
||||||
|
|
||||||
from models.experimental import attempt_load
|
FILE = Path(__file__).resolve()
|
||||||
from utils.datasets import LoadStreams, LoadImages
|
ROOT = FILE.parents[0] # root directory
|
||||||
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
|
if str(ROOT) not in sys.path:
|
||||||
scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
from utils.plots import colors, plot_one_box
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||||
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
|
||||||
|
from models.common import DetectMultiBackend
|
||||||
|
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
|
||||||
|
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
|
||||||
|
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
|
||||||
|
from utils.plots import Annotator, colors, save_one_box
|
||||||
|
from utils.torch_utils import select_device, time_sync
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
def detect(opt):
|
def run(weights=ROOT / 'yolov3.pt', # model.pt path(s)
|
||||||
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
|
||||||
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
|
imgsz=640, # inference size (pixels)
|
||||||
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
|
conf_thres=0.25, # confidence threshold
|
||||||
('rtsp://', 'rtmp://', 'http://', 'https://'))
|
iou_thres=0.45, # NMS IOU threshold
|
||||||
|
max_det=1000, # maximum detections per image
|
||||||
|
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||||
|
view_img=False, # show results
|
||||||
|
save_txt=False, # save results to *.txt
|
||||||
|
save_conf=False, # save confidences in --save-txt labels
|
||||||
|
save_crop=False, # save cropped prediction boxes
|
||||||
|
nosave=False, # do not save images/videos
|
||||||
|
classes=None, # filter by class: --class 0, or --class 0 2 3
|
||||||
|
agnostic_nms=False, # class-agnostic NMS
|
||||||
|
augment=False, # augmented inference
|
||||||
|
visualize=False, # visualize features
|
||||||
|
update=False, # update all models
|
||||||
|
project=ROOT / 'runs/detect', # save results to project/name
|
||||||
|
name='exp', # save results to project/name
|
||||||
|
exist_ok=False, # existing project/name ok, do not increment
|
||||||
|
line_thickness=3, # bounding box thickness (pixels)
|
||||||
|
hide_labels=False, # hide labels
|
||||||
|
hide_conf=False, # hide confidences
|
||||||
|
half=False, # use FP16 half-precision inference
|
||||||
|
dnn=False, # use OpenCV DNN for ONNX inference
|
||||||
|
):
|
||||||
|
source = str(source)
|
||||||
|
save_img = not nosave and not source.endswith('.txt') # save inference images
|
||||||
|
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
|
||||||
|
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
|
||||||
|
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
|
||||||
|
if is_url and is_file:
|
||||||
|
source = check_file(source) # download
|
||||||
|
|
||||||
# Directories
|
# Directories
|
||||||
save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
|
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
||||||
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
||||||
|
|
||||||
# Initialize
|
|
||||||
set_logging()
|
|
||||||
device = select_device(opt.device)
|
|
||||||
half = device.type != 'cpu' # half precision only supported on CUDA
|
|
||||||
|
|
||||||
# Load model
|
# Load model
|
||||||
model = attempt_load(weights, map_location=device) # load FP32 model
|
device = select_device(device)
|
||||||
stride = int(model.stride.max()) # model stride
|
model = DetectMultiBackend(weights, device=device, dnn=dnn)
|
||||||
imgsz = check_img_size(imgsz, s=stride) # check img_size
|
stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
|
||||||
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||||
if half:
|
|
||||||
model.half() # to FP16
|
|
||||||
|
|
||||||
# Second-stage classifier
|
# Half
|
||||||
classify = False
|
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||||
if classify:
|
if pt:
|
||||||
modelc = load_classifier(name='resnet101', n=2) # initialize
|
model.model.half() if half else model.model.float()
|
||||||
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
|
|
||||||
|
|
||||||
# Set Dataloader
|
# Dataloader
|
||||||
vid_path, vid_writer = None, None
|
|
||||||
if webcam:
|
if webcam:
|
||||||
view_img = check_imshow()
|
view_img = check_imshow()
|
||||||
cudnn.benchmark = True # set True to speed up constant image size inference
|
cudnn.benchmark = True # set True to speed up constant image size inference
|
||||||
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
|
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)
|
||||||
|
bs = len(dataset) # batch_size
|
||||||
else:
|
else:
|
||||||
dataset = LoadImages(source, img_size=imgsz, stride=stride)
|
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
|
||||||
|
bs = 1 # batch_size
|
||||||
|
vid_path, vid_writer = [None] * bs, [None] * bs
|
||||||
|
|
||||||
# Run inference
|
# Run inference
|
||||||
if device.type != 'cpu':
|
if pt and device.type != 'cpu':
|
||||||
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
|
model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters()))) # warmup
|
||||||
t0 = time.time()
|
dt, seen = [0.0, 0.0, 0.0], 0
|
||||||
for path, img, im0s, vid_cap in dataset:
|
for path, im, im0s, vid_cap, s in dataset:
|
||||||
img = torch.from_numpy(img).to(device)
|
t1 = time_sync()
|
||||||
img = img.half() if half else img.float() # uint8 to fp16/32
|
im = torch.from_numpy(im).to(device)
|
||||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
im = im.half() if half else im.float() # uint8 to fp16/32
|
||||||
if img.ndimension() == 3:
|
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||||
img = img.unsqueeze(0)
|
if len(im.shape) == 3:
|
||||||
|
im = im[None] # expand for batch dim
|
||||||
|
t2 = time_sync()
|
||||||
|
dt[0] += t2 - t1
|
||||||
|
|
||||||
# Inference
|
# Inference
|
||||||
t1 = time_synchronized()
|
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
|
||||||
pred = model(img, augment=opt.augment)[0]
|
pred = model(im, augment=augment, visualize=visualize)
|
||||||
|
t3 = time_sync()
|
||||||
|
dt[1] += t3 - t2
|
||||||
|
|
||||||
# Apply NMS
|
# NMS
|
||||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
|
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
|
||||||
max_det=opt.max_det)
|
dt[2] += time_sync() - t3
|
||||||
t2 = time_synchronized()
|
|
||||||
|
|
||||||
# Apply Classifier
|
# Second-stage classifier (optional)
|
||||||
if classify:
|
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
|
||||||
pred = apply_classifier(pred, modelc, img, im0s)
|
|
||||||
|
|
||||||
# Process detections
|
# Process predictions
|
||||||
for i, det in enumerate(pred): # detections per image
|
for i, det in enumerate(pred): # per image
|
||||||
|
seen += 1
|
||||||
if webcam: # batch_size >= 1
|
if webcam: # batch_size >= 1
|
||||||
p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
|
p, im0, frame = path[i], im0s[i].copy(), dataset.count
|
||||||
|
s += f'{i}: '
|
||||||
else:
|
else:
|
||||||
p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
|
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
|
||||||
|
|
||||||
p = Path(p) # to Path
|
p = Path(p) # to Path
|
||||||
save_path = str(save_dir / p.name) # img.jpg
|
save_path = str(save_dir / p.name) # im.jpg
|
||||||
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
|
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
|
||||||
s += '%gx%g ' % img.shape[2:] # print string
|
s += '%gx%g ' % im.shape[2:] # print string
|
||||||
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
|
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||||||
imc = im0.copy() if opt.save_crop else im0 # for opt.save_crop
|
imc = im0.copy() if save_crop else im0 # for save_crop
|
||||||
|
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
|
||||||
if len(det):
|
if len(det):
|
||||||
# Rescale boxes from img_size to im0 size
|
# Rescale boxes from img_size to im0 size
|
||||||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
|
||||||
|
|
||||||
# Print results
|
# Print results
|
||||||
for c in det[:, -1].unique():
|
for c in det[:, -1].unique():
|
||||||
@ -103,21 +152,22 @@ def detect(opt):
|
|||||||
for *xyxy, conf, cls in reversed(det):
|
for *xyxy, conf, cls in reversed(det):
|
||||||
if save_txt: # Write to file
|
if save_txt: # Write to file
|
||||||
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||||||
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
|
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
||||||
with open(txt_path + '.txt', 'a') as f:
|
with open(txt_path + '.txt', 'a') as f:
|
||||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
||||||
|
|
||||||
if save_img or opt.save_crop or view_img: # Add bbox to image
|
if save_img or save_crop or view_img: # Add bbox to image
|
||||||
c = int(cls) # integer class
|
c = int(cls) # integer class
|
||||||
label = None if opt.hide_labels else (names[c] if opt.hide_conf else f'{names[c]} {conf:.2f}')
|
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
|
||||||
plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=opt.line_thickness)
|
annotator.box_label(xyxy, label, color=colors(c, True))
|
||||||
if opt.save_crop:
|
if save_crop:
|
||||||
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
|
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
|
||||||
|
|
||||||
# Print time (inference + NMS)
|
# Print time (inference-only)
|
||||||
print(f'{s}Done. ({t2 - t1:.3f}s)')
|
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
|
||||||
|
|
||||||
# Stream results
|
# Stream results
|
||||||
|
im0 = annotator.result()
|
||||||
if view_img:
|
if view_img:
|
||||||
cv2.imshow(str(p), im0)
|
cv2.imshow(str(p), im0)
|
||||||
cv2.waitKey(1) # 1 millisecond
|
cv2.waitKey(1) # 1 millisecond
|
||||||
@ -127,10 +177,10 @@ def detect(opt):
|
|||||||
if dataset.mode == 'image':
|
if dataset.mode == 'image':
|
||||||
cv2.imwrite(save_path, im0)
|
cv2.imwrite(save_path, im0)
|
||||||
else: # 'video' or 'stream'
|
else: # 'video' or 'stream'
|
||||||
if vid_path != save_path: # new video
|
if vid_path[i] != save_path: # new video
|
||||||
vid_path = save_path
|
vid_path[i] = save_path
|
||||||
if isinstance(vid_writer, cv2.VideoWriter):
|
if isinstance(vid_writer[i], cv2.VideoWriter):
|
||||||
vid_writer.release() # release previous video writer
|
vid_writer[i].release() # release previous video writer
|
||||||
if vid_cap: # video
|
if vid_cap: # video
|
||||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||||
@ -138,47 +188,57 @@ def detect(opt):
|
|||||||
else: # stream
|
else: # stream
|
||||||
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
||||||
save_path += '.mp4'
|
save_path += '.mp4'
|
||||||
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
||||||
vid_writer.write(im0)
|
vid_writer[i].write(im0)
|
||||||
|
|
||||||
|
# Print results
|
||||||
|
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
|
||||||
|
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
|
||||||
if save_txt or save_img:
|
if save_txt or save_img:
|
||||||
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
||||||
print(f"Results saved to {save_dir}{s}")
|
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
|
||||||
|
if update:
|
||||||
print(f'Done. ({time.time() - t0:.3f}s)')
|
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
def parse_opt():
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument('--weights', nargs='+', type=str, default='yolov3.pt', help='model.pt path(s)')
|
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov3.pt', help='model path(s)')
|
||||||
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
|
parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')
|
||||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
|
||||||
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
|
||||||
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
|
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
|
||||||
parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
|
parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
|
||||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||||
parser.add_argument('--view-img', action='store_true', help='display results')
|
parser.add_argument('--view-img', action='store_true', help='show results')
|
||||||
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
||||||
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
||||||
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
|
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
|
||||||
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
|
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
|
||||||
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
|
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
|
||||||
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
||||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||||
|
parser.add_argument('--visualize', action='store_true', help='visualize features')
|
||||||
parser.add_argument('--update', action='store_true', help='update all models')
|
parser.add_argument('--update', action='store_true', help='update all models')
|
||||||
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
|
parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
|
||||||
parser.add_argument('--name', default='exp', help='save results to project/name')
|
parser.add_argument('--name', default='exp', help='save results to project/name')
|
||||||
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
||||||
parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
|
parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
|
||||||
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
|
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
|
||||||
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
|
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
|
||||||
|
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
||||||
|
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
||||||
opt = parser.parse_args()
|
opt = parser.parse_args()
|
||||||
print(opt)
|
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
|
||||||
check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
|
print_args(FILE.stem, opt)
|
||||||
|
return opt
|
||||||
|
|
||||||
if opt.update: # update all models (to fix SourceChangeWarning)
|
|
||||||
for opt.weights in ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']:
|
def main(opt):
|
||||||
detect(opt=opt)
|
check_requirements(exclude=('tensorboard', 'thop'))
|
||||||
strip_optimizer(opt.weights)
|
run(**vars(opt))
|
||||||
else:
|
|
||||||
detect(opt=opt)
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
opt = parse_opt()
|
||||||
|
main(opt)
|
||||||
|
|||||||
369
export.py
Normal file
369
export.py
Normal file
@ -0,0 +1,369 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Export a PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats
|
||||||
|
TensorFlow exports authored by https://github.com/zldrobit
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
$ python path/to/export.py --weights yolov3.pt --include torchscript onnx coreml saved_model pb tflite tfjs
|
||||||
|
|
||||||
|
Inference:
|
||||||
|
$ python path/to/detect.py --weights yolov3.pt
|
||||||
|
yolov3.onnx (must export with --dynamic)
|
||||||
|
yolov3_saved_model
|
||||||
|
yolov3.pb
|
||||||
|
yolov3.tflite
|
||||||
|
|
||||||
|
TensorFlow.js:
|
||||||
|
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
||||||
|
$ npm install
|
||||||
|
$ ln -s ../../yolov5/yolov3_web_model public/yolov3_web_model
|
||||||
|
$ npm start
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from torch.utils.mobile_optimizer import optimize_for_mobile
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[0] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||||
|
|
||||||
|
from models.common import Conv
|
||||||
|
from models.experimental import attempt_load
|
||||||
|
from models.yolo import Detect
|
||||||
|
from utils.activations import SiLU
|
||||||
|
from utils.datasets import LoadImages
|
||||||
|
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, colorstr, file_size, print_args,
|
||||||
|
url2file)
|
||||||
|
from utils.torch_utils import select_device
|
||||||
|
|
||||||
|
|
||||||
|
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
|
||||||
|
# TorchScript model export
|
||||||
|
try:
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
|
||||||
|
f = file.with_suffix('.torchscript.pt')
|
||||||
|
|
||||||
|
ts = torch.jit.trace(model, im, strict=False)
|
||||||
|
d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
|
||||||
|
extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
|
||||||
|
(optimize_for_mobile(ts) if optimize else ts).save(f, _extra_files=extra_files)
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
|
||||||
|
# ONNX export
|
||||||
|
try:
|
||||||
|
check_requirements(('onnx',))
|
||||||
|
import onnx
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
|
||||||
|
f = file.with_suffix('.onnx')
|
||||||
|
|
||||||
|
torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
|
||||||
|
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
|
||||||
|
do_constant_folding=not train,
|
||||||
|
input_names=['images'],
|
||||||
|
output_names=['output'],
|
||||||
|
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
|
||||||
|
'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
|
||||||
|
} if dynamic else None)
|
||||||
|
|
||||||
|
# Checks
|
||||||
|
model_onnx = onnx.load(f) # load onnx model
|
||||||
|
onnx.checker.check_model(model_onnx) # check onnx model
|
||||||
|
# LOGGER.info(onnx.helper.printable_graph(model_onnx.graph)) # print
|
||||||
|
|
||||||
|
# Simplify
|
||||||
|
if simplify:
|
||||||
|
try:
|
||||||
|
check_requirements(('onnx-simplifier',))
|
||||||
|
import onnxsim
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
||||||
|
model_onnx, check = onnxsim.simplify(
|
||||||
|
model_onnx,
|
||||||
|
dynamic_input_shape=dynamic,
|
||||||
|
input_shapes={'images': list(im.shape)} if dynamic else None)
|
||||||
|
assert check, 'assert check failed'
|
||||||
|
onnx.save(model_onnx, f)
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'{prefix} simplifier failure: {e}')
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
|
||||||
|
# CoreML export
|
||||||
|
ct_model = None
|
||||||
|
try:
|
||||||
|
check_requirements(('coremltools',))
|
||||||
|
import coremltools as ct
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
|
||||||
|
f = file.with_suffix('.mlmodel')
|
||||||
|
|
||||||
|
model.train() # CoreML exports should be placed in model.train() mode
|
||||||
|
ts = torch.jit.trace(model, im, strict=False) # TorchScript model
|
||||||
|
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
|
||||||
|
ct_model.save(f)
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'\n{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
return ct_model
|
||||||
|
|
||||||
|
|
||||||
|
def export_saved_model(model, im, file, dynamic,
|
||||||
|
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
|
||||||
|
conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
|
||||||
|
# TensorFlow saved_model export
|
||||||
|
keras_model = None
|
||||||
|
try:
|
||||||
|
import tensorflow as tf
|
||||||
|
from tensorflow import keras
|
||||||
|
|
||||||
|
from models.tf import TFDetect, TFModel
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
|
||||||
|
f = str(file).replace('.pt', '_saved_model')
|
||||||
|
batch_size, ch, *imgsz = list(im.shape) # BCHW
|
||||||
|
|
||||||
|
tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
|
||||||
|
im = tf.zeros((batch_size, *imgsz, 3)) # BHWC order for TensorFlow
|
||||||
|
y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
|
||||||
|
inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
|
||||||
|
outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
|
||||||
|
keras_model = keras.Model(inputs=inputs, outputs=outputs)
|
||||||
|
keras_model.trainable = False
|
||||||
|
keras_model.summary()
|
||||||
|
keras_model.save(f, save_format='tf')
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'\n{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
return keras_model
|
||||||
|
|
||||||
|
|
||||||
|
def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
|
||||||
|
# TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
|
||||||
|
try:
|
||||||
|
import tensorflow as tf
|
||||||
|
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
|
||||||
|
f = file.with_suffix('.pb')
|
||||||
|
|
||||||
|
m = tf.function(lambda x: keras_model(x)) # full model
|
||||||
|
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
|
||||||
|
frozen_func = convert_variables_to_constants_v2(m)
|
||||||
|
frozen_func.graph.as_graph_def()
|
||||||
|
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'\n{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
|
||||||
|
# TensorFlow Lite export
|
||||||
|
try:
|
||||||
|
import tensorflow as tf
|
||||||
|
|
||||||
|
from models.tf import representative_dataset_gen
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
|
||||||
|
batch_size, ch, *imgsz = list(im.shape) # BCHW
|
||||||
|
f = str(file).replace('.pt', '-fp16.tflite')
|
||||||
|
|
||||||
|
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
|
||||||
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
|
||||||
|
converter.target_spec.supported_types = [tf.float16]
|
||||||
|
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
||||||
|
if int8:
|
||||||
|
dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
|
||||||
|
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
|
||||||
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
||||||
|
converter.target_spec.supported_types = []
|
||||||
|
converter.inference_input_type = tf.uint8 # or tf.int8
|
||||||
|
converter.inference_output_type = tf.uint8 # or tf.int8
|
||||||
|
converter.experimental_new_quantizer = False
|
||||||
|
f = str(file).replace('.pt', '-int8.tflite')
|
||||||
|
|
||||||
|
tflite_model = converter.convert()
|
||||||
|
open(f, "wb").write(tflite_model)
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'\n{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
|
||||||
|
# TensorFlow.js export
|
||||||
|
try:
|
||||||
|
check_requirements(('tensorflowjs',))
|
||||||
|
import re
|
||||||
|
|
||||||
|
import tensorflowjs as tfjs
|
||||||
|
|
||||||
|
LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
|
||||||
|
f = str(file).replace('.pt', '_web_model') # js dir
|
||||||
|
f_pb = file.with_suffix('.pb') # *.pb path
|
||||||
|
f_json = f + '/model.json' # *.json path
|
||||||
|
|
||||||
|
cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \
|
||||||
|
f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
|
||||||
|
subprocess.run(cmd, shell=True)
|
||||||
|
|
||||||
|
json = open(f_json).read()
|
||||||
|
with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
|
||||||
|
subst = re.sub(
|
||||||
|
r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
|
||||||
|
r'"Identity.?.?": {"name": "Identity.?.?"}, '
|
||||||
|
r'"Identity.?.?": {"name": "Identity.?.?"}, '
|
||||||
|
r'"Identity.?.?": {"name": "Identity.?.?"}}}',
|
||||||
|
r'{"outputs": {"Identity": {"name": "Identity"}, '
|
||||||
|
r'"Identity_1": {"name": "Identity_1"}, '
|
||||||
|
r'"Identity_2": {"name": "Identity_2"}, '
|
||||||
|
r'"Identity_3": {"name": "Identity_3"}}}',
|
||||||
|
json)
|
||||||
|
j.write(subst)
|
||||||
|
|
||||||
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'\n{prefix} export failure: {e}')
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
||||||
|
weights=ROOT / 'yolov3.pt', # weights path
|
||||||
|
imgsz=(640, 640), # image (height, width)
|
||||||
|
batch_size=1, # batch size
|
||||||
|
device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||||
|
include=('torchscript', 'onnx', 'coreml'), # include formats
|
||||||
|
half=False, # FP16 half-precision export
|
||||||
|
inplace=False, # set Detect() inplace=True
|
||||||
|
train=False, # model.train() mode
|
||||||
|
optimize=False, # TorchScript: optimize for mobile
|
||||||
|
int8=False, # CoreML/TF INT8 quantization
|
||||||
|
dynamic=False, # ONNX/TF: dynamic axes
|
||||||
|
simplify=False, # ONNX: simplify model
|
||||||
|
opset=12, # ONNX: opset version
|
||||||
|
topk_per_class=100, # TF.js NMS: topk per class to keep
|
||||||
|
topk_all=100, # TF.js NMS: topk for all classes to keep
|
||||||
|
iou_thres=0.45, # TF.js NMS: IoU threshold
|
||||||
|
conf_thres=0.25 # TF.js NMS: confidence threshold
|
||||||
|
):
|
||||||
|
t = time.time()
|
||||||
|
include = [x.lower() for x in include]
|
||||||
|
tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
|
||||||
|
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
||||||
|
file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
|
||||||
|
|
||||||
|
# Load PyTorch model
|
||||||
|
device = select_device(device)
|
||||||
|
assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
|
||||||
|
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
||||||
|
nc, names = model.nc, model.names # number of classes, class names
|
||||||
|
|
||||||
|
# Input
|
||||||
|
gs = int(max(model.stride)) # grid size (max stride)
|
||||||
|
imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples
|
||||||
|
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
|
||||||
|
|
||||||
|
# Update model
|
||||||
|
if half:
|
||||||
|
im, model = im.half(), model.half() # to FP16
|
||||||
|
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
||||||
|
for k, m in model.named_modules():
|
||||||
|
if isinstance(m, Conv): # assign export-friendly activations
|
||||||
|
if isinstance(m.act, nn.SiLU):
|
||||||
|
m.act = SiLU()
|
||||||
|
elif isinstance(m, Detect):
|
||||||
|
m.inplace = inplace
|
||||||
|
m.onnx_dynamic = dynamic
|
||||||
|
# m.forward = m.forward_export # assign forward (optional)
|
||||||
|
|
||||||
|
for _ in range(2):
|
||||||
|
y = model(im) # dry runs
|
||||||
|
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
|
||||||
|
|
||||||
|
# Exports
|
||||||
|
if 'torchscript' in include:
|
||||||
|
export_torchscript(model, im, file, optimize)
|
||||||
|
if 'onnx' in include:
|
||||||
|
export_onnx(model, im, file, opset, train, dynamic, simplify)
|
||||||
|
if 'coreml' in include:
|
||||||
|
export_coreml(model, im, file)
|
||||||
|
|
||||||
|
# TensorFlow Exports
|
||||||
|
if any(tf_exports):
|
||||||
|
pb, tflite, tfjs = tf_exports[1:]
|
||||||
|
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
|
||||||
|
model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs,
|
||||||
|
topk_per_class=topk_per_class, topk_all=topk_all, conf_thres=conf_thres,
|
||||||
|
iou_thres=iou_thres) # keras model
|
||||||
|
if pb or tfjs: # pb prerequisite to tfjs
|
||||||
|
export_pb(model, im, file)
|
||||||
|
if tflite:
|
||||||
|
export_tflite(model, im, file, int8=int8, data=data, ncalib=100)
|
||||||
|
if tfjs:
|
||||||
|
export_tfjs(model, im, file)
|
||||||
|
|
||||||
|
# Finish
|
||||||
|
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
|
||||||
|
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
|
||||||
|
f'\nVisualize with https://netron.app')
|
||||||
|
|
||||||
|
|
||||||
|
def parse_opt():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
||||||
|
parser.add_argument('--weights', type=str, default=ROOT / 'yolov3.pt', help='weights path')
|
||||||
|
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
|
||||||
|
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
||||||
|
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||||
|
parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
|
||||||
|
parser.add_argument('--inplace', action='store_true', help='set YOLOv3 Detect() inplace=True')
|
||||||
|
parser.add_argument('--train', action='store_true', help='model.train() mode')
|
||||||
|
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
|
||||||
|
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
|
||||||
|
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
|
||||||
|
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
|
||||||
|
parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version')
|
||||||
|
parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
|
||||||
|
parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
|
||||||
|
parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
|
||||||
|
parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
|
||||||
|
parser.add_argument('--include', nargs='+',
|
||||||
|
default=['torchscript', 'onnx'],
|
||||||
|
help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)')
|
||||||
|
opt = parser.parse_args()
|
||||||
|
print_args(FILE.stem, opt)
|
||||||
|
return opt
|
||||||
|
|
||||||
|
|
||||||
|
def main(opt):
|
||||||
|
run(**vars(opt))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
opt = parse_opt()
|
||||||
|
main(opt)
|
||||||
46
hubconf.py
46
hubconf.py
@ -1,56 +1,61 @@
|
|||||||
"""YOLOv3 PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov3/
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5/
|
||||||
|
|
||||||
Usage:
|
Usage:
|
||||||
import torch
|
import torch
|
||||||
model = torch.hub.load('ultralytics/yolov3', 'yolov3_tiny')
|
model = torch.hub.load('ultralytics/yolov3', 'yolov3')
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
|
||||||
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||||
"""Creates a specified YOLOv3 model
|
"""Creates a specified model
|
||||||
|
|
||||||
Arguments:
|
Arguments:
|
||||||
name (str): name of model, i.e. 'yolov3'
|
name (str): name of model, i.e. 'yolov3'
|
||||||
pretrained (bool): load pretrained weights into the model
|
pretrained (bool): load pretrained weights into the model
|
||||||
channels (int): number of input channels
|
channels (int): number of input channels
|
||||||
classes (int): number of model classes
|
classes (int): number of model classes
|
||||||
autoshape (bool): apply YOLOv3 .autoshape() wrapper to model
|
autoshape (bool): apply .autoshape() wrapper to model
|
||||||
verbose (bool): print all information to screen
|
verbose (bool): print all information to screen
|
||||||
device (str, torch.device, None): device to use for model parameters
|
device (str, torch.device, None): device to use for model parameters
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
YOLOv3 pytorch model
|
pytorch model
|
||||||
"""
|
"""
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from models.yolo import Model, attempt_load
|
from models.experimental import attempt_load
|
||||||
from utils.general import check_requirements, set_logging
|
from models.yolo import Model
|
||||||
from utils.google_utils import attempt_download
|
from utils.downloads import attempt_download
|
||||||
|
from utils.general import check_requirements, intersect_dicts, set_logging
|
||||||
from utils.torch_utils import select_device
|
from utils.torch_utils import select_device
|
||||||
|
|
||||||
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop'))
|
file = Path(__file__).resolve()
|
||||||
|
check_requirements(exclude=('tensorboard', 'thop', 'opencv-python'))
|
||||||
set_logging(verbose=verbose)
|
set_logging(verbose=verbose)
|
||||||
|
|
||||||
fname = Path(name).with_suffix('.pt') # checkpoint filename
|
save_dir = Path('') if str(name).endswith('.pt') else file.parent
|
||||||
|
path = (save_dir / name).with_suffix('.pt') # checkpoint path
|
||||||
try:
|
try:
|
||||||
|
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
|
||||||
|
|
||||||
if pretrained and channels == 3 and classes == 80:
|
if pretrained and channels == 3 and classes == 80:
|
||||||
model = attempt_load(fname, map_location=torch.device('cpu')) # download/load FP32 model
|
model = attempt_load(path, map_location=device) # download/load FP32 model
|
||||||
else:
|
else:
|
||||||
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
|
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
|
||||||
model = Model(cfg, channels, classes) # create model
|
model = Model(cfg, channels, classes) # create model
|
||||||
if pretrained:
|
if pretrained:
|
||||||
ckpt = torch.load(attempt_download(fname), map_location=torch.device('cpu')) # load
|
ckpt = torch.load(attempt_download(path), map_location=device) # load
|
||||||
msd = model.state_dict() # model state_dict
|
|
||||||
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
|
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
|
||||||
csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
|
csd = intersect_dicts(csd, model.state_dict(), exclude=['anchors']) # intersect
|
||||||
model.load_state_dict(csd, strict=False) # load
|
model.load_state_dict(csd, strict=False) # load
|
||||||
if len(ckpt['model'].names) == classes:
|
if len(ckpt['model'].names) == classes:
|
||||||
model.names = ckpt['model'].names # set class names attribute
|
model.names = ckpt['model'].names # set class names attribute
|
||||||
if autoshape:
|
if autoshape:
|
||||||
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
|
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
|
||||||
device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device)
|
|
||||||
return model.to(device)
|
return model.to(device)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@ -60,7 +65,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
|
|||||||
|
|
||||||
|
|
||||||
def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None):
|
def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None):
|
||||||
# YOLOv3 custom or local model
|
# custom or local model
|
||||||
return _create(path, autoshape=autoshape, verbose=verbose, device=device)
|
return _create(path, autoshape=autoshape, verbose=verbose, device=device)
|
||||||
|
|
||||||
|
|
||||||
@ -68,26 +73,31 @@ def yolov3(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True
|
|||||||
# YOLOv3 model https://github.com/ultralytics/yolov3
|
# YOLOv3 model https://github.com/ultralytics/yolov3
|
||||||
return _create('yolov3', pretrained, channels, classes, autoshape, verbose, device)
|
return _create('yolov3', pretrained, channels, classes, autoshape, verbose, device)
|
||||||
|
|
||||||
|
|
||||||
def yolov3_spp(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
def yolov3_spp(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||||
# YOLOv3-SPP model https://github.com/ultralytics/yolov3
|
# YOLOv3-SPP model https://github.com/ultralytics/yolov3
|
||||||
return _create('yolov3-spp', pretrained, channels, classes, autoshape, verbose, device)
|
return _create('yolov3-spp', pretrained, channels, classes, autoshape, verbose, device)
|
||||||
|
|
||||||
|
|
||||||
def yolov3_tiny(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
def yolov3_tiny(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
||||||
# YOLOv3-tiny model https://github.com/ultralytics/yolov3
|
# YOLOv3-tiny model https://github.com/ultralytics/yolov3
|
||||||
return _create('yolov3-tiny', pretrained, channels, classes, autoshape, verbose, device)
|
return _create('yolov3-tiny', pretrained, channels, classes, autoshape, verbose, device)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
model = _create(name='yolov3', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
|
model = _create(name='yolov3-tiny', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
|
||||||
# model = custom(path='path/to/model.pt') # custom
|
# model = custom(path='path/to/model.pt') # custom
|
||||||
|
|
||||||
# Verify inference
|
# Verify inference
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
imgs = ['data/images/zidane.jpg', # filename
|
imgs = ['data/images/zidane.jpg', # filename
|
||||||
'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg', # URI
|
Path('data/images/zidane.jpg'), # Path
|
||||||
|
'https://ultralytics.com/images/zidane.jpg', # URI
|
||||||
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
|
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
|
||||||
Image.open('data/images/bus.jpg'), # PIL
|
Image.open('data/images/bus.jpg'), # PIL
|
||||||
np.zeros((320, 640, 3))] # numpy
|
np.zeros((320, 640, 3))] # numpy
|
||||||
|
|||||||
374
models/common.py
374
models/common.py
@ -1,9 +1,16 @@
|
|||||||
# YOLOv3 common modules
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Common modules
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
import math
|
import math
|
||||||
|
import platform
|
||||||
|
import warnings
|
||||||
from copy import copy
|
from copy import copy
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
import requests
|
import requests
|
||||||
@ -12,10 +19,11 @@ import torch.nn as nn
|
|||||||
from PIL import Image
|
from PIL import Image
|
||||||
from torch.cuda import amp
|
from torch.cuda import amp
|
||||||
|
|
||||||
from utils.datasets import letterbox
|
from utils.datasets import exif_transpose, letterbox
|
||||||
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box
|
from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible,
|
||||||
from utils.plots import colors, plot_one_box
|
non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
|
||||||
from utils.torch_utils import time_synchronized
|
from utils.plots import Annotator, colors, save_one_box
|
||||||
|
from utils.torch_utils import time_sync
|
||||||
|
|
||||||
|
|
||||||
def autopad(k, p=None): # kernel, padding
|
def autopad(k, p=None): # kernel, padding
|
||||||
@ -25,26 +33,27 @@ def autopad(k, p=None): # kernel, padding
|
|||||||
return p
|
return p
|
||||||
|
|
||||||
|
|
||||||
def DWConv(c1, c2, k=1, s=1, act=True):
|
|
||||||
# Depthwise convolution
|
|
||||||
return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
|
|
||||||
|
|
||||||
|
|
||||||
class Conv(nn.Module):
|
class Conv(nn.Module):
|
||||||
# Standard convolution
|
# Standard convolution
|
||||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
||||||
super(Conv, self).__init__()
|
super().__init__()
|
||||||
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
|
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
|
||||||
self.bn = nn.BatchNorm2d(c2)
|
self.bn = nn.BatchNorm2d(c2)
|
||||||
self.act = nn.LeakyReLU(0.1) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
|
self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
return self.act(self.bn(self.conv(x)))
|
return self.act(self.bn(self.conv(x)))
|
||||||
|
|
||||||
def fuseforward(self, x):
|
def forward_fuse(self, x):
|
||||||
return self.act(self.conv(x))
|
return self.act(self.conv(x))
|
||||||
|
|
||||||
|
|
||||||
|
class DWConv(Conv):
|
||||||
|
# Depth-wise convolution class
|
||||||
|
def __init__(self, c1, c2, k=1, s=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
||||||
|
super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
|
||||||
|
|
||||||
|
|
||||||
class TransformerLayer(nn.Module):
|
class TransformerLayer(nn.Module):
|
||||||
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
|
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
|
||||||
def __init__(self, c, num_heads):
|
def __init__(self, c, num_heads):
|
||||||
@ -70,31 +79,21 @@ class TransformerBlock(nn.Module):
|
|||||||
if c1 != c2:
|
if c1 != c2:
|
||||||
self.conv = Conv(c1, c2)
|
self.conv = Conv(c1, c2)
|
||||||
self.linear = nn.Linear(c2, c2) # learnable position embedding
|
self.linear = nn.Linear(c2, c2) # learnable position embedding
|
||||||
self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)])
|
self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
|
||||||
self.c2 = c2
|
self.c2 = c2
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
if self.conv is not None:
|
if self.conv is not None:
|
||||||
x = self.conv(x)
|
x = self.conv(x)
|
||||||
b, _, w, h = x.shape
|
b, _, w, h = x.shape
|
||||||
p = x.flatten(2)
|
p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3)
|
||||||
p = p.unsqueeze(0)
|
return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h)
|
||||||
p = p.transpose(0, 3)
|
|
||||||
p = p.squeeze(3)
|
|
||||||
e = self.linear(p)
|
|
||||||
x = p + e
|
|
||||||
|
|
||||||
x = self.tr(x)
|
|
||||||
x = x.unsqueeze(3)
|
|
||||||
x = x.transpose(0, 3)
|
|
||||||
x = x.reshape(b, self.c2, w, h)
|
|
||||||
return x
|
|
||||||
|
|
||||||
|
|
||||||
class Bottleneck(nn.Module):
|
class Bottleneck(nn.Module):
|
||||||
# Standard bottleneck
|
# Standard bottleneck
|
||||||
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
|
||||||
super(Bottleneck, self).__init__()
|
super().__init__()
|
||||||
c_ = int(c2 * e) # hidden channels
|
c_ = int(c2 * e) # hidden channels
|
||||||
self.cv1 = Conv(c1, c_, 1, 1)
|
self.cv1 = Conv(c1, c_, 1, 1)
|
||||||
self.cv2 = Conv(c_, c2, 3, 1, g=g)
|
self.cv2 = Conv(c_, c2, 3, 1, g=g)
|
||||||
@ -107,15 +106,15 @@ class Bottleneck(nn.Module):
|
|||||||
class BottleneckCSP(nn.Module):
|
class BottleneckCSP(nn.Module):
|
||||||
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
|
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
|
||||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||||
super(BottleneckCSP, self).__init__()
|
super().__init__()
|
||||||
c_ = int(c2 * e) # hidden channels
|
c_ = int(c2 * e) # hidden channels
|
||||||
self.cv1 = Conv(c1, c_, 1, 1)
|
self.cv1 = Conv(c1, c_, 1, 1)
|
||||||
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
|
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
|
||||||
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
|
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
|
||||||
self.cv4 = Conv(2 * c_, c2, 1, 1)
|
self.cv4 = Conv(2 * c_, c2, 1, 1)
|
||||||
self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
|
self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
|
||||||
self.act = nn.LeakyReLU(0.1, inplace=True)
|
self.act = nn.SiLU()
|
||||||
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
|
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
y1 = self.cv3(self.m(self.cv1(x)))
|
y1 = self.cv3(self.m(self.cv1(x)))
|
||||||
@ -126,12 +125,12 @@ class BottleneckCSP(nn.Module):
|
|||||||
class C3(nn.Module):
|
class C3(nn.Module):
|
||||||
# CSP Bottleneck with 3 convolutions
|
# CSP Bottleneck with 3 convolutions
|
||||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||||
super(C3, self).__init__()
|
super().__init__()
|
||||||
c_ = int(c2 * e) # hidden channels
|
c_ = int(c2 * e) # hidden channels
|
||||||
self.cv1 = Conv(c1, c_, 1, 1)
|
self.cv1 = Conv(c1, c_, 1, 1)
|
||||||
self.cv2 = Conv(c1, c_, 1, 1)
|
self.cv2 = Conv(c1, c_, 1, 1)
|
||||||
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
|
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
|
||||||
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
|
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
|
||||||
# self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
|
# self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
@ -146,10 +145,26 @@ class C3TR(C3):
|
|||||||
self.m = TransformerBlock(c_, c_, 4, n)
|
self.m = TransformerBlock(c_, c_, 4, n)
|
||||||
|
|
||||||
|
|
||||||
|
class C3SPP(C3):
|
||||||
|
# C3 module with SPP()
|
||||||
|
def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
|
||||||
|
super().__init__(c1, c2, n, shortcut, g, e)
|
||||||
|
c_ = int(c2 * e)
|
||||||
|
self.m = SPP(c_, c_, k)
|
||||||
|
|
||||||
|
|
||||||
|
class C3Ghost(C3):
|
||||||
|
# C3 module with GhostBottleneck()
|
||||||
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
||||||
|
super().__init__(c1, c2, n, shortcut, g, e)
|
||||||
|
c_ = int(c2 * e) # hidden channels
|
||||||
|
self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
|
||||||
|
|
||||||
|
|
||||||
class SPP(nn.Module):
|
class SPP(nn.Module):
|
||||||
# Spatial pyramid pooling layer used in YOLOv3-SPP
|
# Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729
|
||||||
def __init__(self, c1, c2, k=(5, 9, 13)):
|
def __init__(self, c1, c2, k=(5, 9, 13)):
|
||||||
super(SPP, self).__init__()
|
super().__init__()
|
||||||
c_ = c1 // 2 # hidden channels
|
c_ = c1 // 2 # hidden channels
|
||||||
self.cv1 = Conv(c1, c_, 1, 1)
|
self.cv1 = Conv(c1, c_, 1, 1)
|
||||||
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
|
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
|
||||||
@ -157,13 +172,33 @@ class SPP(nn.Module):
|
|||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
x = self.cv1(x)
|
x = self.cv1(x)
|
||||||
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
|
with warnings.catch_warnings():
|
||||||
|
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
|
||||||
|
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
|
||||||
|
|
||||||
|
|
||||||
|
class SPPF(nn.Module):
|
||||||
|
# Spatial Pyramid Pooling - Fast (SPPF) layer for by Glenn Jocher
|
||||||
|
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
|
||||||
|
super().__init__()
|
||||||
|
c_ = c1 // 2 # hidden channels
|
||||||
|
self.cv1 = Conv(c1, c_, 1, 1)
|
||||||
|
self.cv2 = Conv(c_ * 4, c2, 1, 1)
|
||||||
|
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.cv1(x)
|
||||||
|
with warnings.catch_warnings():
|
||||||
|
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
|
||||||
|
y1 = self.m(x)
|
||||||
|
y2 = self.m(y1)
|
||||||
|
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
|
||||||
|
|
||||||
|
|
||||||
class Focus(nn.Module):
|
class Focus(nn.Module):
|
||||||
# Focus wh information into c-space
|
# Focus wh information into c-space
|
||||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
||||||
super(Focus, self).__init__()
|
super().__init__()
|
||||||
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
|
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
|
||||||
# self.contract = Contract(gain=2)
|
# self.contract = Contract(gain=2)
|
||||||
|
|
||||||
@ -172,6 +207,34 @@ class Focus(nn.Module):
|
|||||||
# return self.conv(self.contract(x))
|
# return self.conv(self.contract(x))
|
||||||
|
|
||||||
|
|
||||||
|
class GhostConv(nn.Module):
|
||||||
|
# Ghost Convolution https://github.com/huawei-noah/ghostnet
|
||||||
|
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
||||||
|
super().__init__()
|
||||||
|
c_ = c2 // 2 # hidden channels
|
||||||
|
self.cv1 = Conv(c1, c_, k, s, None, g, act)
|
||||||
|
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
y = self.cv1(x)
|
||||||
|
return torch.cat([y, self.cv2(y)], 1)
|
||||||
|
|
||||||
|
|
||||||
|
class GhostBottleneck(nn.Module):
|
||||||
|
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
|
||||||
|
def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
|
||||||
|
super().__init__()
|
||||||
|
c_ = c2 // 2
|
||||||
|
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
|
||||||
|
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
|
||||||
|
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
|
||||||
|
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
|
||||||
|
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.conv(x) + self.shortcut(x)
|
||||||
|
|
||||||
|
|
||||||
class Contract(nn.Module):
|
class Contract(nn.Module):
|
||||||
# Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
|
# Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
|
||||||
def __init__(self, gain=2):
|
def __init__(self, gain=2):
|
||||||
@ -179,11 +242,11 @@ class Contract(nn.Module):
|
|||||||
self.gain = gain
|
self.gain = gain
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
|
b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
|
||||||
s = self.gain
|
s = self.gain
|
||||||
x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)
|
x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)
|
||||||
x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
|
x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
|
||||||
return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)
|
return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)
|
||||||
|
|
||||||
|
|
||||||
class Expand(nn.Module):
|
class Expand(nn.Module):
|
||||||
@ -193,64 +256,183 @@ class Expand(nn.Module):
|
|||||||
self.gain = gain
|
self.gain = gain
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
|
b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
|
||||||
s = self.gain
|
s = self.gain
|
||||||
x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80)
|
x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80)
|
||||||
x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
|
x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
|
||||||
return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160)
|
return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160)
|
||||||
|
|
||||||
|
|
||||||
class Concat(nn.Module):
|
class Concat(nn.Module):
|
||||||
# Concatenate a list of tensors along dimension
|
# Concatenate a list of tensors along dimension
|
||||||
def __init__(self, dimension=1):
|
def __init__(self, dimension=1):
|
||||||
super(Concat, self).__init__()
|
super().__init__()
|
||||||
self.d = dimension
|
self.d = dimension
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
return torch.cat(x, self.d)
|
return torch.cat(x, self.d)
|
||||||
|
|
||||||
|
|
||||||
class NMS(nn.Module):
|
class DetectMultiBackend(nn.Module):
|
||||||
# Non-Maximum Suppression (NMS) module
|
# MultiBackend class for python inference on various backends
|
||||||
conf = 0.25 # confidence threshold
|
def __init__(self, weights='yolov3.pt', device=None, dnn=True):
|
||||||
iou = 0.45 # IoU threshold
|
# Usage:
|
||||||
classes = None # (optional list) filter by class
|
# PyTorch: weights = *.pt
|
||||||
max_det = 1000 # maximum number of detections per image
|
# TorchScript: *.torchscript.pt
|
||||||
|
# CoreML: *.mlmodel
|
||||||
|
# TensorFlow: *_saved_model
|
||||||
|
# TensorFlow: *.pb
|
||||||
|
# TensorFlow Lite: *.tflite
|
||||||
|
# ONNX Runtime: *.onnx
|
||||||
|
# OpenCV DNN: *.onnx with dnn=True
|
||||||
|
super().__init__()
|
||||||
|
w = str(weights[0] if isinstance(weights, list) else weights)
|
||||||
|
suffix, suffixes = Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '', '.mlmodel']
|
||||||
|
check_suffix(w, suffixes) # check weights have acceptable suffix
|
||||||
|
pt, onnx, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans
|
||||||
|
jit = pt and 'torchscript' in w.lower()
|
||||||
|
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
||||||
|
|
||||||
def __init__(self):
|
if jit: # TorchScript
|
||||||
super(NMS, self).__init__()
|
LOGGER.info(f'Loading {w} for TorchScript inference...')
|
||||||
|
extra_files = {'config.txt': ''} # model metadata
|
||||||
|
model = torch.jit.load(w, _extra_files=extra_files)
|
||||||
|
if extra_files['config.txt']:
|
||||||
|
d = json.loads(extra_files['config.txt']) # extra_files dict
|
||||||
|
stride, names = int(d['stride']), d['names']
|
||||||
|
elif pt: # PyTorch
|
||||||
|
from models.experimental import attempt_load # scoped to avoid circular import
|
||||||
|
model = torch.jit.load(w) if 'torchscript' in w else attempt_load(weights, map_location=device)
|
||||||
|
stride = int(model.stride.max()) # model stride
|
||||||
|
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
||||||
|
elif coreml: # CoreML *.mlmodel
|
||||||
|
import coremltools as ct
|
||||||
|
model = ct.models.MLModel(w)
|
||||||
|
elif dnn: # ONNX OpenCV DNN
|
||||||
|
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
|
||||||
|
check_requirements(('opencv-python>=4.5.4',))
|
||||||
|
net = cv2.dnn.readNetFromONNX(w)
|
||||||
|
elif onnx: # ONNX Runtime
|
||||||
|
LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
|
||||||
|
check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime'))
|
||||||
|
import onnxruntime
|
||||||
|
session = onnxruntime.InferenceSession(w, None)
|
||||||
|
else: # TensorFlow model (TFLite, pb, saved_model)
|
||||||
|
import tensorflow as tf
|
||||||
|
if pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
|
||||||
|
def wrap_frozen_graph(gd, inputs, outputs):
|
||||||
|
x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
|
||||||
|
return x.prune(tf.nest.map_structure(x.graph.as_graph_element, inputs),
|
||||||
|
tf.nest.map_structure(x.graph.as_graph_element, outputs))
|
||||||
|
|
||||||
def forward(self, x):
|
LOGGER.info(f'Loading {w} for TensorFlow *.pb inference...')
|
||||||
return non_max_suppression(x[0], self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det)
|
graph_def = tf.Graph().as_graph_def()
|
||||||
|
graph_def.ParseFromString(open(w, 'rb').read())
|
||||||
|
frozen_func = wrap_frozen_graph(gd=graph_def, inputs="x:0", outputs="Identity:0")
|
||||||
|
elif saved_model:
|
||||||
|
LOGGER.info(f'Loading {w} for TensorFlow saved_model inference...')
|
||||||
|
model = tf.keras.models.load_model(w)
|
||||||
|
elif tflite: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
|
||||||
|
if 'edgetpu' in w.lower():
|
||||||
|
LOGGER.info(f'Loading {w} for TensorFlow Edge TPU inference...')
|
||||||
|
import tflite_runtime.interpreter as tfli
|
||||||
|
delegate = {'Linux': 'libedgetpu.so.1', # install https://coral.ai/software/#edgetpu-runtime
|
||||||
|
'Darwin': 'libedgetpu.1.dylib',
|
||||||
|
'Windows': 'edgetpu.dll'}[platform.system()]
|
||||||
|
interpreter = tfli.Interpreter(model_path=w, experimental_delegates=[tfli.load_delegate(delegate)])
|
||||||
|
else:
|
||||||
|
LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
|
||||||
|
interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model
|
||||||
|
interpreter.allocate_tensors() # allocate
|
||||||
|
input_details = interpreter.get_input_details() # inputs
|
||||||
|
output_details = interpreter.get_output_details() # outputs
|
||||||
|
self.__dict__.update(locals()) # assign all variables to self
|
||||||
|
|
||||||
|
def forward(self, im, augment=False, visualize=False, val=False):
|
||||||
|
# MultiBackend inference
|
||||||
|
b, ch, h, w = im.shape # batch, channel, height, width
|
||||||
|
if self.pt: # PyTorch
|
||||||
|
y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
|
||||||
|
return y if val else y[0]
|
||||||
|
elif self.coreml: # CoreML *.mlmodel
|
||||||
|
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
||||||
|
im = Image.fromarray((im[0] * 255).astype('uint8'))
|
||||||
|
# im = im.resize((192, 320), Image.ANTIALIAS)
|
||||||
|
y = self.model.predict({'image': im}) # coordinates are xywh normalized
|
||||||
|
box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
|
||||||
|
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)
|
||||||
|
elif self.onnx: # ONNX
|
||||||
|
im = im.cpu().numpy() # torch to numpy
|
||||||
|
if self.dnn: # ONNX OpenCV DNN
|
||||||
|
self.net.setInput(im)
|
||||||
|
y = self.net.forward()
|
||||||
|
else: # ONNX Runtime
|
||||||
|
y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
|
||||||
|
else: # TensorFlow model (TFLite, pb, saved_model)
|
||||||
|
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
||||||
|
if self.pb:
|
||||||
|
y = self.frozen_func(x=self.tf.constant(im)).numpy()
|
||||||
|
elif self.saved_model:
|
||||||
|
y = self.model(im, training=False).numpy()
|
||||||
|
elif self.tflite:
|
||||||
|
input, output = self.input_details[0], self.output_details[0]
|
||||||
|
int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
|
||||||
|
if int8:
|
||||||
|
scale, zero_point = input['quantization']
|
||||||
|
im = (im / scale + zero_point).astype(np.uint8) # de-scale
|
||||||
|
self.interpreter.set_tensor(input['index'], im)
|
||||||
|
self.interpreter.invoke()
|
||||||
|
y = self.interpreter.get_tensor(output['index'])
|
||||||
|
if int8:
|
||||||
|
scale, zero_point = output['quantization']
|
||||||
|
y = (y.astype(np.float32) - zero_point) * scale # re-scale
|
||||||
|
y[..., 0] *= w # x
|
||||||
|
y[..., 1] *= h # y
|
||||||
|
y[..., 2] *= w # w
|
||||||
|
y[..., 3] *= h # h
|
||||||
|
y = torch.tensor(y)
|
||||||
|
return (y, []) if val else y
|
||||||
|
|
||||||
|
|
||||||
class AutoShape(nn.Module):
|
class AutoShape(nn.Module):
|
||||||
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
||||||
conf = 0.25 # NMS confidence threshold
|
conf = 0.25 # NMS confidence threshold
|
||||||
iou = 0.45 # NMS IoU threshold
|
iou = 0.45 # NMS IoU threshold
|
||||||
classes = None # (optional list) filter by class
|
classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
|
||||||
|
multi_label = False # NMS multiple labels per box
|
||||||
max_det = 1000 # maximum number of detections per image
|
max_det = 1000 # maximum number of detections per image
|
||||||
|
|
||||||
def __init__(self, model):
|
def __init__(self, model):
|
||||||
super(AutoShape, self).__init__()
|
super().__init__()
|
||||||
self.model = model.eval()
|
self.model = model.eval()
|
||||||
|
|
||||||
def autoshape(self):
|
def autoshape(self):
|
||||||
print('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
|
||||||
|
return self
|
||||||
|
|
||||||
|
def _apply(self, fn):
|
||||||
|
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
|
||||||
|
self = super()._apply(fn)
|
||||||
|
m = self.model.model[-1] # Detect()
|
||||||
|
m.stride = fn(m.stride)
|
||||||
|
m.grid = list(map(fn, m.grid))
|
||||||
|
if isinstance(m.anchor_grid, list):
|
||||||
|
m.anchor_grid = list(map(fn, m.anchor_grid))
|
||||||
return self
|
return self
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
def forward(self, imgs, size=640, augment=False, profile=False):
|
def forward(self, imgs, size=640, augment=False, profile=False):
|
||||||
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
||||||
# filename: imgs = 'data/images/zidane.jpg'
|
# file: imgs = 'data/images/zidane.jpg' # str or PosixPath
|
||||||
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
|
# URI: = 'https://ultralytics.com/images/zidane.jpg'
|
||||||
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
||||||
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
|
# PIL: = Image.open('image.jpg') or ImageGrab.grab() # HWC x(640,1280,3)
|
||||||
# numpy: = np.zeros((640,1280,3)) # HWC
|
# numpy: = np.zeros((640,1280,3)) # HWC
|
||||||
# torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
|
# torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
|
||||||
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
|
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
|
||||||
|
|
||||||
t = [time_synchronized()]
|
t = [time_sync()]
|
||||||
p = next(self.model.parameters()) # for device and type
|
p = next(self.model.parameters()) # for device and type
|
||||||
if isinstance(imgs, torch.Tensor): # torch
|
if isinstance(imgs, torch.Tensor): # torch
|
||||||
with amp.autocast(enabled=p.device.type != 'cpu'):
|
with amp.autocast(enabled=p.device.type != 'cpu'):
|
||||||
@ -261,14 +443,15 @@ class AutoShape(nn.Module):
|
|||||||
shape0, shape1, files = [], [], [] # image and inference shapes, filenames
|
shape0, shape1, files = [], [], [] # image and inference shapes, filenames
|
||||||
for i, im in enumerate(imgs):
|
for i, im in enumerate(imgs):
|
||||||
f = f'image{i}' # filename
|
f = f'image{i}' # filename
|
||||||
if isinstance(im, str): # filename or uri
|
if isinstance(im, (str, Path)): # filename or uri
|
||||||
im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
|
im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
|
||||||
|
im = np.asarray(exif_transpose(im))
|
||||||
elif isinstance(im, Image.Image): # PIL Image
|
elif isinstance(im, Image.Image): # PIL Image
|
||||||
im, f = np.asarray(im), getattr(im, 'filename', f) or f
|
im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f
|
||||||
files.append(Path(f).with_suffix('.jpg').name)
|
files.append(Path(f).with_suffix('.jpg').name)
|
||||||
if im.shape[0] < 5: # image in CHW
|
if im.shape[0] < 5: # image in CHW
|
||||||
im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
|
im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
|
||||||
im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
|
im = im[..., :3] if im.ndim == 3 else np.tile(im[..., None], 3) # enforce 3ch input
|
||||||
s = im.shape[:2] # HWC
|
s = im.shape[:2] # HWC
|
||||||
shape0.append(s) # image shape
|
shape0.append(s) # image shape
|
||||||
g = (size / max(s)) # gain
|
g = (size / max(s)) # gain
|
||||||
@ -278,29 +461,30 @@ class AutoShape(nn.Module):
|
|||||||
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
||||||
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
|
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
|
||||||
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
|
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
|
||||||
x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
|
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
|
||||||
t.append(time_synchronized())
|
t.append(time_sync())
|
||||||
|
|
||||||
with amp.autocast(enabled=p.device.type != 'cpu'):
|
with amp.autocast(enabled=p.device.type != 'cpu'):
|
||||||
# Inference
|
# Inference
|
||||||
y = self.model(x, augment, profile)[0] # forward
|
y = self.model(x, augment, profile)[0] # forward
|
||||||
t.append(time_synchronized())
|
t.append(time_sync())
|
||||||
|
|
||||||
# Post-process
|
# Post-process
|
||||||
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, max_det=self.max_det) # NMS
|
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes,
|
||||||
|
multi_label=self.multi_label, max_det=self.max_det) # NMS
|
||||||
for i in range(n):
|
for i in range(n):
|
||||||
scale_coords(shape1, y[i][:, :4], shape0[i])
|
scale_coords(shape1, y[i][:, :4], shape0[i])
|
||||||
|
|
||||||
t.append(time_synchronized())
|
t.append(time_sync())
|
||||||
return Detections(imgs, y, files, t, self.names, x.shape)
|
return Detections(imgs, y, files, t, self.names, x.shape)
|
||||||
|
|
||||||
|
|
||||||
class Detections:
|
class Detections:
|
||||||
# detections class for YOLOv3 inference results
|
# detections class for inference results
|
||||||
def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
|
def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
|
||||||
super(Detections, self).__init__()
|
super().__init__()
|
||||||
d = pred[0].device # device
|
d = pred[0].device # device
|
||||||
gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
|
gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in imgs] # normalizations
|
||||||
self.imgs = imgs # list of images as numpy arrays
|
self.imgs = imgs # list of images as numpy arrays
|
||||||
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
|
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
|
||||||
self.names = names # class names
|
self.names = names # class names
|
||||||
@ -314,47 +498,59 @@ class Detections:
|
|||||||
self.s = shape # inference BCHW shape
|
self.s = shape # inference BCHW shape
|
||||||
|
|
||||||
def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
|
def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
|
||||||
|
crops = []
|
||||||
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
|
||||||
str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
|
s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string
|
||||||
if pred is not None:
|
if pred.shape[0]:
|
||||||
for c in pred[:, -1].unique():
|
for c in pred[:, -1].unique():
|
||||||
n = (pred[:, -1] == c).sum() # detections per class
|
n = (pred[:, -1] == c).sum() # detections per class
|
||||||
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
||||||
if show or save or render or crop:
|
if show or save or render or crop:
|
||||||
for *box, conf, cls in pred: # xyxy, confidence, class
|
annotator = Annotator(im, example=str(self.names))
|
||||||
|
for *box, conf, cls in reversed(pred): # xyxy, confidence, class
|
||||||
label = f'{self.names[int(cls)]} {conf:.2f}'
|
label = f'{self.names[int(cls)]} {conf:.2f}'
|
||||||
if crop:
|
if crop:
|
||||||
save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i])
|
file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
|
||||||
|
crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
|
||||||
|
'im': save_one_box(box, im, file=file, save=save)})
|
||||||
else: # all others
|
else: # all others
|
||||||
plot_one_box(box, im, label=label, color=colors(cls))
|
annotator.box_label(box, label, color=colors(cls))
|
||||||
|
im = annotator.im
|
||||||
|
else:
|
||||||
|
s += '(no detections)'
|
||||||
|
|
||||||
im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
|
im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
|
||||||
if pprint:
|
if pprint:
|
||||||
print(str.rstrip(', '))
|
LOGGER.info(s.rstrip(', '))
|
||||||
if show:
|
if show:
|
||||||
im.show(self.files[i]) # show
|
im.show(self.files[i]) # show
|
||||||
if save:
|
if save:
|
||||||
f = self.files[i]
|
f = self.files[i]
|
||||||
im.save(save_dir / f) # save
|
im.save(save_dir / f) # save
|
||||||
print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
|
if i == self.n - 1:
|
||||||
|
LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
|
||||||
if render:
|
if render:
|
||||||
self.imgs[i] = np.asarray(im)
|
self.imgs[i] = np.asarray(im)
|
||||||
|
if crop:
|
||||||
|
if save:
|
||||||
|
LOGGER.info(f'Saved results to {save_dir}\n')
|
||||||
|
return crops
|
||||||
|
|
||||||
def print(self):
|
def print(self):
|
||||||
self.display(pprint=True) # print results
|
self.display(pprint=True) # print results
|
||||||
print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)
|
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' %
|
||||||
|
self.t)
|
||||||
|
|
||||||
def show(self):
|
def show(self):
|
||||||
self.display(show=True) # show results
|
self.display(show=True) # show results
|
||||||
|
|
||||||
def save(self, save_dir='runs/hub/exp'):
|
def save(self, save_dir='runs/detect/exp'):
|
||||||
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
|
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
|
||||||
self.display(save=True, save_dir=save_dir) # save results
|
self.display(save=True, save_dir=save_dir) # save results
|
||||||
|
|
||||||
def crop(self, save_dir='runs/hub/exp'):
|
def crop(self, save=True, save_dir='runs/detect/exp'):
|
||||||
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir
|
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
|
||||||
self.display(crop=True, save_dir=save_dir) # crop results
|
return self.display(crop=True, save=save, save_dir=save_dir) # crop results
|
||||||
print(f'Saved results to {save_dir}\n')
|
|
||||||
|
|
||||||
def render(self):
|
def render(self):
|
||||||
self.display(render=True) # render results
|
self.display(render=True) # render results
|
||||||
@ -385,7 +581,7 @@ class Detections:
|
|||||||
class Classify(nn.Module):
|
class Classify(nn.Module):
|
||||||
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
|
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
|
||||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
|
||||||
super(Classify, self).__init__()
|
super().__init__()
|
||||||
self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
|
self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
|
||||||
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
|
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
|
||||||
self.flat = nn.Flatten()
|
self.flat = nn.Flatten()
|
||||||
|
|||||||
@ -1,18 +1,22 @@
|
|||||||
# YOLOv3 experimental modules
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Experimental modules
|
||||||
|
"""
|
||||||
|
import math
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
|
|
||||||
from models.common import Conv, DWConv
|
from models.common import Conv
|
||||||
from utils.google_utils import attempt_download
|
from utils.downloads import attempt_download
|
||||||
|
|
||||||
|
|
||||||
class CrossConv(nn.Module):
|
class CrossConv(nn.Module):
|
||||||
# Cross Convolution Downsample
|
# Cross Convolution Downsample
|
||||||
def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
|
def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
|
||||||
# ch_in, ch_out, kernel, stride, groups, expansion, shortcut
|
# ch_in, ch_out, kernel, stride, groups, expansion, shortcut
|
||||||
super(CrossConv, self).__init__()
|
super().__init__()
|
||||||
c_ = int(c2 * e) # hidden channels
|
c_ = int(c2 * e) # hidden channels
|
||||||
self.cv1 = Conv(c1, c_, (1, k), (1, s))
|
self.cv1 = Conv(c1, c_, (1, k), (1, s))
|
||||||
self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
|
self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
|
||||||
@ -25,11 +29,11 @@ class CrossConv(nn.Module):
|
|||||||
class Sum(nn.Module):
|
class Sum(nn.Module):
|
||||||
# Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
# Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
||||||
def __init__(self, n, weight=False): # n: number of inputs
|
def __init__(self, n, weight=False): # n: number of inputs
|
||||||
super(Sum, self).__init__()
|
super().__init__()
|
||||||
self.weight = weight # apply weights boolean
|
self.weight = weight # apply weights boolean
|
||||||
self.iter = range(n - 1) # iter object
|
self.iter = range(n - 1) # iter object
|
||||||
if weight:
|
if weight:
|
||||||
self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights
|
self.w = nn.Parameter(-torch.arange(1.0, n) / 2, requires_grad=True) # layer weights
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
y = x[0] # no weight
|
y = x[0] # no weight
|
||||||
@ -43,86 +47,66 @@ class Sum(nn.Module):
|
|||||||
return y
|
return y
|
||||||
|
|
||||||
|
|
||||||
class GhostConv(nn.Module):
|
|
||||||
# Ghost Convolution https://github.com/huawei-noah/ghostnet
|
|
||||||
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
|
||||||
super(GhostConv, self).__init__()
|
|
||||||
c_ = c2 // 2 # hidden channels
|
|
||||||
self.cv1 = Conv(c1, c_, k, s, None, g, act)
|
|
||||||
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
y = self.cv1(x)
|
|
||||||
return torch.cat([y, self.cv2(y)], 1)
|
|
||||||
|
|
||||||
|
|
||||||
class GhostBottleneck(nn.Module):
|
|
||||||
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
|
|
||||||
def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
|
|
||||||
super(GhostBottleneck, self).__init__()
|
|
||||||
c_ = c2 // 2
|
|
||||||
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
|
|
||||||
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
|
|
||||||
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
|
|
||||||
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
|
|
||||||
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
return self.conv(x) + self.shortcut(x)
|
|
||||||
|
|
||||||
|
|
||||||
class MixConv2d(nn.Module):
|
class MixConv2d(nn.Module):
|
||||||
# Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
|
# Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595
|
||||||
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
|
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kernel, stride, ch_strategy
|
||||||
super(MixConv2d, self).__init__()
|
super().__init__()
|
||||||
groups = len(k)
|
n = len(k) # number of convolutions
|
||||||
if equal_ch: # equal c_ per group
|
if equal_ch: # equal c_ per group
|
||||||
i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
|
i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices
|
||||||
c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
|
c_ = [(i == g).sum() for g in range(n)] # intermediate channels
|
||||||
else: # equal weight.numel() per group
|
else: # equal weight.numel() per group
|
||||||
b = [c2] + [0] * groups
|
b = [c2] + [0] * n
|
||||||
a = np.eye(groups + 1, groups, k=-1)
|
a = np.eye(n + 1, n, k=-1)
|
||||||
a -= np.roll(a, 1, axis=1)
|
a -= np.roll(a, 1, axis=1)
|
||||||
a *= np.array(k) ** 2
|
a *= np.array(k) ** 2
|
||||||
a[0] = 1
|
a[0] = 1
|
||||||
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
||||||
|
|
||||||
self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
|
self.m = nn.ModuleList(
|
||||||
|
[nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)])
|
||||||
self.bn = nn.BatchNorm2d(c2)
|
self.bn = nn.BatchNorm2d(c2)
|
||||||
self.act = nn.LeakyReLU(0.1, inplace=True)
|
self.act = nn.SiLU()
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
return self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
||||||
|
|
||||||
|
|
||||||
class Ensemble(nn.ModuleList):
|
class Ensemble(nn.ModuleList):
|
||||||
# Ensemble of models
|
# Ensemble of models
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
super(Ensemble, self).__init__()
|
super().__init__()
|
||||||
|
|
||||||
def forward(self, x, augment=False):
|
def forward(self, x, augment=False, profile=False, visualize=False):
|
||||||
y = []
|
y = []
|
||||||
for module in self:
|
for module in self:
|
||||||
y.append(module(x, augment)[0])
|
y.append(module(x, augment, profile, visualize)[0])
|
||||||
# y = torch.stack(y).max(0)[0] # max ensemble
|
# y = torch.stack(y).max(0)[0] # max ensemble
|
||||||
# y = torch.stack(y).mean(0) # mean ensemble
|
# y = torch.stack(y).mean(0) # mean ensemble
|
||||||
y = torch.cat(y, 1) # nms ensemble
|
y = torch.cat(y, 1) # nms ensemble
|
||||||
return y, None # inference, train output
|
return y, None # inference, train output
|
||||||
|
|
||||||
|
|
||||||
def attempt_load(weights, map_location=None, inplace=True):
|
def attempt_load(weights, map_location=None, inplace=True, fuse=True):
|
||||||
from models.yolo import Detect, Model
|
from models.yolo import Detect, Model
|
||||||
|
|
||||||
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
||||||
model = Ensemble()
|
model = Ensemble()
|
||||||
for w in weights if isinstance(weights, list) else [weights]:
|
for w in weights if isinstance(weights, list) else [weights]:
|
||||||
ckpt = torch.load(attempt_download(w), map_location=map_location) # load
|
ckpt = torch.load(attempt_download(w), map_location=map_location) # load
|
||||||
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
|
if fuse:
|
||||||
|
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
|
||||||
|
else:
|
||||||
|
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval()) # without layer fuse
|
||||||
|
|
||||||
# Compatibility updates
|
# Compatibility updates
|
||||||
for m in model.modules():
|
for m in model.modules():
|
||||||
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]:
|
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]:
|
||||||
m.inplace = inplace # pytorch 1.7.0 compatibility
|
m.inplace = inplace # pytorch 1.7.0 compatibility
|
||||||
|
if type(m) is Detect:
|
||||||
|
if not isinstance(m.anchor_grid, list): # new Detect Layer compatibility
|
||||||
|
delattr(m, 'anchor_grid')
|
||||||
|
setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl)
|
||||||
elif type(m) is Conv:
|
elif type(m) is Conv:
|
||||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
||||||
|
|
||||||
|
|||||||
145
models/export.py
145
models/export.py
@ -1,145 +0,0 @@
|
|||||||
"""Exports a YOLOv3 *.pt model to TorchScript, ONNX, CoreML formats
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
$ python path/to/models/export.py --weights yolov3.pt --img 640 --batch 1
|
|
||||||
"""
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import sys
|
|
||||||
import time
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
from torch.utils.mobile_optimizer import optimize_for_mobile
|
|
||||||
|
|
||||||
import models
|
|
||||||
from models.experimental import attempt_load
|
|
||||||
from utils.activations import Hardswish, SiLU
|
|
||||||
from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
|
|
||||||
from utils.torch_utils import select_device
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path')
|
|
||||||
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
|
|
||||||
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
|
||||||
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
|
||||||
parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
|
|
||||||
parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
|
|
||||||
parser.add_argument('--inplace', action='store_true', help='set YOLOv3 Detect() inplace=True')
|
|
||||||
parser.add_argument('--train', action='store_true', help='model.train() mode')
|
|
||||||
parser.add_argument('--optimize', action='store_true', help='optimize TorchScript for mobile') # TorchScript-only
|
|
||||||
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
|
|
||||||
parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
|
|
||||||
parser.add_argument('--opset-version', type=int, default=12, help='ONNX opset version') # ONNX-only
|
|
||||||
opt = parser.parse_args()
|
|
||||||
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
|
|
||||||
opt.include = [x.lower() for x in opt.include]
|
|
||||||
print(opt)
|
|
||||||
set_logging()
|
|
||||||
t = time.time()
|
|
||||||
|
|
||||||
# Load PyTorch model
|
|
||||||
device = select_device(opt.device)
|
|
||||||
model = attempt_load(opt.weights, map_location=device) # load FP32 model
|
|
||||||
labels = model.names
|
|
||||||
|
|
||||||
# Checks
|
|
||||||
gs = int(max(model.stride)) # grid size (max stride)
|
|
||||||
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
|
|
||||||
assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
|
|
||||||
|
|
||||||
# Input
|
|
||||||
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
|
|
||||||
|
|
||||||
# Update model
|
|
||||||
if opt.half:
|
|
||||||
img, model = img.half(), model.half() # to FP16
|
|
||||||
if opt.train:
|
|
||||||
model.train() # training mode (no grid construction in Detect layer)
|
|
||||||
for k, m in model.named_modules():
|
|
||||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
|
||||||
if isinstance(m, models.common.Conv): # assign export-friendly activations
|
|
||||||
if isinstance(m.act, nn.Hardswish):
|
|
||||||
m.act = Hardswish()
|
|
||||||
elif isinstance(m.act, nn.SiLU):
|
|
||||||
m.act = SiLU()
|
|
||||||
elif isinstance(m, models.yolo.Detect):
|
|
||||||
m.inplace = opt.inplace
|
|
||||||
m.onnx_dynamic = opt.dynamic
|
|
||||||
# m.forward = m.forward_export # assign forward (optional)
|
|
||||||
|
|
||||||
for _ in range(2):
|
|
||||||
y = model(img) # dry runs
|
|
||||||
print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
|
|
||||||
|
|
||||||
# TorchScript export -----------------------------------------------------------------------------------------------
|
|
||||||
if 'torchscript' in opt.include or 'coreml' in opt.include:
|
|
||||||
prefix = colorstr('TorchScript:')
|
|
||||||
try:
|
|
||||||
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
|
||||||
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
|
||||||
ts = torch.jit.trace(model, img, strict=False)
|
|
||||||
(optimize_for_mobile(ts) if opt.optimize else ts).save(f)
|
|
||||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
|
||||||
except Exception as e:
|
|
||||||
print(f'{prefix} export failure: {e}')
|
|
||||||
|
|
||||||
# ONNX export ------------------------------------------------------------------------------------------------------
|
|
||||||
if 'onnx' in opt.include:
|
|
||||||
prefix = colorstr('ONNX:')
|
|
||||||
try:
|
|
||||||
import onnx
|
|
||||||
|
|
||||||
print(f'{prefix} starting export with onnx {onnx.__version__}...')
|
|
||||||
f = opt.weights.replace('.pt', '.onnx') # filename
|
|
||||||
torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
|
|
||||||
training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL,
|
|
||||||
do_constant_folding=not opt.train,
|
|
||||||
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
|
|
||||||
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
|
|
||||||
|
|
||||||
# Checks
|
|
||||||
model_onnx = onnx.load(f) # load onnx model
|
|
||||||
onnx.checker.check_model(model_onnx) # check onnx model
|
|
||||||
# print(onnx.helper.printable_graph(model_onnx.graph)) # print
|
|
||||||
|
|
||||||
# Simplify
|
|
||||||
if opt.simplify:
|
|
||||||
try:
|
|
||||||
check_requirements(['onnx-simplifier'])
|
|
||||||
import onnxsim
|
|
||||||
|
|
||||||
print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
|
||||||
model_onnx, check = onnxsim.simplify(
|
|
||||||
model_onnx,
|
|
||||||
dynamic_input_shape=opt.dynamic,
|
|
||||||
input_shapes={'images': list(img.shape)} if opt.dynamic else None)
|
|
||||||
assert check, 'assert check failed'
|
|
||||||
onnx.save(model_onnx, f)
|
|
||||||
except Exception as e:
|
|
||||||
print(f'{prefix} simplifier failure: {e}')
|
|
||||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
|
||||||
except Exception as e:
|
|
||||||
print(f'{prefix} export failure: {e}')
|
|
||||||
|
|
||||||
# CoreML export ----------------------------------------------------------------------------------------------------
|
|
||||||
if 'coreml' in opt.include:
|
|
||||||
prefix = colorstr('CoreML:')
|
|
||||||
try:
|
|
||||||
import coremltools as ct
|
|
||||||
|
|
||||||
print(f'{prefix} starting export with coremltools {ct.__version__}...')
|
|
||||||
assert opt.train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`'
|
|
||||||
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
|
||||||
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
|
||||||
model.save(f)
|
|
||||||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
|
||||||
except Exception as e:
|
|
||||||
print(f'{prefix} export failure: {e}')
|
|
||||||
|
|
||||||
# Finish
|
|
||||||
print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.')
|
|
||||||
465
models/tf.py
Normal file
465
models/tf.py
Normal file
@ -0,0 +1,465 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
TensorFlow, Keras and TFLite versions of
|
||||||
|
Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
$ python models/tf.py --weights yolov3.pt
|
||||||
|
|
||||||
|
Export:
|
||||||
|
$ python path/to/export.py --weights yolov3.pt --include saved_model pb tflite tfjs
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import logging
|
||||||
|
import sys
|
||||||
|
from copy import deepcopy
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[1] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
# ROOT = ROOT.relative_to(Path.cwd()) # relative
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import tensorflow as tf
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from tensorflow import keras
|
||||||
|
|
||||||
|
from models.common import C3, SPP, SPPF, Bottleneck, BottleneckCSP, Concat, Conv, DWConv, Focus, autopad
|
||||||
|
from models.experimental import CrossConv, MixConv2d, attempt_load
|
||||||
|
from models.yolo import Detect
|
||||||
|
from utils.activations import SiLU
|
||||||
|
from utils.general import LOGGER, make_divisible, print_args
|
||||||
|
|
||||||
|
|
||||||
|
class TFBN(keras.layers.Layer):
|
||||||
|
# TensorFlow BatchNormalization wrapper
|
||||||
|
def __init__(self, w=None):
|
||||||
|
super().__init__()
|
||||||
|
self.bn = keras.layers.BatchNormalization(
|
||||||
|
beta_initializer=keras.initializers.Constant(w.bias.numpy()),
|
||||||
|
gamma_initializer=keras.initializers.Constant(w.weight.numpy()),
|
||||||
|
moving_mean_initializer=keras.initializers.Constant(w.running_mean.numpy()),
|
||||||
|
moving_variance_initializer=keras.initializers.Constant(w.running_var.numpy()),
|
||||||
|
epsilon=w.eps)
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return self.bn(inputs)
|
||||||
|
|
||||||
|
|
||||||
|
class TFPad(keras.layers.Layer):
|
||||||
|
def __init__(self, pad):
|
||||||
|
super().__init__()
|
||||||
|
self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]])
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return tf.pad(inputs, self.pad, mode='constant', constant_values=0)
|
||||||
|
|
||||||
|
|
||||||
|
class TFConv(keras.layers.Layer):
|
||||||
|
# Standard convolution
|
||||||
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
|
||||||
|
# ch_in, ch_out, weights, kernel, stride, padding, groups
|
||||||
|
super().__init__()
|
||||||
|
assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
|
||||||
|
assert isinstance(k, int), "Convolution with multiple kernels are not allowed."
|
||||||
|
# TensorFlow convolution padding is inconsistent with PyTorch (e.g. k=3 s=2 'SAME' padding)
|
||||||
|
# see https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch
|
||||||
|
|
||||||
|
conv = keras.layers.Conv2D(
|
||||||
|
c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False if hasattr(w, 'bn') else True,
|
||||||
|
kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()),
|
||||||
|
bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy()))
|
||||||
|
self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
|
||||||
|
self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity
|
||||||
|
|
||||||
|
# activations
|
||||||
|
if isinstance(w.act, nn.LeakyReLU):
|
||||||
|
self.act = (lambda x: keras.activations.relu(x, alpha=0.1)) if act else tf.identity
|
||||||
|
elif isinstance(w.act, nn.Hardswish):
|
||||||
|
self.act = (lambda x: x * tf.nn.relu6(x + 3) * 0.166666667) if act else tf.identity
|
||||||
|
elif isinstance(w.act, (nn.SiLU, SiLU)):
|
||||||
|
self.act = (lambda x: keras.activations.swish(x)) if act else tf.identity
|
||||||
|
else:
|
||||||
|
raise Exception(f'no matching TensorFlow activation found for {w.act}')
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return self.act(self.bn(self.conv(inputs)))
|
||||||
|
|
||||||
|
|
||||||
|
class TFFocus(keras.layers.Layer):
|
||||||
|
# Focus wh information into c-space
|
||||||
|
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
|
||||||
|
# ch_in, ch_out, kernel, stride, padding, groups
|
||||||
|
super().__init__()
|
||||||
|
self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv)
|
||||||
|
|
||||||
|
def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c)
|
||||||
|
# inputs = inputs / 255 # normalize 0-255 to 0-1
|
||||||
|
return self.conv(tf.concat([inputs[:, ::2, ::2, :],
|
||||||
|
inputs[:, 1::2, ::2, :],
|
||||||
|
inputs[:, ::2, 1::2, :],
|
||||||
|
inputs[:, 1::2, 1::2, :]], 3))
|
||||||
|
|
||||||
|
|
||||||
|
class TFBottleneck(keras.layers.Layer):
|
||||||
|
# Standard bottleneck
|
||||||
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None): # ch_in, ch_out, shortcut, groups, expansion
|
||||||
|
super().__init__()
|
||||||
|
c_ = int(c2 * e) # hidden channels
|
||||||
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
|
||||||
|
self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2)
|
||||||
|
self.add = shortcut and c1 == c2
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs))
|
||||||
|
|
||||||
|
|
||||||
|
class TFConv2d(keras.layers.Layer):
|
||||||
|
# Substitution for PyTorch nn.Conv2D
|
||||||
|
def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None):
|
||||||
|
super().__init__()
|
||||||
|
assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
|
||||||
|
self.conv = keras.layers.Conv2D(
|
||||||
|
c2, k, s, 'VALID', use_bias=bias,
|
||||||
|
kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()),
|
||||||
|
bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None, )
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return self.conv(inputs)
|
||||||
|
|
||||||
|
|
||||||
|
class TFBottleneckCSP(keras.layers.Layer):
|
||||||
|
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
|
||||||
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
|
||||||
|
# ch_in, ch_out, number, shortcut, groups, expansion
|
||||||
|
super().__init__()
|
||||||
|
c_ = int(c2 * e) # hidden channels
|
||||||
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
|
||||||
|
self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2)
|
||||||
|
self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3)
|
||||||
|
self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4)
|
||||||
|
self.bn = TFBN(w.bn)
|
||||||
|
self.act = lambda x: keras.activations.relu(x, alpha=0.1)
|
||||||
|
self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
y1 = self.cv3(self.m(self.cv1(inputs)))
|
||||||
|
y2 = self.cv2(inputs)
|
||||||
|
return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3))))
|
||||||
|
|
||||||
|
|
||||||
|
class TFC3(keras.layers.Layer):
|
||||||
|
# CSP Bottleneck with 3 convolutions
|
||||||
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
|
||||||
|
# ch_in, ch_out, number, shortcut, groups, expansion
|
||||||
|
super().__init__()
|
||||||
|
c_ = int(c2 * e) # hidden channels
|
||||||
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
|
||||||
|
self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
|
||||||
|
self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
|
||||||
|
self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3))
|
||||||
|
|
||||||
|
|
||||||
|
class TFSPP(keras.layers.Layer):
|
||||||
|
# Spatial pyramid pooling layer used in YOLOv3-SPP
|
||||||
|
def __init__(self, c1, c2, k=(5, 9, 13), w=None):
|
||||||
|
super().__init__()
|
||||||
|
c_ = c1 // 2 # hidden channels
|
||||||
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
|
||||||
|
self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
|
||||||
|
self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k]
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
x = self.cv1(inputs)
|
||||||
|
return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3))
|
||||||
|
|
||||||
|
|
||||||
|
class TFSPPF(keras.layers.Layer):
|
||||||
|
# Spatial pyramid pooling-Fast layer
|
||||||
|
def __init__(self, c1, c2, k=5, w=None):
|
||||||
|
super().__init__()
|
||||||
|
c_ = c1 // 2 # hidden channels
|
||||||
|
self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
|
||||||
|
self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2)
|
||||||
|
self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding='SAME')
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
x = self.cv1(inputs)
|
||||||
|
y1 = self.m(x)
|
||||||
|
y2 = self.m(y1)
|
||||||
|
return self.cv2(tf.concat([x, y1, y2, self.m(y2)], 3))
|
||||||
|
|
||||||
|
|
||||||
|
class TFDetect(keras.layers.Layer):
|
||||||
|
def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None): # detection layer
|
||||||
|
super().__init__()
|
||||||
|
self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32)
|
||||||
|
self.nc = nc # number of classes
|
||||||
|
self.no = nc + 5 # number of outputs per anchor
|
||||||
|
self.nl = len(anchors) # number of detection layers
|
||||||
|
self.na = len(anchors[0]) // 2 # number of anchors
|
||||||
|
self.grid = [tf.zeros(1)] * self.nl # init grid
|
||||||
|
self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32)
|
||||||
|
self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]),
|
||||||
|
[self.nl, 1, -1, 1, 2])
|
||||||
|
self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]
|
||||||
|
self.training = False # set to False after building model
|
||||||
|
self.imgsz = imgsz
|
||||||
|
for i in range(self.nl):
|
||||||
|
ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
|
||||||
|
self.grid[i] = self._make_grid(nx, ny)
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
z = [] # inference output
|
||||||
|
x = []
|
||||||
|
for i in range(self.nl):
|
||||||
|
x.append(self.m[i](inputs[i]))
|
||||||
|
# x(bs,20,20,255) to x(bs,3,20,20,85)
|
||||||
|
ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
|
||||||
|
x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na, self.no]), [0, 2, 1, 3])
|
||||||
|
|
||||||
|
if not self.training: # inference
|
||||||
|
y = tf.sigmoid(x[i])
|
||||||
|
xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
|
||||||
|
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]
|
||||||
|
# Normalize xywh to 0-1 to reduce calibration error
|
||||||
|
xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
|
||||||
|
wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
|
||||||
|
y = tf.concat([xy, wh, y[..., 4:]], -1)
|
||||||
|
z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no]))
|
||||||
|
|
||||||
|
return x if self.training else (tf.concat(z, 1), x)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _make_grid(nx=20, ny=20):
|
||||||
|
# yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
||||||
|
# return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
|
||||||
|
xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
|
||||||
|
return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32)
|
||||||
|
|
||||||
|
|
||||||
|
class TFUpsample(keras.layers.Layer):
|
||||||
|
def __init__(self, size, scale_factor, mode, w=None): # warning: all arguments needed including 'w'
|
||||||
|
super().__init__()
|
||||||
|
assert scale_factor == 2, "scale_factor must be 2"
|
||||||
|
self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode)
|
||||||
|
# self.upsample = keras.layers.UpSampling2D(size=scale_factor, interpolation=mode)
|
||||||
|
# with default arguments: align_corners=False, half_pixel_centers=False
|
||||||
|
# self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x,
|
||||||
|
# size=(x.shape[1] * 2, x.shape[2] * 2))
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return self.upsample(inputs)
|
||||||
|
|
||||||
|
|
||||||
|
class TFConcat(keras.layers.Layer):
|
||||||
|
def __init__(self, dimension=1, w=None):
|
||||||
|
super().__init__()
|
||||||
|
assert dimension == 1, "convert only NCHW to NHWC concat"
|
||||||
|
self.d = 3
|
||||||
|
|
||||||
|
def call(self, inputs):
|
||||||
|
return tf.concat(inputs, self.d)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3)
|
||||||
|
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
|
||||||
|
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
|
||||||
|
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
|
||||||
|
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
|
||||||
|
|
||||||
|
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
|
||||||
|
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
|
||||||
|
m_str = m
|
||||||
|
m = eval(m) if isinstance(m, str) else m # eval strings
|
||||||
|
for j, a in enumerate(args):
|
||||||
|
try:
|
||||||
|
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
||||||
|
except NameError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
n = max(round(n * gd), 1) if n > 1 else n # depth gain
|
||||||
|
if m in [nn.Conv2d, Conv, Bottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:
|
||||||
|
c1, c2 = ch[f], args[0]
|
||||||
|
c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
|
||||||
|
|
||||||
|
args = [c1, c2, *args[1:]]
|
||||||
|
if m in [BottleneckCSP, C3]:
|
||||||
|
args.insert(2, n)
|
||||||
|
n = 1
|
||||||
|
elif m is nn.BatchNorm2d:
|
||||||
|
args = [ch[f]]
|
||||||
|
elif m is Concat:
|
||||||
|
c2 = sum(ch[-1 if x == -1 else x + 1] for x in f)
|
||||||
|
elif m is Detect:
|
||||||
|
args.append([ch[x + 1] for x in f])
|
||||||
|
if isinstance(args[1], int): # number of anchors
|
||||||
|
args[1] = [list(range(args[1] * 2))] * len(f)
|
||||||
|
args.append(imgsz)
|
||||||
|
else:
|
||||||
|
c2 = ch[f]
|
||||||
|
|
||||||
|
tf_m = eval('TF' + m_str.replace('nn.', ''))
|
||||||
|
m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \
|
||||||
|
else tf_m(*args, w=model.model[i]) # module
|
||||||
|
|
||||||
|
torch_m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
|
||||||
|
t = str(m)[8:-2].replace('__main__.', '') # module type
|
||||||
|
np = sum(x.numel() for x in torch_m_.parameters()) # number params
|
||||||
|
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
|
||||||
|
LOGGER.info(f'{i:>3}{str(f):>18}{str(n):>3}{np:>10} {t:<40}{str(args):<30}') # print
|
||||||
|
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
||||||
|
layers.append(m_)
|
||||||
|
ch.append(c2)
|
||||||
|
return keras.Sequential(layers), sorted(save)
|
||||||
|
|
||||||
|
|
||||||
|
class TFModel:
|
||||||
|
def __init__(self, cfg='yolov3.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)): # model, channels, classes
|
||||||
|
super().__init__()
|
||||||
|
if isinstance(cfg, dict):
|
||||||
|
self.yaml = cfg # model dict
|
||||||
|
else: # is *.yaml
|
||||||
|
import yaml # for torch hub
|
||||||
|
self.yaml_file = Path(cfg).name
|
||||||
|
with open(cfg) as f:
|
||||||
|
self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||||
|
|
||||||
|
# Define model
|
||||||
|
if nc and nc != self.yaml['nc']:
|
||||||
|
LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}")
|
||||||
|
self.yaml['nc'] = nc # override yaml value
|
||||||
|
self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)
|
||||||
|
|
||||||
|
def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
|
||||||
|
conf_thres=0.25):
|
||||||
|
y = [] # outputs
|
||||||
|
x = inputs
|
||||||
|
for i, m in enumerate(self.model.layers):
|
||||||
|
if m.f != -1: # if not from previous layer
|
||||||
|
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
|
||||||
|
|
||||||
|
x = m(x) # run
|
||||||
|
y.append(x if m.i in self.savelist else None) # save output
|
||||||
|
|
||||||
|
# Add TensorFlow NMS
|
||||||
|
if tf_nms:
|
||||||
|
boxes = self._xywh2xyxy(x[0][..., :4])
|
||||||
|
probs = x[0][:, :, 4:5]
|
||||||
|
classes = x[0][:, :, 5:]
|
||||||
|
scores = probs * classes
|
||||||
|
if agnostic_nms:
|
||||||
|
nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres)
|
||||||
|
return nms, x[1]
|
||||||
|
else:
|
||||||
|
boxes = tf.expand_dims(boxes, 2)
|
||||||
|
nms = tf.image.combined_non_max_suppression(
|
||||||
|
boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False)
|
||||||
|
return nms, x[1]
|
||||||
|
|
||||||
|
return x[0] # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...]
|
||||||
|
# x = x[0][0] # [x(1,6300,85), ...] to x(6300,85)
|
||||||
|
# xywh = x[..., :4] # x(6300,4) boxes
|
||||||
|
# conf = x[..., 4:5] # x(6300,1) confidences
|
||||||
|
# cls = tf.reshape(tf.cast(tf.argmax(x[..., 5:], axis=1), tf.float32), (-1, 1)) # x(6300,1) classes
|
||||||
|
# return tf.concat([conf, cls, xywh], 1)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _xywh2xyxy(xywh):
|
||||||
|
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
||||||
|
x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1)
|
||||||
|
return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)
|
||||||
|
|
||||||
|
|
||||||
|
class AgnosticNMS(keras.layers.Layer):
|
||||||
|
# TF Agnostic NMS
|
||||||
|
def call(self, input, topk_all, iou_thres, conf_thres):
|
||||||
|
# wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450
|
||||||
|
return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres, conf_thres), input,
|
||||||
|
fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32),
|
||||||
|
name='agnostic_nms')
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25): # agnostic NMS
|
||||||
|
boxes, classes, scores = x
|
||||||
|
class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
|
||||||
|
scores_inp = tf.reduce_max(scores, -1)
|
||||||
|
selected_inds = tf.image.non_max_suppression(
|
||||||
|
boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres)
|
||||||
|
selected_boxes = tf.gather(boxes, selected_inds)
|
||||||
|
padded_boxes = tf.pad(selected_boxes,
|
||||||
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
|
||||||
|
mode="CONSTANT", constant_values=0.0)
|
||||||
|
selected_scores = tf.gather(scores_inp, selected_inds)
|
||||||
|
padded_scores = tf.pad(selected_scores,
|
||||||
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
|
||||||
|
mode="CONSTANT", constant_values=-1.0)
|
||||||
|
selected_classes = tf.gather(class_inds, selected_inds)
|
||||||
|
padded_classes = tf.pad(selected_classes,
|
||||||
|
paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
|
||||||
|
mode="CONSTANT", constant_values=-1.0)
|
||||||
|
valid_detections = tf.shape(selected_inds)[0]
|
||||||
|
return padded_boxes, padded_scores, padded_classes, valid_detections
|
||||||
|
|
||||||
|
|
||||||
|
def representative_dataset_gen(dataset, ncalib=100):
|
||||||
|
# Representative dataset generator for use with converter.representative_dataset, returns a generator of np arrays
|
||||||
|
for n, (path, img, im0s, vid_cap, string) in enumerate(dataset):
|
||||||
|
input = np.transpose(img, [1, 2, 0])
|
||||||
|
input = np.expand_dims(input, axis=0).astype(np.float32)
|
||||||
|
input /= 255
|
||||||
|
yield [input]
|
||||||
|
if n >= ncalib:
|
||||||
|
break
|
||||||
|
|
||||||
|
|
||||||
|
def run(weights=ROOT / 'yolov3.pt', # weights path
|
||||||
|
imgsz=(640, 640), # inference size h,w
|
||||||
|
batch_size=1, # batch size
|
||||||
|
dynamic=False, # dynamic batch size
|
||||||
|
):
|
||||||
|
# PyTorch model
|
||||||
|
im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image
|
||||||
|
model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False)
|
||||||
|
y = model(im) # inference
|
||||||
|
model.info()
|
||||||
|
|
||||||
|
# TensorFlow model
|
||||||
|
im = tf.zeros((batch_size, *imgsz, 3)) # BHWC image
|
||||||
|
tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
|
||||||
|
y = tf_model.predict(im) # inference
|
||||||
|
|
||||||
|
# Keras model
|
||||||
|
im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
|
||||||
|
keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im))
|
||||||
|
keras_model.summary()
|
||||||
|
|
||||||
|
LOGGER.info('PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.')
|
||||||
|
|
||||||
|
|
||||||
|
def parse_opt():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('--weights', type=str, default=ROOT / 'yolov3.pt', help='weights path')
|
||||||
|
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
|
||||||
|
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
||||||
|
parser.add_argument('--dynamic', action='store_true', help='dynamic batch size')
|
||||||
|
opt = parser.parse_args()
|
||||||
|
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
|
||||||
|
print_args(FILE.stem, opt)
|
||||||
|
return opt
|
||||||
|
|
||||||
|
|
||||||
|
def main(opt):
|
||||||
|
run(**vars(opt))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
opt = parse_opt()
|
||||||
|
main(opt)
|
||||||
204
models/yolo.py
204
models/yolo.py
@ -1,27 +1,32 @@
|
|||||||
"""YOLOv3-specific modules
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
YOLO-specific modules
|
||||||
|
|
||||||
Usage:
|
Usage:
|
||||||
$ python path/to/models/yolo.py --cfg yolov3.yaml
|
$ python path/to/models/yolo.py --cfg yolov3.yaml
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import logging
|
|
||||||
import sys
|
import sys
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories
|
FILE = Path(__file__).resolve()
|
||||||
logger = logging.getLogger(__name__)
|
ROOT = FILE.parents[1] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
# ROOT = ROOT.relative_to(Path.cwd()) # relative
|
||||||
|
|
||||||
from models.common import *
|
from models.common import *
|
||||||
from models.experimental import *
|
from models.experimental import *
|
||||||
from utils.autoanchor import check_anchor_order
|
from utils.autoanchor import check_anchor_order
|
||||||
from utils.general import make_divisible, check_file, set_logging
|
from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args
|
||||||
from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
|
from utils.plots import feature_visualization
|
||||||
select_device, copy_attr
|
from utils.torch_utils import (copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device,
|
||||||
|
time_sync)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import thop # for FLOPS computation
|
import thop # for FLOPs computation
|
||||||
except ImportError:
|
except ImportError:
|
||||||
thop = None
|
thop = None
|
||||||
|
|
||||||
@ -31,20 +36,18 @@ class Detect(nn.Module):
|
|||||||
onnx_dynamic = False # ONNX export parameter
|
onnx_dynamic = False # ONNX export parameter
|
||||||
|
|
||||||
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
|
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
|
||||||
super(Detect, self).__init__()
|
super().__init__()
|
||||||
self.nc = nc # number of classes
|
self.nc = nc # number of classes
|
||||||
self.no = nc + 5 # number of outputs per anchor
|
self.no = nc + 5 # number of outputs per anchor
|
||||||
self.nl = len(anchors) # number of detection layers
|
self.nl = len(anchors) # number of detection layers
|
||||||
self.na = len(anchors[0]) // 2 # number of anchors
|
self.na = len(anchors[0]) // 2 # number of anchors
|
||||||
self.grid = [torch.zeros(1)] * self.nl # init grid
|
self.grid = [torch.zeros(1)] * self.nl # init grid
|
||||||
a = torch.tensor(anchors).float().view(self.nl, -1, 2)
|
self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid
|
||||||
self.register_buffer('anchors', a) # shape(nl,na,2)
|
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
|
||||||
self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)
|
|
||||||
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
|
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
|
||||||
self.inplace = inplace # use in-place ops (e.g. slice assignment)
|
self.inplace = inplace # use in-place ops (e.g. slice assignment)
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
# x = x.copy() # for profiling
|
|
||||||
z = [] # inference output
|
z = [] # inference output
|
||||||
for i in range(self.nl):
|
for i in range(self.nl):
|
||||||
x[i] = self.m[i](x[i]) # conv
|
x[i] = self.m[i](x[i]) # conv
|
||||||
@ -52,50 +55,55 @@ class Detect(nn.Module):
|
|||||||
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
||||||
|
|
||||||
if not self.training: # inference
|
if not self.training: # inference
|
||||||
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
|
if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
|
||||||
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
|
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
|
||||||
|
|
||||||
y = x[i].sigmoid()
|
y = x[i].sigmoid()
|
||||||
if self.inplace:
|
if self.inplace:
|
||||||
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
|
y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
|
||||||
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
||||||
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
|
else: # for on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
|
||||||
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
|
xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
|
||||||
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
|
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
|
||||||
y = torch.cat((xy, wh, y[..., 4:]), -1)
|
y = torch.cat((xy, wh, y[..., 4:]), -1)
|
||||||
z.append(y.view(bs, -1, self.no))
|
z.append(y.view(bs, -1, self.no))
|
||||||
|
|
||||||
return x if self.training else (torch.cat(z, 1), x)
|
return x if self.training else (torch.cat(z, 1), x)
|
||||||
|
|
||||||
@staticmethod
|
def _make_grid(self, nx=20, ny=20, i=0):
|
||||||
def _make_grid(nx=20, ny=20):
|
d = self.anchors[i].device
|
||||||
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
|
if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
|
||||||
return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
|
yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing='ij')
|
||||||
|
else:
|
||||||
|
yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])
|
||||||
|
grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
|
||||||
|
anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
|
||||||
|
.view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
|
||||||
|
return grid, anchor_grid
|
||||||
|
|
||||||
|
|
||||||
class Model(nn.Module):
|
class Model(nn.Module):
|
||||||
def __init__(self, cfg='yolov3.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
|
def __init__(self, cfg='yolov3.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
|
||||||
super(Model, self).__init__()
|
super().__init__()
|
||||||
if isinstance(cfg, dict):
|
if isinstance(cfg, dict):
|
||||||
self.yaml = cfg # model dict
|
self.yaml = cfg # model dict
|
||||||
else: # is *.yaml
|
else: # is *.yaml
|
||||||
import yaml # for torch hub
|
import yaml # for torch hub
|
||||||
self.yaml_file = Path(cfg).name
|
self.yaml_file = Path(cfg).name
|
||||||
with open(cfg) as f:
|
with open(cfg, encoding='ascii', errors='ignore') as f:
|
||||||
self.yaml = yaml.safe_load(f) # model dict
|
self.yaml = yaml.safe_load(f) # model dict
|
||||||
|
|
||||||
# Define model
|
# Define model
|
||||||
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
||||||
if nc and nc != self.yaml['nc']:
|
if nc and nc != self.yaml['nc']:
|
||||||
logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
|
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
|
||||||
self.yaml['nc'] = nc # override yaml value
|
self.yaml['nc'] = nc # override yaml value
|
||||||
if anchors:
|
if anchors:
|
||||||
logger.info(f'Overriding model.yaml anchors with anchors={anchors}')
|
LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
|
||||||
self.yaml['anchors'] = round(anchors) # override yaml value
|
self.yaml['anchors'] = round(anchors) # override yaml value
|
||||||
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
|
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
|
||||||
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
|
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
|
||||||
self.inplace = self.yaml.get('inplace', True)
|
self.inplace = self.yaml.get('inplace', True)
|
||||||
# logger.info([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
|
||||||
|
|
||||||
# Build strides, anchors
|
# Build strides, anchors
|
||||||
m = self.model[-1] # Detect()
|
m = self.model[-1] # Detect()
|
||||||
@ -107,53 +115,42 @@ class Model(nn.Module):
|
|||||||
check_anchor_order(m)
|
check_anchor_order(m)
|
||||||
self.stride = m.stride
|
self.stride = m.stride
|
||||||
self._initialize_biases() # only run once
|
self._initialize_biases() # only run once
|
||||||
# logger.info('Strides: %s' % m.stride.tolist())
|
|
||||||
|
|
||||||
# Init weights, biases
|
# Init weights, biases
|
||||||
initialize_weights(self)
|
initialize_weights(self)
|
||||||
self.info()
|
self.info()
|
||||||
logger.info('')
|
LOGGER.info('')
|
||||||
|
|
||||||
def forward(self, x, augment=False, profile=False):
|
def forward(self, x, augment=False, profile=False, visualize=False):
|
||||||
if augment:
|
if augment:
|
||||||
return self.forward_augment(x) # augmented inference, None
|
return self._forward_augment(x) # augmented inference, None
|
||||||
else:
|
return self._forward_once(x, profile, visualize) # single-scale inference, train
|
||||||
return self.forward_once(x, profile) # single-scale inference, train
|
|
||||||
|
|
||||||
def forward_augment(self, x):
|
def _forward_augment(self, x):
|
||||||
img_size = x.shape[-2:] # height, width
|
img_size = x.shape[-2:] # height, width
|
||||||
s = [1, 0.83, 0.67] # scales
|
s = [1, 0.83, 0.67] # scales
|
||||||
f = [None, 3, None] # flips (2-ud, 3-lr)
|
f = [None, 3, None] # flips (2-ud, 3-lr)
|
||||||
y = [] # outputs
|
y = [] # outputs
|
||||||
for si, fi in zip(s, f):
|
for si, fi in zip(s, f):
|
||||||
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
|
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
|
||||||
yi = self.forward_once(xi)[0] # forward
|
yi = self._forward_once(xi)[0] # forward
|
||||||
# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
|
# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
|
||||||
yi = self._descale_pred(yi, fi, si, img_size)
|
yi = self._descale_pred(yi, fi, si, img_size)
|
||||||
y.append(yi)
|
y.append(yi)
|
||||||
|
y = self._clip_augmented(y) # clip augmented tails
|
||||||
return torch.cat(y, 1), None # augmented inference, train
|
return torch.cat(y, 1), None # augmented inference, train
|
||||||
|
|
||||||
def forward_once(self, x, profile=False):
|
def _forward_once(self, x, profile=False, visualize=False):
|
||||||
y, dt = [], [] # outputs
|
y, dt = [], [] # outputs
|
||||||
for m in self.model:
|
for m in self.model:
|
||||||
if m.f != -1: # if not from previous layer
|
if m.f != -1: # if not from previous layer
|
||||||
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
|
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
|
||||||
|
|
||||||
if profile:
|
if profile:
|
||||||
o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPS
|
self._profile_one_layer(m, x, dt)
|
||||||
t = time_synchronized()
|
|
||||||
for _ in range(10):
|
|
||||||
_ = m(x)
|
|
||||||
dt.append((time_synchronized() - t) * 100)
|
|
||||||
if m == self.model[0]:
|
|
||||||
logger.info(f"{'time (ms)':>10s} {'GFLOPS':>10s} {'params':>10s} {'module'}")
|
|
||||||
logger.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}')
|
|
||||||
|
|
||||||
x = m(x) # run
|
x = m(x) # run
|
||||||
y.append(x if m.i in self.save else None) # save output
|
y.append(x if m.i in self.save else None) # save output
|
||||||
|
if visualize:
|
||||||
if profile:
|
feature_visualization(x, m.type, m.i, save_dir=visualize)
|
||||||
logger.info('%.1fms total' % sum(dt))
|
|
||||||
return x
|
return x
|
||||||
|
|
||||||
def _descale_pred(self, p, flips, scale, img_size):
|
def _descale_pred(self, p, flips, scale, img_size):
|
||||||
@ -173,6 +170,30 @@ class Model(nn.Module):
|
|||||||
p = torch.cat((x, y, wh, p[..., 4:]), -1)
|
p = torch.cat((x, y, wh, p[..., 4:]), -1)
|
||||||
return p
|
return p
|
||||||
|
|
||||||
|
def _clip_augmented(self, y):
|
||||||
|
# Clip augmented inference tails
|
||||||
|
nl = self.model[-1].nl # number of detection layers (P3-P5)
|
||||||
|
g = sum(4 ** x for x in range(nl)) # grid points
|
||||||
|
e = 1 # exclude layer count
|
||||||
|
i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e)) # indices
|
||||||
|
y[0] = y[0][:, :-i] # large
|
||||||
|
i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # indices
|
||||||
|
y[-1] = y[-1][:, i:] # small
|
||||||
|
return y
|
||||||
|
|
||||||
|
def _profile_one_layer(self, m, x, dt):
|
||||||
|
c = isinstance(m, Detect) # is final layer, copy input as inplace fix
|
||||||
|
o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs
|
||||||
|
t = time_sync()
|
||||||
|
for _ in range(10):
|
||||||
|
m(x.copy() if c else x)
|
||||||
|
dt.append((time_sync() - t) * 100)
|
||||||
|
if m == self.model[0]:
|
||||||
|
LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s} {'module'}")
|
||||||
|
LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}')
|
||||||
|
if c:
|
||||||
|
LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total")
|
||||||
|
|
||||||
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
|
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
|
||||||
# https://arxiv.org/abs/1708.02002 section 3.3
|
# https://arxiv.org/abs/1708.02002 section 3.3
|
||||||
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
||||||
@ -180,47 +201,33 @@ class Model(nn.Module):
|
|||||||
for mi, s in zip(m.m, m.stride): # from
|
for mi, s in zip(m.m, m.stride): # from
|
||||||
b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
|
b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
|
||||||
b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
||||||
b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls
|
b.data[:, 5:] += math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum()) # cls
|
||||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||||
|
|
||||||
def _print_biases(self):
|
def _print_biases(self):
|
||||||
m = self.model[-1] # Detect() module
|
m = self.model[-1] # Detect() module
|
||||||
for mi in m.m: # from
|
for mi in m.m: # from
|
||||||
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
|
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
|
||||||
logger.info(
|
LOGGER.info(
|
||||||
('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
|
('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
|
||||||
|
|
||||||
# def _print_weights(self):
|
# def _print_weights(self):
|
||||||
# for m in self.model.modules():
|
# for m in self.model.modules():
|
||||||
# if type(m) is Bottleneck:
|
# if type(m) is Bottleneck:
|
||||||
# logger.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
# LOGGER.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
||||||
|
|
||||||
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
||||||
logger.info('Fusing layers... ')
|
LOGGER.info('Fusing layers... ')
|
||||||
for m in self.model.modules():
|
for m in self.model.modules():
|
||||||
if type(m) is Conv and hasattr(m, 'bn'):
|
if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
|
||||||
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
||||||
delattr(m, 'bn') # remove batchnorm
|
delattr(m, 'bn') # remove batchnorm
|
||||||
m.forward = m.fuseforward # update forward
|
m.forward = m.forward_fuse # update forward
|
||||||
self.info()
|
self.info()
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def nms(self, mode=True): # add or remove NMS module
|
|
||||||
present = type(self.model[-1]) is NMS # last layer is NMS
|
|
||||||
if mode and not present:
|
|
||||||
logger.info('Adding NMS... ')
|
|
||||||
m = NMS() # module
|
|
||||||
m.f = -1 # from
|
|
||||||
m.i = self.model[-1].i + 1 # index
|
|
||||||
self.model.add_module(name='%s' % m.i, module=m) # add
|
|
||||||
self.eval()
|
|
||||||
elif not mode and present:
|
|
||||||
logger.info('Removing NMS... ')
|
|
||||||
self.model = self.model[:-1] # remove
|
|
||||||
return self
|
|
||||||
|
|
||||||
def autoshape(self): # add AutoShape module
|
def autoshape(self): # add AutoShape module
|
||||||
logger.info('Adding AutoShape... ')
|
LOGGER.info('Adding AutoShape... ')
|
||||||
m = AutoShape(self) # wrap model
|
m = AutoShape(self) # wrap model
|
||||||
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
|
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
|
||||||
return m
|
return m
|
||||||
@ -228,9 +235,20 @@ class Model(nn.Module):
|
|||||||
def info(self, verbose=False, img_size=640): # print model information
|
def info(self, verbose=False, img_size=640): # print model information
|
||||||
model_info(self, verbose, img_size)
|
model_info(self, verbose, img_size)
|
||||||
|
|
||||||
|
def _apply(self, fn):
|
||||||
|
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
|
||||||
|
self = super()._apply(fn)
|
||||||
|
m = self.model[-1] # Detect()
|
||||||
|
if isinstance(m, Detect):
|
||||||
|
m.stride = fn(m.stride)
|
||||||
|
m.grid = list(map(fn, m.grid))
|
||||||
|
if isinstance(m.anchor_grid, list):
|
||||||
|
m.anchor_grid = list(map(fn, m.anchor_grid))
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
def parse_model(d, ch): # model_dict, input_channels(3)
|
def parse_model(d, ch): # model_dict, input_channels(3)
|
||||||
logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
|
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
|
||||||
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
|
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
|
||||||
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
|
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
|
||||||
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
|
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
|
||||||
@ -241,24 +259,24 @@ def parse_model(d, ch): # model_dict, input_channels(3)
|
|||||||
for j, a in enumerate(args):
|
for j, a in enumerate(args):
|
||||||
try:
|
try:
|
||||||
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
||||||
except:
|
except NameError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
n = max(round(n * gd), 1) if n > 1 else n # depth gain
|
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
|
||||||
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
|
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
|
||||||
C3, C3TR]:
|
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
|
||||||
c1, c2 = ch[f], args[0]
|
c1, c2 = ch[f], args[0]
|
||||||
if c2 != no: # if not output
|
if c2 != no: # if not output
|
||||||
c2 = make_divisible(c2 * gw, 8)
|
c2 = make_divisible(c2 * gw, 8)
|
||||||
|
|
||||||
args = [c1, c2, *args[1:]]
|
args = [c1, c2, *args[1:]]
|
||||||
if m in [BottleneckCSP, C3, C3TR]:
|
if m in [BottleneckCSP, C3, C3TR, C3Ghost]:
|
||||||
args.insert(2, n) # number of repeats
|
args.insert(2, n) # number of repeats
|
||||||
n = 1
|
n = 1
|
||||||
elif m is nn.BatchNorm2d:
|
elif m is nn.BatchNorm2d:
|
||||||
args = [ch[f]]
|
args = [ch[f]]
|
||||||
elif m is Concat:
|
elif m is Concat:
|
||||||
c2 = sum([ch[x] for x in f])
|
c2 = sum(ch[x] for x in f)
|
||||||
elif m is Detect:
|
elif m is Detect:
|
||||||
args.append([ch[x] for x in f])
|
args.append([ch[x] for x in f])
|
||||||
if isinstance(args[1], int): # number of anchors
|
if isinstance(args[1], int): # number of anchors
|
||||||
@ -270,11 +288,11 @@ def parse_model(d, ch): # model_dict, input_channels(3)
|
|||||||
else:
|
else:
|
||||||
c2 = ch[f]
|
c2 = ch[f]
|
||||||
|
|
||||||
m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module
|
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
|
||||||
t = str(m)[8:-2].replace('__main__.', '') # module type
|
t = str(m)[8:-2].replace('__main__.', '') # module type
|
||||||
np = sum([x.numel() for x in m_.parameters()]) # number params
|
np = sum(x.numel() for x in m_.parameters()) # number params
|
||||||
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
|
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
|
||||||
logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
|
LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}') # print
|
||||||
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
||||||
layers.append(m_)
|
layers.append(m_)
|
||||||
if i == 0:
|
if i == 0:
|
||||||
@ -285,11 +303,13 @@ def parse_model(d, ch): # model_dict, input_channels(3)
|
|||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument('--cfg', type=str, default='yolov3.yaml', help='model.yaml')
|
parser.add_argument('--cfg', type=str, default='yolov3yaml', help='model.yaml')
|
||||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||||
|
parser.add_argument('--profile', action='store_true', help='profile model speed')
|
||||||
|
parser.add_argument('--test', action='store_true', help='test all yolo*.yaml')
|
||||||
opt = parser.parse_args()
|
opt = parser.parse_args()
|
||||||
opt.cfg = check_file(opt.cfg) # check file
|
opt.cfg = check_yaml(opt.cfg) # check YAML
|
||||||
set_logging()
|
print_args(FILE.stem, opt)
|
||||||
device = select_device(opt.device)
|
device = select_device(opt.device)
|
||||||
|
|
||||||
# Create model
|
# Create model
|
||||||
@ -297,12 +317,20 @@ if __name__ == '__main__':
|
|||||||
model.train()
|
model.train()
|
||||||
|
|
||||||
# Profile
|
# Profile
|
||||||
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 320, 320).to(device)
|
if opt.profile:
|
||||||
# y = model(img, profile=True)
|
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
|
||||||
|
y = model(img, profile=True)
|
||||||
|
|
||||||
|
# Test all models
|
||||||
|
if opt.test:
|
||||||
|
for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
|
||||||
|
try:
|
||||||
|
_ = Model(cfg)
|
||||||
|
except Exception as e:
|
||||||
|
print(f'Error in {cfg}: {e}')
|
||||||
|
|
||||||
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
|
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
|
||||||
# from torch.utils.tensorboard import SummaryWriter
|
# from torch.utils.tensorboard import SummaryWriter
|
||||||
# tb_writer = SummaryWriter('.')
|
# tb_writer = SummaryWriter('.')
|
||||||
# logger.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/")
|
# LOGGER.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/")
|
||||||
# tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph
|
# tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph
|
||||||
# tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard
|
|
||||||
|
|||||||
@ -1,9 +1,9 @@
|
|||||||
# parameters
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
|
# Parameters
|
||||||
nc: 80 # number of classes
|
nc: 80 # number of classes
|
||||||
depth_multiple: 1.0 # model depth multiple
|
depth_multiple: 1.0 # model depth multiple
|
||||||
width_multiple: 1.0 # layer channel multiple
|
width_multiple: 1.0 # layer channel multiple
|
||||||
|
|
||||||
# anchors
|
|
||||||
anchors:
|
anchors:
|
||||||
- [10,13, 16,30, 33,23] # P3/8
|
- [10,13, 16,30, 33,23] # P3/8
|
||||||
- [30,61, 62,45, 59,119] # P4/16
|
- [30,61, 62,45, 59,119] # P4/16
|
||||||
|
|||||||
@ -1,9 +1,9 @@
|
|||||||
# parameters
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
|
# Parameters
|
||||||
nc: 80 # number of classes
|
nc: 80 # number of classes
|
||||||
depth_multiple: 1.0 # model depth multiple
|
depth_multiple: 1.0 # model depth multiple
|
||||||
width_multiple: 1.0 # layer channel multiple
|
width_multiple: 1.0 # layer channel multiple
|
||||||
|
|
||||||
# anchors
|
|
||||||
anchors:
|
anchors:
|
||||||
- [10,14, 23,27, 37,58] # P4/16
|
- [10,14, 23,27, 37,58] # P4/16
|
||||||
- [81,82, 135,169, 344,319] # P5/32
|
- [81,82, 135,169, 344,319] # P5/32
|
||||||
|
|||||||
@ -1,9 +1,9 @@
|
|||||||
# parameters
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
|
||||||
|
# Parameters
|
||||||
nc: 80 # number of classes
|
nc: 80 # number of classes
|
||||||
depth_multiple: 1.0 # model depth multiple
|
depth_multiple: 1.0 # model depth multiple
|
||||||
width_multiple: 1.0 # layer channel multiple
|
width_multiple: 1.0 # layer channel multiple
|
||||||
|
|
||||||
# anchors
|
|
||||||
anchors:
|
anchors:
|
||||||
- [10,13, 16,30, 33,23] # P3/8
|
- [10,13, 16,30, 33,23] # P3/8
|
||||||
- [30,61, 62,45, 59,119] # P4/16
|
- [30,61, 62,45, 59,119] # P4/16
|
||||||
|
|||||||
@ -1,30 +1,36 @@
|
|||||||
# pip install -r requirements.txt
|
# pip install -r requirements.txt
|
||||||
|
|
||||||
# base ----------------------------------------
|
# Base ----------------------------------------
|
||||||
matplotlib>=3.2.2
|
matplotlib>=3.2.2
|
||||||
numpy>=1.18.5
|
numpy>=1.18.5
|
||||||
opencv-python>=4.1.2
|
opencv-python>=4.1.2
|
||||||
Pillow
|
Pillow>=7.1.2
|
||||||
PyYAML>=5.3.1
|
PyYAML>=5.3.1
|
||||||
|
requests>=2.23.0
|
||||||
scipy>=1.4.1
|
scipy>=1.4.1
|
||||||
torch>=1.7.0
|
torch>=1.7.0
|
||||||
torchvision>=0.8.1
|
torchvision>=0.8.1
|
||||||
tqdm>=4.41.0
|
tqdm>=4.41.0
|
||||||
|
|
||||||
# logging -------------------------------------
|
# Logging -------------------------------------
|
||||||
tensorboard>=2.4.1
|
tensorboard>=2.4.1
|
||||||
# wandb
|
# wandb
|
||||||
|
|
||||||
# plotting ------------------------------------
|
# Plotting ------------------------------------
|
||||||
|
pandas>=1.1.4
|
||||||
seaborn>=0.11.0
|
seaborn>=0.11.0
|
||||||
pandas
|
|
||||||
|
|
||||||
# export --------------------------------------
|
# Export --------------------------------------
|
||||||
# coremltools>=4.1
|
# coremltools>=4.1 # CoreML export
|
||||||
# onnx>=1.9.0
|
# onnx>=1.9.0 # ONNX export
|
||||||
# scikit-learn==0.19.2 # for coreml quantization
|
# onnx-simplifier>=0.3.6 # ONNX simplifier
|
||||||
|
# scikit-learn==0.19.2 # CoreML quantization
|
||||||
|
# tensorflow>=2.4.1 # TFLite export
|
||||||
|
# tensorflowjs>=3.9.0 # TF.js export
|
||||||
|
|
||||||
# extras --------------------------------------
|
# Extras --------------------------------------
|
||||||
|
# albumentations>=1.0.3
|
||||||
# Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
|
# Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
|
||||||
pycocotools>=2.0 # COCO mAP
|
# pycocotools>=2.0 # COCO mAP
|
||||||
thop # FLOPS computation
|
# roboflow
|
||||||
|
thop # FLOPs computation
|
||||||
|
|||||||
51
setup.cfg
Normal file
51
setup.cfg
Normal file
@ -0,0 +1,51 @@
|
|||||||
|
# Project-wide configuration file, can be used for package metadata and other toll configurations
|
||||||
|
# Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments
|
||||||
|
|
||||||
|
[metadata]
|
||||||
|
license_file = LICENSE
|
||||||
|
description-file = README.md
|
||||||
|
|
||||||
|
|
||||||
|
[tool:pytest]
|
||||||
|
norecursedirs =
|
||||||
|
.git
|
||||||
|
dist
|
||||||
|
build
|
||||||
|
addopts =
|
||||||
|
--doctest-modules
|
||||||
|
--durations=25
|
||||||
|
--color=yes
|
||||||
|
|
||||||
|
|
||||||
|
[flake8]
|
||||||
|
max-line-length = 120
|
||||||
|
exclude = .tox,*.egg,build,temp
|
||||||
|
select = E,W,F
|
||||||
|
doctests = True
|
||||||
|
verbose = 2
|
||||||
|
# https://pep8.readthedocs.io/en/latest/intro.html#error-codes
|
||||||
|
format = pylint
|
||||||
|
# see: https://www.flake8rules.com/
|
||||||
|
ignore =
|
||||||
|
E731 # Do not assign a lambda expression, use a def
|
||||||
|
F405
|
||||||
|
E402
|
||||||
|
F841
|
||||||
|
E741
|
||||||
|
F821
|
||||||
|
E722
|
||||||
|
F401
|
||||||
|
W504
|
||||||
|
E127
|
||||||
|
W504
|
||||||
|
E231
|
||||||
|
E501
|
||||||
|
F403
|
||||||
|
E302
|
||||||
|
F541
|
||||||
|
|
||||||
|
|
||||||
|
[isort]
|
||||||
|
# https://pycqa.github.io/isort/docs/configuration/options.html
|
||||||
|
line_length = 120
|
||||||
|
multi_line_output = 0
|
||||||
349
test.py
349
test.py
@ -1,349 +0,0 @@
|
|||||||
import argparse
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
|
||||||
from threading import Thread
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
import yaml
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
from models.experimental import attempt_load
|
|
||||||
from utils.datasets import create_dataloader
|
|
||||||
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
|
|
||||||
box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
|
|
||||||
from utils.metrics import ap_per_class, ConfusionMatrix
|
|
||||||
from utils.plots import plot_images, output_to_target, plot_study_txt
|
|
||||||
from utils.torch_utils import select_device, time_synchronized
|
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
|
||||||
def test(data,
|
|
||||||
weights=None,
|
|
||||||
batch_size=32,
|
|
||||||
imgsz=640,
|
|
||||||
conf_thres=0.001,
|
|
||||||
iou_thres=0.6, # for NMS
|
|
||||||
save_json=False,
|
|
||||||
single_cls=False,
|
|
||||||
augment=False,
|
|
||||||
verbose=False,
|
|
||||||
model=None,
|
|
||||||
dataloader=None,
|
|
||||||
save_dir=Path(''), # for saving images
|
|
||||||
save_txt=False, # for auto-labelling
|
|
||||||
save_hybrid=False, # for hybrid auto-labelling
|
|
||||||
save_conf=False, # save auto-label confidences
|
|
||||||
plots=True,
|
|
||||||
wandb_logger=None,
|
|
||||||
compute_loss=None,
|
|
||||||
half_precision=True,
|
|
||||||
is_coco=False,
|
|
||||||
opt=None):
|
|
||||||
# Initialize/load model and set device
|
|
||||||
training = model is not None
|
|
||||||
if training: # called by train.py
|
|
||||||
device = next(model.parameters()).device # get model device
|
|
||||||
|
|
||||||
else: # called directly
|
|
||||||
set_logging()
|
|
||||||
device = select_device(opt.device, batch_size=batch_size)
|
|
||||||
|
|
||||||
# Directories
|
|
||||||
save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
|
|
||||||
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
|
||||||
|
|
||||||
# Load model
|
|
||||||
model = attempt_load(weights, map_location=device) # load FP32 model
|
|
||||||
gs = max(int(model.stride.max()), 32) # grid size (max stride)
|
|
||||||
imgsz = check_img_size(imgsz, s=gs) # check img_size
|
|
||||||
|
|
||||||
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
|
||||||
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
|
||||||
# model = nn.DataParallel(model)
|
|
||||||
|
|
||||||
# Half
|
|
||||||
half = device.type != 'cpu' and half_precision # half precision only supported on CUDA
|
|
||||||
if half:
|
|
||||||
model.half()
|
|
||||||
|
|
||||||
# Configure
|
|
||||||
model.eval()
|
|
||||||
if isinstance(data, str):
|
|
||||||
is_coco = data.endswith('coco.yaml')
|
|
||||||
with open(data) as f:
|
|
||||||
data = yaml.safe_load(f)
|
|
||||||
check_dataset(data) # check
|
|
||||||
nc = 1 if single_cls else int(data['nc']) # number of classes
|
|
||||||
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
|
||||||
niou = iouv.numel()
|
|
||||||
|
|
||||||
# Logging
|
|
||||||
log_imgs = 0
|
|
||||||
if wandb_logger and wandb_logger.wandb:
|
|
||||||
log_imgs = min(wandb_logger.log_imgs, 100)
|
|
||||||
# Dataloader
|
|
||||||
if not training:
|
|
||||||
if device.type != 'cpu':
|
|
||||||
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
|
|
||||||
task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images
|
|
||||||
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,
|
|
||||||
prefix=colorstr(f'{task}: '))[0]
|
|
||||||
|
|
||||||
seen = 0
|
|
||||||
confusion_matrix = ConfusionMatrix(nc=nc)
|
|
||||||
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
|
|
||||||
coco91class = coco80_to_coco91_class()
|
|
||||||
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
|
|
||||||
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
|
|
||||||
loss = torch.zeros(3, device=device)
|
|
||||||
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
|
|
||||||
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
|
|
||||||
img = img.to(device, non_blocking=True)
|
|
||||||
img = img.half() if half else img.float() # uint8 to fp16/32
|
|
||||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
||||||
targets = targets.to(device)
|
|
||||||
nb, _, height, width = img.shape # batch size, channels, height, width
|
|
||||||
|
|
||||||
# Run model
|
|
||||||
t = time_synchronized()
|
|
||||||
out, train_out = model(img, augment=augment) # inference and training outputs
|
|
||||||
t0 += time_synchronized() - t
|
|
||||||
|
|
||||||
# Compute loss
|
|
||||||
if compute_loss:
|
|
||||||
loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
|
|
||||||
|
|
||||||
# Run NMS
|
|
||||||
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
|
||||||
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
|
||||||
t = time_synchronized()
|
|
||||||
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
|
|
||||||
t1 += time_synchronized() - t
|
|
||||||
|
|
||||||
# Statistics per image
|
|
||||||
for si, pred in enumerate(out):
|
|
||||||
labels = targets[targets[:, 0] == si, 1:]
|
|
||||||
nl = len(labels)
|
|
||||||
tcls = labels[:, 0].tolist() if nl else [] # target class
|
|
||||||
path = Path(paths[si])
|
|
||||||
seen += 1
|
|
||||||
|
|
||||||
if len(pred) == 0:
|
|
||||||
if nl:
|
|
||||||
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
|
||||||
continue
|
|
||||||
|
|
||||||
# Predictions
|
|
||||||
if single_cls:
|
|
||||||
pred[:, 5] = 0
|
|
||||||
predn = pred.clone()
|
|
||||||
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
|
|
||||||
|
|
||||||
# Append to text file
|
|
||||||
if save_txt:
|
|
||||||
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
|
|
||||||
for *xyxy, conf, cls in predn.tolist():
|
|
||||||
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
|
||||||
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
|
||||||
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
|
|
||||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
|
||||||
|
|
||||||
# W&B logging - Media Panel Plots
|
|
||||||
if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0: # Check for test operation
|
|
||||||
if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
|
|
||||||
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
|
||||||
"class_id": int(cls),
|
|
||||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
|
||||||
"scores": {"class_score": conf},
|
|
||||||
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
|
||||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
|
||||||
wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))
|
|
||||||
wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None
|
|
||||||
|
|
||||||
# Append to pycocotools JSON dictionary
|
|
||||||
if save_json:
|
|
||||||
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
|
|
||||||
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
|
|
||||||
box = xyxy2xywh(predn[:, :4]) # xywh
|
|
||||||
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
|
||||||
for p, b in zip(pred.tolist(), box.tolist()):
|
|
||||||
jdict.append({'image_id': image_id,
|
|
||||||
'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
|
|
||||||
'bbox': [round(x, 3) for x in b],
|
|
||||||
'score': round(p[4], 5)})
|
|
||||||
|
|
||||||
# Assign all predictions as incorrect
|
|
||||||
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
|
|
||||||
if nl:
|
|
||||||
detected = [] # target indices
|
|
||||||
tcls_tensor = labels[:, 0]
|
|
||||||
|
|
||||||
# target boxes
|
|
||||||
tbox = xywh2xyxy(labels[:, 1:5])
|
|
||||||
scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
|
|
||||||
if plots:
|
|
||||||
confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
|
|
||||||
|
|
||||||
# Per target class
|
|
||||||
for cls in torch.unique(tcls_tensor):
|
|
||||||
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # target indices
|
|
||||||
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # prediction indices
|
|
||||||
|
|
||||||
# Search for detections
|
|
||||||
if pi.shape[0]:
|
|
||||||
# Prediction to target ious
|
|
||||||
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
|
|
||||||
|
|
||||||
# Append detections
|
|
||||||
detected_set = set()
|
|
||||||
for j in (ious > iouv[0]).nonzero(as_tuple=False):
|
|
||||||
d = ti[i[j]] # detected target
|
|
||||||
if d.item() not in detected_set:
|
|
||||||
detected_set.add(d.item())
|
|
||||||
detected.append(d)
|
|
||||||
correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
|
|
||||||
if len(detected) == nl: # all targets already located in image
|
|
||||||
break
|
|
||||||
|
|
||||||
# Append statistics (correct, conf, pcls, tcls)
|
|
||||||
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
|
|
||||||
|
|
||||||
# Plot images
|
|
||||||
if plots and batch_i < 3:
|
|
||||||
f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
|
|
||||||
Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
|
|
||||||
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
|
|
||||||
Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start()
|
|
||||||
|
|
||||||
# Compute statistics
|
|
||||||
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
|
|
||||||
if len(stats) and stats[0].any():
|
|
||||||
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
|
|
||||||
ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
|
|
||||||
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
|
|
||||||
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
|
|
||||||
else:
|
|
||||||
nt = torch.zeros(1)
|
|
||||||
|
|
||||||
# Print results
|
|
||||||
pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
|
|
||||||
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
|
|
||||||
|
|
||||||
# Print results per class
|
|
||||||
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
|
|
||||||
for i, c in enumerate(ap_class):
|
|
||||||
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
|
|
||||||
|
|
||||||
# Print speeds
|
|
||||||
t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
|
|
||||||
if not training:
|
|
||||||
print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
|
|
||||||
|
|
||||||
# Plots
|
|
||||||
if plots:
|
|
||||||
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
|
|
||||||
if wandb_logger and wandb_logger.wandb:
|
|
||||||
val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
|
|
||||||
wandb_logger.log({"Validation": val_batches})
|
|
||||||
if wandb_images:
|
|
||||||
wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
|
|
||||||
|
|
||||||
# Save JSON
|
|
||||||
if save_json and len(jdict):
|
|
||||||
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
|
|
||||||
anno_json = '../coco/annotations/instances_val2017.json' # annotations json
|
|
||||||
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
|
|
||||||
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
|
|
||||||
with open(pred_json, 'w') as f:
|
|
||||||
json.dump(jdict, f)
|
|
||||||
|
|
||||||
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
|
|
||||||
from pycocotools.coco import COCO
|
|
||||||
from pycocotools.cocoeval import COCOeval
|
|
||||||
|
|
||||||
anno = COCO(anno_json) # init annotations api
|
|
||||||
pred = anno.loadRes(pred_json) # init predictions api
|
|
||||||
eval = COCOeval(anno, pred, 'bbox')
|
|
||||||
if is_coco:
|
|
||||||
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
|
|
||||||
eval.evaluate()
|
|
||||||
eval.accumulate()
|
|
||||||
eval.summarize()
|
|
||||||
map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
|
|
||||||
except Exception as e:
|
|
||||||
print(f'pycocotools unable to run: {e}')
|
|
||||||
|
|
||||||
# Return results
|
|
||||||
model.float() # for training
|
|
||||||
if not training:
|
|
||||||
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
|
||||||
print(f"Results saved to {save_dir}{s}")
|
|
||||||
maps = np.zeros(nc) + map
|
|
||||||
for i, c in enumerate(ap_class):
|
|
||||||
maps[c] = ap[i]
|
|
||||||
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
parser = argparse.ArgumentParser(prog='test.py')
|
|
||||||
parser.add_argument('--weights', nargs='+', type=str, default='yolov3.pt', help='model.pt path(s)')
|
|
||||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
|
|
||||||
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
|
||||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
|
||||||
parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
|
|
||||||
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
|
|
||||||
parser.add_argument('--task', default='val', help='train, val, test, speed or study')
|
|
||||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
|
||||||
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
|
|
||||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
|
||||||
parser.add_argument('--verbose', action='store_true', help='report mAP by class')
|
|
||||||
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
|
||||||
parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
|
|
||||||
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
|
||||||
parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
|
|
||||||
parser.add_argument('--project', default='runs/test', help='save to project/name')
|
|
||||||
parser.add_argument('--name', default='exp', help='save to project/name')
|
|
||||||
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
|
||||||
opt = parser.parse_args()
|
|
||||||
opt.save_json |= opt.data.endswith('coco.yaml')
|
|
||||||
opt.data = check_file(opt.data) # check file
|
|
||||||
print(opt)
|
|
||||||
check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
|
|
||||||
|
|
||||||
if opt.task in ('train', 'val', 'test'): # run normally
|
|
||||||
test(opt.data,
|
|
||||||
opt.weights,
|
|
||||||
opt.batch_size,
|
|
||||||
opt.img_size,
|
|
||||||
opt.conf_thres,
|
|
||||||
opt.iou_thres,
|
|
||||||
opt.save_json,
|
|
||||||
opt.single_cls,
|
|
||||||
opt.augment,
|
|
||||||
opt.verbose,
|
|
||||||
save_txt=opt.save_txt | opt.save_hybrid,
|
|
||||||
save_hybrid=opt.save_hybrid,
|
|
||||||
save_conf=opt.save_conf,
|
|
||||||
opt=opt
|
|
||||||
)
|
|
||||||
|
|
||||||
elif opt.task == 'speed': # speed benchmarks
|
|
||||||
for w in opt.weights:
|
|
||||||
test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False, opt=opt)
|
|
||||||
|
|
||||||
elif opt.task == 'study': # run over a range of settings and save/plot
|
|
||||||
# python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt
|
|
||||||
x = list(range(256, 1536 + 128, 128)) # x axis (image sizes)
|
|
||||||
for w in opt.weights:
|
|
||||||
f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt' # filename to save to
|
|
||||||
y = [] # y axis
|
|
||||||
for i in x: # img-size
|
|
||||||
print(f'\nRunning {f} point {i}...')
|
|
||||||
r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
|
|
||||||
plots=False, opt=opt)
|
|
||||||
y.append(r + t) # results and times
|
|
||||||
np.savetxt(f, y, fmt='%10.4g') # save
|
|
||||||
os.system('zip -r study.zip study_*.txt')
|
|
||||||
plot_study_txt(x=x) # plot
|
|
||||||
656
train.py
656
train.py
@ -1,147 +1,178 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Train a model on a custom dataset
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
$ python path/to/train.py --data coco128.yaml --weights yolov3.pt --img 640
|
||||||
|
"""
|
||||||
import argparse
|
import argparse
|
||||||
import logging
|
|
||||||
import math
|
import math
|
||||||
import os
|
import os
|
||||||
import random
|
import random
|
||||||
|
import sys
|
||||||
import time
|
import time
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from threading import Thread
|
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import torch
|
||||||
import torch.distributed as dist
|
import torch.distributed as dist
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.nn.functional as F
|
|
||||||
import torch.optim as optim
|
|
||||||
import torch.optim.lr_scheduler as lr_scheduler
|
|
||||||
import torch.utils.data
|
|
||||||
import yaml
|
import yaml
|
||||||
from torch.cuda import amp
|
from torch.cuda import amp
|
||||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||||
from torch.utils.tensorboard import SummaryWriter
|
from torch.optim import SGD, Adam, lr_scheduler
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
import test # import test.py to get mAP after each epoch
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[0] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||||
|
|
||||||
|
import val # for end-of-epoch mAP
|
||||||
from models.experimental import attempt_load
|
from models.experimental import attempt_load
|
||||||
from models.yolo import Model
|
from models.yolo import Model
|
||||||
from utils.autoanchor import check_anchors
|
from utils.autoanchor import check_anchors
|
||||||
|
from utils.autobatch import check_train_batch_size
|
||||||
|
from utils.callbacks import Callbacks
|
||||||
from utils.datasets import create_dataloader
|
from utils.datasets import create_dataloader
|
||||||
from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
|
from utils.downloads import attempt_download
|
||||||
fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
|
from utils.general import (LOGGER, NCOLS, check_dataset, check_file, check_git_status, check_img_size,
|
||||||
check_requirements, print_mutation, set_logging, one_cycle, colorstr
|
check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,
|
||||||
from utils.google_utils import attempt_download
|
init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods,
|
||||||
|
one_cycle, print_args, print_mutation, strip_optimizer)
|
||||||
|
from utils.loggers import Loggers
|
||||||
|
from utils.loggers.wandb.wandb_utils import check_wandb_resume
|
||||||
from utils.loss import ComputeLoss
|
from utils.loss import ComputeLoss
|
||||||
from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
|
from utils.metrics import fitness
|
||||||
from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, de_parallel
|
from utils.plots import plot_evolve, plot_labels
|
||||||
from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
|
from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, select_device, torch_distributed_zero_first
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
|
||||||
|
RANK = int(os.getenv('RANK', -1))
|
||||||
|
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
|
||||||
|
|
||||||
|
|
||||||
def train(hyp, opt, device, tb_writer=None):
|
def train(hyp, # path/to/hyp.yaml or hyp dictionary
|
||||||
logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
|
opt,
|
||||||
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
|
device,
|
||||||
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
|
callbacks
|
||||||
|
):
|
||||||
|
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \
|
||||||
|
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
|
||||||
|
opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
|
||||||
|
|
||||||
# Directories
|
# Directories
|
||||||
wdir = save_dir / 'weights'
|
w = save_dir / 'weights' # weights dir
|
||||||
wdir.mkdir(parents=True, exist_ok=True) # make dir
|
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
|
||||||
last = wdir / 'last.pt'
|
last, best = w / 'last.pt', w / 'best.pt'
|
||||||
best = wdir / 'best.pt'
|
|
||||||
results_file = save_dir / 'results.txt'
|
# Hyperparameters
|
||||||
|
if isinstance(hyp, str):
|
||||||
|
with open(hyp, errors='ignore') as f:
|
||||||
|
hyp = yaml.safe_load(f) # load hyps dict
|
||||||
|
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
|
||||||
|
|
||||||
# Save run settings
|
# Save run settings
|
||||||
with open(save_dir / 'hyp.yaml', 'w') as f:
|
with open(save_dir / 'hyp.yaml', 'w') as f:
|
||||||
yaml.safe_dump(hyp, f, sort_keys=False)
|
yaml.safe_dump(hyp, f, sort_keys=False)
|
||||||
with open(save_dir / 'opt.yaml', 'w') as f:
|
with open(save_dir / 'opt.yaml', 'w') as f:
|
||||||
yaml.safe_dump(vars(opt), f, sort_keys=False)
|
yaml.safe_dump(vars(opt), f, sort_keys=False)
|
||||||
|
data_dict = None
|
||||||
|
|
||||||
# Configure
|
# Loggers
|
||||||
plots = not opt.evolve # create plots
|
if RANK in [-1, 0]:
|
||||||
|
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
|
||||||
|
if loggers.wandb:
|
||||||
|
data_dict = loggers.wandb.data_dict
|
||||||
|
if resume:
|
||||||
|
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp
|
||||||
|
|
||||||
|
# Register actions
|
||||||
|
for k in methods(loggers):
|
||||||
|
callbacks.register_action(k, callback=getattr(loggers, k))
|
||||||
|
|
||||||
|
# Config
|
||||||
|
plots = not evolve # create plots
|
||||||
cuda = device.type != 'cpu'
|
cuda = device.type != 'cpu'
|
||||||
init_seeds(2 + rank)
|
init_seeds(1 + RANK)
|
||||||
with open(opt.data) as f:
|
with torch_distributed_zero_first(LOCAL_RANK):
|
||||||
data_dict = yaml.safe_load(f) # data dict
|
data_dict = data_dict or check_dataset(data) # check if None
|
||||||
|
train_path, val_path = data_dict['train'], data_dict['val']
|
||||||
# Logging- Doing this before checking the dataset. Might update data_dict
|
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
|
||||||
loggers = {'wandb': None} # loggers dict
|
names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
|
||||||
if rank in [-1, 0]:
|
assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check
|
||||||
opt.hyp = hyp # add hyperparameters
|
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
|
||||||
run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
|
|
||||||
wandb_logger = WandbLogger(opt, save_dir.stem, run_id, data_dict)
|
|
||||||
loggers['wandb'] = wandb_logger.wandb
|
|
||||||
data_dict = wandb_logger.data_dict
|
|
||||||
if wandb_logger.wandb:
|
|
||||||
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming
|
|
||||||
|
|
||||||
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
|
|
||||||
names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
|
|
||||||
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
|
|
||||||
is_coco = opt.data.endswith('coco.yaml') and nc == 80 # COCO dataset
|
|
||||||
|
|
||||||
# Model
|
# Model
|
||||||
|
check_suffix(weights, '.pt') # check weights
|
||||||
pretrained = weights.endswith('.pt')
|
pretrained = weights.endswith('.pt')
|
||||||
if pretrained:
|
if pretrained:
|
||||||
with torch_distributed_zero_first(rank):
|
with torch_distributed_zero_first(LOCAL_RANK):
|
||||||
weights = attempt_download(weights) # download if not found locally
|
weights = attempt_download(weights) # download if not found locally
|
||||||
ckpt = torch.load(weights, map_location=device) # load checkpoint
|
ckpt = torch.load(weights, map_location=device) # load checkpoint
|
||||||
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
||||||
exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys
|
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
|
||||||
state_dict = ckpt['model'].float().state_dict() # to FP32
|
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
|
||||||
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect
|
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
|
||||||
model.load_state_dict(state_dict, strict=False) # load
|
model.load_state_dict(csd, strict=False) # load
|
||||||
logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report
|
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
|
||||||
else:
|
else:
|
||||||
model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
||||||
with torch_distributed_zero_first(rank):
|
|
||||||
check_dataset(data_dict) # check
|
|
||||||
train_path = data_dict['train']
|
|
||||||
test_path = data_dict['val']
|
|
||||||
|
|
||||||
# Freeze
|
# Freeze
|
||||||
freeze = [] # parameter names to freeze (full or partial)
|
freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze
|
||||||
for k, v in model.named_parameters():
|
for k, v in model.named_parameters():
|
||||||
v.requires_grad = True # train all layers
|
v.requires_grad = True # train all layers
|
||||||
if any(x in k for x in freeze):
|
if any(x in k for x in freeze):
|
||||||
print('freezing %s' % k)
|
LOGGER.info(f'freezing {k}')
|
||||||
v.requires_grad = False
|
v.requires_grad = False
|
||||||
|
|
||||||
|
# Image size
|
||||||
|
gs = max(int(model.stride.max()), 32) # grid size (max stride)
|
||||||
|
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
|
||||||
|
|
||||||
|
# Batch size
|
||||||
|
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
|
||||||
|
batch_size = check_train_batch_size(model, imgsz)
|
||||||
|
|
||||||
# Optimizer
|
# Optimizer
|
||||||
nbs = 64 # nominal batch size
|
nbs = 64 # nominal batch size
|
||||||
accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing
|
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
|
||||||
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
|
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
|
||||||
logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
|
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
|
||||||
|
|
||||||
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
|
g0, g1, g2 = [], [], [] # optimizer parameter groups
|
||||||
for k, v in model.named_modules():
|
for v in model.modules():
|
||||||
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
|
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
|
||||||
pg2.append(v.bias) # biases
|
g2.append(v.bias)
|
||||||
if isinstance(v, nn.BatchNorm2d):
|
if isinstance(v, nn.BatchNorm2d): # weight (no decay)
|
||||||
pg0.append(v.weight) # no decay
|
g0.append(v.weight)
|
||||||
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
|
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay)
|
||||||
pg1.append(v.weight) # apply decay
|
g1.append(v.weight)
|
||||||
|
|
||||||
if opt.adam:
|
if opt.adam:
|
||||||
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
|
optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
|
||||||
else:
|
else:
|
||||||
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
|
optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
|
||||||
|
|
||||||
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
|
optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']}) # add g1 with weight_decay
|
||||||
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
|
optimizer.add_param_group({'params': g2}) # add g2 (biases)
|
||||||
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
|
LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "
|
||||||
del pg0, pg1, pg2
|
f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias")
|
||||||
|
del g0, g1, g2
|
||||||
|
|
||||||
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
|
# Scheduler
|
||||||
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
|
|
||||||
if opt.linear_lr:
|
if opt.linear_lr:
|
||||||
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
|
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
|
||||||
else:
|
else:
|
||||||
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
|
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
|
||||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
|
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
|
||||||
# plot_lr_scheduler(optimizer, scheduler, epochs)
|
|
||||||
|
|
||||||
# EMA
|
# EMA
|
||||||
ema = ModelEMA(model) if rank in [-1, 0] else None
|
ema = ModelEMA(model) if RANK in [-1, 0] else None
|
||||||
|
|
||||||
# Resume
|
# Resume
|
||||||
start_epoch, best_fitness = 0, 0.0
|
start_epoch, best_fitness = 0, 0.0
|
||||||
@ -156,80 +187,70 @@ def train(hyp, opt, device, tb_writer=None):
|
|||||||
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
|
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
|
||||||
ema.updates = ckpt['updates']
|
ema.updates = ckpt['updates']
|
||||||
|
|
||||||
# Results
|
|
||||||
if ckpt.get('training_results') is not None:
|
|
||||||
results_file.write_text(ckpt['training_results']) # write results.txt
|
|
||||||
|
|
||||||
# Epochs
|
# Epochs
|
||||||
start_epoch = ckpt['epoch'] + 1
|
start_epoch = ckpt['epoch'] + 1
|
||||||
if opt.resume:
|
if resume:
|
||||||
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
|
assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.'
|
||||||
if epochs < start_epoch:
|
if epochs < start_epoch:
|
||||||
logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
|
LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")
|
||||||
(weights, ckpt['epoch'], epochs))
|
|
||||||
epochs += ckpt['epoch'] # finetune additional epochs
|
epochs += ckpt['epoch'] # finetune additional epochs
|
||||||
|
|
||||||
del ckpt, state_dict
|
del ckpt, csd
|
||||||
|
|
||||||
# Image sizes
|
|
||||||
gs = max(int(model.stride.max()), 32) # grid size (max stride)
|
|
||||||
nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj'])
|
|
||||||
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
|
|
||||||
|
|
||||||
# DP mode
|
# DP mode
|
||||||
if cuda and rank == -1 and torch.cuda.device_count() > 1:
|
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
|
||||||
|
LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
|
||||||
|
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
|
||||||
model = torch.nn.DataParallel(model)
|
model = torch.nn.DataParallel(model)
|
||||||
|
|
||||||
# SyncBatchNorm
|
# SyncBatchNorm
|
||||||
if opt.sync_bn and cuda and rank != -1:
|
if opt.sync_bn and cuda and RANK != -1:
|
||||||
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
|
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
|
||||||
logger.info('Using SyncBatchNorm()')
|
LOGGER.info('Using SyncBatchNorm()')
|
||||||
|
|
||||||
# Trainloader
|
# Trainloader
|
||||||
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
|
train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
|
||||||
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
|
hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK,
|
||||||
world_size=opt.world_size, workers=opt.workers,
|
workers=workers, image_weights=opt.image_weights, quad=opt.quad,
|
||||||
image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
|
prefix=colorstr('train: '), shuffle=True)
|
||||||
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
|
mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class
|
||||||
nb = len(dataloader) # number of batches
|
nb = len(train_loader) # number of batches
|
||||||
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
|
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
|
||||||
|
|
||||||
# Process 0
|
# Process 0
|
||||||
if rank in [-1, 0]:
|
if RANK in [-1, 0]:
|
||||||
testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader
|
val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls,
|
||||||
hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
|
hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1,
|
||||||
world_size=opt.world_size, workers=opt.workers,
|
workers=workers, pad=0.5,
|
||||||
pad=0.5, prefix=colorstr('val: '))[0]
|
prefix=colorstr('val: '))[0]
|
||||||
|
|
||||||
if not opt.resume:
|
if not resume:
|
||||||
labels = np.concatenate(dataset.labels, 0)
|
labels = np.concatenate(dataset.labels, 0)
|
||||||
c = torch.tensor(labels[:, 0]) # classes
|
# c = torch.tensor(labels[:, 0]) # classes
|
||||||
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
|
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
|
||||||
# model._initialize_biases(cf.to(device))
|
# model._initialize_biases(cf.to(device))
|
||||||
if plots:
|
if plots:
|
||||||
plot_labels(labels, names, save_dir, loggers)
|
plot_labels(labels, names, save_dir)
|
||||||
if tb_writer:
|
|
||||||
tb_writer.add_histogram('classes', c, 0)
|
|
||||||
|
|
||||||
# Anchors
|
# Anchors
|
||||||
if not opt.noautoanchor:
|
if not opt.noautoanchor:
|
||||||
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
|
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
|
||||||
model.half().float() # pre-reduce anchor precision
|
model.half().float() # pre-reduce anchor precision
|
||||||
|
|
||||||
|
callbacks.run('on_pretrain_routine_end')
|
||||||
|
|
||||||
# DDP mode
|
# DDP mode
|
||||||
if cuda and rank != -1:
|
if cuda and RANK != -1:
|
||||||
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
|
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
|
||||||
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
|
|
||||||
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
|
|
||||||
|
|
||||||
# Model parameters
|
# Model parameters
|
||||||
hyp['box'] *= 3. / nl # scale to layers
|
nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
|
||||||
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
|
hyp['box'] *= 3 / nl # scale to layers
|
||||||
hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers
|
hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers
|
||||||
|
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
|
||||||
hyp['label_smoothing'] = opt.label_smoothing
|
hyp['label_smoothing'] = opt.label_smoothing
|
||||||
model.nc = nc # attach number of classes to model
|
model.nc = nc # attach number of classes to model
|
||||||
model.hyp = hyp # attach hyperparameters to model
|
model.hyp = hyp # attach hyperparameters to model
|
||||||
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
|
|
||||||
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
|
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
|
||||||
model.names = names
|
model.names = names
|
||||||
|
|
||||||
@ -237,53 +258,47 @@ def train(hyp, opt, device, tb_writer=None):
|
|||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
|
nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
|
||||||
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
|
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
|
||||||
|
last_opt_step = -1
|
||||||
maps = np.zeros(nc) # mAP per class
|
maps = np.zeros(nc) # mAP per class
|
||||||
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
|
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
|
||||||
scheduler.last_epoch = start_epoch - 1 # do not move
|
scheduler.last_epoch = start_epoch - 1 # do not move
|
||||||
scaler = amp.GradScaler(enabled=cuda)
|
scaler = amp.GradScaler(enabled=cuda)
|
||||||
|
stopper = EarlyStopping(patience=opt.patience)
|
||||||
compute_loss = ComputeLoss(model) # init loss class
|
compute_loss = ComputeLoss(model) # init loss class
|
||||||
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
|
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
|
||||||
f'Using {dataloader.num_workers} dataloader workers\n'
|
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
|
||||||
f'Logging results to {save_dir}\n'
|
f"Logging results to {colorstr('bold', save_dir)}\n"
|
||||||
f'Starting training for {epochs} epochs...')
|
f'Starting training for {epochs} epochs...')
|
||||||
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
|
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
|
||||||
model.train()
|
model.train()
|
||||||
|
|
||||||
# Update image weights (optional)
|
# Update image weights (optional, single-GPU only)
|
||||||
if opt.image_weights:
|
if opt.image_weights:
|
||||||
# Generate indices
|
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
|
||||||
if rank in [-1, 0]:
|
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
|
||||||
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
|
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
|
||||||
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
|
|
||||||
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
|
|
||||||
# Broadcast if DDP
|
|
||||||
if rank != -1:
|
|
||||||
indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
|
|
||||||
dist.broadcast(indices, 0)
|
|
||||||
if rank != 0:
|
|
||||||
dataset.indices = indices.cpu().numpy()
|
|
||||||
|
|
||||||
# Update mosaic border
|
# Update mosaic border (optional)
|
||||||
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
|
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
|
||||||
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
|
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
|
||||||
|
|
||||||
mloss = torch.zeros(4, device=device) # mean losses
|
mloss = torch.zeros(3, device=device) # mean losses
|
||||||
if rank != -1:
|
if RANK != -1:
|
||||||
dataloader.sampler.set_epoch(epoch)
|
train_loader.sampler.set_epoch(epoch)
|
||||||
pbar = enumerate(dataloader)
|
pbar = enumerate(train_loader)
|
||||||
logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
|
LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))
|
||||||
if rank in [-1, 0]:
|
if RANK in [-1, 0]:
|
||||||
pbar = tqdm(pbar, total=nb) # progress bar
|
pbar = tqdm(pbar, total=nb, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
|
||||||
optimizer.zero_grad()
|
optimizer.zero_grad()
|
||||||
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
|
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
|
||||||
ni = i + nb * epoch # number integrated batches (since train start)
|
ni = i + nb * epoch # number integrated batches (since train start)
|
||||||
imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
|
imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0
|
||||||
|
|
||||||
# Warmup
|
# Warmup
|
||||||
if ni <= nw:
|
if ni <= nw:
|
||||||
xi = [0, nw] # x interp
|
xi = [0, nw] # x interp
|
||||||
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
|
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
|
||||||
accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
|
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
|
||||||
for j, x in enumerate(optimizer.param_groups):
|
for j, x in enumerate(optimizer.param_groups):
|
||||||
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
|
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
|
||||||
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
|
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
|
||||||
@ -296,14 +311,14 @@ def train(hyp, opt, device, tb_writer=None):
|
|||||||
sf = sz / max(imgs.shape[2:]) # scale factor
|
sf = sz / max(imgs.shape[2:]) # scale factor
|
||||||
if sf != 1:
|
if sf != 1:
|
||||||
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
|
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
|
||||||
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
|
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
|
||||||
|
|
||||||
# Forward
|
# Forward
|
||||||
with amp.autocast(enabled=cuda):
|
with amp.autocast(enabled=cuda):
|
||||||
pred = model(imgs) # forward
|
pred = model(imgs) # forward
|
||||||
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
|
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
|
||||||
if rank != -1:
|
if RANK != -1:
|
||||||
loss *= opt.world_size # gradient averaged between devices in DDP mode
|
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
|
||||||
if opt.quad:
|
if opt.quad:
|
||||||
loss *= 4.
|
loss *= 4.
|
||||||
|
|
||||||
@ -311,234 +326,203 @@ def train(hyp, opt, device, tb_writer=None):
|
|||||||
scaler.scale(loss).backward()
|
scaler.scale(loss).backward()
|
||||||
|
|
||||||
# Optimize
|
# Optimize
|
||||||
if ni % accumulate == 0:
|
if ni - last_opt_step >= accumulate:
|
||||||
scaler.step(optimizer) # optimizer.step
|
scaler.step(optimizer) # optimizer.step
|
||||||
scaler.update()
|
scaler.update()
|
||||||
optimizer.zero_grad()
|
optimizer.zero_grad()
|
||||||
if ema:
|
if ema:
|
||||||
ema.update(model)
|
ema.update(model)
|
||||||
|
last_opt_step = ni
|
||||||
# Print
|
|
||||||
if rank in [-1, 0]:
|
|
||||||
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
|
|
||||||
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
|
|
||||||
s = ('%10s' * 2 + '%10.4g' * 6) % (
|
|
||||||
'%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
|
|
||||||
pbar.set_description(s)
|
|
||||||
|
|
||||||
# Plot
|
|
||||||
if plots and ni < 3:
|
|
||||||
f = save_dir / f'train_batch{ni}.jpg' # filename
|
|
||||||
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
|
||||||
if tb_writer:
|
|
||||||
tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs, strict=False), []) # model graph
|
|
||||||
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
|
|
||||||
elif plots and ni == 10 and wandb_logger.wandb:
|
|
||||||
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
|
|
||||||
save_dir.glob('train*.jpg') if x.exists()]})
|
|
||||||
|
|
||||||
# end batch ------------------------------------------------------------------------------------------------
|
|
||||||
# end epoch ----------------------------------------------------------------------------------------------------
|
|
||||||
|
|
||||||
# Scheduler
|
|
||||||
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
|
|
||||||
scheduler.step()
|
|
||||||
|
|
||||||
# DDP process 0 or single-GPU
|
|
||||||
if rank in [-1, 0]:
|
|
||||||
# mAP
|
|
||||||
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
|
|
||||||
final_epoch = epoch + 1 == epochs
|
|
||||||
if not opt.notest or final_epoch: # Calculate mAP
|
|
||||||
wandb_logger.current_epoch = epoch + 1
|
|
||||||
results, maps, times = test.test(data_dict,
|
|
||||||
batch_size=batch_size * 2,
|
|
||||||
imgsz=imgsz_test,
|
|
||||||
model=ema.ema,
|
|
||||||
single_cls=opt.single_cls,
|
|
||||||
dataloader=testloader,
|
|
||||||
save_dir=save_dir,
|
|
||||||
save_json=is_coco and final_epoch,
|
|
||||||
verbose=nc < 50 and final_epoch,
|
|
||||||
plots=plots and final_epoch,
|
|
||||||
wandb_logger=wandb_logger,
|
|
||||||
compute_loss=compute_loss,
|
|
||||||
is_coco=is_coco)
|
|
||||||
|
|
||||||
# Write
|
|
||||||
with open(results_file, 'a') as f:
|
|
||||||
f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
|
|
||||||
|
|
||||||
# Log
|
# Log
|
||||||
tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
|
if RANK in [-1, 0]:
|
||||||
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
|
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
|
||||||
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
|
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
|
||||||
'x/lr0', 'x/lr1', 'x/lr2'] # params
|
pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (
|
||||||
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
|
f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
|
||||||
if tb_writer:
|
callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)
|
||||||
tb_writer.add_scalar(tag, x, epoch) # tensorboard
|
# end batch ------------------------------------------------------------------------------------------------
|
||||||
if wandb_logger.wandb:
|
|
||||||
wandb_logger.log({tag: x}) # W&B
|
# Scheduler
|
||||||
|
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
|
||||||
|
scheduler.step()
|
||||||
|
|
||||||
|
if RANK in [-1, 0]:
|
||||||
|
# mAP
|
||||||
|
callbacks.run('on_train_epoch_end', epoch=epoch)
|
||||||
|
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
|
||||||
|
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
|
||||||
|
if not noval or final_epoch: # Calculate mAP
|
||||||
|
results, maps, _ = val.run(data_dict,
|
||||||
|
batch_size=batch_size // WORLD_SIZE * 2,
|
||||||
|
imgsz=imgsz,
|
||||||
|
model=ema.ema,
|
||||||
|
single_cls=single_cls,
|
||||||
|
dataloader=val_loader,
|
||||||
|
save_dir=save_dir,
|
||||||
|
plots=False,
|
||||||
|
callbacks=callbacks,
|
||||||
|
compute_loss=compute_loss)
|
||||||
|
|
||||||
# Update best mAP
|
# Update best mAP
|
||||||
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
|
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
|
||||||
if fi > best_fitness:
|
if fi > best_fitness:
|
||||||
best_fitness = fi
|
best_fitness = fi
|
||||||
wandb_logger.end_epoch(best_result=best_fitness == fi)
|
log_vals = list(mloss) + list(results) + lr
|
||||||
|
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
|
||||||
|
|
||||||
# Save model
|
# Save model
|
||||||
if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
|
if (not nosave) or (final_epoch and not evolve): # if save
|
||||||
ckpt = {'epoch': epoch,
|
ckpt = {'epoch': epoch,
|
||||||
'best_fitness': best_fitness,
|
'best_fitness': best_fitness,
|
||||||
'training_results': results_file.read_text(),
|
|
||||||
'model': deepcopy(de_parallel(model)).half(),
|
'model': deepcopy(de_parallel(model)).half(),
|
||||||
'ema': deepcopy(ema.ema).half(),
|
'ema': deepcopy(ema.ema).half(),
|
||||||
'updates': ema.updates,
|
'updates': ema.updates,
|
||||||
'optimizer': optimizer.state_dict(),
|
'optimizer': optimizer.state_dict(),
|
||||||
'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None}
|
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
|
||||||
|
'date': datetime.now().isoformat()}
|
||||||
|
|
||||||
# Save last, best and delete
|
# Save last, best and delete
|
||||||
torch.save(ckpt, last)
|
torch.save(ckpt, last)
|
||||||
if best_fitness == fi:
|
if best_fitness == fi:
|
||||||
torch.save(ckpt, best)
|
torch.save(ckpt, best)
|
||||||
if wandb_logger.wandb:
|
if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):
|
||||||
if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
|
torch.save(ckpt, w / f'epoch{epoch}.pt')
|
||||||
wandb_logger.log_model(
|
|
||||||
last.parent, opt, epoch, fi, best_model=best_fitness == fi)
|
|
||||||
del ckpt
|
del ckpt
|
||||||
|
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
|
||||||
|
|
||||||
|
# Stop Single-GPU
|
||||||
|
if RANK == -1 and stopper(epoch=epoch, fitness=fi):
|
||||||
|
break
|
||||||
|
|
||||||
|
# Stop DDP TODO: known issues shttps://github.com/ultralytics/yolov5/pull/4576
|
||||||
|
# stop = stopper(epoch=epoch, fitness=fi)
|
||||||
|
# if RANK == 0:
|
||||||
|
# dist.broadcast_object_list([stop], 0) # broadcast 'stop' to all ranks
|
||||||
|
|
||||||
|
# Stop DPP
|
||||||
|
# with torch_distributed_zero_first(RANK):
|
||||||
|
# if stop:
|
||||||
|
# break # must break all DDP ranks
|
||||||
|
|
||||||
# end epoch ----------------------------------------------------------------------------------------------------
|
# end epoch ----------------------------------------------------------------------------------------------------
|
||||||
# end training
|
# end training -----------------------------------------------------------------------------------------------------
|
||||||
if rank in [-1, 0]:
|
if RANK in [-1, 0]:
|
||||||
logger.info(f'{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.\n')
|
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
|
||||||
if plots:
|
for f in last, best:
|
||||||
plot_results(save_dir=save_dir) # save as results.png
|
if f.exists():
|
||||||
if wandb_logger.wandb:
|
strip_optimizer(f) # strip optimizers
|
||||||
files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
|
if f is best:
|
||||||
wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
|
LOGGER.info(f'\nValidating {f}...')
|
||||||
if (save_dir / f).exists()]})
|
results, _, _ = val.run(data_dict,
|
||||||
|
batch_size=batch_size // WORLD_SIZE * 2,
|
||||||
|
imgsz=imgsz,
|
||||||
|
model=attempt_load(f, device).half(),
|
||||||
|
iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65
|
||||||
|
single_cls=single_cls,
|
||||||
|
dataloader=val_loader,
|
||||||
|
save_dir=save_dir,
|
||||||
|
save_json=is_coco,
|
||||||
|
verbose=True,
|
||||||
|
plots=True,
|
||||||
|
callbacks=callbacks,
|
||||||
|
compute_loss=compute_loss) # val best model with plots
|
||||||
|
if is_coco:
|
||||||
|
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
|
||||||
|
|
||||||
if not opt.evolve:
|
callbacks.run('on_train_end', last, best, plots, epoch, results)
|
||||||
if is_coco: # COCO dataset
|
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
|
||||||
for m in [last, best] if best.exists() else [last]: # speed, mAP tests
|
|
||||||
results, _, _ = test.test(opt.data,
|
|
||||||
batch_size=batch_size * 2,
|
|
||||||
imgsz=imgsz_test,
|
|
||||||
conf_thres=0.001,
|
|
||||||
iou_thres=0.7,
|
|
||||||
model=attempt_load(m, device).half(),
|
|
||||||
single_cls=opt.single_cls,
|
|
||||||
dataloader=testloader,
|
|
||||||
save_dir=save_dir,
|
|
||||||
save_json=True,
|
|
||||||
plots=False,
|
|
||||||
is_coco=is_coco)
|
|
||||||
|
|
||||||
# Strip optimizers
|
|
||||||
for f in last, best:
|
|
||||||
if f.exists():
|
|
||||||
strip_optimizer(f) # strip optimizers
|
|
||||||
if wandb_logger.wandb: # Log the stripped model
|
|
||||||
wandb_logger.wandb.log_artifact(str(best if best.exists() else last), type='model',
|
|
||||||
name='run_' + wandb_logger.wandb_run.id + '_model',
|
|
||||||
aliases=['latest', 'best', 'stripped'])
|
|
||||||
wandb_logger.finish_run()
|
|
||||||
else:
|
|
||||||
dist.destroy_process_group()
|
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
return results
|
return results
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
def parse_opt(known=False):
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument('--weights', type=str, default='yolov3.pt', help='initial weights path')
|
parser.add_argument('--weights', type=str, default=ROOT / 'yolov3.pt', help='initial weights path')
|
||||||
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
|
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
|
||||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
||||||
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
|
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch.yaml', help='hyperparameters path')
|
||||||
parser.add_argument('--epochs', type=int, default=300)
|
parser.add_argument('--epochs', type=int, default=300)
|
||||||
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
|
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
|
||||||
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
|
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
|
||||||
parser.add_argument('--rect', action='store_true', help='rectangular training')
|
parser.add_argument('--rect', action='store_true', help='rectangular training')
|
||||||
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
|
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
|
||||||
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
||||||
parser.add_argument('--notest', action='store_true', help='only test final epoch')
|
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
|
||||||
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
|
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
|
||||||
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
|
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
|
||||||
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
|
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
|
||||||
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
|
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
|
||||||
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
|
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
|
||||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||||
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
|
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
|
||||||
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
|
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
|
||||||
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
|
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
|
||||||
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
|
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
|
||||||
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
|
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
|
||||||
parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
|
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
|
||||||
parser.add_argument('--project', default='runs/train', help='save to project/name')
|
|
||||||
parser.add_argument('--entity', default=None, help='W&B entity')
|
|
||||||
parser.add_argument('--name', default='exp', help='save to project/name')
|
parser.add_argument('--name', default='exp', help='save to project/name')
|
||||||
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
||||||
parser.add_argument('--quad', action='store_true', help='quad dataloader')
|
parser.add_argument('--quad', action='store_true', help='quad dataloader')
|
||||||
parser.add_argument('--linear-lr', action='store_true', help='linear LR')
|
parser.add_argument('--linear-lr', action='store_true', help='linear LR')
|
||||||
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
|
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
|
||||||
parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
|
parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
|
||||||
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
|
parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')
|
||||||
parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
|
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
|
||||||
parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
|
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
|
||||||
opt = parser.parse_args()
|
|
||||||
|
|
||||||
# Set DDP variables
|
# Weights & Biases arguments
|
||||||
opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
|
parser.add_argument('--entity', default=None, help='W&B: Entity')
|
||||||
opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
|
parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as artifact table')
|
||||||
set_logging(opt.global_rank)
|
parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
|
||||||
if opt.global_rank in [-1, 0]:
|
parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')
|
||||||
|
|
||||||
|
opt = parser.parse_known_args()[0] if known else parser.parse_args()
|
||||||
|
return opt
|
||||||
|
|
||||||
|
|
||||||
|
def main(opt, callbacks=Callbacks()):
|
||||||
|
# Checks
|
||||||
|
if RANK in [-1, 0]:
|
||||||
|
print_args(FILE.stem, opt)
|
||||||
check_git_status()
|
check_git_status()
|
||||||
check_requirements(exclude=('pycocotools', 'thop'))
|
check_requirements(exclude=['thop'])
|
||||||
|
|
||||||
# Resume
|
# Resume
|
||||||
wandb_run = check_wandb_resume(opt)
|
if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run
|
||||||
if opt.resume and not wandb_run: # resume an interrupted run
|
|
||||||
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
|
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
|
||||||
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
|
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
|
||||||
apriori = opt.global_rank, opt.local_rank
|
with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f:
|
||||||
with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
|
|
||||||
opt = argparse.Namespace(**yaml.safe_load(f)) # replace
|
opt = argparse.Namespace(**yaml.safe_load(f)) # replace
|
||||||
opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = \
|
opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate
|
||||||
'', ckpt, True, opt.total_batch_size, *apriori # reinstate
|
LOGGER.info(f'Resuming training from {ckpt}')
|
||||||
logger.info('Resuming training from %s' % ckpt)
|
|
||||||
else:
|
else:
|
||||||
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
|
opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
|
||||||
opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp) # check files
|
check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
|
||||||
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
|
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
|
||||||
opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test)
|
if opt.evolve:
|
||||||
opt.name = 'evolve' if opt.evolve else opt.name
|
opt.project = str(ROOT / 'runs/evolve')
|
||||||
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve))
|
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
|
||||||
|
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
|
||||||
|
|
||||||
# DDP mode
|
# DDP mode
|
||||||
opt.total_batch_size = opt.batch_size
|
|
||||||
device = select_device(opt.device, batch_size=opt.batch_size)
|
device = select_device(opt.device, batch_size=opt.batch_size)
|
||||||
if opt.local_rank != -1:
|
if LOCAL_RANK != -1:
|
||||||
assert torch.cuda.device_count() > opt.local_rank
|
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
|
||||||
torch.cuda.set_device(opt.local_rank)
|
assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count'
|
||||||
device = torch.device('cuda', opt.local_rank)
|
|
||||||
dist.init_process_group(backend='nccl', init_method='env://') # distributed backend
|
|
||||||
assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
|
|
||||||
assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'
|
assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'
|
||||||
opt.batch_size = opt.total_batch_size // opt.world_size
|
assert not opt.evolve, '--evolve argument is not compatible with DDP training'
|
||||||
|
torch.cuda.set_device(LOCAL_RANK)
|
||||||
# Hyperparameters
|
device = torch.device('cuda', LOCAL_RANK)
|
||||||
with open(opt.hyp) as f:
|
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
|
||||||
hyp = yaml.safe_load(f) # load hyps
|
|
||||||
|
|
||||||
# Train
|
# Train
|
||||||
logger.info(opt)
|
|
||||||
if not opt.evolve:
|
if not opt.evolve:
|
||||||
tb_writer = None # init loggers
|
train(opt.hyp, opt, device, callbacks)
|
||||||
if opt.global_rank in [-1, 0]:
|
if WORLD_SIZE > 1 and RANK == 0:
|
||||||
prefix = colorstr('tensorboard: ')
|
LOGGER.info('Destroying process group... ')
|
||||||
logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
|
dist.destroy_process_group()
|
||||||
tb_writer = SummaryWriter(opt.save_dir) # Tensorboard
|
|
||||||
train(hyp, opt, device, tb_writer)
|
|
||||||
|
|
||||||
# Evolve hyperparameters (optional)
|
# Evolve hyperparameters (optional)
|
||||||
else:
|
else:
|
||||||
@ -570,23 +554,27 @@ if __name__ == '__main__':
|
|||||||
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
|
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
|
||||||
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
|
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
|
||||||
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
|
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
|
||||||
'mixup': (1, 0.0, 1.0)} # image mixup (probability)
|
'mixup': (1, 0.0, 1.0), # image mixup (probability)
|
||||||
|
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
|
||||||
|
|
||||||
assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
|
with open(opt.hyp, errors='ignore') as f:
|
||||||
opt.notest, opt.nosave = True, True # only test/save final epoch
|
hyp = yaml.safe_load(f) # load hyps dict
|
||||||
|
if 'anchors' not in hyp: # anchors commented in hyp.yaml
|
||||||
|
hyp['anchors'] = 3
|
||||||
|
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch
|
||||||
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
|
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
|
||||||
yaml_file = Path(opt.save_dir) / 'hyp_evolved.yaml' # save best result here
|
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
|
||||||
if opt.bucket:
|
if opt.bucket:
|
||||||
os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket) # download evolve.txt if exists
|
os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {save_dir}') # download evolve.csv if exists
|
||||||
|
|
||||||
for _ in range(300): # generations to evolve
|
for _ in range(opt.evolve): # generations to evolve
|
||||||
if Path('evolve.txt').exists(): # if evolve.txt exists: select best hyps and mutate
|
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
|
||||||
# Select parent(s)
|
# Select parent(s)
|
||||||
parent = 'single' # parent selection method: 'single' or 'weighted'
|
parent = 'single' # parent selection method: 'single' or 'weighted'
|
||||||
x = np.loadtxt('evolve.txt', ndmin=2)
|
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
|
||||||
n = min(5, len(x)) # number of previous results to consider
|
n = min(5, len(x)) # number of previous results to consider
|
||||||
x = x[np.argsort(-fitness(x))][:n] # top n mutations
|
x = x[np.argsort(-fitness(x))][:n] # top n mutations
|
||||||
w = fitness(x) - fitness(x).min() # weights
|
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
|
||||||
if parent == 'single' or len(x) == 1:
|
if parent == 'single' or len(x) == 1:
|
||||||
# x = x[random.randint(0, n - 1)] # random selection
|
# x = x[random.randint(0, n - 1)] # random selection
|
||||||
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
|
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
|
||||||
@ -597,7 +585,7 @@ if __name__ == '__main__':
|
|||||||
mp, s = 0.8, 0.2 # mutation probability, sigma
|
mp, s = 0.8, 0.2 # mutation probability, sigma
|
||||||
npr = np.random
|
npr = np.random
|
||||||
npr.seed(int(time.time()))
|
npr.seed(int(time.time()))
|
||||||
g = np.array([x[0] for x in meta.values()]) # gains 0-1
|
g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
|
||||||
ng = len(meta)
|
ng = len(meta)
|
||||||
v = np.ones(ng)
|
v = np.ones(ng)
|
||||||
while all(v == 1): # mutate until a change occurs (prevent duplicates)
|
while all(v == 1): # mutate until a change occurs (prevent duplicates)
|
||||||
@ -612,12 +600,26 @@ if __name__ == '__main__':
|
|||||||
hyp[k] = round(hyp[k], 5) # significant digits
|
hyp[k] = round(hyp[k], 5) # significant digits
|
||||||
|
|
||||||
# Train mutation
|
# Train mutation
|
||||||
results = train(hyp.copy(), opt, device)
|
results = train(hyp.copy(), opt, device, callbacks)
|
||||||
|
|
||||||
# Write mutation results
|
# Write mutation results
|
||||||
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
|
print_mutation(results, hyp.copy(), save_dir, opt.bucket)
|
||||||
|
|
||||||
# Plot results
|
# Plot results
|
||||||
plot_evolution(yaml_file)
|
plot_evolve(evolve_csv)
|
||||||
print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n'
|
LOGGER.info(f'Hyperparameter evolution finished\n'
|
||||||
f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}')
|
f"Results saved to {colorstr('bold', save_dir)}\n"
|
||||||
|
f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}')
|
||||||
|
|
||||||
|
|
||||||
|
def run(**kwargs):
|
||||||
|
# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov3.pt')
|
||||||
|
opt = parse_opt(True)
|
||||||
|
for k, v in kwargs.items():
|
||||||
|
setattr(opt, k, v)
|
||||||
|
main(opt)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
opt = parse_opt()
|
||||||
|
main(opt)
|
||||||
|
|||||||
861
tutorial.ipynb
vendored
861
tutorial.ipynb
vendored
File diff suppressed because one or more lines are too long
@ -0,0 +1,18 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
utils/initialization
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def notebook_init():
|
||||||
|
# For notebooks
|
||||||
|
print('Checking setup...')
|
||||||
|
from IPython import display # to display images and clear console output
|
||||||
|
|
||||||
|
from utils.general import emojis
|
||||||
|
from utils.torch_utils import select_device # imports
|
||||||
|
|
||||||
|
display.clear_output()
|
||||||
|
select_device(newline=False)
|
||||||
|
print(emojis('Setup complete ✅'))
|
||||||
|
return display
|
||||||
@ -1,4 +1,7 @@
|
|||||||
# Activation functions
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Activation functions
|
||||||
|
"""
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
@ -16,7 +19,7 @@ class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
|
|||||||
@staticmethod
|
@staticmethod
|
||||||
def forward(x):
|
def forward(x):
|
||||||
# return x * F.hardsigmoid(x) # for torchscript and CoreML
|
# return x * F.hardsigmoid(x) # for torchscript and CoreML
|
||||||
return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX
|
return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for torchscript, CoreML and ONNX
|
||||||
|
|
||||||
|
|
||||||
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
||||||
|
|||||||
277
utils/augmentations.py
Normal file
277
utils/augmentations.py
Normal file
@ -0,0 +1,277 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Image augmentation functions
|
||||||
|
"""
|
||||||
|
|
||||||
|
import math
|
||||||
|
import random
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from utils.general import LOGGER, check_version, colorstr, resample_segments, segment2box
|
||||||
|
from utils.metrics import bbox_ioa
|
||||||
|
|
||||||
|
|
||||||
|
class Albumentations:
|
||||||
|
# Albumentations class (optional, only used if package is installed)
|
||||||
|
def __init__(self):
|
||||||
|
self.transform = None
|
||||||
|
try:
|
||||||
|
import albumentations as A
|
||||||
|
check_version(A.__version__, '1.0.3', hard=True) # version requirement
|
||||||
|
|
||||||
|
self.transform = A.Compose([
|
||||||
|
A.Blur(p=0.01),
|
||||||
|
A.MedianBlur(p=0.01),
|
||||||
|
A.ToGray(p=0.01),
|
||||||
|
A.CLAHE(p=0.01),
|
||||||
|
A.RandomBrightnessContrast(p=0.0),
|
||||||
|
A.RandomGamma(p=0.0),
|
||||||
|
A.ImageCompression(quality_lower=75, p=0.0)],
|
||||||
|
bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
|
||||||
|
|
||||||
|
LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p))
|
||||||
|
except ImportError: # package not installed, skip
|
||||||
|
pass
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(colorstr('albumentations: ') + f'{e}')
|
||||||
|
|
||||||
|
def __call__(self, im, labels, p=1.0):
|
||||||
|
if self.transform and random.random() < p:
|
||||||
|
new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0]) # transformed
|
||||||
|
im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])])
|
||||||
|
return im, labels
|
||||||
|
|
||||||
|
|
||||||
|
def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
|
||||||
|
# HSV color-space augmentation
|
||||||
|
if hgain or sgain or vgain:
|
||||||
|
r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains
|
||||||
|
hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))
|
||||||
|
dtype = im.dtype # uint8
|
||||||
|
|
||||||
|
x = np.arange(0, 256, dtype=r.dtype)
|
||||||
|
lut_hue = ((x * r[0]) % 180).astype(dtype)
|
||||||
|
lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
|
||||||
|
lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
|
||||||
|
|
||||||
|
im_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
|
||||||
|
cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR, dst=im) # no return needed
|
||||||
|
|
||||||
|
|
||||||
|
def hist_equalize(im, clahe=True, bgr=False):
|
||||||
|
# Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255
|
||||||
|
yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
|
||||||
|
if clahe:
|
||||||
|
c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||||
|
yuv[:, :, 0] = c.apply(yuv[:, :, 0])
|
||||||
|
else:
|
||||||
|
yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0]) # equalize Y channel histogram
|
||||||
|
return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB) # convert YUV image to RGB
|
||||||
|
|
||||||
|
|
||||||
|
def replicate(im, labels):
|
||||||
|
# Replicate labels
|
||||||
|
h, w = im.shape[:2]
|
||||||
|
boxes = labels[:, 1:].astype(int)
|
||||||
|
x1, y1, x2, y2 = boxes.T
|
||||||
|
s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels)
|
||||||
|
for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices
|
||||||
|
x1b, y1b, x2b, y2b = boxes[i]
|
||||||
|
bh, bw = y2b - y1b, x2b - x1b
|
||||||
|
yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y
|
||||||
|
x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]
|
||||||
|
im[y1a:y2a, x1a:x2a] = im[y1b:y2b, x1b:x2b] # im4[ymin:ymax, xmin:xmax]
|
||||||
|
labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)
|
||||||
|
|
||||||
|
return im, labels
|
||||||
|
|
||||||
|
|
||||||
|
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
||||||
|
# Resize and pad image while meeting stride-multiple constraints
|
||||||
|
shape = im.shape[:2] # current shape [height, width]
|
||||||
|
if isinstance(new_shape, int):
|
||||||
|
new_shape = (new_shape, new_shape)
|
||||||
|
|
||||||
|
# Scale ratio (new / old)
|
||||||
|
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||||
|
if not scaleup: # only scale down, do not scale up (for better val mAP)
|
||||||
|
r = min(r, 1.0)
|
||||||
|
|
||||||
|
# Compute padding
|
||||||
|
ratio = r, r # width, height ratios
|
||||||
|
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
||||||
|
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
||||||
|
if auto: # minimum rectangle
|
||||||
|
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
||||||
|
elif scaleFill: # stretch
|
||||||
|
dw, dh = 0.0, 0.0
|
||||||
|
new_unpad = (new_shape[1], new_shape[0])
|
||||||
|
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
|
||||||
|
|
||||||
|
dw /= 2 # divide padding into 2 sides
|
||||||
|
dh /= 2
|
||||||
|
|
||||||
|
if shape[::-1] != new_unpad: # resize
|
||||||
|
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||||
|
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
||||||
|
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
||||||
|
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
||||||
|
return im, ratio, (dw, dh)
|
||||||
|
|
||||||
|
|
||||||
|
def random_perspective(im, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
|
||||||
|
border=(0, 0)):
|
||||||
|
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
|
||||||
|
# targets = [cls, xyxy]
|
||||||
|
|
||||||
|
height = im.shape[0] + border[0] * 2 # shape(h,w,c)
|
||||||
|
width = im.shape[1] + border[1] * 2
|
||||||
|
|
||||||
|
# Center
|
||||||
|
C = np.eye(3)
|
||||||
|
C[0, 2] = -im.shape[1] / 2 # x translation (pixels)
|
||||||
|
C[1, 2] = -im.shape[0] / 2 # y translation (pixels)
|
||||||
|
|
||||||
|
# Perspective
|
||||||
|
P = np.eye(3)
|
||||||
|
P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)
|
||||||
|
P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)
|
||||||
|
|
||||||
|
# Rotation and Scale
|
||||||
|
R = np.eye(3)
|
||||||
|
a = random.uniform(-degrees, degrees)
|
||||||
|
# a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations
|
||||||
|
s = random.uniform(1 - scale, 1 + scale)
|
||||||
|
# s = 2 ** random.uniform(-scale, scale)
|
||||||
|
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
|
||||||
|
|
||||||
|
# Shear
|
||||||
|
S = np.eye(3)
|
||||||
|
S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)
|
||||||
|
S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)
|
||||||
|
|
||||||
|
# Translation
|
||||||
|
T = np.eye(3)
|
||||||
|
T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width # x translation (pixels)
|
||||||
|
T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height # y translation (pixels)
|
||||||
|
|
||||||
|
# Combined rotation matrix
|
||||||
|
M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
|
||||||
|
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
|
||||||
|
if perspective:
|
||||||
|
im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
|
||||||
|
else: # affine
|
||||||
|
im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
|
||||||
|
|
||||||
|
# Visualize
|
||||||
|
# import matplotlib.pyplot as plt
|
||||||
|
# ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
|
||||||
|
# ax[0].imshow(im[:, :, ::-1]) # base
|
||||||
|
# ax[1].imshow(im2[:, :, ::-1]) # warped
|
||||||
|
|
||||||
|
# Transform label coordinates
|
||||||
|
n = len(targets)
|
||||||
|
if n:
|
||||||
|
use_segments = any(x.any() for x in segments)
|
||||||
|
new = np.zeros((n, 4))
|
||||||
|
if use_segments: # warp segments
|
||||||
|
segments = resample_segments(segments) # upsample
|
||||||
|
for i, segment in enumerate(segments):
|
||||||
|
xy = np.ones((len(segment), 3))
|
||||||
|
xy[:, :2] = segment
|
||||||
|
xy = xy @ M.T # transform
|
||||||
|
xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2] # perspective rescale or affine
|
||||||
|
|
||||||
|
# clip
|
||||||
|
new[i] = segment2box(xy, width, height)
|
||||||
|
|
||||||
|
else: # warp boxes
|
||||||
|
xy = np.ones((n * 4, 3))
|
||||||
|
xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
|
||||||
|
xy = xy @ M.T # transform
|
||||||
|
xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
|
||||||
|
|
||||||
|
# create new boxes
|
||||||
|
x = xy[:, [0, 2, 4, 6]]
|
||||||
|
y = xy[:, [1, 3, 5, 7]]
|
||||||
|
new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||||
|
|
||||||
|
# clip
|
||||||
|
new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
|
||||||
|
new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
|
||||||
|
|
||||||
|
# filter candidates
|
||||||
|
i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)
|
||||||
|
targets = targets[i]
|
||||||
|
targets[:, 1:5] = new[i]
|
||||||
|
|
||||||
|
return im, targets
|
||||||
|
|
||||||
|
|
||||||
|
def copy_paste(im, labels, segments, p=0.5):
|
||||||
|
# Implement Copy-Paste augmentation https://arxiv.org/abs/2012.07177, labels as nx5 np.array(cls, xyxy)
|
||||||
|
n = len(segments)
|
||||||
|
if p and n:
|
||||||
|
h, w, c = im.shape # height, width, channels
|
||||||
|
im_new = np.zeros(im.shape, np.uint8)
|
||||||
|
for j in random.sample(range(n), k=round(p * n)):
|
||||||
|
l, s = labels[j], segments[j]
|
||||||
|
box = w - l[3], l[2], w - l[1], l[4]
|
||||||
|
ioa = bbox_ioa(box, labels[:, 1:5]) # intersection over area
|
||||||
|
if (ioa < 0.30).all(): # allow 30% obscuration of existing labels
|
||||||
|
labels = np.concatenate((labels, [[l[0], *box]]), 0)
|
||||||
|
segments.append(np.concatenate((w - s[:, 0:1], s[:, 1:2]), 1))
|
||||||
|
cv2.drawContours(im_new, [segments[j].astype(np.int32)], -1, (255, 255, 255), cv2.FILLED)
|
||||||
|
|
||||||
|
result = cv2.bitwise_and(src1=im, src2=im_new)
|
||||||
|
result = cv2.flip(result, 1) # augment segments (flip left-right)
|
||||||
|
i = result > 0 # pixels to replace
|
||||||
|
# i[:, :] = result.max(2).reshape(h, w, 1) # act over ch
|
||||||
|
im[i] = result[i] # cv2.imwrite('debug.jpg', im) # debug
|
||||||
|
|
||||||
|
return im, labels, segments
|
||||||
|
|
||||||
|
|
||||||
|
def cutout(im, labels, p=0.5):
|
||||||
|
# Applies image cutout augmentation https://arxiv.org/abs/1708.04552
|
||||||
|
if random.random() < p:
|
||||||
|
h, w = im.shape[:2]
|
||||||
|
scales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16 # image size fraction
|
||||||
|
for s in scales:
|
||||||
|
mask_h = random.randint(1, int(h * s)) # create random masks
|
||||||
|
mask_w = random.randint(1, int(w * s))
|
||||||
|
|
||||||
|
# box
|
||||||
|
xmin = max(0, random.randint(0, w) - mask_w // 2)
|
||||||
|
ymin = max(0, random.randint(0, h) - mask_h // 2)
|
||||||
|
xmax = min(w, xmin + mask_w)
|
||||||
|
ymax = min(h, ymin + mask_h)
|
||||||
|
|
||||||
|
# apply random color mask
|
||||||
|
im[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]
|
||||||
|
|
||||||
|
# return unobscured labels
|
||||||
|
if len(labels) and s > 0.03:
|
||||||
|
box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
|
||||||
|
ioa = bbox_ioa(box, labels[:, 1:5]) # intersection over area
|
||||||
|
labels = labels[ioa < 0.60] # remove >60% obscured labels
|
||||||
|
|
||||||
|
return labels
|
||||||
|
|
||||||
|
|
||||||
|
def mixup(im, labels, im2, labels2):
|
||||||
|
# Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
|
||||||
|
r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0
|
||||||
|
im = (im * r + im2 * (1 - r)).astype(np.uint8)
|
||||||
|
labels = np.concatenate((labels, labels2), 0)
|
||||||
|
return im, labels
|
||||||
|
|
||||||
|
|
||||||
|
def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16): # box1(4,n), box2(4,n)
|
||||||
|
# Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio
|
||||||
|
w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
|
||||||
|
w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
|
||||||
|
ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps)) # aspect ratio
|
||||||
|
return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr) # candidates
|
||||||
@ -1,28 +1,32 @@
|
|||||||
# Auto-anchor utils
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Auto-anchor utils
|
||||||
|
"""
|
||||||
|
|
||||||
|
import random
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
import yaml
|
import yaml
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from utils.general import colorstr
|
from utils.general import LOGGER, colorstr, emojis
|
||||||
|
|
||||||
|
PREFIX = colorstr('AutoAnchor: ')
|
||||||
|
|
||||||
|
|
||||||
def check_anchor_order(m):
|
def check_anchor_order(m):
|
||||||
# Check anchor order against stride order for YOLOv3 Detect() module m, and correct if necessary
|
# Check anchor order against stride order for Detect() module m, and correct if necessary
|
||||||
a = m.anchor_grid.prod(-1).view(-1) # anchor area
|
a = m.anchors.prod(-1).view(-1) # anchor area
|
||||||
da = a[-1] - a[0] # delta a
|
da = a[-1] - a[0] # delta a
|
||||||
ds = m.stride[-1] - m.stride[0] # delta s
|
ds = m.stride[-1] - m.stride[0] # delta s
|
||||||
if da.sign() != ds.sign(): # same order
|
if da.sign() != ds.sign(): # same order
|
||||||
print('Reversing anchor order')
|
LOGGER.info(f'{PREFIX}Reversing anchor order')
|
||||||
m.anchors[:] = m.anchors.flip(0)
|
m.anchors[:] = m.anchors.flip(0)
|
||||||
m.anchor_grid[:] = m.anchor_grid.flip(0)
|
|
||||||
|
|
||||||
|
|
||||||
def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
||||||
# Check anchor fit to data, recompute if necessary
|
# Check anchor fit to data, recompute if necessary
|
||||||
prefix = colorstr('autoanchor: ')
|
|
||||||
print(f'\n{prefix}Analyzing anchors... ', end='')
|
|
||||||
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
|
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
|
||||||
shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
||||||
scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
|
scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
|
||||||
@ -30,39 +34,39 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
|||||||
|
|
||||||
def metric(k): # compute metric
|
def metric(k): # compute metric
|
||||||
r = wh[:, None] / k[None]
|
r = wh[:, None] / k[None]
|
||||||
x = torch.min(r, 1. / r).min(2)[0] # ratio metric
|
x = torch.min(r, 1 / r).min(2)[0] # ratio metric
|
||||||
best = x.max(1)[0] # best_x
|
best = x.max(1)[0] # best_x
|
||||||
aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold
|
aat = (x > 1 / thr).float().sum(1).mean() # anchors above threshold
|
||||||
bpr = (best > 1. / thr).float().mean() # best possible recall
|
bpr = (best > 1 / thr).float().mean() # best possible recall
|
||||||
return bpr, aat
|
return bpr, aat
|
||||||
|
|
||||||
anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
|
anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors
|
||||||
bpr, aat = metric(anchors)
|
bpr, aat = metric(anchors.cpu().view(-1, 2))
|
||||||
print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
|
s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
|
||||||
if bpr < 0.98: # threshold to recompute
|
if bpr > 0.98: # threshold to recompute
|
||||||
print('. Attempting to improve anchors, please wait...')
|
LOGGER.info(emojis(f'{s}Current anchors are a good fit to dataset ✅'))
|
||||||
na = m.anchor_grid.numel() // 2 # number of anchors
|
else:
|
||||||
|
LOGGER.info(emojis(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...'))
|
||||||
|
na = m.anchors.numel() // 2 # number of anchors
|
||||||
try:
|
try:
|
||||||
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f'{prefix}ERROR: {e}')
|
LOGGER.info(f'{PREFIX}ERROR: {e}')
|
||||||
new_bpr = metric(anchors)[0]
|
new_bpr = metric(anchors)[0]
|
||||||
if new_bpr > bpr: # replace anchors
|
if new_bpr > bpr: # replace anchors
|
||||||
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
||||||
m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
|
|
||||||
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
||||||
check_anchor_order(m)
|
check_anchor_order(m)
|
||||||
print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
LOGGER.info(f'{PREFIX}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
||||||
else:
|
else:
|
||||||
print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.')
|
LOGGER.info(f'{PREFIX}Original anchors better than new anchors. Proceeding with original anchors.')
|
||||||
print('') # newline
|
|
||||||
|
|
||||||
|
|
||||||
def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
|
def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
|
||||||
""" Creates kmeans-evolved anchors from training dataset
|
""" Creates kmeans-evolved anchors from training dataset
|
||||||
|
|
||||||
Arguments:
|
Arguments:
|
||||||
path: path to dataset *.yaml, or a loaded dataset
|
dataset: path to data.yaml, or a loaded dataset
|
||||||
n: number of anchors
|
n: number of anchors
|
||||||
img_size: image size used for training
|
img_size: image size used for training
|
||||||
thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
|
thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
|
||||||
@ -77,12 +81,11 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
|||||||
"""
|
"""
|
||||||
from scipy.cluster.vq import kmeans
|
from scipy.cluster.vq import kmeans
|
||||||
|
|
||||||
thr = 1. / thr
|
thr = 1 / thr
|
||||||
prefix = colorstr('autoanchor: ')
|
|
||||||
|
|
||||||
def metric(k, wh): # compute metrics
|
def metric(k, wh): # compute metrics
|
||||||
r = wh[:, None] / k[None]
|
r = wh[:, None] / k[None]
|
||||||
x = torch.min(r, 1. / r).min(2)[0] # ratio metric
|
x = torch.min(r, 1 / r).min(2)[0] # ratio metric
|
||||||
# x = wh_iou(wh, torch.tensor(k)) # iou metric
|
# x = wh_iou(wh, torch.tensor(k)) # iou metric
|
||||||
return x, x.max(1)[0] # x, best_x
|
return x, x.max(1)[0] # x, best_x
|
||||||
|
|
||||||
@ -90,24 +93,24 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
|||||||
_, best = metric(torch.tensor(k, dtype=torch.float32), wh)
|
_, best = metric(torch.tensor(k, dtype=torch.float32), wh)
|
||||||
return (best * (best > thr).float()).mean() # fitness
|
return (best * (best > thr).float()).mean() # fitness
|
||||||
|
|
||||||
def print_results(k):
|
def print_results(k, verbose=True):
|
||||||
k = k[np.argsort(k.prod(1))] # sort small to large
|
k = k[np.argsort(k.prod(1))] # sort small to large
|
||||||
x, best = metric(k, wh0)
|
x, best = metric(k, wh0)
|
||||||
bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
|
bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
|
||||||
print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr')
|
s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \
|
||||||
print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, '
|
f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \
|
||||||
f'past_thr={x[x > thr].mean():.3f}-mean: ', end='')
|
f'past_thr={x[x > thr].mean():.3f}-mean: '
|
||||||
for i, x in enumerate(k):
|
for i, x in enumerate(k):
|
||||||
print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg
|
s += '%i,%i, ' % (round(x[0]), round(x[1]))
|
||||||
|
if verbose:
|
||||||
|
LOGGER.info(s[:-2])
|
||||||
return k
|
return k
|
||||||
|
|
||||||
if isinstance(path, str): # *.yaml file
|
if isinstance(dataset, str): # *.yaml file
|
||||||
with open(path) as f:
|
with open(dataset, errors='ignore') as f:
|
||||||
data_dict = yaml.safe_load(f) # model dict
|
data_dict = yaml.safe_load(f) # model dict
|
||||||
from utils.datasets import LoadImagesAndLabels
|
from utils.datasets import LoadImagesAndLabels
|
||||||
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
||||||
else:
|
|
||||||
dataset = path # dataset
|
|
||||||
|
|
||||||
# Get label wh
|
# Get label wh
|
||||||
shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
|
||||||
@ -116,19 +119,19 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
|||||||
# Filter
|
# Filter
|
||||||
i = (wh0 < 3.0).any(1).sum()
|
i = (wh0 < 3.0).any(1).sum()
|
||||||
if i:
|
if i:
|
||||||
print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
|
LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
|
||||||
wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
|
wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
|
||||||
# wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
|
# wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
|
||||||
|
|
||||||
# Kmeans calculation
|
# Kmeans calculation
|
||||||
print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
|
LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
|
||||||
s = wh.std(0) # sigmas for whitening
|
s = wh.std(0) # sigmas for whitening
|
||||||
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
|
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
|
||||||
assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}')
|
assert len(k) == n, f'{PREFIX}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}'
|
||||||
k *= s
|
k *= s
|
||||||
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
||||||
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
||||||
k = print_results(k)
|
k = print_results(k, verbose=False)
|
||||||
|
|
||||||
# Plot
|
# Plot
|
||||||
# k, d = [None] * 20, [None] * 20
|
# k, d = [None] * 20, [None] * 20
|
||||||
@ -145,17 +148,17 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
|||||||
# Evolve
|
# Evolve
|
||||||
npr = np.random
|
npr = np.random
|
||||||
f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
|
f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
|
||||||
pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar
|
pbar = tqdm(range(gen), desc=f'{PREFIX}Evolving anchors with Genetic Algorithm:') # progress bar
|
||||||
for _ in pbar:
|
for _ in pbar:
|
||||||
v = np.ones(sh)
|
v = np.ones(sh)
|
||||||
while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
|
while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
|
||||||
v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
|
v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
|
||||||
kg = (k.copy() * v).clip(min=2.0)
|
kg = (k.copy() * v).clip(min=2.0)
|
||||||
fg = anchor_fitness(kg)
|
fg = anchor_fitness(kg)
|
||||||
if fg > f:
|
if fg > f:
|
||||||
f, k = fg, kg.copy()
|
f, k = fg, kg.copy()
|
||||||
pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
|
pbar.desc = f'{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
|
||||||
if verbose:
|
if verbose:
|
||||||
print_results(k)
|
print_results(k, verbose)
|
||||||
|
|
||||||
return print_results(k)
|
return print_results(k)
|
||||||
|
|||||||
57
utils/autobatch.py
Normal file
57
utils/autobatch.py
Normal file
@ -0,0 +1,57 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Auto-batch utils
|
||||||
|
"""
|
||||||
|
|
||||||
|
from copy import deepcopy
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from torch.cuda import amp
|
||||||
|
|
||||||
|
from utils.general import LOGGER, colorstr
|
||||||
|
from utils.torch_utils import profile
|
||||||
|
|
||||||
|
|
||||||
|
def check_train_batch_size(model, imgsz=640):
|
||||||
|
# Check training batch size
|
||||||
|
with amp.autocast():
|
||||||
|
return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size
|
||||||
|
|
||||||
|
|
||||||
|
def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
|
||||||
|
# Automatically estimate best batch size to use `fraction` of available CUDA memory
|
||||||
|
# Usage:
|
||||||
|
# import torch
|
||||||
|
# from utils.autobatch import autobatch
|
||||||
|
# model = torch.hub.load('ultralytics/yolov3', 'yolov3', autoshape=False)
|
||||||
|
# print(autobatch(model))
|
||||||
|
|
||||||
|
prefix = colorstr('AutoBatch: ')
|
||||||
|
LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}')
|
||||||
|
device = next(model.parameters()).device # get model device
|
||||||
|
if device.type == 'cpu':
|
||||||
|
LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}')
|
||||||
|
return batch_size
|
||||||
|
|
||||||
|
d = str(device).upper() # 'CUDA:0'
|
||||||
|
properties = torch.cuda.get_device_properties(device) # device properties
|
||||||
|
t = properties.total_memory / 1024 ** 3 # (GiB)
|
||||||
|
r = torch.cuda.memory_reserved(device) / 1024 ** 3 # (GiB)
|
||||||
|
a = torch.cuda.memory_allocated(device) / 1024 ** 3 # (GiB)
|
||||||
|
f = t - (r + a) # free inside reserved
|
||||||
|
LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free')
|
||||||
|
|
||||||
|
batch_sizes = [1, 2, 4, 8, 16]
|
||||||
|
try:
|
||||||
|
img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes]
|
||||||
|
y = profile(img, model, n=3, device=device)
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.warning(f'{prefix}{e}')
|
||||||
|
|
||||||
|
y = [x[2] for x in y if x] # memory [2]
|
||||||
|
batch_sizes = batch_sizes[:len(y)]
|
||||||
|
p = np.polyfit(batch_sizes, y, deg=1) # first degree polynomial fit
|
||||||
|
b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size)
|
||||||
|
LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)')
|
||||||
|
return b
|
||||||
@ -1,26 +0,0 @@
|
|||||||
# AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/
|
|
||||||
# This script will run on every instance restart, not only on first start
|
|
||||||
# --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA ---
|
|
||||||
|
|
||||||
Content-Type: multipart/mixed; boundary="//"
|
|
||||||
MIME-Version: 1.0
|
|
||||||
|
|
||||||
--//
|
|
||||||
Content-Type: text/cloud-config; charset="us-ascii"
|
|
||||||
MIME-Version: 1.0
|
|
||||||
Content-Transfer-Encoding: 7bit
|
|
||||||
Content-Disposition: attachment; filename="cloud-config.txt"
|
|
||||||
|
|
||||||
#cloud-config
|
|
||||||
cloud_final_modules:
|
|
||||||
- [scripts-user, always]
|
|
||||||
|
|
||||||
--//
|
|
||||||
Content-Type: text/x-shellscript; charset="us-ascii"
|
|
||||||
MIME-Version: 1.0
|
|
||||||
Content-Transfer-Encoding: 7bit
|
|
||||||
Content-Disposition: attachment; filename="userdata.txt"
|
|
||||||
|
|
||||||
#!/bin/bash
|
|
||||||
# --- paste contents of userdata.sh here ---
|
|
||||||
--//
|
|
||||||
@ -1,37 +0,0 @@
|
|||||||
# Resume all interrupted trainings in yolov5/ dir including DDP trainings
|
|
||||||
# Usage: $ python utils/aws/resume.py
|
|
||||||
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import yaml
|
|
||||||
|
|
||||||
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
|
||||||
|
|
||||||
port = 0 # --master_port
|
|
||||||
path = Path('').resolve()
|
|
||||||
for last in path.rglob('*/**/last.pt'):
|
|
||||||
ckpt = torch.load(last)
|
|
||||||
if ckpt['optimizer'] is None:
|
|
||||||
continue
|
|
||||||
|
|
||||||
# Load opt.yaml
|
|
||||||
with open(last.parent.parent / 'opt.yaml') as f:
|
|
||||||
opt = yaml.safe_load(f)
|
|
||||||
|
|
||||||
# Get device count
|
|
||||||
d = opt['device'].split(',') # devices
|
|
||||||
nd = len(d) # number of devices
|
|
||||||
ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel
|
|
||||||
|
|
||||||
if ddp: # multi-GPU
|
|
||||||
port += 1
|
|
||||||
cmd = f'python -m torch.distributed.launch --nproc_per_node {nd} --master_port {port} train.py --resume {last}'
|
|
||||||
else: # single-GPU
|
|
||||||
cmd = f'python train.py --resume {last}'
|
|
||||||
|
|
||||||
cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread
|
|
||||||
print(cmd)
|
|
||||||
os.system(cmd)
|
|
||||||
@ -1,27 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html
|
|
||||||
# This script will run only once on first instance start (for a re-start script see mime.sh)
|
|
||||||
# /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir
|
|
||||||
# Use >300 GB SSD
|
|
||||||
|
|
||||||
cd home/ubuntu
|
|
||||||
if [ ! -d yolov5 ]; then
|
|
||||||
echo "Running first-time script." # install dependencies, download COCO, pull Docker
|
|
||||||
git clone https://github.com/ultralytics/yolov5 -b master && sudo chmod -R 777 yolov5
|
|
||||||
cd yolov5
|
|
||||||
bash data/scripts/get_coco.sh && echo "Data done." &
|
|
||||||
sudo docker pull ultralytics/yolov5:latest && echo "Docker done." &
|
|
||||||
python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." &
|
|
||||||
wait && echo "All tasks done." # finish background tasks
|
|
||||||
else
|
|
||||||
echo "Running re-start script." # resume interrupted runs
|
|
||||||
i=0
|
|
||||||
list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour'
|
|
||||||
while IFS= read -r id; do
|
|
||||||
((i++))
|
|
||||||
echo "restarting container $i: $id"
|
|
||||||
sudo docker start $id
|
|
||||||
# sudo docker exec -it $id python train.py --resume # single-GPU
|
|
||||||
sudo docker exec -d $id python utils/aws/resume.py # multi-scenario
|
|
||||||
done <<<"$list"
|
|
||||||
fi
|
|
||||||
76
utils/callbacks.py
Normal file
76
utils/callbacks.py
Normal file
@ -0,0 +1,76 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Callback utils
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class Callbacks:
|
||||||
|
""""
|
||||||
|
Handles all registered callbacks for Hooks
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Define the available callbacks
|
||||||
|
_callbacks = {
|
||||||
|
'on_pretrain_routine_start': [],
|
||||||
|
'on_pretrain_routine_end': [],
|
||||||
|
|
||||||
|
'on_train_start': [],
|
||||||
|
'on_train_epoch_start': [],
|
||||||
|
'on_train_batch_start': [],
|
||||||
|
'optimizer_step': [],
|
||||||
|
'on_before_zero_grad': [],
|
||||||
|
'on_train_batch_end': [],
|
||||||
|
'on_train_epoch_end': [],
|
||||||
|
|
||||||
|
'on_val_start': [],
|
||||||
|
'on_val_batch_start': [],
|
||||||
|
'on_val_image_end': [],
|
||||||
|
'on_val_batch_end': [],
|
||||||
|
'on_val_end': [],
|
||||||
|
|
||||||
|
'on_fit_epoch_end': [], # fit = train + val
|
||||||
|
'on_model_save': [],
|
||||||
|
'on_train_end': [],
|
||||||
|
|
||||||
|
'teardown': [],
|
||||||
|
}
|
||||||
|
|
||||||
|
def register_action(self, hook, name='', callback=None):
|
||||||
|
"""
|
||||||
|
Register a new action to a callback hook
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hook The callback hook name to register the action to
|
||||||
|
name The name of the action for later reference
|
||||||
|
callback The callback to fire
|
||||||
|
"""
|
||||||
|
assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
|
||||||
|
assert callable(callback), f"callback '{callback}' is not callable"
|
||||||
|
self._callbacks[hook].append({'name': name, 'callback': callback})
|
||||||
|
|
||||||
|
def get_registered_actions(self, hook=None):
|
||||||
|
""""
|
||||||
|
Returns all the registered actions by callback hook
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hook The name of the hook to check, defaults to all
|
||||||
|
"""
|
||||||
|
if hook:
|
||||||
|
return self._callbacks[hook]
|
||||||
|
else:
|
||||||
|
return self._callbacks
|
||||||
|
|
||||||
|
def run(self, hook, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Loop through the registered actions and fire all callbacks
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hook The name of the hook to check, defaults to all
|
||||||
|
args Arguments to receive from
|
||||||
|
kwargs Keyword Arguments to receive from
|
||||||
|
"""
|
||||||
|
|
||||||
|
assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
|
||||||
|
|
||||||
|
for logger in self._callbacks[hook]:
|
||||||
|
logger['callback'](*args, **kwargs)
|
||||||
File diff suppressed because it is too large
Load Diff
@ -1,10 +1,15 @@
|
|||||||
# Google utils: https://cloud.google.com/storage/docs/reference/libraries
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Download utils
|
||||||
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import platform
|
import platform
|
||||||
import subprocess
|
import subprocess
|
||||||
import time
|
import time
|
||||||
|
import urllib
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from zipfile import ZipFile
|
||||||
|
|
||||||
import requests
|
import requests
|
||||||
import torch
|
import torch
|
||||||
@ -19,30 +24,32 @@ def gsutil_getsize(url=''):
|
|||||||
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
|
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
|
||||||
# Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
|
# Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
|
||||||
file = Path(file)
|
file = Path(file)
|
||||||
try: # GitHub
|
assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
|
||||||
|
try: # url1
|
||||||
print(f'Downloading {url} to {file}...')
|
print(f'Downloading {url} to {file}...')
|
||||||
torch.hub.download_url_to_file(url, str(file))
|
torch.hub.download_url_to_file(url, str(file))
|
||||||
assert file.exists() and file.stat().st_size > min_bytes # check
|
assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check
|
||||||
except Exception as e: # GCP
|
except Exception as e: # url2
|
||||||
file.unlink(missing_ok=True) # remove partial downloads
|
file.unlink(missing_ok=True) # remove partial downloads
|
||||||
print(f'Download error: {e}\nRe-attempting {url2 or url} to {file}...')
|
print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
|
||||||
os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
|
os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
|
||||||
finally:
|
finally:
|
||||||
if not file.exists() or file.stat().st_size < min_bytes: # check
|
if not file.exists() or file.stat().st_size < min_bytes: # check
|
||||||
file.unlink(missing_ok=True) # remove partial downloads
|
file.unlink(missing_ok=True) # remove partial downloads
|
||||||
print(f'ERROR: Download failure: {error_msg or url}')
|
print(f"ERROR: {assert_msg}\n{error_msg}")
|
||||||
print('')
|
print('')
|
||||||
|
|
||||||
|
|
||||||
def attempt_download(file, repo='ultralytics/yolov3'):
|
def attempt_download(file, repo='ultralytics/yolov3'): # from utils.downloads import *; attempt_download()
|
||||||
# Attempt file download if does not exist
|
# Attempt file download if does not exist
|
||||||
file = Path(str(file).strip().replace("'", ''))
|
file = Path(str(file).strip().replace("'", ''))
|
||||||
|
|
||||||
if not file.exists():
|
if not file.exists():
|
||||||
# URL specified
|
# URL specified
|
||||||
name = file.name
|
name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc.
|
||||||
if str(file).startswith(('http:/', 'https:/')): # download
|
if str(file).startswith(('http:/', 'https:/')): # download
|
||||||
url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
|
url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
|
||||||
|
name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth...
|
||||||
safe_download(file=name, url=url, min_bytes=1E5)
|
safe_download(file=name, url=url, min_bytes=1E5)
|
||||||
return name
|
return name
|
||||||
|
|
||||||
@ -50,7 +57,7 @@ def attempt_download(file, repo='ultralytics/yolov3'):
|
|||||||
file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required)
|
file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required)
|
||||||
try:
|
try:
|
||||||
response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
|
response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
|
||||||
assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
|
assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov3.pt'...]
|
||||||
tag = response['tag_name'] # i.e. 'v1.0'
|
tag = response['tag_name'] # i.e. 'v1.0'
|
||||||
except: # fallback plan
|
except: # fallback plan
|
||||||
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
||||||
@ -70,7 +77,7 @@ def attempt_download(file, repo='ultralytics/yolov3'):
|
|||||||
|
|
||||||
|
|
||||||
def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
|
def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
|
||||||
# Downloads a file from Google Drive. from yolov3.utils.google_utils import *; gdrive_download()
|
# Downloads a file from Google Drive. from yolov3.utils.downloads import *; gdrive_download()
|
||||||
t = time.time()
|
t = time.time()
|
||||||
file = Path(file)
|
file = Path(file)
|
||||||
cookie = Path('cookie') # gdrive cookie
|
cookie = Path('cookie') # gdrive cookie
|
||||||
@ -97,8 +104,8 @@ def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
|
|||||||
# Unzip if archive
|
# Unzip if archive
|
||||||
if file.suffix == '.zip':
|
if file.suffix == '.zip':
|
||||||
print('unzipping... ', end='')
|
print('unzipping... ', end='')
|
||||||
os.system(f'unzip -q {file}') # unzip
|
ZipFile(file).extractall(path=file.parent) # unzip
|
||||||
file.unlink() # remove zip to free space
|
file.unlink() # remove zip
|
||||||
|
|
||||||
print(f'Done ({time.time() - t:.1f}s)')
|
print(f'Done ({time.time() - t:.1f}s)')
|
||||||
return r
|
return r
|
||||||
@ -111,6 +118,9 @@ def get_token(cookie="./cookie"):
|
|||||||
return line.split()[-1]
|
return line.split()[-1]
|
||||||
return ""
|
return ""
|
||||||
|
|
||||||
|
# Google utils: https://cloud.google.com/storage/docs/reference/libraries ----------------------------------------------
|
||||||
|
#
|
||||||
|
#
|
||||||
# def upload_blob(bucket_name, source_file_name, destination_blob_name):
|
# def upload_blob(bucket_name, source_file_name, destination_blob_name):
|
||||||
# # Uploads a file to a bucket
|
# # Uploads a file to a bucket
|
||||||
# # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
|
# # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
|
||||||
@ -1,68 +0,0 @@
|
|||||||
# Flask REST API
|
|
||||||
[REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/).
|
|
||||||
|
|
||||||
## Requirements
|
|
||||||
|
|
||||||
[Flask](https://palletsprojects.com/p/flask/) is required. Install with:
|
|
||||||
```shell
|
|
||||||
$ pip install Flask
|
|
||||||
```
|
|
||||||
|
|
||||||
## Run
|
|
||||||
|
|
||||||
After Flask installation run:
|
|
||||||
|
|
||||||
```shell
|
|
||||||
$ python3 restapi.py --port 5000
|
|
||||||
```
|
|
||||||
|
|
||||||
Then use [curl](https://curl.se/) to perform a request:
|
|
||||||
|
|
||||||
```shell
|
|
||||||
$ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'`
|
|
||||||
```
|
|
||||||
|
|
||||||
The model inference results are returned as a JSON response:
|
|
||||||
|
|
||||||
```json
|
|
||||||
[
|
|
||||||
{
|
|
||||||
"class": 0,
|
|
||||||
"confidence": 0.8900438547,
|
|
||||||
"height": 0.9318675399,
|
|
||||||
"name": "person",
|
|
||||||
"width": 0.3264600933,
|
|
||||||
"xcenter": 0.7438579798,
|
|
||||||
"ycenter": 0.5207948685
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"class": 0,
|
|
||||||
"confidence": 0.8440024257,
|
|
||||||
"height": 0.7155083418,
|
|
||||||
"name": "person",
|
|
||||||
"width": 0.6546785235,
|
|
||||||
"xcenter": 0.427829951,
|
|
||||||
"ycenter": 0.6334488392
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"class": 27,
|
|
||||||
"confidence": 0.3771208823,
|
|
||||||
"height": 0.3902671337,
|
|
||||||
"name": "tie",
|
|
||||||
"width": 0.0696444362,
|
|
||||||
"xcenter": 0.3675483763,
|
|
||||||
"ycenter": 0.7991207838
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"class": 27,
|
|
||||||
"confidence": 0.3527112305,
|
|
||||||
"height": 0.1540903747,
|
|
||||||
"name": "tie",
|
|
||||||
"width": 0.0336618312,
|
|
||||||
"xcenter": 0.7814827561,
|
|
||||||
"ycenter": 0.5065554976
|
|
||||||
}
|
|
||||||
]
|
|
||||||
```
|
|
||||||
|
|
||||||
An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py`
|
|
||||||
@ -1,13 +0,0 @@
|
|||||||
"""Perform test request"""
|
|
||||||
import pprint
|
|
||||||
|
|
||||||
import requests
|
|
||||||
|
|
||||||
DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
|
|
||||||
TEST_IMAGE = "zidane.jpg"
|
|
||||||
|
|
||||||
image_data = open(TEST_IMAGE, "rb").read()
|
|
||||||
|
|
||||||
response = requests.post(DETECTION_URL, files={"image": image_data}).json()
|
|
||||||
|
|
||||||
pprint.pprint(response)
|
|
||||||
@ -1,37 +0,0 @@
|
|||||||
"""
|
|
||||||
Run a rest API exposing the yolov5s object detection model
|
|
||||||
"""
|
|
||||||
import argparse
|
|
||||||
import io
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from PIL import Image
|
|
||||||
from flask import Flask, request
|
|
||||||
|
|
||||||
app = Flask(__name__)
|
|
||||||
|
|
||||||
DETECTION_URL = "/v1/object-detection/yolov5s"
|
|
||||||
|
|
||||||
|
|
||||||
@app.route(DETECTION_URL, methods=["POST"])
|
|
||||||
def predict():
|
|
||||||
if not request.method == "POST":
|
|
||||||
return
|
|
||||||
|
|
||||||
if request.files.get("image"):
|
|
||||||
image_file = request.files["image"]
|
|
||||||
image_bytes = image_file.read()
|
|
||||||
|
|
||||||
img = Image.open(io.BytesIO(image_bytes))
|
|
||||||
|
|
||||||
results = model(img, size=640) # reduce size=320 for faster inference
|
|
||||||
return results.pandas().xyxy[0].to_json(orient="records")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
parser = argparse.ArgumentParser(description="Flask API exposing YOLOv3 model")
|
|
||||||
parser.add_argument("--port", default=5000, type=int, help="port number")
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache
|
|
||||||
app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat
|
|
||||||
562
utils/general.py
562
utils/general.py
@ -1,5 +1,9 @@
|
|||||||
# YOLOv3 general utils
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
General utils
|
||||||
|
"""
|
||||||
|
|
||||||
|
import contextlib
|
||||||
import glob
|
import glob
|
||||||
import logging
|
import logging
|
||||||
import math
|
import math
|
||||||
@ -7,11 +11,15 @@ import os
|
|||||||
import platform
|
import platform
|
||||||
import random
|
import random
|
||||||
import re
|
import re
|
||||||
import subprocess
|
import shutil
|
||||||
|
import signal
|
||||||
import time
|
import time
|
||||||
|
import urllib
|
||||||
from itertools import repeat
|
from itertools import repeat
|
||||||
from multiprocessing.pool import ThreadPool
|
from multiprocessing.pool import ThreadPool
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from subprocess import check_output
|
||||||
|
from zipfile import ZipFile
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@ -21,9 +29,8 @@ import torch
|
|||||||
import torchvision
|
import torchvision
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
from utils.google_utils import gsutil_getsize
|
from utils.downloads import gsutil_getsize
|
||||||
from utils.metrics import fitness
|
from utils.metrics import box_iou, fitness
|
||||||
from utils.torch_utils import init_torch_seeds
|
|
||||||
|
|
||||||
# Settings
|
# Settings
|
||||||
torch.set_printoptions(linewidth=320, precision=5, profile='long')
|
torch.set_printoptions(linewidth=320, precision=5, profile='long')
|
||||||
@ -32,18 +39,96 @@ pd.options.display.max_columns = 10
|
|||||||
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(min(os.cpu_count(), 8)) # NumExpr max threads
|
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[1] # root directory
|
||||||
|
|
||||||
def set_logging(rank=-1, verbose=True):
|
|
||||||
logging.basicConfig(
|
def set_logging(name=None, verbose=True):
|
||||||
format="%(message)s",
|
# Sets level and returns logger
|
||||||
level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
|
rank = int(os.getenv('RANK', -1)) # rank in world for Multi-GPU trainings
|
||||||
|
logging.basicConfig(format="%(message)s", level=logging.INFO if (verbose and rank in (-1, 0)) else logging.WARNING)
|
||||||
|
return logging.getLogger(name)
|
||||||
|
|
||||||
|
|
||||||
|
LOGGER = set_logging(__name__) # define globally (used in train.py, val.py, detect.py, etc.)
|
||||||
|
|
||||||
|
|
||||||
|
class Profile(contextlib.ContextDecorator):
|
||||||
|
# Usage: @Profile() decorator or 'with Profile():' context manager
|
||||||
|
def __enter__(self):
|
||||||
|
self.start = time.time()
|
||||||
|
|
||||||
|
def __exit__(self, type, value, traceback):
|
||||||
|
print(f'Profile results: {time.time() - self.start:.5f}s')
|
||||||
|
|
||||||
|
|
||||||
|
class Timeout(contextlib.ContextDecorator):
|
||||||
|
# Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager
|
||||||
|
def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True):
|
||||||
|
self.seconds = int(seconds)
|
||||||
|
self.timeout_message = timeout_msg
|
||||||
|
self.suppress = bool(suppress_timeout_errors)
|
||||||
|
|
||||||
|
def _timeout_handler(self, signum, frame):
|
||||||
|
raise TimeoutError(self.timeout_message)
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
signal.signal(signal.SIGALRM, self._timeout_handler) # Set handler for SIGALRM
|
||||||
|
signal.alarm(self.seconds) # start countdown for SIGALRM to be raised
|
||||||
|
|
||||||
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||||
|
signal.alarm(0) # Cancel SIGALRM if it's scheduled
|
||||||
|
if self.suppress and exc_type is TimeoutError: # Suppress TimeoutError
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
class WorkingDirectory(contextlib.ContextDecorator):
|
||||||
|
# Usage: @WorkingDirectory(dir) decorator or 'with WorkingDirectory(dir):' context manager
|
||||||
|
def __init__(self, new_dir):
|
||||||
|
self.dir = new_dir # new dir
|
||||||
|
self.cwd = Path.cwd().resolve() # current dir
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
os.chdir(self.dir)
|
||||||
|
|
||||||
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||||
|
os.chdir(self.cwd)
|
||||||
|
|
||||||
|
|
||||||
|
def try_except(func):
|
||||||
|
# try-except function. Usage: @try_except decorator
|
||||||
|
def handler(*args, **kwargs):
|
||||||
|
try:
|
||||||
|
func(*args, **kwargs)
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
|
||||||
|
return handler
|
||||||
|
|
||||||
|
|
||||||
|
def methods(instance):
|
||||||
|
# Get class/instance methods
|
||||||
|
return [f for f in dir(instance) if callable(getattr(instance, f)) and not f.startswith("__")]
|
||||||
|
|
||||||
|
|
||||||
|
def print_args(name, opt):
|
||||||
|
# Print argparser arguments
|
||||||
|
LOGGER.info(colorstr(f'{name}: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
|
||||||
|
|
||||||
|
|
||||||
def init_seeds(seed=0):
|
def init_seeds(seed=0):
|
||||||
# Initialize random number generator (RNG) seeds
|
# Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html
|
||||||
|
# cudnn seed 0 settings are slower and more reproducible, else faster and less reproducible
|
||||||
|
import torch.backends.cudnn as cudnn
|
||||||
random.seed(seed)
|
random.seed(seed)
|
||||||
np.random.seed(seed)
|
np.random.seed(seed)
|
||||||
init_torch_seeds(seed)
|
torch.manual_seed(seed)
|
||||||
|
cudnn.benchmark, cudnn.deterministic = (False, True) if seed == 0 else (True, False)
|
||||||
|
|
||||||
|
|
||||||
|
def intersect_dicts(da, db, exclude=()):
|
||||||
|
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
||||||
|
return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
|
||||||
|
|
||||||
|
|
||||||
def get_latest_run(search_dir='.'):
|
def get_latest_run(search_dir='.'):
|
||||||
@ -52,81 +137,136 @@ def get_latest_run(search_dir='.'):
|
|||||||
return max(last_list, key=os.path.getctime) if last_list else ''
|
return max(last_list, key=os.path.getctime) if last_list else ''
|
||||||
|
|
||||||
|
|
||||||
|
def user_config_dir(dir='Ultralytics', env_var='YOLOV3_CONFIG_DIR'):
|
||||||
|
# Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
|
||||||
|
env = os.getenv(env_var)
|
||||||
|
if env:
|
||||||
|
path = Path(env) # use environment variable
|
||||||
|
else:
|
||||||
|
cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs
|
||||||
|
path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir
|
||||||
|
path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable
|
||||||
|
path.mkdir(exist_ok=True) # make if required
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
def is_writeable(dir, test=False):
|
||||||
|
# Return True if directory has write permissions, test opening a file with write permissions if test=True
|
||||||
|
if test: # method 1
|
||||||
|
file = Path(dir) / 'tmp.txt'
|
||||||
|
try:
|
||||||
|
with open(file, 'w'): # open file with write permissions
|
||||||
|
pass
|
||||||
|
file.unlink() # remove file
|
||||||
|
return True
|
||||||
|
except OSError:
|
||||||
|
return False
|
||||||
|
else: # method 2
|
||||||
|
return os.access(dir, os.R_OK) # possible issues on Windows
|
||||||
|
|
||||||
|
|
||||||
def is_docker():
|
def is_docker():
|
||||||
# Is environment a Docker container
|
# Is environment a Docker container?
|
||||||
return Path('/workspace').exists() # or Path('/.dockerenv').exists()
|
return Path('/workspace').exists() # or Path('/.dockerenv').exists()
|
||||||
|
|
||||||
|
|
||||||
def is_colab():
|
def is_colab():
|
||||||
# Is environment a Google Colab instance
|
# Is environment a Google Colab instance?
|
||||||
try:
|
try:
|
||||||
import google.colab
|
import google.colab
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except ImportError:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def is_pip():
|
||||||
|
# Is file in a pip package?
|
||||||
|
return 'site-packages' in Path(__file__).resolve().parts
|
||||||
|
|
||||||
|
|
||||||
|
def is_ascii(s=''):
|
||||||
|
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
|
||||||
|
s = str(s) # convert list, tuple, None, etc. to str
|
||||||
|
return len(s.encode().decode('ascii', 'ignore')) == len(s)
|
||||||
|
|
||||||
|
|
||||||
|
def is_chinese(s='人工智能'):
|
||||||
|
# Is string composed of any Chinese characters?
|
||||||
|
return re.search('[\u4e00-\u9fff]', s)
|
||||||
|
|
||||||
|
|
||||||
def emojis(str=''):
|
def emojis(str=''):
|
||||||
# Return platform-dependent emoji-safe version of string
|
# Return platform-dependent emoji-safe version of string
|
||||||
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
||||||
|
|
||||||
|
|
||||||
def file_size(file):
|
def file_size(path):
|
||||||
# Return file size in MB
|
# Return file/dir size (MB)
|
||||||
return Path(file).stat().st_size / 1e6
|
path = Path(path)
|
||||||
|
if path.is_file():
|
||||||
|
return path.stat().st_size / 1E6
|
||||||
|
elif path.is_dir():
|
||||||
|
return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6
|
||||||
|
else:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
def check_online():
|
def check_online():
|
||||||
# Check internet connectivity
|
# Check internet connectivity
|
||||||
import socket
|
import socket
|
||||||
try:
|
try:
|
||||||
socket.create_connection(("1.1.1.1", 443), 5) # check host accesability
|
socket.create_connection(("1.1.1.1", 443), 5) # check host accessibility
|
||||||
return True
|
return True
|
||||||
except OSError:
|
except OSError:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
@try_except
|
||||||
|
@WorkingDirectory(ROOT)
|
||||||
def check_git_status():
|
def check_git_status():
|
||||||
# Recommend 'git pull' if code is out of date
|
# Recommend 'git pull' if code is out of date
|
||||||
|
msg = ', for updates see https://github.com/ultralytics/yolov3'
|
||||||
print(colorstr('github: '), end='')
|
print(colorstr('github: '), end='')
|
||||||
try:
|
assert Path('.git').exists(), 'skipping check (not a git repository)' + msg
|
||||||
assert Path('.git').exists(), 'skipping check (not a git repository)'
|
assert not is_docker(), 'skipping check (Docker image)' + msg
|
||||||
assert not is_docker(), 'skipping check (Docker image)'
|
assert check_online(), 'skipping check (offline)' + msg
|
||||||
assert check_online(), 'skipping check (offline)'
|
|
||||||
|
|
||||||
cmd = 'git fetch && git config --get remote.origin.url'
|
cmd = 'git fetch && git config --get remote.origin.url'
|
||||||
url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git') # github repo url
|
url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip('.git') # git fetch
|
||||||
branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
|
branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
|
||||||
n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
|
n = int(check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
|
||||||
if n > 0:
|
if n > 0:
|
||||||
s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
|
s = f"⚠️ YOLOv3 is out of date by {n} commit{'s' * (n > 1)}. Use `git pull` or `git clone {url}` to update."
|
||||||
f"Use 'git pull' to update or 'git clone {url}' to download latest."
|
else:
|
||||||
else:
|
s = f'up to date with {url} ✅'
|
||||||
s = f'up to date with {url} ✅'
|
print(emojis(s)) # emoji-safe
|
||||||
print(emojis(s)) # emoji-safe
|
|
||||||
except Exception as e:
|
|
||||||
print(e)
|
|
||||||
|
|
||||||
|
|
||||||
def check_python(minimum='3.7.0', required=True):
|
def check_python(minimum='3.6.2'):
|
||||||
# Check current python version vs. required python version
|
# Check current python version vs. required python version
|
||||||
current = platform.python_version()
|
check_version(platform.python_version(), minimum, name='Python ', hard=True)
|
||||||
result = pkg.parse_version(current) >= pkg.parse_version(minimum)
|
|
||||||
if required:
|
|
||||||
assert result, f'Python {minimum} required by YOLOv3, but Python {current} is currently installed'
|
|
||||||
return result
|
|
||||||
|
|
||||||
|
|
||||||
def check_requirements(requirements='requirements.txt', exclude=()):
|
def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False):
|
||||||
|
# Check version vs. required version
|
||||||
|
current, minimum = (pkg.parse_version(x) for x in (current, minimum))
|
||||||
|
result = (current == minimum) if pinned else (current >= minimum) # bool
|
||||||
|
if hard: # assert min requirements met
|
||||||
|
assert result, f'{name}{minimum} required by YOLOv3, but {name}{current} is currently installed'
|
||||||
|
else:
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
@try_except
|
||||||
|
def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), install=True):
|
||||||
# Check installed dependencies meet requirements (pass *.txt file or list of packages)
|
# Check installed dependencies meet requirements (pass *.txt file or list of packages)
|
||||||
prefix = colorstr('red', 'bold', 'requirements:')
|
prefix = colorstr('red', 'bold', 'requirements:')
|
||||||
check_python() # check python version
|
check_python() # check python version
|
||||||
if isinstance(requirements, (str, Path)): # requirements.txt file
|
if isinstance(requirements, (str, Path)): # requirements.txt file
|
||||||
file = Path(requirements)
|
file = Path(requirements)
|
||||||
if not file.exists():
|
assert file.exists(), f"{prefix} {file.resolve()} not found, check failed."
|
||||||
print(f"{prefix} {file.resolve()} not found, check failed.")
|
with file.open() as f:
|
||||||
return
|
requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(f) if x.name not in exclude]
|
||||||
requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
|
|
||||||
else: # list or tuple of packages
|
else: # list or tuple of packages
|
||||||
requirements = [x for x in requirements if x not in exclude]
|
requirements = [x for x in requirements if x not in exclude]
|
||||||
|
|
||||||
@ -135,25 +275,33 @@ def check_requirements(requirements='requirements.txt', exclude=()):
|
|||||||
try:
|
try:
|
||||||
pkg.require(r)
|
pkg.require(r)
|
||||||
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
|
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
|
||||||
n += 1
|
s = f"{prefix} {r} not found and is required by YOLOv3"
|
||||||
print(f"{prefix} {r} not found and is required by YOLOv3, attempting auto-update...")
|
if install:
|
||||||
try:
|
print(f"{s}, attempting auto-update...")
|
||||||
print(subprocess.check_output(f"pip install '{r}'", shell=True).decode())
|
try:
|
||||||
except Exception as e:
|
assert check_online(), f"'pip install {r}' skipped (offline)"
|
||||||
print(f'{prefix} {e}')
|
print(check_output(f"pip install '{r}'", shell=True).decode())
|
||||||
|
n += 1
|
||||||
|
except Exception as e:
|
||||||
|
print(f'{prefix} {e}')
|
||||||
|
else:
|
||||||
|
print(f'{s}. Please install and rerun your command.')
|
||||||
|
|
||||||
if n: # if packages updated
|
if n: # if packages updated
|
||||||
source = file.resolve() if 'file' in locals() else requirements
|
source = file.resolve() if 'file' in locals() else requirements
|
||||||
s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
|
s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
|
||||||
f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
|
f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
|
||||||
print(emojis(s)) # emoji-safe
|
print(emojis(s))
|
||||||
|
|
||||||
|
|
||||||
def check_img_size(img_size, s=32):
|
def check_img_size(imgsz, s=32, floor=0):
|
||||||
# Verify img_size is a multiple of stride s
|
# Verify image size is a multiple of stride s in each dimension
|
||||||
new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
|
if isinstance(imgsz, int): # integer i.e. img_size=640
|
||||||
if new_size != img_size:
|
new_size = max(make_divisible(imgsz, int(s)), floor)
|
||||||
print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
|
else: # list i.e. img_size=[640, 480]
|
||||||
|
new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
|
||||||
|
if new_size != imgsz:
|
||||||
|
print(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
|
||||||
return new_size
|
return new_size
|
||||||
|
|
||||||
|
|
||||||
@ -172,53 +320,114 @@ def check_imshow():
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
def check_file(file):
|
def check_suffix(file='yolov3.pt', suffix=('.pt',), msg=''):
|
||||||
|
# Check file(s) for acceptable suffix
|
||||||
|
if file and suffix:
|
||||||
|
if isinstance(suffix, str):
|
||||||
|
suffix = [suffix]
|
||||||
|
for f in file if isinstance(file, (list, tuple)) else [file]:
|
||||||
|
s = Path(f).suffix.lower() # file suffix
|
||||||
|
if len(s):
|
||||||
|
assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
|
||||||
|
|
||||||
|
|
||||||
|
def check_yaml(file, suffix=('.yaml', '.yml')):
|
||||||
|
# Search/download YAML file (if necessary) and return path, checking suffix
|
||||||
|
return check_file(file, suffix)
|
||||||
|
|
||||||
|
|
||||||
|
def check_file(file, suffix=''):
|
||||||
# Search/download file (if necessary) and return path
|
# Search/download file (if necessary) and return path
|
||||||
|
check_suffix(file, suffix) # optional
|
||||||
file = str(file) # convert to str()
|
file = str(file) # convert to str()
|
||||||
if Path(file).is_file() or file == '': # exists
|
if Path(file).is_file() or file == '': # exists
|
||||||
return file
|
return file
|
||||||
elif file.startswith(('http://', 'https://')): # download
|
elif file.startswith(('http:/', 'https:/')): # download
|
||||||
url, file = file, Path(file).name
|
url = str(Path(file)).replace(':/', '://') # Pathlib turns :// -> :/
|
||||||
print(f'Downloading {url} to {file}...')
|
file = Path(urllib.parse.unquote(file).split('?')[0]).name # '%2F' to '/', split https://url.com/file.txt?auth
|
||||||
torch.hub.download_url_to_file(url, file)
|
if Path(file).is_file():
|
||||||
assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
|
print(f'Found {url} locally at {file}') # file already exists
|
||||||
|
else:
|
||||||
|
print(f'Downloading {url} to {file}...')
|
||||||
|
torch.hub.download_url_to_file(url, file)
|
||||||
|
assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
|
||||||
return file
|
return file
|
||||||
else: # search
|
else: # search
|
||||||
files = glob.glob('./**/' + file, recursive=True) # find file
|
files = []
|
||||||
|
for d in 'data', 'models', 'utils': # search directories
|
||||||
|
files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file
|
||||||
assert len(files), f'File not found: {file}' # assert file was found
|
assert len(files), f'File not found: {file}' # assert file was found
|
||||||
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
|
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
|
||||||
return files[0] # return file
|
return files[0] # return file
|
||||||
|
|
||||||
|
|
||||||
def check_dataset(dict):
|
def check_dataset(data, autodownload=True):
|
||||||
# Download dataset if not found locally
|
# Download and/or unzip dataset if not found locally
|
||||||
val, s = dict.get('val'), dict.get('download')
|
# Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip
|
||||||
if val and len(val):
|
|
||||||
|
# Download (optional)
|
||||||
|
extract_dir = ''
|
||||||
|
if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip
|
||||||
|
download(data, dir='../datasets', unzip=True, delete=False, curl=False, threads=1)
|
||||||
|
data = next((Path('../datasets') / Path(data).stem).rglob('*.yaml'))
|
||||||
|
extract_dir, autodownload = data.parent, False
|
||||||
|
|
||||||
|
# Read yaml (optional)
|
||||||
|
if isinstance(data, (str, Path)):
|
||||||
|
with open(data, errors='ignore') as f:
|
||||||
|
data = yaml.safe_load(f) # dictionary
|
||||||
|
|
||||||
|
# Parse yaml
|
||||||
|
path = extract_dir or Path(data.get('path') or '') # optional 'path' default to '.'
|
||||||
|
for k in 'train', 'val', 'test':
|
||||||
|
if data.get(k): # prepend path
|
||||||
|
data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
|
||||||
|
|
||||||
|
assert 'nc' in data, "Dataset 'nc' key missing."
|
||||||
|
if 'names' not in data:
|
||||||
|
data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing
|
||||||
|
train, val, test, s = (data.get(x) for x in ('train', 'val', 'test', 'download'))
|
||||||
|
if val:
|
||||||
val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
|
val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
|
||||||
if not all(x.exists() for x in val):
|
if not all(x.exists() for x in val):
|
||||||
print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
|
print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
|
||||||
if s and len(s): # download script
|
if s and autodownload: # download script
|
||||||
|
root = path.parent if 'path' in data else '..' # unzip directory i.e. '../'
|
||||||
if s.startswith('http') and s.endswith('.zip'): # URL
|
if s.startswith('http') and s.endswith('.zip'): # URL
|
||||||
f = Path(s).name # filename
|
f = Path(s).name # filename
|
||||||
print(f'Downloading {s} ...')
|
print(f'Downloading {s} to {f}...')
|
||||||
torch.hub.download_url_to_file(s, f)
|
torch.hub.download_url_to_file(s, f)
|
||||||
r = os.system(f'unzip -q {f} -d ../ && rm {f}') # unzip
|
Path(root).mkdir(parents=True, exist_ok=True) # create root
|
||||||
|
ZipFile(f).extractall(path=root) # unzip
|
||||||
|
Path(f).unlink() # remove zip
|
||||||
|
r = None # success
|
||||||
elif s.startswith('bash '): # bash script
|
elif s.startswith('bash '): # bash script
|
||||||
print(f'Running {s} ...')
|
print(f'Running {s} ...')
|
||||||
r = os.system(s)
|
r = os.system(s)
|
||||||
else: # python script
|
else: # python script
|
||||||
r = exec(s) # return None
|
r = exec(s, {'yaml': data}) # return None
|
||||||
print('Dataset autodownload %s\n' % ('success' if r in (0, None) else 'failure')) # print result
|
print(f"Dataset autodownload {f'success, saved to {root}' if r in (0, None) else 'failure'}\n")
|
||||||
else:
|
else:
|
||||||
raise Exception('Dataset not found.')
|
raise Exception('Dataset not found.')
|
||||||
|
|
||||||
|
return data # dictionary
|
||||||
|
|
||||||
|
|
||||||
|
def url2file(url):
|
||||||
|
# Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt
|
||||||
|
url = str(Path(url)).replace(':/', '://') # Pathlib turns :// -> :/
|
||||||
|
file = Path(urllib.parse.unquote(url)).name.split('?')[0] # '%2F' to '/', split https://url.com/file.txt?auth
|
||||||
|
return file
|
||||||
|
|
||||||
|
|
||||||
def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
|
def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
|
||||||
# Multi-threaded file download and unzip function
|
# Multi-threaded file download and unzip function, used in data.yaml for autodownload
|
||||||
def download_one(url, dir):
|
def download_one(url, dir):
|
||||||
# Download 1 file
|
# Download 1 file
|
||||||
f = dir / Path(url).name # filename
|
f = dir / Path(url).name # filename
|
||||||
if not f.exists():
|
if Path(url).is_file(): # exists in current path
|
||||||
|
Path(url).rename(f) # move to dir
|
||||||
|
elif not f.exists():
|
||||||
print(f'Downloading {url} to {f}...')
|
print(f'Downloading {url} to {f}...')
|
||||||
if curl:
|
if curl:
|
||||||
os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -") # curl download, retry and resume on fail
|
os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -") # curl download, retry and resume on fail
|
||||||
@ -227,12 +436,11 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
|
|||||||
if unzip and f.suffix in ('.zip', '.gz'):
|
if unzip and f.suffix in ('.zip', '.gz'):
|
||||||
print(f'Unzipping {f}...')
|
print(f'Unzipping {f}...')
|
||||||
if f.suffix == '.zip':
|
if f.suffix == '.zip':
|
||||||
s = f'unzip -qo {f} -d {dir} && rm {f}' # unzip -quiet -overwrite
|
ZipFile(f).extractall(path=dir) # unzip
|
||||||
elif f.suffix == '.gz':
|
elif f.suffix == '.gz':
|
||||||
s = f'tar xfz {f} --directory {f.parent}' # unzip
|
os.system(f'tar xfz {f} --directory {f.parent}') # unzip
|
||||||
if delete: # delete zip file after unzip
|
if delete:
|
||||||
s += f' && rm {f}'
|
f.unlink() # remove zip
|
||||||
os.system(s)
|
|
||||||
|
|
||||||
dir = Path(dir)
|
dir = Path(dir)
|
||||||
dir.mkdir(parents=True, exist_ok=True) # make directory
|
dir.mkdir(parents=True, exist_ok=True) # make directory
|
||||||
@ -242,7 +450,7 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
|
|||||||
pool.close()
|
pool.close()
|
||||||
pool.join()
|
pool.join()
|
||||||
else:
|
else:
|
||||||
for u in tuple(url) if isinstance(url, str) else url:
|
for u in [url] if isinstance(url, (str, Path)) else url:
|
||||||
download_one(u, dir)
|
download_one(u, dir)
|
||||||
|
|
||||||
|
|
||||||
@ -257,7 +465,7 @@ def clean_str(s):
|
|||||||
|
|
||||||
|
|
||||||
def one_cycle(y1=0.0, y2=1.0, steps=100):
|
def one_cycle(y1=0.0, y2=1.0, steps=100):
|
||||||
# lambda function for sinusoidal ramp from y1 to y2
|
# lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf
|
||||||
return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
|
return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
|
||||||
|
|
||||||
|
|
||||||
@ -355,6 +563,18 @@ def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
|||||||
return y
|
return y
|
||||||
|
|
||||||
|
|
||||||
|
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
||||||
|
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
|
||||||
|
if clip:
|
||||||
|
clip_coords(x, (h - eps, w - eps)) # warning: inplace clip
|
||||||
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
|
y[:, 0] = ((x[:, 0] + x[:, 2]) / 2) / w # x center
|
||||||
|
y[:, 1] = ((x[:, 1] + x[:, 3]) / 2) / h # y center
|
||||||
|
y[:, 2] = (x[:, 2] - x[:, 0]) / w # width
|
||||||
|
y[:, 3] = (x[:, 3] - x[:, 1]) / h # height
|
||||||
|
return y
|
||||||
|
|
||||||
|
|
||||||
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
||||||
# Convert normalized segments into pixel segments, shape (n,2)
|
# Convert normalized segments into pixel segments, shape (n,2)
|
||||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||||
@ -405,90 +625,16 @@ def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
|||||||
return coords
|
return coords
|
||||||
|
|
||||||
|
|
||||||
def clip_coords(boxes, img_shape):
|
def clip_coords(boxes, shape):
|
||||||
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
||||||
boxes[:, 0].clamp_(0, img_shape[1]) # x1
|
if isinstance(boxes, torch.Tensor): # faster individually
|
||||||
boxes[:, 1].clamp_(0, img_shape[0]) # y1
|
boxes[:, 0].clamp_(0, shape[1]) # x1
|
||||||
boxes[:, 2].clamp_(0, img_shape[1]) # x2
|
boxes[:, 1].clamp_(0, shape[0]) # y1
|
||||||
boxes[:, 3].clamp_(0, img_shape[0]) # y2
|
boxes[:, 2].clamp_(0, shape[1]) # x2
|
||||||
|
boxes[:, 3].clamp_(0, shape[0]) # y2
|
||||||
|
else: # np.array (faster grouped)
|
||||||
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
|
||||||
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
||||||
box2 = box2.T
|
|
||||||
|
|
||||||
# Get the coordinates of bounding boxes
|
|
||||||
if x1y1x2y2: # x1, y1, x2, y2 = box1
|
|
||||||
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
|
|
||||||
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
|
|
||||||
else: # transform from xywh to xyxy
|
|
||||||
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
|
|
||||||
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
|
|
||||||
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
|
|
||||||
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
|
|
||||||
|
|
||||||
# Intersection area
|
|
||||||
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
|
||||||
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
|
||||||
|
|
||||||
# Union Area
|
|
||||||
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
|
||||||
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
|
||||||
union = w1 * h1 + w2 * h2 - inter + eps
|
|
||||||
|
|
||||||
iou = inter / union
|
|
||||||
if GIoU or DIoU or CIoU:
|
|
||||||
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
|
||||||
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
|
||||||
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
|
||||||
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
|
|
||||||
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
|
|
||||||
(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
|
|
||||||
if DIoU:
|
|
||||||
return iou - rho2 / c2 # DIoU
|
|
||||||
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
|
||||||
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
|
||||||
with torch.no_grad():
|
|
||||||
alpha = v / (v - iou + (1 + eps))
|
|
||||||
return iou - (rho2 / c2 + v * alpha) # CIoU
|
|
||||||
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
|
||||||
c_area = cw * ch + eps # convex area
|
|
||||||
return iou - (c_area - union) / c_area # GIoU
|
|
||||||
else:
|
|
||||||
return iou # IoU
|
|
||||||
|
|
||||||
|
|
||||||
def box_iou(box1, box2):
|
|
||||||
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
|
||||||
"""
|
|
||||||
Return intersection-over-union (Jaccard index) of boxes.
|
|
||||||
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
|
||||||
Arguments:
|
|
||||||
box1 (Tensor[N, 4])
|
|
||||||
box2 (Tensor[M, 4])
|
|
||||||
Returns:
|
|
||||||
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
|
||||||
IoU values for every element in boxes1 and boxes2
|
|
||||||
"""
|
|
||||||
|
|
||||||
def box_area(box):
|
|
||||||
# box = 4xn
|
|
||||||
return (box[2] - box[0]) * (box[3] - box[1])
|
|
||||||
|
|
||||||
area1 = box_area(box1.T)
|
|
||||||
area2 = box_area(box2.T)
|
|
||||||
|
|
||||||
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
|
||||||
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
||||||
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
|
||||||
|
|
||||||
|
|
||||||
def wh_iou(wh1, wh2):
|
|
||||||
# Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
|
|
||||||
wh1 = wh1[:, None] # [N,1,2]
|
|
||||||
wh2 = wh2[None] # [1,M,2]
|
|
||||||
inter = torch.min(wh1, wh2).prod(2) # [N,M]
|
|
||||||
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
|
|
||||||
|
|
||||||
|
|
||||||
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
||||||
@ -601,39 +747,48 @@ def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_op
|
|||||||
print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
|
print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
|
||||||
|
|
||||||
|
|
||||||
def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
|
def print_mutation(results, hyp, save_dir, bucket):
|
||||||
# Print mutation results to evolve.txt (for use with train.py --evolve)
|
evolve_csv, results_csv, evolve_yaml = save_dir / 'evolve.csv', save_dir / 'results.csv', save_dir / 'hyp_evolve.yaml'
|
||||||
a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys
|
keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
|
||||||
b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values
|
'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps]
|
||||||
c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
|
keys = tuple(x.strip() for x in keys)
|
||||||
print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
|
vals = results + tuple(hyp.values())
|
||||||
|
n = len(keys)
|
||||||
|
|
||||||
|
# Download (optional)
|
||||||
if bucket:
|
if bucket:
|
||||||
url = 'gs://%s/evolve.txt' % bucket
|
url = f'gs://{bucket}/evolve.csv'
|
||||||
if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0):
|
if gsutil_getsize(url) > (os.path.getsize(evolve_csv) if os.path.exists(evolve_csv) else 0):
|
||||||
os.system('gsutil cp %s .' % url) # download evolve.txt if larger than local
|
os.system(f'gsutil cp {url} {save_dir}') # download evolve.csv if larger than local
|
||||||
|
|
||||||
with open('evolve.txt', 'a') as f: # append result
|
# Log to evolve.csv
|
||||||
f.write(c + b + '\n')
|
s = '' if evolve_csv.exists() else (('%20s,' * n % keys).rstrip(',') + '\n') # add header
|
||||||
x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows
|
with open(evolve_csv, 'a') as f:
|
||||||
x = x[np.argsort(-fitness(x))] # sort
|
f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n')
|
||||||
np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness
|
|
||||||
|
# Print to screen
|
||||||
|
print(colorstr('evolve: ') + ', '.join(f'{x.strip():>20s}' for x in keys))
|
||||||
|
print(colorstr('evolve: ') + ', '.join(f'{x:20.5g}' for x in vals), end='\n\n\n')
|
||||||
|
|
||||||
# Save yaml
|
# Save yaml
|
||||||
for i, k in enumerate(hyp.keys()):
|
with open(evolve_yaml, 'w') as f:
|
||||||
hyp[k] = float(x[0, i + 7])
|
data = pd.read_csv(evolve_csv)
|
||||||
with open(yaml_file, 'w') as f:
|
data = data.rename(columns=lambda x: x.strip()) # strip keys
|
||||||
results = tuple(x[0, :7])
|
i = np.argmax(fitness(data.values[:, :7])) #
|
||||||
c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
|
f.write('# YOLOv3 Hyperparameter Evolution Results\n' +
|
||||||
f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
|
f'# Best generation: {i}\n' +
|
||||||
|
f'# Last generation: {len(data)}\n' +
|
||||||
|
'# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) + '\n' +
|
||||||
|
'# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n')
|
||||||
yaml.safe_dump(hyp, f, sort_keys=False)
|
yaml.safe_dump(hyp, f, sort_keys=False)
|
||||||
|
|
||||||
if bucket:
|
if bucket:
|
||||||
os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload
|
os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}') # upload
|
||||||
|
|
||||||
|
|
||||||
def apply_classifier(x, model, img, im0):
|
def apply_classifier(x, model, img, im0):
|
||||||
# Apply a second stage classifier to yolo outputs
|
# Apply a second stage classifier to YOLO outputs
|
||||||
|
# Example model = torchvision.models.__dict__['efficientnet_b0'](pretrained=True).to(device).eval()
|
||||||
im0 = [im0] if isinstance(im0, np.ndarray) else im0
|
im0 = [im0] if isinstance(im0, np.ndarray) else im0
|
||||||
for i, d in enumerate(x): # per image
|
for i, d in enumerate(x): # per image
|
||||||
if d is not None and len(d):
|
if d is not None and len(d):
|
||||||
@ -654,11 +809,11 @@ def apply_classifier(x, model, img, im0):
|
|||||||
for j, a in enumerate(d): # per item
|
for j, a in enumerate(d): # per item
|
||||||
cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
|
cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
|
||||||
im = cv2.resize(cutout, (224, 224)) # BGR
|
im = cv2.resize(cutout, (224, 224)) # BGR
|
||||||
# cv2.imwrite('test%i.jpg' % j, cutout)
|
# cv2.imwrite('example%i.jpg' % j, cutout)
|
||||||
|
|
||||||
im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||||
im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
|
im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
|
||||||
im /= 255.0 # 0 - 255 to 0.0 - 1.0
|
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||||
ims.append(im)
|
ims.append(im)
|
||||||
|
|
||||||
pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
|
pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
|
||||||
@ -667,33 +822,20 @@ def apply_classifier(x, model, img, im0):
|
|||||||
return x
|
return x
|
||||||
|
|
||||||
|
|
||||||
def save_one_box(xyxy, im, file='image.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
|
|
||||||
xyxy = torch.tensor(xyxy).view(-1, 4)
|
|
||||||
b = xyxy2xywh(xyxy) # boxes
|
|
||||||
if square:
|
|
||||||
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
|
||||||
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
|
||||||
xyxy = xywh2xyxy(b).long()
|
|
||||||
clip_coords(xyxy, im.shape)
|
|
||||||
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
|
||||||
if save:
|
|
||||||
cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
|
|
||||||
return crop
|
|
||||||
|
|
||||||
|
|
||||||
def increment_path(path, exist_ok=False, sep='', mkdir=False):
|
def increment_path(path, exist_ok=False, sep='', mkdir=False):
|
||||||
# Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
# Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
||||||
path = Path(path) # os-agnostic
|
path = Path(path) # os-agnostic
|
||||||
if path.exists() and not exist_ok:
|
if path.exists() and not exist_ok:
|
||||||
suffix = path.suffix
|
path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '')
|
||||||
path = path.with_suffix('')
|
|
||||||
dirs = glob.glob(f"{path}{sep}*") # similar paths
|
dirs = glob.glob(f"{path}{sep}*") # similar paths
|
||||||
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
|
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
|
||||||
i = [int(m.groups()[0]) for m in matches if m] # indices
|
i = [int(m.groups()[0]) for m in matches if m] # indices
|
||||||
n = max(i) + 1 if i else 2 # increment number
|
n = max(i) + 1 if i else 2 # increment number
|
||||||
path = Path(f"{path}{sep}{n}{suffix}") # update path
|
path = Path(f"{path}{sep}{n}{suffix}") # increment path
|
||||||
dir = path if path.suffix == '' else path.parent # directory
|
if mkdir:
|
||||||
if not dir.exists() and mkdir:
|
path.mkdir(parents=True, exist_ok=True) # make directory
|
||||||
dir.mkdir(parents=True, exist_ok=True) # make directory
|
|
||||||
return path
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
# Variables
|
||||||
|
NCOLS = 0 if is_docker() else shutil.get_terminal_size().columns # terminal window size
|
||||||
|
|||||||
@ -1,25 +0,0 @@
|
|||||||
FROM gcr.io/google-appengine/python
|
|
||||||
|
|
||||||
# Create a virtualenv for dependencies. This isolates these packages from
|
|
||||||
# system-level packages.
|
|
||||||
# Use -p python3 or -p python3.7 to select python version. Default is version 2.
|
|
||||||
RUN virtualenv /env -p python3
|
|
||||||
|
|
||||||
# Setting these environment variables are the same as running
|
|
||||||
# source /env/bin/activate.
|
|
||||||
ENV VIRTUAL_ENV /env
|
|
||||||
ENV PATH /env/bin:$PATH
|
|
||||||
|
|
||||||
RUN apt-get update && apt-get install -y python-opencv
|
|
||||||
|
|
||||||
# Copy the application's requirements.txt and run pip to install all
|
|
||||||
# dependencies into the virtualenv.
|
|
||||||
ADD requirements.txt /app/requirements.txt
|
|
||||||
RUN pip install -r /app/requirements.txt
|
|
||||||
|
|
||||||
# Add the application source code.
|
|
||||||
ADD . /app
|
|
||||||
|
|
||||||
# Run a WSGI server to serve the application. gunicorn must be declared as
|
|
||||||
# a dependency in requirements.txt.
|
|
||||||
CMD gunicorn -b :$PORT main:app
|
|
||||||
@ -1,4 +0,0 @@
|
|||||||
# add these requirements in your app on top of the existing ones
|
|
||||||
pip==19.2
|
|
||||||
Flask==1.0.2
|
|
||||||
gunicorn==19.9.0
|
|
||||||
@ -1,14 +0,0 @@
|
|||||||
runtime: custom
|
|
||||||
env: flex
|
|
||||||
|
|
||||||
service: yolov3app
|
|
||||||
|
|
||||||
liveness_check:
|
|
||||||
initial_delay_sec: 600
|
|
||||||
|
|
||||||
manual_scaling:
|
|
||||||
instances: 1
|
|
||||||
resources:
|
|
||||||
cpu: 1
|
|
||||||
memory_gb: 4
|
|
||||||
disk_size_gb: 20
|
|
||||||
156
utils/loggers/__init__.py
Normal file
156
utils/loggers/__init__.py
Normal file
@ -0,0 +1,156 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Logging utils
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import warnings
|
||||||
|
from threading import Thread
|
||||||
|
|
||||||
|
import pkg_resources as pkg
|
||||||
|
import torch
|
||||||
|
from torch.utils.tensorboard import SummaryWriter
|
||||||
|
|
||||||
|
from utils.general import colorstr, emojis
|
||||||
|
from utils.loggers.wandb.wandb_utils import WandbLogger
|
||||||
|
from utils.plots import plot_images, plot_results
|
||||||
|
from utils.torch_utils import de_parallel
|
||||||
|
|
||||||
|
LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases
|
||||||
|
RANK = int(os.getenv('RANK', -1))
|
||||||
|
|
||||||
|
try:
|
||||||
|
import wandb
|
||||||
|
|
||||||
|
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||||
|
if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]:
|
||||||
|
wandb_login_success = wandb.login(timeout=30)
|
||||||
|
if not wandb_login_success:
|
||||||
|
wandb = None
|
||||||
|
except (ImportError, AssertionError):
|
||||||
|
wandb = None
|
||||||
|
|
||||||
|
|
||||||
|
class Loggers():
|
||||||
|
# Loggers class
|
||||||
|
def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
|
||||||
|
self.save_dir = save_dir
|
||||||
|
self.weights = weights
|
||||||
|
self.opt = opt
|
||||||
|
self.hyp = hyp
|
||||||
|
self.logger = logger # for printing results to console
|
||||||
|
self.include = include
|
||||||
|
self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
|
||||||
|
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics
|
||||||
|
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
|
||||||
|
'x/lr0', 'x/lr1', 'x/lr2'] # params
|
||||||
|
for k in LOGGERS:
|
||||||
|
setattr(self, k, None) # init empty logger dictionary
|
||||||
|
self.csv = True # always log to csv
|
||||||
|
|
||||||
|
# Message
|
||||||
|
if not wandb:
|
||||||
|
prefix = colorstr('Weights & Biases: ')
|
||||||
|
s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv3 🚀 runs (RECOMMENDED)"
|
||||||
|
print(emojis(s))
|
||||||
|
|
||||||
|
# TensorBoard
|
||||||
|
s = self.save_dir
|
||||||
|
if 'tb' in self.include and not self.opt.evolve:
|
||||||
|
prefix = colorstr('TensorBoard: ')
|
||||||
|
self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/")
|
||||||
|
self.tb = SummaryWriter(str(s))
|
||||||
|
|
||||||
|
# W&B
|
||||||
|
if wandb and 'wandb' in self.include:
|
||||||
|
wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://')
|
||||||
|
run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None
|
||||||
|
self.opt.hyp = self.hyp # add hyperparameters
|
||||||
|
self.wandb = WandbLogger(self.opt, run_id)
|
||||||
|
else:
|
||||||
|
self.wandb = None
|
||||||
|
|
||||||
|
def on_pretrain_routine_end(self):
|
||||||
|
# Callback runs on pre-train routine end
|
||||||
|
paths = self.save_dir.glob('*labels*.jpg') # training labels
|
||||||
|
if self.wandb:
|
||||||
|
self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
|
||||||
|
|
||||||
|
def on_train_batch_end(self, ni, model, imgs, targets, paths, plots, sync_bn):
|
||||||
|
# Callback runs on train batch end
|
||||||
|
if plots:
|
||||||
|
if ni == 0:
|
||||||
|
if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754
|
||||||
|
with warnings.catch_warnings():
|
||||||
|
warnings.simplefilter('ignore') # suppress jit trace warning
|
||||||
|
self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
|
||||||
|
if ni < 3:
|
||||||
|
f = self.save_dir / f'train_batch{ni}.jpg' # filename
|
||||||
|
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
|
||||||
|
if self.wandb and ni == 10:
|
||||||
|
files = sorted(self.save_dir.glob('train*.jpg'))
|
||||||
|
self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
|
||||||
|
|
||||||
|
def on_train_epoch_end(self, epoch):
|
||||||
|
# Callback runs on train epoch end
|
||||||
|
if self.wandb:
|
||||||
|
self.wandb.current_epoch = epoch + 1
|
||||||
|
|
||||||
|
def on_val_image_end(self, pred, predn, path, names, im):
|
||||||
|
# Callback runs on val image end
|
||||||
|
if self.wandb:
|
||||||
|
self.wandb.val_one_image(pred, predn, path, names, im)
|
||||||
|
|
||||||
|
def on_val_end(self):
|
||||||
|
# Callback runs on val end
|
||||||
|
if self.wandb:
|
||||||
|
files = sorted(self.save_dir.glob('val*.jpg'))
|
||||||
|
self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||||
|
|
||||||
|
def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
|
||||||
|
# Callback runs at the end of each fit (train+val) epoch
|
||||||
|
x = {k: v for k, v in zip(self.keys, vals)} # dict
|
||||||
|
if self.csv:
|
||||||
|
file = self.save_dir / 'results.csv'
|
||||||
|
n = len(x) + 1 # number of cols
|
||||||
|
s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n') # add header
|
||||||
|
with open(file, 'a') as f:
|
||||||
|
f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
|
||||||
|
|
||||||
|
if self.tb:
|
||||||
|
for k, v in x.items():
|
||||||
|
self.tb.add_scalar(k, v, epoch)
|
||||||
|
|
||||||
|
if self.wandb:
|
||||||
|
self.wandb.log(x)
|
||||||
|
self.wandb.end_epoch(best_result=best_fitness == fi)
|
||||||
|
|
||||||
|
def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
|
||||||
|
# Callback runs on model save event
|
||||||
|
if self.wandb:
|
||||||
|
if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
|
||||||
|
self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
|
||||||
|
|
||||||
|
def on_train_end(self, last, best, plots, epoch, results):
|
||||||
|
# Callback runs on training end
|
||||||
|
if plots:
|
||||||
|
plot_results(file=self.save_dir / 'results.csv') # save results.png
|
||||||
|
files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
|
||||||
|
files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
|
||||||
|
|
||||||
|
if self.tb:
|
||||||
|
import cv2
|
||||||
|
for f in files:
|
||||||
|
self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
|
||||||
|
|
||||||
|
if self.wandb:
|
||||||
|
self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
|
||||||
|
# Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
|
||||||
|
if not self.opt.evolve:
|
||||||
|
wandb.log_artifact(str(best if best.exists() else last), type='model',
|
||||||
|
name='run_' + self.wandb.wandb_run.id + '_model',
|
||||||
|
aliases=['latest', 'best', 'stripped'])
|
||||||
|
self.wandb.finish_run()
|
||||||
|
else:
|
||||||
|
self.wandb.finish_run()
|
||||||
|
self.wandb = WandbLogger(self.opt)
|
||||||
147
utils/loggers/wandb/README.md
Normal file
147
utils/loggers/wandb/README.md
Normal file
@ -0,0 +1,147 @@
|
|||||||
|
📚 This guide explains how to use **Weights & Biases** (W&B) with YOLOv3 🚀. UPDATED 29 September 2021.
|
||||||
|
* [About Weights & Biases](#about-weights-&-biases)
|
||||||
|
* [First-Time Setup](#first-time-setup)
|
||||||
|
* [Viewing runs](#viewing-runs)
|
||||||
|
* [Advanced Usage: Dataset Versioning and Evaluation](#advanced-usage)
|
||||||
|
* [Reports: Share your work with the world!](#reports)
|
||||||
|
|
||||||
|
## About Weights & Biases
|
||||||
|
Think of [W&B](https://wandb.ai/site?utm_campaign=repo_yolo_wandbtutorial) like GitHub for machine learning models. With a few lines of code, save everything you need to debug, compare and reproduce your models — architecture, hyperparameters, git commits, model weights, GPU usage, and even datasets and predictions.
|
||||||
|
|
||||||
|
Used by top researchers including teams at OpenAI, Lyft, Github, and MILA, W&B is part of the new standard of best practices for machine learning. How W&B can help you optimize your machine learning workflows:
|
||||||
|
|
||||||
|
* [Debug](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Free-2) model performance in real time
|
||||||
|
* [GPU usage](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#System-4) visualized automatically
|
||||||
|
* [Custom charts](https://wandb.ai/wandb/customizable-charts/reports/Powerful-Custom-Charts-To-Debug-Model-Peformance--VmlldzoyNzY4ODI) for powerful, extensible visualization
|
||||||
|
* [Share insights](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Share-8) interactively with collaborators
|
||||||
|
* [Optimize hyperparameters](https://docs.wandb.com/sweeps) efficiently
|
||||||
|
* [Track](https://docs.wandb.com/artifacts) datasets, pipelines, and production models
|
||||||
|
|
||||||
|
## First-Time Setup
|
||||||
|
<details open>
|
||||||
|
<summary> Toggle Details </summary>
|
||||||
|
When you first train, W&B will prompt you to create a new account and will generate an **API key** for you. If you are an existing user you can retrieve your key from https://wandb.ai/authorize. This key is used to tell W&B where to log your data. You only need to supply your key once, and then it is remembered on the same device.
|
||||||
|
|
||||||
|
W&B will create a cloud **project** (default is 'YOLOv3') for your training runs, and each new training run will be provided a unique run **name** within that project as project/name. You can also manually set your project and run name as:
|
||||||
|
|
||||||
|
```shell
|
||||||
|
$ python train.py --project ... --name ...
|
||||||
|
```
|
||||||
|
|
||||||
|
YOLOv3 notebook example: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||||
|
<img width="960" alt="Screen Shot 2021-09-29 at 10 23 13 PM" src="https://user-images.githubusercontent.com/26833433/135392431-1ab7920a-c49d-450a-b0b0-0c86ec86100e.png">
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
## Viewing Runs
|
||||||
|
<details open>
|
||||||
|
<summary> Toggle Details </summary>
|
||||||
|
Run information streams from your environment to the W&B cloud console as you train. This allows you to monitor and even cancel runs in <b>realtime</b> . All important information is logged:
|
||||||
|
|
||||||
|
* Training & Validation losses
|
||||||
|
* Metrics: Precision, Recall, mAP@0.5, mAP@0.5:0.95
|
||||||
|
* Learning Rate over time
|
||||||
|
* A bounding box debugging panel, showing the training progress over time
|
||||||
|
* GPU: Type, **GPU Utilization**, power, temperature, **CUDA memory usage**
|
||||||
|
* System: Disk I/0, CPU utilization, RAM memory usage
|
||||||
|
* Your trained model as W&B Artifact
|
||||||
|
* Environment: OS and Python types, Git repository and state, **training command**
|
||||||
|
|
||||||
|
<p align="center"><img width="900" alt="Weights & Biases dashboard" src="https://user-images.githubusercontent.com/26833433/135390767-c28b050f-8455-4004-adb0-3b730386e2b2.png"></p>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
## Advanced Usage
|
||||||
|
You can leverage W&B artifacts and Tables integration to easily visualize and manage your datasets, models and training evaluations. Here are some quick examples to get you started.
|
||||||
|
<details open>
|
||||||
|
<h3>1. Visualize and Version Datasets</h3>
|
||||||
|
Log, visualize, dynamically query, and understand your data with <a href='https://docs.wandb.ai/guides/data-vis/tables'>W&B Tables</a>. You can use the following command to log your dataset as a W&B Table. This will generate a <code>{dataset}_wandb.yaml</code> file which can be used to train from dataset artifact.
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python utils/logger/wandb/log_dataset.py --project ... --name ... --data .. </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<h3> 2: Train and Log Evaluation simultaneousy </h3>
|
||||||
|
This is an extension of the previous section, but it'll also training after uploading the dataset. <b> This also evaluation Table</b>
|
||||||
|
Evaluation table compares your predictions and ground truths across the validation set for each epoch. It uses the references to the already uploaded datasets,
|
||||||
|
so no images will be uploaded from your system more than once.
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python utils/logger/wandb/log_dataset.py --data .. --upload_data </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<h3> 3: Train using dataset artifact </h3>
|
||||||
|
When you upload a dataset as described in the first section, you get a new config file with an added `_wandb` to its name. This file contains the information that
|
||||||
|
can be used to train a model directly from the dataset artifact. <b> This also logs evaluation </b>
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python utils/logger/wandb/log_dataset.py --data {data}_wandb.yaml </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<h3> 4: Save model checkpoints as artifacts </h3>
|
||||||
|
To enable saving and versioning checkpoints of your experiment, pass `--save_period n` with the base cammand, where `n` represents checkpoint interval.
|
||||||
|
You can also log both the dataset and model checkpoints simultaneously. If not passed, only the final model will be logged
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python train.py --save_period 1 </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<h3> 5: Resume runs from checkpoint artifacts. </h3>
|
||||||
|
Any run can be resumed using artifacts if the <code>--resume</code> argument starts with <code>wandb-artifact://</code> prefix followed by the run path, i.e, <code>wandb-artifact://username/project/runid </code>. This doesn't require the model checkpoint to be present on the local system.
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python train.py --resume wandb-artifact://{run_path} </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<h3> 6: Resume runs from dataset artifact & checkpoint artifacts. </h3>
|
||||||
|
<b> Local dataset or model checkpoints are not required. This can be used to resume runs directly on a different device </b>
|
||||||
|
The syntax is same as the previous section, but you'll need to lof both the dataset and model checkpoints as artifacts, i.e, set bot <code>--upload_dataset</code> or
|
||||||
|
train from <code>_wandb.yaml</code> file and set <code>--save_period</code>
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary> <b>Usage</b> </summary>
|
||||||
|
<b>Code</b> <code> $ python train.py --resume wandb-artifact://{run_path} </code>
|
||||||
|
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
|
||||||
|
<h3> Reports </h3>
|
||||||
|
W&B Reports can be created from your saved runs for sharing online. Once a report is created you will receive a link you can use to publically share your results. Here is an example report created from the COCO128 tutorial trainings of all YOLOv5 models ([link](https://wandb.ai/glenn-jocher/yolov5_tutorial/reports/YOLOv5-COCO128-Tutorial-Results--VmlldzozMDI5OTY)).
|
||||||
|
|
||||||
|
<img width="900" alt="Weights & Biases Reports" src="https://user-images.githubusercontent.com/26833433/135394029-a17eaf86-c6c1-4b1d-bb80-b90e83aaffa7.png">
|
||||||
|
|
||||||
|
|
||||||
|
## Environments
|
||||||
|
|
||||||
|
YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
|
||||||
|
|
||||||
|
- **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||||
|
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
|
||||||
|
- **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
|
||||||
|
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
|
||||||
|
|
||||||
|
|
||||||
|
## Status
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
If this badge is green, all [YOLOv3 GitHub Actions](https://github.com/ultralytics/yolov3/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv3 training ([train.py](https://github.com/ultralytics/yolov3/blob/master/train.py)), validation ([val.py](https://github.com/ultralytics/yolov3/blob/master/val.py)), inference ([detect.py](https://github.com/ultralytics/yolov3/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov3/blob/master/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.
|
||||||
@ -1,16 +1,16 @@
|
|||||||
import argparse
|
import argparse
|
||||||
|
|
||||||
import yaml
|
|
||||||
|
|
||||||
from wandb_utils import WandbLogger
|
from wandb_utils import WandbLogger
|
||||||
|
|
||||||
|
from utils.general import LOGGER
|
||||||
|
|
||||||
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||||
|
|
||||||
|
|
||||||
def create_dataset_artifact(opt):
|
def create_dataset_artifact(opt):
|
||||||
with open(opt.data) as f:
|
logger = WandbLogger(opt, None, job_type='Dataset Creation') # TODO: return value unused
|
||||||
data = yaml.safe_load(f) # data dict
|
if not logger.wandb:
|
||||||
logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
|
LOGGER.info("install wandb using `pip install wandb` to log the dataset")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
@ -18,6 +18,9 @@ if __name__ == '__main__':
|
|||||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
||||||
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
||||||
parser.add_argument('--project', type=str, default='YOLOv3', help='name of W&B Project')
|
parser.add_argument('--project', type=str, default='YOLOv3', help='name of W&B Project')
|
||||||
|
parser.add_argument('--entity', default=None, help='W&B entity')
|
||||||
|
parser.add_argument('--name', type=str, default='log dataset', help='name of W&B run')
|
||||||
|
|
||||||
opt = parser.parse_args()
|
opt = parser.parse_args()
|
||||||
opt.resume = False # Explicitly disallow resume check for dataset upload job
|
opt.resume = False # Explicitly disallow resume check for dataset upload job
|
||||||
|
|
||||||
41
utils/loggers/wandb/sweep.py
Normal file
41
utils/loggers/wandb/sweep.py
Normal file
@ -0,0 +1,41 @@
|
|||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import wandb
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[3] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
|
||||||
|
from train import parse_opt, train
|
||||||
|
from utils.callbacks import Callbacks
|
||||||
|
from utils.general import increment_path
|
||||||
|
from utils.torch_utils import select_device
|
||||||
|
|
||||||
|
|
||||||
|
def sweep():
|
||||||
|
wandb.init()
|
||||||
|
# Get hyp dict from sweep agent
|
||||||
|
hyp_dict = vars(wandb.config).get("_items")
|
||||||
|
|
||||||
|
# Workaround: get necessary opt args
|
||||||
|
opt = parse_opt(known=True)
|
||||||
|
opt.batch_size = hyp_dict.get("batch_size")
|
||||||
|
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve))
|
||||||
|
opt.epochs = hyp_dict.get("epochs")
|
||||||
|
opt.nosave = True
|
||||||
|
opt.data = hyp_dict.get("data")
|
||||||
|
opt.weights = str(opt.weights)
|
||||||
|
opt.cfg = str(opt.cfg)
|
||||||
|
opt.data = str(opt.data)
|
||||||
|
opt.hyp = str(opt.hyp)
|
||||||
|
opt.project = str(opt.project)
|
||||||
|
device = select_device(opt.device, batch_size=opt.batch_size)
|
||||||
|
|
||||||
|
# train
|
||||||
|
train(hyp_dict, opt, device, callbacks=Callbacks())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
sweep()
|
||||||
143
utils/loggers/wandb/sweep.yaml
Normal file
143
utils/loggers/wandb/sweep.yaml
Normal file
@ -0,0 +1,143 @@
|
|||||||
|
# Hyperparameters for training
|
||||||
|
# To set range-
|
||||||
|
# Provide min and max values as:
|
||||||
|
# parameter:
|
||||||
|
#
|
||||||
|
# min: scalar
|
||||||
|
# max: scalar
|
||||||
|
# OR
|
||||||
|
#
|
||||||
|
# Set a specific list of search space-
|
||||||
|
# parameter:
|
||||||
|
# values: [scalar1, scalar2, scalar3...]
|
||||||
|
#
|
||||||
|
# You can use grid, bayesian and hyperopt search strategy
|
||||||
|
# For more info on configuring sweeps visit - https://docs.wandb.ai/guides/sweeps/configuration
|
||||||
|
|
||||||
|
program: utils/loggers/wandb/sweep.py
|
||||||
|
method: random
|
||||||
|
metric:
|
||||||
|
name: metrics/mAP_0.5
|
||||||
|
goal: maximize
|
||||||
|
|
||||||
|
parameters:
|
||||||
|
# hyperparameters: set either min, max range or values list
|
||||||
|
data:
|
||||||
|
value: "data/coco128.yaml"
|
||||||
|
batch_size:
|
||||||
|
values: [64]
|
||||||
|
epochs:
|
||||||
|
values: [10]
|
||||||
|
|
||||||
|
lr0:
|
||||||
|
distribution: uniform
|
||||||
|
min: 1e-5
|
||||||
|
max: 1e-1
|
||||||
|
lrf:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.01
|
||||||
|
max: 1.0
|
||||||
|
momentum:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.6
|
||||||
|
max: 0.98
|
||||||
|
weight_decay:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.001
|
||||||
|
warmup_epochs:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 5.0
|
||||||
|
warmup_momentum:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.95
|
||||||
|
warmup_bias_lr:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.2
|
||||||
|
box:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.02
|
||||||
|
max: 0.2
|
||||||
|
cls:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.2
|
||||||
|
max: 4.0
|
||||||
|
cls_pw:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.5
|
||||||
|
max: 2.0
|
||||||
|
obj:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.2
|
||||||
|
max: 4.0
|
||||||
|
obj_pw:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.5
|
||||||
|
max: 2.0
|
||||||
|
iou_t:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.1
|
||||||
|
max: 0.7
|
||||||
|
anchor_t:
|
||||||
|
distribution: uniform
|
||||||
|
min: 2.0
|
||||||
|
max: 8.0
|
||||||
|
fl_gamma:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.1
|
||||||
|
hsv_h:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.1
|
||||||
|
hsv_s:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.9
|
||||||
|
hsv_v:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.9
|
||||||
|
degrees:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 45.0
|
||||||
|
translate:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.9
|
||||||
|
scale:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.9
|
||||||
|
shear:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 10.0
|
||||||
|
perspective:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 0.001
|
||||||
|
flipud:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 1.0
|
||||||
|
fliplr:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 1.0
|
||||||
|
mosaic:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 1.0
|
||||||
|
mixup:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 1.0
|
||||||
|
copy_paste:
|
||||||
|
distribution: uniform
|
||||||
|
min: 0.0
|
||||||
|
max: 1.0
|
||||||
532
utils/loggers/wandb/wandb_utils.py
Normal file
532
utils/loggers/wandb/wandb_utils.py
Normal file
@ -0,0 +1,532 @@
|
|||||||
|
"""Utilities and tools for tracking runs with Weights & Biases."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Dict
|
||||||
|
|
||||||
|
import pkg_resources as pkg
|
||||||
|
import yaml
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[3] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
|
||||||
|
from utils.datasets import LoadImagesAndLabels, img2label_paths
|
||||||
|
from utils.general import LOGGER, check_dataset, check_file
|
||||||
|
|
||||||
|
try:
|
||||||
|
import wandb
|
||||||
|
|
||||||
|
assert hasattr(wandb, '__version__') # verify package import not local dir
|
||||||
|
except (ImportError, AssertionError):
|
||||||
|
wandb = None
|
||||||
|
|
||||||
|
RANK = int(os.getenv('RANK', -1))
|
||||||
|
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
||||||
|
|
||||||
|
|
||||||
|
def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
|
||||||
|
return from_string[len(prefix):]
|
||||||
|
|
||||||
|
|
||||||
|
def check_wandb_config_file(data_config_file):
|
||||||
|
wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
|
||||||
|
if Path(wandb_config).is_file():
|
||||||
|
return wandb_config
|
||||||
|
return data_config_file
|
||||||
|
|
||||||
|
|
||||||
|
def check_wandb_dataset(data_file):
|
||||||
|
is_trainset_wandb_artifact = False
|
||||||
|
is_valset_wandb_artifact = False
|
||||||
|
if check_file(data_file) and data_file.endswith('.yaml'):
|
||||||
|
with open(data_file, errors='ignore') as f:
|
||||||
|
data_dict = yaml.safe_load(f)
|
||||||
|
is_trainset_wandb_artifact = (isinstance(data_dict['train'], str) and
|
||||||
|
data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX))
|
||||||
|
is_valset_wandb_artifact = (isinstance(data_dict['val'], str) and
|
||||||
|
data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX))
|
||||||
|
if is_trainset_wandb_artifact or is_valset_wandb_artifact:
|
||||||
|
return data_dict
|
||||||
|
else:
|
||||||
|
return check_dataset(data_file)
|
||||||
|
|
||||||
|
|
||||||
|
def get_run_info(run_path):
|
||||||
|
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
|
||||||
|
run_id = run_path.stem
|
||||||
|
project = run_path.parent.stem
|
||||||
|
entity = run_path.parent.parent.stem
|
||||||
|
model_artifact_name = 'run_' + run_id + '_model'
|
||||||
|
return entity, project, run_id, model_artifact_name
|
||||||
|
|
||||||
|
|
||||||
|
def check_wandb_resume(opt):
|
||||||
|
process_wandb_config_ddp_mode(opt) if RANK not in [-1, 0] else None
|
||||||
|
if isinstance(opt.resume, str):
|
||||||
|
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
if RANK not in [-1, 0]: # For resuming DDP runs
|
||||||
|
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||||
|
api = wandb.Api()
|
||||||
|
artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
|
||||||
|
modeldir = artifact.download()
|
||||||
|
opt.weights = str(Path(modeldir) / "last.pt")
|
||||||
|
return True
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def process_wandb_config_ddp_mode(opt):
|
||||||
|
with open(check_file(opt.data), errors='ignore') as f:
|
||||||
|
data_dict = yaml.safe_load(f) # data dict
|
||||||
|
train_dir, val_dir = None, None
|
||||||
|
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
api = wandb.Api()
|
||||||
|
train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
|
||||||
|
train_dir = train_artifact.download()
|
||||||
|
train_path = Path(train_dir) / 'data/images/'
|
||||||
|
data_dict['train'] = str(train_path)
|
||||||
|
|
||||||
|
if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
api = wandb.Api()
|
||||||
|
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
|
||||||
|
val_dir = val_artifact.download()
|
||||||
|
val_path = Path(val_dir) / 'data/images/'
|
||||||
|
data_dict['val'] = str(val_path)
|
||||||
|
if train_dir or val_dir:
|
||||||
|
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
|
||||||
|
with open(ddp_data_path, 'w') as f:
|
||||||
|
yaml.safe_dump(data_dict, f)
|
||||||
|
opt.data = ddp_data_path
|
||||||
|
|
||||||
|
|
||||||
|
class WandbLogger():
|
||||||
|
"""Log training runs, datasets, models, and predictions to Weights & Biases.
|
||||||
|
|
||||||
|
This logger sends information to W&B at wandb.ai. By default, this information
|
||||||
|
includes hyperparameters, system configuration and metrics, model metrics,
|
||||||
|
and basic data metrics and analyses.
|
||||||
|
|
||||||
|
By providing additional command line arguments to train.py, datasets,
|
||||||
|
models and predictions can also be logged.
|
||||||
|
|
||||||
|
For more on how this logger is used, see the Weights & Biases documentation:
|
||||||
|
https://docs.wandb.com/guides/integrations/yolov5
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, opt, run_id=None, job_type='Training'):
|
||||||
|
"""
|
||||||
|
- Initialize WandbLogger instance
|
||||||
|
- Upload dataset if opt.upload_dataset is True
|
||||||
|
- Setup trainig processes if job_type is 'Training'
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
opt (namespace) -- Commandline arguments for this run
|
||||||
|
run_id (str) -- Run ID of W&B run to be resumed
|
||||||
|
job_type (str) -- To set the job_type for this run
|
||||||
|
|
||||||
|
"""
|
||||||
|
# Pre-training routine --
|
||||||
|
self.job_type = job_type
|
||||||
|
self.wandb, self.wandb_run = wandb, None if not wandb else wandb.run
|
||||||
|
self.val_artifact, self.train_artifact = None, None
|
||||||
|
self.train_artifact_path, self.val_artifact_path = None, None
|
||||||
|
self.result_artifact = None
|
||||||
|
self.val_table, self.result_table = None, None
|
||||||
|
self.bbox_media_panel_images = []
|
||||||
|
self.val_table_path_map = None
|
||||||
|
self.max_imgs_to_log = 16
|
||||||
|
self.wandb_artifact_data_dict = None
|
||||||
|
self.data_dict = None
|
||||||
|
# It's more elegant to stick to 1 wandb.init call,
|
||||||
|
# but useful config data is overwritten in the WandbLogger's wandb.init call
|
||||||
|
if isinstance(opt.resume, str): # checks resume from artifact
|
||||||
|
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
||||||
|
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
|
||||||
|
assert wandb, 'install wandb to resume wandb runs'
|
||||||
|
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
|
||||||
|
self.wandb_run = wandb.init(id=run_id,
|
||||||
|
project=project,
|
||||||
|
entity=entity,
|
||||||
|
resume='allow',
|
||||||
|
allow_val_change=True)
|
||||||
|
opt.resume = model_artifact_name
|
||||||
|
elif self.wandb:
|
||||||
|
self.wandb_run = wandb.init(config=opt,
|
||||||
|
resume="allow",
|
||||||
|
project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem,
|
||||||
|
entity=opt.entity,
|
||||||
|
name=opt.name if opt.name != 'exp' else None,
|
||||||
|
job_type=job_type,
|
||||||
|
id=run_id,
|
||||||
|
allow_val_change=True) if not wandb.run else wandb.run
|
||||||
|
if self.wandb_run:
|
||||||
|
if self.job_type == 'Training':
|
||||||
|
if opt.upload_dataset:
|
||||||
|
if not opt.resume:
|
||||||
|
self.wandb_artifact_data_dict = self.check_and_upload_dataset(opt)
|
||||||
|
|
||||||
|
if opt.resume:
|
||||||
|
# resume from artifact
|
||||||
|
if isinstance(opt.resume, str) and opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
self.data_dict = dict(self.wandb_run.config.data_dict)
|
||||||
|
else: # local resume
|
||||||
|
self.data_dict = check_wandb_dataset(opt.data)
|
||||||
|
else:
|
||||||
|
self.data_dict = check_wandb_dataset(opt.data)
|
||||||
|
self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict
|
||||||
|
|
||||||
|
# write data_dict to config. useful for resuming from artifacts. Do this only when not resuming.
|
||||||
|
self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict},
|
||||||
|
allow_val_change=True)
|
||||||
|
self.setup_training(opt)
|
||||||
|
|
||||||
|
if self.job_type == 'Dataset Creation':
|
||||||
|
self.data_dict = self.check_and_upload_dataset(opt)
|
||||||
|
|
||||||
|
def check_and_upload_dataset(self, opt):
|
||||||
|
"""
|
||||||
|
Check if the dataset format is compatible and upload it as W&B artifact
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
opt (namespace)-- Commandline arguments for current run
|
||||||
|
|
||||||
|
returns:
|
||||||
|
Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links.
|
||||||
|
"""
|
||||||
|
assert wandb, 'Install wandb to upload dataset'
|
||||||
|
config_path = self.log_dataset_artifact(opt.data,
|
||||||
|
opt.single_cls,
|
||||||
|
'YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem)
|
||||||
|
LOGGER.info(f"Created dataset config file {config_path}")
|
||||||
|
with open(config_path, errors='ignore') as f:
|
||||||
|
wandb_data_dict = yaml.safe_load(f)
|
||||||
|
return wandb_data_dict
|
||||||
|
|
||||||
|
def setup_training(self, opt):
|
||||||
|
"""
|
||||||
|
Setup the necessary processes for training YOLO models:
|
||||||
|
- Attempt to download model checkpoint and dataset artifacts if opt.resume stats with WANDB_ARTIFACT_PREFIX
|
||||||
|
- Update data_dict, to contain info of previous run if resumed and the paths of dataset artifact if downloaded
|
||||||
|
- Setup log_dict, initialize bbox_interval
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
opt (namespace) -- commandline arguments for this run
|
||||||
|
|
||||||
|
"""
|
||||||
|
self.log_dict, self.current_epoch = {}, 0
|
||||||
|
self.bbox_interval = opt.bbox_interval
|
||||||
|
if isinstance(opt.resume, str):
|
||||||
|
modeldir, _ = self.download_model_artifact(opt)
|
||||||
|
if modeldir:
|
||||||
|
self.weights = Path(modeldir) / "last.pt"
|
||||||
|
config = self.wandb_run.config
|
||||||
|
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str(
|
||||||
|
self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \
|
||||||
|
config.hyp
|
||||||
|
data_dict = self.data_dict
|
||||||
|
if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download
|
||||||
|
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),
|
||||||
|
opt.artifact_alias)
|
||||||
|
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
|
||||||
|
opt.artifact_alias)
|
||||||
|
|
||||||
|
if self.train_artifact_path is not None:
|
||||||
|
train_path = Path(self.train_artifact_path) / 'data/images/'
|
||||||
|
data_dict['train'] = str(train_path)
|
||||||
|
if self.val_artifact_path is not None:
|
||||||
|
val_path = Path(self.val_artifact_path) / 'data/images/'
|
||||||
|
data_dict['val'] = str(val_path)
|
||||||
|
|
||||||
|
if self.val_artifact is not None:
|
||||||
|
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||||
|
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||||
|
self.val_table = self.val_artifact.get("val")
|
||||||
|
if self.val_table_path_map is None:
|
||||||
|
self.map_val_table_path()
|
||||||
|
if opt.bbox_interval == -1:
|
||||||
|
self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
|
||||||
|
train_from_artifact = self.train_artifact_path is not None and self.val_artifact_path is not None
|
||||||
|
# Update the the data_dict to point to local artifacts dir
|
||||||
|
if train_from_artifact:
|
||||||
|
self.data_dict = data_dict
|
||||||
|
|
||||||
|
def download_dataset_artifact(self, path, alias):
|
||||||
|
"""
|
||||||
|
download the model checkpoint artifact if the path starts with WANDB_ARTIFACT_PREFIX
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
path -- path of the dataset to be used for training
|
||||||
|
alias (str)-- alias of the artifact to be download/used for training
|
||||||
|
|
||||||
|
returns:
|
||||||
|
(str, wandb.Artifact) -- path of the downladed dataset and it's corresponding artifact object if dataset
|
||||||
|
is found otherwise returns (None, None)
|
||||||
|
"""
|
||||||
|
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
||||||
|
dataset_artifact = wandb.use_artifact(artifact_path.as_posix().replace("\\", "/"))
|
||||||
|
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
|
||||||
|
datadir = dataset_artifact.download()
|
||||||
|
return datadir, dataset_artifact
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
def download_model_artifact(self, opt):
|
||||||
|
"""
|
||||||
|
download the model checkpoint artifact if the resume path starts with WANDB_ARTIFACT_PREFIX
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
opt (namespace) -- Commandline arguments for this run
|
||||||
|
"""
|
||||||
|
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
||||||
|
model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
|
||||||
|
assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
|
||||||
|
modeldir = model_artifact.download()
|
||||||
|
epochs_trained = model_artifact.metadata.get('epochs_trained')
|
||||||
|
total_epochs = model_artifact.metadata.get('total_epochs')
|
||||||
|
is_finished = total_epochs is None
|
||||||
|
assert not is_finished, 'training is finished, can only resume incomplete runs.'
|
||||||
|
return modeldir, model_artifact
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
def log_model(self, path, opt, epoch, fitness_score, best_model=False):
|
||||||
|
"""
|
||||||
|
Log the model checkpoint as W&B artifact
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
path (Path) -- Path of directory containing the checkpoints
|
||||||
|
opt (namespace) -- Command line arguments for this run
|
||||||
|
epoch (int) -- Current epoch number
|
||||||
|
fitness_score (float) -- fitness score for current epoch
|
||||||
|
best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
|
||||||
|
"""
|
||||||
|
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
|
||||||
|
'original_url': str(path),
|
||||||
|
'epochs_trained': epoch + 1,
|
||||||
|
'save period': opt.save_period,
|
||||||
|
'project': opt.project,
|
||||||
|
'total_epochs': opt.epochs,
|
||||||
|
'fitness_score': fitness_score
|
||||||
|
})
|
||||||
|
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
|
||||||
|
wandb.log_artifact(model_artifact,
|
||||||
|
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
||||||
|
LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")
|
||||||
|
|
||||||
|
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
|
||||||
|
"""
|
||||||
|
Log the dataset as W&B artifact and return the new data file with W&B links
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
data_file (str) -- the .yaml file with information about the dataset like - path, classes etc.
|
||||||
|
single_class (boolean) -- train multi-class data as single-class
|
||||||
|
project (str) -- project name. Used to construct the artifact path
|
||||||
|
overwrite_config (boolean) -- overwrites the data.yaml file if set to true otherwise creates a new
|
||||||
|
file with _wandb postfix. Eg -> data_wandb.yaml
|
||||||
|
|
||||||
|
returns:
|
||||||
|
the new .yaml file with artifact links. it can be used to start training directly from artifacts
|
||||||
|
"""
|
||||||
|
self.data_dict = check_dataset(data_file) # parse and check
|
||||||
|
data = dict(self.data_dict)
|
||||||
|
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
|
||||||
|
names = {k: v for k, v in enumerate(names)} # to index dictionary
|
||||||
|
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||||
|
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
|
||||||
|
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
||||||
|
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
|
||||||
|
if data.get('train'):
|
||||||
|
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
|
||||||
|
if data.get('val'):
|
||||||
|
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
|
||||||
|
path = Path(data_file).stem
|
||||||
|
path = (path if overwrite_config else path + '_wandb') + '.yaml' # updated data.yaml path
|
||||||
|
data.pop('download', None)
|
||||||
|
data.pop('path', None)
|
||||||
|
with open(path, 'w') as f:
|
||||||
|
yaml.safe_dump(data, f)
|
||||||
|
|
||||||
|
if self.job_type == 'Training': # builds correct artifact pipeline graph
|
||||||
|
self.wandb_run.use_artifact(self.val_artifact)
|
||||||
|
self.wandb_run.use_artifact(self.train_artifact)
|
||||||
|
self.val_artifact.wait()
|
||||||
|
self.val_table = self.val_artifact.get('val')
|
||||||
|
self.map_val_table_path()
|
||||||
|
else:
|
||||||
|
self.wandb_run.log_artifact(self.train_artifact)
|
||||||
|
self.wandb_run.log_artifact(self.val_artifact)
|
||||||
|
return path
|
||||||
|
|
||||||
|
def map_val_table_path(self):
|
||||||
|
"""
|
||||||
|
Map the validation dataset Table like name of file -> it's id in the W&B Table.
|
||||||
|
Useful for - referencing artifacts for evaluation.
|
||||||
|
"""
|
||||||
|
self.val_table_path_map = {}
|
||||||
|
LOGGER.info("Mapping dataset")
|
||||||
|
for i, data in enumerate(tqdm(self.val_table.data)):
|
||||||
|
self.val_table_path_map[data[3]] = data[0]
|
||||||
|
|
||||||
|
def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[int,str], name: str = 'dataset'):
|
||||||
|
"""
|
||||||
|
Create and return W&B artifact containing W&B Table of the dataset.
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
dataset -- instance of LoadImagesAndLabels class used to iterate over the data to build Table
|
||||||
|
class_to_id -- hash map that maps class ids to labels
|
||||||
|
name -- name of the artifact
|
||||||
|
|
||||||
|
returns:
|
||||||
|
dataset artifact to be logged or used
|
||||||
|
"""
|
||||||
|
# TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
|
||||||
|
artifact = wandb.Artifact(name=name, type="dataset")
|
||||||
|
img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
|
||||||
|
img_files = tqdm(dataset.img_files) if not img_files else img_files
|
||||||
|
for img_file in img_files:
|
||||||
|
if Path(img_file).is_dir():
|
||||||
|
artifact.add_dir(img_file, name='data/images')
|
||||||
|
labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
|
||||||
|
artifact.add_dir(labels_path, name='data/labels')
|
||||||
|
else:
|
||||||
|
artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
|
||||||
|
label_file = Path(img2label_paths([img_file])[0])
|
||||||
|
artifact.add_file(str(label_file),
|
||||||
|
name='data/labels/' + label_file.name) if label_file.exists() else None
|
||||||
|
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
|
||||||
|
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
|
||||||
|
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
|
||||||
|
box_data, img_classes = [], {}
|
||||||
|
for cls, *xywh in labels[:, 1:].tolist():
|
||||||
|
cls = int(cls)
|
||||||
|
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
|
||||||
|
"class_id": cls,
|
||||||
|
"box_caption": "%s" % (class_to_id[cls])})
|
||||||
|
img_classes[cls] = class_to_id[cls]
|
||||||
|
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
|
||||||
|
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
|
||||||
|
Path(paths).name)
|
||||||
|
artifact.add(table, name)
|
||||||
|
return artifact
|
||||||
|
|
||||||
|
def log_training_progress(self, predn, path, names):
|
||||||
|
"""
|
||||||
|
Build evaluation Table. Uses reference from validation dataset table.
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
predn (list): list of predictions in the native space in the format - [xmin, ymin, xmax, ymax, confidence, class]
|
||||||
|
path (str): local path of the current evaluation image
|
||||||
|
names (dict(int, str)): hash map that maps class ids to labels
|
||||||
|
"""
|
||||||
|
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
|
||||||
|
box_data = []
|
||||||
|
total_conf = 0
|
||||||
|
for *xyxy, conf, cls in predn.tolist():
|
||||||
|
if conf >= 0.25:
|
||||||
|
box_data.append(
|
||||||
|
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||||
|
"class_id": int(cls),
|
||||||
|
"box_caption": f"{names[cls]} {conf:.3f}",
|
||||||
|
"scores": {"class_score": conf},
|
||||||
|
"domain": "pixel"})
|
||||||
|
total_conf += conf
|
||||||
|
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||||
|
id = self.val_table_path_map[Path(path).name]
|
||||||
|
self.result_table.add_data(self.current_epoch,
|
||||||
|
id,
|
||||||
|
self.val_table.data[id][1],
|
||||||
|
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
|
||||||
|
total_conf / max(1, len(box_data))
|
||||||
|
)
|
||||||
|
|
||||||
|
def val_one_image(self, pred, predn, path, names, im):
|
||||||
|
"""
|
||||||
|
Log validation data for one image. updates the result Table if validation dataset is uploaded and log bbox media panel
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
pred (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class]
|
||||||
|
predn (list): list of predictions in the native space - [xmin, ymin, xmax, ymax, confidence, class]
|
||||||
|
path (str): local path of the current evaluation image
|
||||||
|
"""
|
||||||
|
if self.val_table and self.result_table: # Log Table if Val dataset is uploaded as artifact
|
||||||
|
self.log_training_progress(predn, path, names)
|
||||||
|
|
||||||
|
if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0:
|
||||||
|
if self.current_epoch % self.bbox_interval == 0:
|
||||||
|
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||||
|
"class_id": int(cls),
|
||||||
|
"box_caption": f"{names[cls]} {conf:.3f}",
|
||||||
|
"scores": {"class_score": conf},
|
||||||
|
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
||||||
|
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||||
|
self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name))
|
||||||
|
|
||||||
|
def log(self, log_dict):
|
||||||
|
"""
|
||||||
|
save the metrics to the logging dictionary
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
log_dict (Dict) -- metrics/media to be logged in current step
|
||||||
|
"""
|
||||||
|
if self.wandb_run:
|
||||||
|
for key, value in log_dict.items():
|
||||||
|
self.log_dict[key] = value
|
||||||
|
|
||||||
|
def end_epoch(self, best_result=False):
|
||||||
|
"""
|
||||||
|
commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
|
||||||
|
|
||||||
|
arguments:
|
||||||
|
best_result (boolean): Boolean representing if the result of this evaluation is best or not
|
||||||
|
"""
|
||||||
|
if self.wandb_run:
|
||||||
|
with all_logging_disabled():
|
||||||
|
if self.bbox_media_panel_images:
|
||||||
|
self.log_dict["BoundingBoxDebugger"] = self.bbox_media_panel_images
|
||||||
|
try:
|
||||||
|
wandb.log(self.log_dict)
|
||||||
|
except BaseException as e:
|
||||||
|
LOGGER.info(f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}")
|
||||||
|
self.wandb_run.finish()
|
||||||
|
self.wandb_run = None
|
||||||
|
|
||||||
|
self.log_dict = {}
|
||||||
|
self.bbox_media_panel_images = []
|
||||||
|
if self.result_artifact:
|
||||||
|
self.result_artifact.add(self.result_table, 'result')
|
||||||
|
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
|
||||||
|
('best' if best_result else '')])
|
||||||
|
|
||||||
|
wandb.log({"evaluation": self.result_table})
|
||||||
|
self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"])
|
||||||
|
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
||||||
|
|
||||||
|
def finish_run(self):
|
||||||
|
"""
|
||||||
|
Log metrics if any and finish the current W&B run
|
||||||
|
"""
|
||||||
|
if self.wandb_run:
|
||||||
|
if self.log_dict:
|
||||||
|
with all_logging_disabled():
|
||||||
|
wandb.log(self.log_dict)
|
||||||
|
wandb.run.finish()
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def all_logging_disabled(highest_level=logging.CRITICAL):
|
||||||
|
""" source - https://gist.github.com/simon-weber/7853144
|
||||||
|
A context manager that will prevent any logging messages triggered during the body from being processed.
|
||||||
|
:param highest_level: the maximum logging level in use.
|
||||||
|
This would only need to be changed if a custom level greater than CRITICAL is defined.
|
||||||
|
"""
|
||||||
|
previous_level = logging.root.manager.disable
|
||||||
|
logging.disable(highest_level)
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
logging.disable(previous_level)
|
||||||
@ -1,9 +1,12 @@
|
|||||||
# Loss functions
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Loss functions
|
||||||
|
"""
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
|
|
||||||
from utils.general import bbox_iou
|
from utils.metrics import bbox_iou
|
||||||
from utils.torch_utils import is_parallel
|
from utils.torch_utils import is_parallel
|
||||||
|
|
||||||
|
|
||||||
@ -15,7 +18,7 @@ def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#iss
|
|||||||
class BCEBlurWithLogitsLoss(nn.Module):
|
class BCEBlurWithLogitsLoss(nn.Module):
|
||||||
# BCEwithLogitLoss() with reduced missing label effects.
|
# BCEwithLogitLoss() with reduced missing label effects.
|
||||||
def __init__(self, alpha=0.05):
|
def __init__(self, alpha=0.05):
|
||||||
super(BCEBlurWithLogitsLoss, self).__init__()
|
super().__init__()
|
||||||
self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()
|
self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()
|
||||||
self.alpha = alpha
|
self.alpha = alpha
|
||||||
|
|
||||||
@ -32,7 +35,7 @@ class BCEBlurWithLogitsLoss(nn.Module):
|
|||||||
class FocalLoss(nn.Module):
|
class FocalLoss(nn.Module):
|
||||||
# Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
|
# Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
|
||||||
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
||||||
super(FocalLoss, self).__init__()
|
super().__init__()
|
||||||
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
||||||
self.gamma = gamma
|
self.gamma = gamma
|
||||||
self.alpha = alpha
|
self.alpha = alpha
|
||||||
@ -62,7 +65,7 @@ class FocalLoss(nn.Module):
|
|||||||
class QFocalLoss(nn.Module):
|
class QFocalLoss(nn.Module):
|
||||||
# Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
|
# Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
|
||||||
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
||||||
super(QFocalLoss, self).__init__()
|
super().__init__()
|
||||||
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
||||||
self.gamma = gamma
|
self.gamma = gamma
|
||||||
self.alpha = alpha
|
self.alpha = alpha
|
||||||
@ -88,7 +91,7 @@ class QFocalLoss(nn.Module):
|
|||||||
class ComputeLoss:
|
class ComputeLoss:
|
||||||
# Compute losses
|
# Compute losses
|
||||||
def __init__(self, model, autobalance=False):
|
def __init__(self, model, autobalance=False):
|
||||||
super(ComputeLoss, self).__init__()
|
self.sort_obj_iou = False
|
||||||
device = next(model.parameters()).device # get model device
|
device = next(model.parameters()).device # get model device
|
||||||
h = model.hyp # hyperparameters
|
h = model.hyp # hyperparameters
|
||||||
|
|
||||||
@ -105,9 +108,9 @@ class ComputeLoss:
|
|||||||
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
|
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
|
||||||
|
|
||||||
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
|
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
|
||||||
self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7
|
self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
|
||||||
self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index
|
self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index
|
||||||
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
|
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
|
||||||
for k in 'na', 'nc', 'nl', 'anchors':
|
for k in 'na', 'nc', 'nl', 'anchors':
|
||||||
setattr(self, k, getattr(det, k))
|
setattr(self, k, getattr(det, k))
|
||||||
|
|
||||||
@ -126,14 +129,18 @@ class ComputeLoss:
|
|||||||
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
|
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
|
||||||
|
|
||||||
# Regression
|
# Regression
|
||||||
pxy = ps[:, :2].sigmoid() * 2. - 0.5
|
pxy = ps[:, :2].sigmoid() * 2 - 0.5
|
||||||
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
|
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
|
||||||
pbox = torch.cat((pxy, pwh), 1) # predicted box
|
pbox = torch.cat((pxy, pwh), 1) # predicted box
|
||||||
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
|
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
|
||||||
lbox += (1.0 - iou).mean() # iou loss
|
lbox += (1.0 - iou).mean() # iou loss
|
||||||
|
|
||||||
# Objectness
|
# Objectness
|
||||||
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
|
score_iou = iou.detach().clamp(0).type(tobj.dtype)
|
||||||
|
if self.sort_obj_iou:
|
||||||
|
sort_id = torch.argsort(score_iou)
|
||||||
|
b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id]
|
||||||
|
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio
|
||||||
|
|
||||||
# Classification
|
# Classification
|
||||||
if self.nc > 1: # cls loss (only if multiple classes)
|
if self.nc > 1: # cls loss (only if multiple classes)
|
||||||
@ -157,8 +164,7 @@ class ComputeLoss:
|
|||||||
lcls *= self.hyp['cls']
|
lcls *= self.hyp['cls']
|
||||||
bs = tobj.shape[0] # batch size
|
bs = tobj.shape[0] # batch size
|
||||||
|
|
||||||
loss = lbox + lobj + lcls
|
return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()
|
||||||
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
|
|
||||||
|
|
||||||
def build_targets(self, p, targets):
|
def build_targets(self, p, targets):
|
||||||
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
|
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
|
||||||
@ -170,7 +176,7 @@ class ComputeLoss:
|
|||||||
|
|
||||||
g = 0.5 # bias
|
g = 0.5 # bias
|
||||||
off = torch.tensor([[0, 0],
|
off = torch.tensor([[0, 0],
|
||||||
# [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
|
[1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
|
||||||
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
|
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
|
||||||
], device=targets.device).float() * g # offsets
|
], device=targets.device).float() * g # offsets
|
||||||
|
|
||||||
@ -183,17 +189,17 @@ class ComputeLoss:
|
|||||||
if nt:
|
if nt:
|
||||||
# Matches
|
# Matches
|
||||||
r = t[:, :, 4:6] / anchors[:, None] # wh ratio
|
r = t[:, :, 4:6] / anchors[:, None] # wh ratio
|
||||||
j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare
|
j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare
|
||||||
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
|
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
|
||||||
t = t[j] # filter
|
t = t[j] # filter
|
||||||
|
|
||||||
# Offsets
|
# Offsets
|
||||||
gxy = t[:, 2:4] # grid xy
|
gxy = t[:, 2:4] # grid xy
|
||||||
gxi = gain[[2, 3]] - gxy # inverse
|
gxi = gain[[2, 3]] - gxy # inverse
|
||||||
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
|
j, k = ((gxy % 1 < g) & (gxy > 1)).T
|
||||||
l, m = ((gxi % 1. < g) & (gxi > 1.)).T
|
l, m = ((gxi % 1 < g) & (gxi > 1)).T
|
||||||
j = torch.stack((torch.ones_like(j),))
|
j = torch.stack((torch.ones_like(j), j, k, l, m))
|
||||||
t = t.repeat((off.shape[0], 1, 1))[j]
|
t = t.repeat((5, 1, 1))[j]
|
||||||
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
|
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
|
||||||
else:
|
else:
|
||||||
t = targets[0]
|
t = targets[0]
|
||||||
|
|||||||
136
utils/metrics.py
136
utils/metrics.py
@ -1,13 +1,16 @@
|
|||||||
# Model validation metrics
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Model validation metrics
|
||||||
|
"""
|
||||||
|
|
||||||
|
import math
|
||||||
|
import warnings
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from . import general
|
|
||||||
|
|
||||||
|
|
||||||
def fitness(x):
|
def fitness(x):
|
||||||
# Model fitness as a weighted combination of metrics
|
# Model fitness as a weighted combination of metrics
|
||||||
@ -68,6 +71,8 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
|
|||||||
|
|
||||||
# Compute F1 (harmonic mean of precision and recall)
|
# Compute F1 (harmonic mean of precision and recall)
|
||||||
f1 = 2 * p * r / (p + r + 1e-16)
|
f1 = 2 * p * r / (p + r + 1e-16)
|
||||||
|
names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
|
||||||
|
names = {i: v for i, v in enumerate(names)} # to dict
|
||||||
if plot:
|
if plot:
|
||||||
plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
|
plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
|
||||||
plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
|
plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
|
||||||
@ -88,8 +93,8 @@ def compute_ap(recall, precision):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
# Append sentinel values to beginning and end
|
# Append sentinel values to beginning and end
|
||||||
mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01]))
|
mrec = np.concatenate(([0.0], recall, [1.0]))
|
||||||
mpre = np.concatenate(([1.], precision, [0.]))
|
mpre = np.concatenate(([1.0], precision, [0.0]))
|
||||||
|
|
||||||
# Compute the precision envelope
|
# Compute the precision envelope
|
||||||
mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
|
mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
|
||||||
@ -127,7 +132,7 @@ class ConfusionMatrix:
|
|||||||
detections = detections[detections[:, 4] > self.conf]
|
detections = detections[detections[:, 4] > self.conf]
|
||||||
gt_classes = labels[:, 0].int()
|
gt_classes = labels[:, 0].int()
|
||||||
detection_classes = detections[:, 5].int()
|
detection_classes = detections[:, 5].int()
|
||||||
iou = general.box_iou(labels[:, 1:], detections[:, :4])
|
iou = box_iou(labels[:, 1:], detections[:, :4])
|
||||||
|
|
||||||
x = torch.where(iou > self.iou_thres)
|
x = torch.where(iou > self.iou_thres)
|
||||||
if x[0].shape[0]:
|
if x[0].shape[0]:
|
||||||
@ -157,30 +162,135 @@ class ConfusionMatrix:
|
|||||||
def matrix(self):
|
def matrix(self):
|
||||||
return self.matrix
|
return self.matrix
|
||||||
|
|
||||||
def plot(self, save_dir='', names=()):
|
def plot(self, normalize=True, save_dir='', names=()):
|
||||||
try:
|
try:
|
||||||
import seaborn as sn
|
import seaborn as sn
|
||||||
|
|
||||||
array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize
|
array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns
|
||||||
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
|
array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
|
||||||
|
|
||||||
fig = plt.figure(figsize=(12, 9), tight_layout=True)
|
fig = plt.figure(figsize=(12, 9), tight_layout=True)
|
||||||
sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
|
sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
|
||||||
labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
|
labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
|
||||||
sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
|
with warnings.catch_warnings():
|
||||||
xticklabels=names + ['background FP'] if labels else "auto",
|
warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
|
||||||
yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
|
sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
|
||||||
|
xticklabels=names + ['background FP'] if labels else "auto",
|
||||||
|
yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
|
||||||
fig.axes[0].set_xlabel('True')
|
fig.axes[0].set_xlabel('True')
|
||||||
fig.axes[0].set_ylabel('Predicted')
|
fig.axes[0].set_ylabel('Predicted')
|
||||||
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
||||||
|
plt.close()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
pass
|
print(f'WARNING: ConfusionMatrix plot failure: {e}')
|
||||||
|
|
||||||
def print(self):
|
def print(self):
|
||||||
for i in range(self.nc + 1):
|
for i in range(self.nc + 1):
|
||||||
print(' '.join(map(str, self.matrix[i])))
|
print(' '.join(map(str, self.matrix[i])))
|
||||||
|
|
||||||
|
|
||||||
|
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
|
||||||
|
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
|
||||||
|
box2 = box2.T
|
||||||
|
|
||||||
|
# Get the coordinates of bounding boxes
|
||||||
|
if x1y1x2y2: # x1, y1, x2, y2 = box1
|
||||||
|
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
|
||||||
|
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
|
||||||
|
else: # transform from xywh to xyxy
|
||||||
|
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
|
||||||
|
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
|
||||||
|
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
|
||||||
|
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
|
||||||
|
|
||||||
|
# Intersection area
|
||||||
|
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
||||||
|
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
||||||
|
|
||||||
|
# Union Area
|
||||||
|
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
||||||
|
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
||||||
|
union = w1 * h1 + w2 * h2 - inter + eps
|
||||||
|
|
||||||
|
iou = inter / union
|
||||||
|
if GIoU or DIoU or CIoU:
|
||||||
|
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
||||||
|
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
||||||
|
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
||||||
|
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
|
||||||
|
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
|
||||||
|
(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
|
||||||
|
if DIoU:
|
||||||
|
return iou - rho2 / c2 # DIoU
|
||||||
|
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
||||||
|
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
||||||
|
with torch.no_grad():
|
||||||
|
alpha = v / (v - iou + (1 + eps))
|
||||||
|
return iou - (rho2 / c2 + v * alpha) # CIoU
|
||||||
|
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
||||||
|
c_area = cw * ch + eps # convex area
|
||||||
|
return iou - (c_area - union) / c_area # GIoU
|
||||||
|
else:
|
||||||
|
return iou # IoU
|
||||||
|
|
||||||
|
|
||||||
|
def box_iou(box1, box2):
|
||||||
|
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
||||||
|
"""
|
||||||
|
Return intersection-over-union (Jaccard index) of boxes.
|
||||||
|
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||||||
|
Arguments:
|
||||||
|
box1 (Tensor[N, 4])
|
||||||
|
box2 (Tensor[M, 4])
|
||||||
|
Returns:
|
||||||
|
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
||||||
|
IoU values for every element in boxes1 and boxes2
|
||||||
|
"""
|
||||||
|
|
||||||
|
def box_area(box):
|
||||||
|
# box = 4xn
|
||||||
|
return (box[2] - box[0]) * (box[3] - box[1])
|
||||||
|
|
||||||
|
area1 = box_area(box1.T)
|
||||||
|
area2 = box_area(box2.T)
|
||||||
|
|
||||||
|
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||||||
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
||||||
|
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
||||||
|
|
||||||
|
|
||||||
|
def bbox_ioa(box1, box2, eps=1E-7):
|
||||||
|
""" Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2
|
||||||
|
box1: np.array of shape(4)
|
||||||
|
box2: np.array of shape(nx4)
|
||||||
|
returns: np.array of shape(n)
|
||||||
|
"""
|
||||||
|
|
||||||
|
box2 = box2.transpose()
|
||||||
|
|
||||||
|
# Get the coordinates of bounding boxes
|
||||||
|
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
|
||||||
|
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
|
||||||
|
|
||||||
|
# Intersection area
|
||||||
|
inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
|
||||||
|
(np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)
|
||||||
|
|
||||||
|
# box2 area
|
||||||
|
box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps
|
||||||
|
|
||||||
|
# Intersection over box2 area
|
||||||
|
return inter_area / box2_area
|
||||||
|
|
||||||
|
|
||||||
|
def wh_iou(wh1, wh2):
|
||||||
|
# Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
|
||||||
|
wh1 = wh1[:, None] # [N,1,2]
|
||||||
|
wh2 = wh2[None] # [1,M,2]
|
||||||
|
inter = torch.min(wh1, wh2).prod(2) # [N,M]
|
||||||
|
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
|
||||||
|
|
||||||
|
|
||||||
# Plots ----------------------------------------------------------------------------------------------------------------
|
# Plots ----------------------------------------------------------------------------------------------------------------
|
||||||
|
|
||||||
def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
|
def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
|
||||||
@ -201,6 +311,7 @@ def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
|
|||||||
ax.set_ylim(0, 1)
|
ax.set_ylim(0, 1)
|
||||||
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
||||||
fig.savefig(Path(save_dir), dpi=250)
|
fig.savefig(Path(save_dir), dpi=250)
|
||||||
|
plt.close()
|
||||||
|
|
||||||
|
|
||||||
def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
|
def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
|
||||||
@ -221,3 +332,4 @@ def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence'
|
|||||||
ax.set_ylim(0, 1)
|
ax.set_ylim(0, 1)
|
||||||
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
|
||||||
fig.savefig(Path(save_dir), dpi=250)
|
fig.savefig(Path(save_dir), dpi=250)
|
||||||
|
plt.close()
|
||||||
|
|||||||
437
utils/plots.py
437
utils/plots.py
@ -1,9 +1,10 @@
|
|||||||
# Plotting utils
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Plotting utils
|
||||||
|
"""
|
||||||
|
|
||||||
import glob
|
|
||||||
import math
|
import math
|
||||||
import os
|
import os
|
||||||
import random
|
|
||||||
from copy import copy
|
from copy import copy
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@ -12,15 +13,17 @@ import matplotlib
|
|||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
import seaborn as sns
|
import seaborn as sn
|
||||||
import torch
|
import torch
|
||||||
import yaml
|
|
||||||
from PIL import Image, ImageDraw, ImageFont
|
from PIL import Image, ImageDraw, ImageFont
|
||||||
|
|
||||||
from utils.general import xywh2xyxy, xyxy2xywh
|
from utils.general import (LOGGER, Timeout, check_requirements, clip_coords, increment_path, is_ascii, is_chinese,
|
||||||
|
try_except, user_config_dir, xywh2xyxy, xyxy2xywh)
|
||||||
from utils.metrics import fitness
|
from utils.metrics import fitness
|
||||||
|
|
||||||
# Settings
|
# Settings
|
||||||
|
CONFIG_DIR = user_config_dir() # Ultralytics settings dir
|
||||||
|
RANK = int(os.getenv('RANK', -1))
|
||||||
matplotlib.rc('font', **{'size': 11})
|
matplotlib.rc('font', **{'size': 11})
|
||||||
matplotlib.use('Agg') # for writing to files only
|
matplotlib.use('Agg') # for writing to files only
|
||||||
|
|
||||||
@ -46,6 +49,105 @@ class Colors:
|
|||||||
colors = Colors() # create instance for 'from utils.plots import colors'
|
colors = Colors() # create instance for 'from utils.plots import colors'
|
||||||
|
|
||||||
|
|
||||||
|
def check_font(font='Arial.ttf', size=10):
|
||||||
|
# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
|
||||||
|
font = Path(font)
|
||||||
|
font = font if font.exists() else (CONFIG_DIR / font.name)
|
||||||
|
try:
|
||||||
|
return ImageFont.truetype(str(font) if font.exists() else font.name, size)
|
||||||
|
except Exception as e: # download if missing
|
||||||
|
url = "https://ultralytics.com/assets/" + font.name
|
||||||
|
print(f'Downloading {url} to {font}...')
|
||||||
|
torch.hub.download_url_to_file(url, str(font), progress=False)
|
||||||
|
try:
|
||||||
|
return ImageFont.truetype(str(font), size)
|
||||||
|
except TypeError:
|
||||||
|
check_requirements('Pillow>=8.4.0') # known issue https://github.com/ultralytics/yolov5/issues/5374
|
||||||
|
|
||||||
|
|
||||||
|
class Annotator:
|
||||||
|
if RANK in (-1, 0):
|
||||||
|
check_font() # download TTF if necessary
|
||||||
|
|
||||||
|
# Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
||||||
|
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
||||||
|
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
||||||
|
self.pil = pil or not is_ascii(example) or is_chinese(example)
|
||||||
|
if self.pil: # use PIL
|
||||||
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
||||||
|
self.draw = ImageDraw.Draw(self.im)
|
||||||
|
self.font = check_font(font='Arial.Unicode.ttf' if is_chinese(example) else font,
|
||||||
|
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
||||||
|
else: # use cv2
|
||||||
|
self.im = im
|
||||||
|
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
||||||
|
|
||||||
|
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
||||||
|
# Add one xyxy box to image with label
|
||||||
|
if self.pil or not is_ascii(label):
|
||||||
|
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
||||||
|
if label:
|
||||||
|
w, h = self.font.getsize(label) # text width, height
|
||||||
|
outside = box[1] - h >= 0 # label fits outside box
|
||||||
|
self.draw.rectangle([box[0],
|
||||||
|
box[1] - h if outside else box[1],
|
||||||
|
box[0] + w + 1,
|
||||||
|
box[1] + 1 if outside else box[1] + h + 1], fill=color)
|
||||||
|
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
||||||
|
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
||||||
|
else: # cv2
|
||||||
|
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
||||||
|
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
|
||||||
|
if label:
|
||||||
|
tf = max(self.lw - 1, 1) # font thickness
|
||||||
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
||||||
|
outside = p1[1] - h - 3 >= 0 # label fits outside box
|
||||||
|
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
||||||
|
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
||||||
|
cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color,
|
||||||
|
thickness=tf, lineType=cv2.LINE_AA)
|
||||||
|
|
||||||
|
def rectangle(self, xy, fill=None, outline=None, width=1):
|
||||||
|
# Add rectangle to image (PIL-only)
|
||||||
|
self.draw.rectangle(xy, fill, outline, width)
|
||||||
|
|
||||||
|
def text(self, xy, text, txt_color=(255, 255, 255)):
|
||||||
|
# Add text to image (PIL-only)
|
||||||
|
w, h = self.font.getsize(text) # text width, height
|
||||||
|
self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)
|
||||||
|
|
||||||
|
def result(self):
|
||||||
|
# Return annotated image as array
|
||||||
|
return np.asarray(self.im)
|
||||||
|
|
||||||
|
|
||||||
|
def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')):
|
||||||
|
"""
|
||||||
|
x: Features to be visualized
|
||||||
|
module_type: Module type
|
||||||
|
stage: Module stage within model
|
||||||
|
n: Maximum number of feature maps to plot
|
||||||
|
save_dir: Directory to save results
|
||||||
|
"""
|
||||||
|
if 'Detect' not in module_type:
|
||||||
|
batch, channels, height, width = x.shape # batch, channels, height, width
|
||||||
|
if height > 1 and width > 1:
|
||||||
|
f = f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
|
||||||
|
|
||||||
|
blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
|
||||||
|
n = min(n, channels) # number of plots
|
||||||
|
fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols
|
||||||
|
ax = ax.ravel()
|
||||||
|
plt.subplots_adjust(wspace=0.05, hspace=0.05)
|
||||||
|
for i in range(n):
|
||||||
|
ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
|
||||||
|
ax[i].axis('off')
|
||||||
|
|
||||||
|
print(f'Saving {save_dir / f}... ({n}/{channels})')
|
||||||
|
plt.savefig(save_dir / f, dpi=300, bbox_inches='tight')
|
||||||
|
plt.close()
|
||||||
|
|
||||||
|
|
||||||
def hist2d(x, y, n=100):
|
def hist2d(x, y, n=100):
|
||||||
# 2d histogram used in labels.png and evolve.png
|
# 2d histogram used in labels.png and evolve.png
|
||||||
xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
|
xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
|
||||||
@ -68,54 +170,6 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
|||||||
return filtfilt(b, a, data) # forward-backward filter
|
return filtfilt(b, a, data) # forward-backward filter
|
||||||
|
|
||||||
|
|
||||||
def plot_one_box(x, im, color=(128, 128, 128), label=None, line_thickness=3):
|
|
||||||
# Plots one bounding box on image 'im' using OpenCV
|
|
||||||
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
|
|
||||||
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness
|
|
||||||
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
|
||||||
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
|
|
||||||
if label:
|
|
||||||
tf = max(tl - 1, 1) # font thickness
|
|
||||||
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
|
||||||
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
|
||||||
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled
|
|
||||||
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
|
||||||
|
|
||||||
|
|
||||||
def plot_one_box_PIL(box, im, color=(128, 128, 128), label=None, line_thickness=None):
|
|
||||||
# Plots one bounding box on image 'im' using PIL
|
|
||||||
im = Image.fromarray(im)
|
|
||||||
draw = ImageDraw.Draw(im)
|
|
||||||
line_thickness = line_thickness or max(int(min(im.size) / 200), 2)
|
|
||||||
draw.rectangle(box, width=line_thickness, outline=color) # plot
|
|
||||||
if label:
|
|
||||||
font = ImageFont.truetype("Arial.ttf", size=max(round(max(im.size) / 40), 12))
|
|
||||||
txt_width, txt_height = font.getsize(label)
|
|
||||||
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=color)
|
|
||||||
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
|
|
||||||
return np.asarray(im)
|
|
||||||
|
|
||||||
|
|
||||||
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
|
|
||||||
# Compares the two methods for width-height anchor multiplication
|
|
||||||
# https://github.com/ultralytics/yolov3/issues/168
|
|
||||||
x = np.arange(-4.0, 4.0, .1)
|
|
||||||
ya = np.exp(x)
|
|
||||||
yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2
|
|
||||||
|
|
||||||
fig = plt.figure(figsize=(6, 3), tight_layout=True)
|
|
||||||
plt.plot(x, ya, '.-', label='YOLOv3')
|
|
||||||
plt.plot(x, yb ** 2, '.-', label='YOLOv5 ^2')
|
|
||||||
plt.plot(x, yb ** 1.6, '.-', label='YOLOv5 ^1.6')
|
|
||||||
plt.xlim(left=-4, right=4)
|
|
||||||
plt.ylim(bottom=0, top=6)
|
|
||||||
plt.xlabel('input')
|
|
||||||
plt.ylabel('output')
|
|
||||||
plt.grid()
|
|
||||||
plt.legend()
|
|
||||||
fig.savefig('comparison.png', dpi=200)
|
|
||||||
|
|
||||||
|
|
||||||
def output_to_target(output):
|
def output_to_target(output):
|
||||||
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
|
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
|
||||||
targets = []
|
targets = []
|
||||||
@ -125,82 +179,65 @@ def output_to_target(output):
|
|||||||
return np.array(targets)
|
return np.array(targets)
|
||||||
|
|
||||||
|
|
||||||
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
|
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=1920, max_subplots=16):
|
||||||
# Plot image grid with labels
|
# Plot image grid with labels
|
||||||
|
|
||||||
if isinstance(images, torch.Tensor):
|
if isinstance(images, torch.Tensor):
|
||||||
images = images.cpu().float().numpy()
|
images = images.cpu().float().numpy()
|
||||||
if isinstance(targets, torch.Tensor):
|
if isinstance(targets, torch.Tensor):
|
||||||
targets = targets.cpu().numpy()
|
targets = targets.cpu().numpy()
|
||||||
|
|
||||||
# un-normalise
|
|
||||||
if np.max(images[0]) <= 1:
|
if np.max(images[0]) <= 1:
|
||||||
images *= 255
|
images *= 255 # de-normalise (optional)
|
||||||
|
|
||||||
tl = 3 # line thickness
|
|
||||||
tf = max(tl - 1, 1) # font thickness
|
|
||||||
bs, _, h, w = images.shape # batch size, _, height, width
|
bs, _, h, w = images.shape # batch size, _, height, width
|
||||||
bs = min(bs, max_subplots) # limit plot images
|
bs = min(bs, max_subplots) # limit plot images
|
||||||
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
||||||
|
|
||||||
# Check if we should resize
|
# Build Image
|
||||||
scale_factor = max_size / max(h, w)
|
|
||||||
if scale_factor < 1:
|
|
||||||
h = math.ceil(scale_factor * h)
|
|
||||||
w = math.ceil(scale_factor * w)
|
|
||||||
|
|
||||||
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
||||||
for i, img in enumerate(images):
|
for i, im in enumerate(images):
|
||||||
if i == max_subplots: # if last batch has fewer images than we expect
|
if i == max_subplots: # if last batch has fewer images than we expect
|
||||||
break
|
break
|
||||||
|
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
||||||
|
im = im.transpose(1, 2, 0)
|
||||||
|
mosaic[y:y + h, x:x + w, :] = im
|
||||||
|
|
||||||
block_x = int(w * (i // ns))
|
# Resize (optional)
|
||||||
block_y = int(h * (i % ns))
|
scale = max_size / ns / max(h, w)
|
||||||
|
if scale < 1:
|
||||||
|
h = math.ceil(scale * h)
|
||||||
|
w = math.ceil(scale * w)
|
||||||
|
mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
|
||||||
|
|
||||||
img = img.transpose(1, 2, 0)
|
# Annotate
|
||||||
if scale_factor < 1:
|
fs = int((h + w) * ns * 0.01) # font size
|
||||||
img = cv2.resize(img, (w, h))
|
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True)
|
||||||
|
for i in range(i + 1):
|
||||||
mosaic[block_y:block_y + h, block_x:block_x + w, :] = img
|
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
||||||
|
annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
|
||||||
|
if paths:
|
||||||
|
annotator.text((x + 5, y + 5 + h), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
|
||||||
if len(targets) > 0:
|
if len(targets) > 0:
|
||||||
image_targets = targets[targets[:, 0] == i]
|
ti = targets[targets[:, 0] == i] # image targets
|
||||||
boxes = xywh2xyxy(image_targets[:, 2:6]).T
|
boxes = xywh2xyxy(ti[:, 2:6]).T
|
||||||
classes = image_targets[:, 1].astype('int')
|
classes = ti[:, 1].astype('int')
|
||||||
labels = image_targets.shape[1] == 6 # labels if no conf column
|
labels = ti.shape[1] == 6 # labels if no conf column
|
||||||
conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred)
|
conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred)
|
||||||
|
|
||||||
if boxes.shape[1]:
|
if boxes.shape[1]:
|
||||||
if boxes.max() <= 1.01: # if normalized with tolerance 0.01
|
if boxes.max() <= 1.01: # if normalized with tolerance 0.01
|
||||||
boxes[[0, 2]] *= w # scale to pixels
|
boxes[[0, 2]] *= w # scale to pixels
|
||||||
boxes[[1, 3]] *= h
|
boxes[[1, 3]] *= h
|
||||||
elif scale_factor < 1: # absolute coords need scale if image scales
|
elif scale < 1: # absolute coords need scale if image scales
|
||||||
boxes *= scale_factor
|
boxes *= scale
|
||||||
boxes[[0, 2]] += block_x
|
boxes[[0, 2]] += x
|
||||||
boxes[[1, 3]] += block_y
|
boxes[[1, 3]] += y
|
||||||
for j, box in enumerate(boxes.T):
|
for j, box in enumerate(boxes.T.tolist()):
|
||||||
cls = int(classes[j])
|
cls = classes[j]
|
||||||
color = colors(cls)
|
color = colors(cls)
|
||||||
cls = names[cls] if names else cls
|
cls = names[cls] if names else cls
|
||||||
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
||||||
label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
|
label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}'
|
||||||
plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
|
annotator.box_label(box, label, color=color)
|
||||||
|
annotator.im.save(fname) # save
|
||||||
# Draw image filename labels
|
|
||||||
if paths:
|
|
||||||
label = Path(paths[i]).name[:40] # trim to 40 char
|
|
||||||
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
|
||||||
cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf,
|
|
||||||
lineType=cv2.LINE_AA)
|
|
||||||
|
|
||||||
# Image border
|
|
||||||
cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3)
|
|
||||||
|
|
||||||
if fname:
|
|
||||||
r = min(1280. / max(h, w) / ns, 1.0) # ratio to limit image size
|
|
||||||
mosaic = cv2.resize(mosaic, (int(ns * w * r), int(ns * h * r)), interpolation=cv2.INTER_AREA)
|
|
||||||
# cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB)) # cv2 save
|
|
||||||
Image.fromarray(mosaic).save(fname) # PIL save
|
|
||||||
return mosaic
|
|
||||||
|
|
||||||
|
|
||||||
def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
|
def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
|
||||||
@ -220,9 +257,9 @@ def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
|
|||||||
plt.close()
|
plt.close()
|
||||||
|
|
||||||
|
|
||||||
def plot_test_txt(): # from utils.plots import *; plot_test()
|
def plot_val_txt(): # from utils.plots import *; plot_val()
|
||||||
# Plot test.txt histograms
|
# Plot val.txt histograms
|
||||||
x = np.loadtxt('test.txt', dtype=np.float32)
|
x = np.loadtxt('val.txt', dtype=np.float32)
|
||||||
box = xyxy2xywh(x[:, :4])
|
box = xyxy2xywh(x[:, :4])
|
||||||
cx, cy = box[:, 0], box[:, 1]
|
cx, cy = box[:, 0], box[:, 1]
|
||||||
|
|
||||||
@ -244,29 +281,32 @@ def plot_targets_txt(): # from utils.plots import *; plot_targets_txt()
|
|||||||
fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
|
fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
|
||||||
ax = ax.ravel()
|
ax = ax.ravel()
|
||||||
for i in range(4):
|
for i in range(4):
|
||||||
ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std()))
|
ax[i].hist(x[i], bins=100, label=f'{x[i].mean():.3g} +/- {x[i].std():.3g}')
|
||||||
ax[i].legend()
|
ax[i].legend()
|
||||||
ax[i].set_title(s[i])
|
ax[i].set_title(s[i])
|
||||||
plt.savefig('targets.jpg', dpi=200)
|
plt.savefig('targets.jpg', dpi=200)
|
||||||
|
|
||||||
|
|
||||||
def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_txt()
|
def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_val_study()
|
||||||
# Plot study.txt generated by test.py
|
# Plot file=study.txt generated by val.py (or plot all study*.txt in dir)
|
||||||
fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
|
save_dir = Path(file).parent if file else Path(dir)
|
||||||
# ax = ax.ravel()
|
plot2 = False # plot additional results
|
||||||
|
if plot2:
|
||||||
|
ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)[1].ravel()
|
||||||
|
|
||||||
fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
|
fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
|
||||||
# for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov3-tiny', 'yolov3', 'yolov3-spp', 'yolov5l']]:
|
# for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov3', 'yolov3-spp', 'yolov3-tiny']]:
|
||||||
for f in sorted(Path(path).glob('study*.txt')):
|
for f in sorted(save_dir.glob('study*.txt')):
|
||||||
y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
|
y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
|
||||||
x = np.arange(y.shape[1]) if x is None else np.array(x)
|
x = np.arange(y.shape[1]) if x is None else np.array(x)
|
||||||
s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
|
if plot2:
|
||||||
# for i in range(7):
|
s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_preprocess (ms/img)', 't_inference (ms/img)', 't_NMS (ms/img)']
|
||||||
# ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
|
for i in range(7):
|
||||||
# ax[i].set_title(s[i])
|
ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
|
||||||
|
ax[i].set_title(s[i])
|
||||||
|
|
||||||
j = y[3].argmax() + 1
|
j = y[3].argmax() + 1
|
||||||
ax2.plot(y[6, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8,
|
ax2.plot(y[5, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8,
|
||||||
label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
|
label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
|
||||||
|
|
||||||
ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
|
ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
|
||||||
@ -275,22 +315,26 @@ def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_tx
|
|||||||
ax2.grid(alpha=0.2)
|
ax2.grid(alpha=0.2)
|
||||||
ax2.set_yticks(np.arange(20, 60, 5))
|
ax2.set_yticks(np.arange(20, 60, 5))
|
||||||
ax2.set_xlim(0, 57)
|
ax2.set_xlim(0, 57)
|
||||||
ax2.set_ylim(15, 55)
|
ax2.set_ylim(25, 55)
|
||||||
ax2.set_xlabel('GPU Speed (ms/img)')
|
ax2.set_xlabel('GPU Speed (ms/img)')
|
||||||
ax2.set_ylabel('COCO AP val')
|
ax2.set_ylabel('COCO AP val')
|
||||||
ax2.legend(loc='lower right')
|
ax2.legend(loc='lower right')
|
||||||
plt.savefig(str(Path(path).name) + '.png', dpi=300)
|
f = save_dir / 'study.png'
|
||||||
|
print(f'Saving {f}...')
|
||||||
|
plt.savefig(f, dpi=300)
|
||||||
|
|
||||||
|
|
||||||
def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
@try_except # known issue https://github.com/ultralytics/yolov5/issues/5395
|
||||||
|
@Timeout(30) # known issue https://github.com/ultralytics/yolov5/issues/5611
|
||||||
|
def plot_labels(labels, names=(), save_dir=Path('')):
|
||||||
# plot dataset labels
|
# plot dataset labels
|
||||||
print('Plotting labels... ')
|
LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
|
||||||
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
|
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
|
||||||
nc = int(c.max() + 1) # number of classes
|
nc = int(c.max() + 1) # number of classes
|
||||||
x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
|
x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
|
||||||
|
|
||||||
# seaborn correlogram
|
# seaborn correlogram
|
||||||
sns.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
|
sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
|
||||||
plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
|
plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
|
||||||
plt.close()
|
plt.close()
|
||||||
|
|
||||||
@ -305,8 +349,8 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
|||||||
ax[0].set_xticklabels(names, rotation=90, fontsize=10)
|
ax[0].set_xticklabels(names, rotation=90, fontsize=10)
|
||||||
else:
|
else:
|
||||||
ax[0].set_xlabel('classes')
|
ax[0].set_xlabel('classes')
|
||||||
sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
|
sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
|
||||||
sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
|
sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
|
||||||
|
|
||||||
# rectangles
|
# rectangles
|
||||||
labels[:, 1:3] = 0.5 # center
|
labels[:, 1:3] = 0.5 # center
|
||||||
@ -325,34 +369,57 @@ def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
|
|||||||
matplotlib.use('Agg')
|
matplotlib.use('Agg')
|
||||||
plt.close()
|
plt.close()
|
||||||
|
|
||||||
# loggers
|
|
||||||
for k, v in loggers.items() or {}:
|
|
||||||
if k == 'wandb' and v:
|
|
||||||
v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}, commit=False)
|
|
||||||
|
|
||||||
|
def plot_evolve(evolve_csv='path/to/evolve.csv'): # from utils.plots import *; plot_evolve()
|
||||||
def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution()
|
# Plot evolve.csv hyp evolution results
|
||||||
# Plot hyperparameter evolution results in evolve.txt
|
evolve_csv = Path(evolve_csv)
|
||||||
with open(yaml_file) as f:
|
data = pd.read_csv(evolve_csv)
|
||||||
hyp = yaml.safe_load(f)
|
keys = [x.strip() for x in data.columns]
|
||||||
x = np.loadtxt('evolve.txt', ndmin=2)
|
x = data.values
|
||||||
f = fitness(x)
|
f = fitness(x)
|
||||||
# weights = (f - f.min()) ** 2 # for weighted results
|
j = np.argmax(f) # max fitness index
|
||||||
plt.figure(figsize=(10, 12), tight_layout=True)
|
plt.figure(figsize=(10, 12), tight_layout=True)
|
||||||
matplotlib.rc('font', **{'size': 8})
|
matplotlib.rc('font', **{'size': 8})
|
||||||
for i, (k, v) in enumerate(hyp.items()):
|
for i, k in enumerate(keys[7:]):
|
||||||
y = x[:, i + 7]
|
v = x[:, 7 + i]
|
||||||
# mu = (y * weights).sum() / weights.sum() # best weighted result
|
mu = v[j] # best single result
|
||||||
mu = y[f.argmax()] # best single result
|
|
||||||
plt.subplot(6, 5, i + 1)
|
plt.subplot(6, 5, i + 1)
|
||||||
plt.scatter(y, f, c=hist2d(y, f, 20), cmap='viridis', alpha=.8, edgecolors='none')
|
plt.scatter(v, f, c=hist2d(v, f, 20), cmap='viridis', alpha=.8, edgecolors='none')
|
||||||
plt.plot(mu, f.max(), 'k+', markersize=15)
|
plt.plot(mu, f.max(), 'k+', markersize=15)
|
||||||
plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9}) # limit to 40 characters
|
plt.title(f'{k} = {mu:.3g}', fontdict={'size': 9}) # limit to 40 characters
|
||||||
if i % 5 != 0:
|
if i % 5 != 0:
|
||||||
plt.yticks([])
|
plt.yticks([])
|
||||||
print('%15s: %.3g' % (k, mu))
|
print(f'{k:>15}: {mu:.3g}')
|
||||||
plt.savefig('evolve.png', dpi=200)
|
f = evolve_csv.with_suffix('.png') # filename
|
||||||
print('\nPlot saved as evolve.png')
|
plt.savefig(f, dpi=200)
|
||||||
|
plt.close()
|
||||||
|
print(f'Saved {f}')
|
||||||
|
|
||||||
|
|
||||||
|
def plot_results(file='path/to/results.csv', dir=''):
|
||||||
|
# Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
|
||||||
|
save_dir = Path(file).parent if file else Path(dir)
|
||||||
|
fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
|
||||||
|
ax = ax.ravel()
|
||||||
|
files = list(save_dir.glob('results*.csv'))
|
||||||
|
assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
|
||||||
|
for fi, f in enumerate(files):
|
||||||
|
try:
|
||||||
|
data = pd.read_csv(f)
|
||||||
|
s = [x.strip() for x in data.columns]
|
||||||
|
x = data.values[:, 0]
|
||||||
|
for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]):
|
||||||
|
y = data.values[:, j]
|
||||||
|
# y[y == 0] = np.nan # don't show zero values
|
||||||
|
ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8)
|
||||||
|
ax[i].set_title(s[j], fontsize=12)
|
||||||
|
# if j in [8, 9, 10]: # share train and val loss y axes
|
||||||
|
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
|
||||||
|
except Exception as e:
|
||||||
|
print(f'Warning: Plotting error for {f}: {e}')
|
||||||
|
ax[1].legend()
|
||||||
|
fig.savefig(save_dir / 'results.png', dpi=200)
|
||||||
|
plt.close()
|
||||||
|
|
||||||
|
|
||||||
def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
|
def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
|
||||||
@ -381,66 +448,22 @@ def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
|
|||||||
else:
|
else:
|
||||||
a.remove()
|
a.remove()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print('Warning: Plotting error for %s; %s' % (f, e))
|
print(f'Warning: Plotting error for {f}; {e}')
|
||||||
|
|
||||||
ax[1].legend()
|
ax[1].legend()
|
||||||
plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
|
plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
|
||||||
|
|
||||||
|
|
||||||
def plot_results_overlay(start=0, stop=0): # from utils.plots import *; plot_results_overlay()
|
def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
|
||||||
# Plot training 'results*.txt', overlaying train and val losses
|
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
||||||
s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95'] # legends
|
xyxy = torch.tensor(xyxy).view(-1, 4)
|
||||||
t = ['Box', 'Objectness', 'Classification', 'P-R', 'mAP-F1'] # titles
|
b = xyxy2xywh(xyxy) # boxes
|
||||||
for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')):
|
if square:
|
||||||
results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
|
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
||||||
n = results.shape[1] # number of rows
|
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
||||||
x = range(start, min(stop, n) if stop else n)
|
xyxy = xywh2xyxy(b).long()
|
||||||
fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True)
|
clip_coords(xyxy, im.shape)
|
||||||
ax = ax.ravel()
|
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
||||||
for i in range(5):
|
if save:
|
||||||
for j in [i, i + 5]:
|
file.parent.mkdir(parents=True, exist_ok=True) # make directory
|
||||||
y = results[j, x]
|
cv2.imwrite(str(increment_path(file).with_suffix('.jpg')), crop)
|
||||||
ax[i].plot(x, y, marker='.', label=s[j])
|
return crop
|
||||||
# y_smooth = butter_lowpass_filtfilt(y)
|
|
||||||
# ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j])
|
|
||||||
|
|
||||||
ax[i].set_title(t[i])
|
|
||||||
ax[i].legend()
|
|
||||||
ax[i].set_ylabel(f) if i == 0 else None # add filename
|
|
||||||
fig.savefig(f.replace('.txt', '.png'), dpi=200)
|
|
||||||
|
|
||||||
|
|
||||||
def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''):
|
|
||||||
# Plot training 'results*.txt'. from utils.plots import *; plot_results(save_dir='runs/train/exp')
|
|
||||||
fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
|
|
||||||
ax = ax.ravel()
|
|
||||||
s = ['Box', 'Objectness', 'Classification', 'Precision', 'Recall',
|
|
||||||
'val Box', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95']
|
|
||||||
if bucket:
|
|
||||||
# files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id]
|
|
||||||
files = ['results%g.txt' % x for x in id]
|
|
||||||
c = ('gsutil cp ' + '%s ' * len(files) + '.') % tuple('gs://%s/results%g.txt' % (bucket, x) for x in id)
|
|
||||||
os.system(c)
|
|
||||||
else:
|
|
||||||
files = list(Path(save_dir).glob('results*.txt'))
|
|
||||||
assert len(files), 'No results.txt files found in %s, nothing to plot.' % os.path.abspath(save_dir)
|
|
||||||
for fi, f in enumerate(files):
|
|
||||||
try:
|
|
||||||
results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
|
|
||||||
n = results.shape[1] # number of rows
|
|
||||||
x = range(start, min(stop, n) if stop else n)
|
|
||||||
for i in range(10):
|
|
||||||
y = results[i, x]
|
|
||||||
if i in [0, 1, 2, 5, 6, 7]:
|
|
||||||
y[y == 0] = np.nan # don't show zero loss values
|
|
||||||
# y /= y[0] # normalize
|
|
||||||
label = labels[fi] if len(labels) else f.stem
|
|
||||||
ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8)
|
|
||||||
ax[i].set_title(s[i])
|
|
||||||
# if i in [5, 6, 7]: # share train and val loss y axes
|
|
||||||
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
|
|
||||||
except Exception as e:
|
|
||||||
print('Warning: Plotting error for %s; %s' % (f, e))
|
|
||||||
|
|
||||||
ax[1].legend()
|
|
||||||
fig.savefig(Path(save_dir) / 'results.png', dpi=200)
|
|
||||||
|
|||||||
@ -1,7 +1,9 @@
|
|||||||
# YOLOv3 PyTorch utils
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
PyTorch utils
|
||||||
|
"""
|
||||||
|
|
||||||
import datetime
|
import datetime
|
||||||
import logging
|
|
||||||
import math
|
import math
|
||||||
import os
|
import os
|
||||||
import platform
|
import platform
|
||||||
@ -12,16 +14,16 @@ from copy import deepcopy
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.backends.cudnn as cudnn
|
import torch.distributed as dist
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
import torchvision
|
|
||||||
|
from utils.general import LOGGER
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import thop # for FLOPS computation
|
import thop # for FLOPs computation
|
||||||
except ImportError:
|
except ImportError:
|
||||||
thop = None
|
thop = None
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
@ -30,19 +32,10 @@ def torch_distributed_zero_first(local_rank: int):
|
|||||||
Decorator to make all processes in distributed training wait for each local_master to do something.
|
Decorator to make all processes in distributed training wait for each local_master to do something.
|
||||||
"""
|
"""
|
||||||
if local_rank not in [-1, 0]:
|
if local_rank not in [-1, 0]:
|
||||||
torch.distributed.barrier()
|
dist.barrier(device_ids=[local_rank])
|
||||||
yield
|
yield
|
||||||
if local_rank == 0:
|
if local_rank == 0:
|
||||||
torch.distributed.barrier()
|
dist.barrier(device_ids=[0])
|
||||||
|
|
||||||
|
|
||||||
def init_torch_seeds(seed=0):
|
|
||||||
# Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
|
|
||||||
torch.manual_seed(seed)
|
|
||||||
if seed == 0: # slower, more reproducible
|
|
||||||
cudnn.benchmark, cudnn.deterministic = False, True
|
|
||||||
else: # faster, less reproducible
|
|
||||||
cudnn.benchmark, cudnn.deterministic = True, False
|
|
||||||
|
|
||||||
|
|
||||||
def date_modified(path=__file__):
|
def date_modified(path=__file__):
|
||||||
@ -60,10 +53,11 @@ def git_describe(path=Path(__file__).parent): # path must be a directory
|
|||||||
return '' # not a git repository
|
return '' # not a git repository
|
||||||
|
|
||||||
|
|
||||||
def select_device(device='', batch_size=None):
|
def select_device(device='', batch_size=None, newline=True):
|
||||||
# device = 'cpu' or '0' or '0,1,2,3'
|
# device = 'cpu' or '0' or '0,1,2,3'
|
||||||
s = f'YOLOv3 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
|
s = f'YOLOv3 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
|
||||||
cpu = device.lower() == 'cpu'
|
device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0'
|
||||||
|
cpu = device == 'cpu'
|
||||||
if cpu:
|
if cpu:
|
||||||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
|
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
|
||||||
elif device: # non-cpu device requested
|
elif device: # non-cpu device requested
|
||||||
@ -72,65 +66,80 @@ def select_device(device='', batch_size=None):
|
|||||||
|
|
||||||
cuda = not cpu and torch.cuda.is_available()
|
cuda = not cpu and torch.cuda.is_available()
|
||||||
if cuda:
|
if cuda:
|
||||||
devices = device.split(',') if device else range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
||||||
n = len(devices) # device count
|
n = len(devices) # device count
|
||||||
if n > 1 and batch_size: # check batch_size is divisible by device_count
|
if n > 1 and batch_size: # check batch_size is divisible by device_count
|
||||||
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
|
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
|
||||||
space = ' ' * len(s)
|
space = ' ' * (len(s) + 1)
|
||||||
for i, d in enumerate(devices):
|
for i, d in enumerate(devices):
|
||||||
p = torch.cuda.get_device_properties(i)
|
p = torch.cuda.get_device_properties(i)
|
||||||
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
|
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2:.0f}MiB)\n" # bytes to MB
|
||||||
else:
|
else:
|
||||||
s += 'CPU\n'
|
s += 'CPU\n'
|
||||||
|
|
||||||
logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
|
if not newline:
|
||||||
|
s = s.rstrip()
|
||||||
|
LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
|
||||||
return torch.device('cuda:0' if cuda else 'cpu')
|
return torch.device('cuda:0' if cuda else 'cpu')
|
||||||
|
|
||||||
|
|
||||||
def time_synchronized():
|
def time_sync():
|
||||||
# pytorch-accurate time
|
# pytorch-accurate time
|
||||||
if torch.cuda.is_available():
|
if torch.cuda.is_available():
|
||||||
torch.cuda.synchronize()
|
torch.cuda.synchronize()
|
||||||
return time.time()
|
return time.time()
|
||||||
|
|
||||||
|
|
||||||
def profile(x, ops, n=100, device=None):
|
def profile(input, ops, n=10, device=None):
|
||||||
# profile a pytorch module or list of modules. Example usage:
|
# speed/memory/FLOPs profiler
|
||||||
# x = torch.randn(16, 3, 640, 640) # input
|
#
|
||||||
|
# Usage:
|
||||||
|
# input = torch.randn(16, 3, 640, 640)
|
||||||
# m1 = lambda x: x * torch.sigmoid(x)
|
# m1 = lambda x: x * torch.sigmoid(x)
|
||||||
# m2 = nn.SiLU()
|
# m2 = nn.SiLU()
|
||||||
# profile(x, [m1, m2], n=100) # profile speed over 100 iterations
|
# profile(input, [m1, m2], n=100) # profile over 100 iterations
|
||||||
|
|
||||||
device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
results = []
|
||||||
x = x.to(device)
|
device = device or select_device()
|
||||||
x.requires_grad = True
|
print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
|
||||||
print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '')
|
f"{'input':>24s}{'output':>24s}")
|
||||||
print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}")
|
|
||||||
for m in ops if isinstance(ops, list) else [ops]:
|
|
||||||
m = m.to(device) if hasattr(m, 'to') else m # device
|
|
||||||
m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m # type
|
|
||||||
dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward
|
|
||||||
try:
|
|
||||||
flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPS
|
|
||||||
except:
|
|
||||||
flops = 0
|
|
||||||
|
|
||||||
for _ in range(n):
|
for x in input if isinstance(input, list) else [input]:
|
||||||
t[0] = time_synchronized()
|
x = x.to(device)
|
||||||
y = m(x)
|
x.requires_grad = True
|
||||||
t[1] = time_synchronized()
|
for m in ops if isinstance(ops, list) else [ops]:
|
||||||
|
m = m.to(device) if hasattr(m, 'to') else m # device
|
||||||
|
m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
|
||||||
|
tf, tb, t = 0, 0, [0, 0, 0] # dt forward, backward
|
||||||
try:
|
try:
|
||||||
_ = y.sum().backward()
|
flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPs
|
||||||
t[2] = time_synchronized()
|
except:
|
||||||
except: # no backward method
|
flops = 0
|
||||||
t[2] = float('nan')
|
|
||||||
dtf += (t[1] - t[0]) * 1000 / n # ms per op forward
|
|
||||||
dtb += (t[2] - t[1]) * 1000 / n # ms per op backward
|
|
||||||
|
|
||||||
s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
|
try:
|
||||||
s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
|
for _ in range(n):
|
||||||
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
|
t[0] = time_sync()
|
||||||
print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
|
y = m(x)
|
||||||
|
t[1] = time_sync()
|
||||||
|
try:
|
||||||
|
_ = (sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()
|
||||||
|
t[2] = time_sync()
|
||||||
|
except Exception as e: # no backward method
|
||||||
|
# print(e) # for debug
|
||||||
|
t[2] = float('nan')
|
||||||
|
tf += (t[1] - t[0]) * 1000 / n # ms per op forward
|
||||||
|
tb += (t[2] - t[1]) * 1000 / n # ms per op backward
|
||||||
|
mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0 # (GB)
|
||||||
|
s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
|
||||||
|
s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
|
||||||
|
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
|
||||||
|
print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')
|
||||||
|
results.append([p, flops, mem, tf, tb, s_in, s_out])
|
||||||
|
except Exception as e:
|
||||||
|
print(e)
|
||||||
|
results.append(None)
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
def is_parallel(model):
|
def is_parallel(model):
|
||||||
@ -143,11 +152,6 @@ def de_parallel(model):
|
|||||||
return model.module if is_parallel(model) else model
|
return model.module if is_parallel(model) else model
|
||||||
|
|
||||||
|
|
||||||
def intersect_dicts(da, db, exclude=()):
|
|
||||||
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
|
||||||
return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
|
|
||||||
|
|
||||||
|
|
||||||
def initialize_weights(model):
|
def initialize_weights(model):
|
||||||
for m in model.modules():
|
for m in model.modules():
|
||||||
t = type(m)
|
t = type(m)
|
||||||
@ -156,7 +160,7 @@ def initialize_weights(model):
|
|||||||
elif t is nn.BatchNorm2d:
|
elif t is nn.BatchNorm2d:
|
||||||
m.eps = 1e-3
|
m.eps = 1e-3
|
||||||
m.momentum = 0.03
|
m.momentum = 0.03
|
||||||
elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
|
elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
|
||||||
m.inplace = True
|
m.inplace = True
|
||||||
|
|
||||||
|
|
||||||
@ -167,7 +171,7 @@ def find_modules(model, mclass=nn.Conv2d):
|
|||||||
|
|
||||||
def sparsity(model):
|
def sparsity(model):
|
||||||
# Return global model sparsity
|
# Return global model sparsity
|
||||||
a, b = 0., 0.
|
a, b = 0, 0
|
||||||
for p in model.parameters():
|
for p in model.parameters():
|
||||||
a += p.numel()
|
a += p.numel()
|
||||||
b += (p == 0).sum()
|
b += (p == 0).sum()
|
||||||
@ -213,42 +217,23 @@ def model_info(model, verbose=False, img_size=640):
|
|||||||
n_p = sum(x.numel() for x in model.parameters()) # number parameters
|
n_p = sum(x.numel() for x in model.parameters()) # number parameters
|
||||||
n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
|
n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
|
||||||
if verbose:
|
if verbose:
|
||||||
print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
|
print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")
|
||||||
for i, (name, p) in enumerate(model.named_parameters()):
|
for i, (name, p) in enumerate(model.named_parameters()):
|
||||||
name = name.replace('module_list.', '')
|
name = name.replace('module_list.', '')
|
||||||
print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
|
print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
|
||||||
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
|
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
|
||||||
|
|
||||||
try: # FLOPS
|
try: # FLOPs
|
||||||
from thop import profile
|
from thop import profile
|
||||||
stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
|
stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
|
||||||
img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
|
img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
|
||||||
flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS
|
flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs
|
||||||
img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
|
img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
|
||||||
fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS
|
fs = ', %.1f GFLOPs' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs
|
||||||
except (ImportError, Exception):
|
except (ImportError, Exception):
|
||||||
fs = ''
|
fs = ''
|
||||||
|
|
||||||
logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
||||||
|
|
||||||
|
|
||||||
def load_classifier(name='resnet101', n=2):
|
|
||||||
# Loads a pretrained model reshaped to n-class output
|
|
||||||
model = torchvision.models.__dict__[name](pretrained=True)
|
|
||||||
|
|
||||||
# ResNet model properties
|
|
||||||
# input_size = [3, 224, 224]
|
|
||||||
# input_space = 'RGB'
|
|
||||||
# input_range = [0, 1]
|
|
||||||
# mean = [0.485, 0.456, 0.406]
|
|
||||||
# std = [0.229, 0.224, 0.225]
|
|
||||||
|
|
||||||
# Reshape output to n classes
|
|
||||||
filters = model.fc.weight.shape[1]
|
|
||||||
model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
|
|
||||||
model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
|
|
||||||
model.fc.out_features = n
|
|
||||||
return model
|
|
||||||
|
|
||||||
|
|
||||||
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
||||||
@ -260,7 +245,7 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
|||||||
s = (int(h * ratio), int(w * ratio)) # new size
|
s = (int(h * ratio), int(w * ratio)) # new size
|
||||||
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
||||||
if not same_shape: # pad/crop img
|
if not same_shape: # pad/crop img
|
||||||
h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
|
h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
|
||||||
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
||||||
|
|
||||||
|
|
||||||
@ -273,6 +258,29 @@ def copy_attr(a, b, include=(), exclude=()):
|
|||||||
setattr(a, k, v)
|
setattr(a, k, v)
|
||||||
|
|
||||||
|
|
||||||
|
class EarlyStopping:
|
||||||
|
# simple early stopper
|
||||||
|
def __init__(self, patience=30):
|
||||||
|
self.best_fitness = 0.0 # i.e. mAP
|
||||||
|
self.best_epoch = 0
|
||||||
|
self.patience = patience or float('inf') # epochs to wait after fitness stops improving to stop
|
||||||
|
self.possible_stop = False # possible stop may occur next epoch
|
||||||
|
|
||||||
|
def __call__(self, epoch, fitness):
|
||||||
|
if fitness >= self.best_fitness: # >= 0 to allow for early zero-fitness stage of training
|
||||||
|
self.best_epoch = epoch
|
||||||
|
self.best_fitness = fitness
|
||||||
|
delta = epoch - self.best_epoch # epochs without improvement
|
||||||
|
self.possible_stop = delta >= (self.patience - 1) # possible stop may occur next epoch
|
||||||
|
stop = delta >= self.patience # stop training if patience exceeded
|
||||||
|
if stop:
|
||||||
|
LOGGER.info(f'Stopping training early as no improvement observed in last {self.patience} epochs. '
|
||||||
|
f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n'
|
||||||
|
f'To update EarlyStopping(patience={self.patience}) pass a new patience value, '
|
||||||
|
f'i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping.')
|
||||||
|
return stop
|
||||||
|
|
||||||
|
|
||||||
class ModelEMA:
|
class ModelEMA:
|
||||||
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
|
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
|
||||||
Keep a moving average of everything in the model state_dict (parameters and buffers).
|
Keep a moving average of everything in the model state_dict (parameters and buffers).
|
||||||
@ -303,7 +311,7 @@ class ModelEMA:
|
|||||||
for k, v in self.ema.state_dict().items():
|
for k, v in self.ema.state_dict().items():
|
||||||
if v.dtype.is_floating_point:
|
if v.dtype.is_floating_point:
|
||||||
v *= d
|
v *= d
|
||||||
v += (1. - d) * msd[k].detach()
|
v += (1 - d) * msd[k].detach()
|
||||||
|
|
||||||
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
||||||
# Update EMA attributes
|
# Update EMA attributes
|
||||||
|
|||||||
@ -1,318 +0,0 @@
|
|||||||
"""Utilities and tools for tracking runs with Weights & Biases."""
|
|
||||||
import json
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import yaml
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
|
|
||||||
from utils.datasets import LoadImagesAndLabels
|
|
||||||
from utils.datasets import img2label_paths
|
|
||||||
from utils.general import colorstr, xywh2xyxy, check_dataset, check_file
|
|
||||||
|
|
||||||
try:
|
|
||||||
import wandb
|
|
||||||
from wandb import init, finish
|
|
||||||
except ImportError:
|
|
||||||
wandb = None
|
|
||||||
|
|
||||||
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
|
|
||||||
|
|
||||||
|
|
||||||
def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
|
|
||||||
return from_string[len(prefix):]
|
|
||||||
|
|
||||||
|
|
||||||
def check_wandb_config_file(data_config_file):
|
|
||||||
wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
|
|
||||||
if Path(wandb_config).is_file():
|
|
||||||
return wandb_config
|
|
||||||
return data_config_file
|
|
||||||
|
|
||||||
|
|
||||||
def get_run_info(run_path):
|
|
||||||
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
|
|
||||||
run_id = run_path.stem
|
|
||||||
project = run_path.parent.stem
|
|
||||||
entity = run_path.parent.parent.stem
|
|
||||||
model_artifact_name = 'run_' + run_id + '_model'
|
|
||||||
return entity, project, run_id, model_artifact_name
|
|
||||||
|
|
||||||
|
|
||||||
def check_wandb_resume(opt):
|
|
||||||
process_wandb_config_ddp_mode(opt) if opt.global_rank not in [-1, 0] else None
|
|
||||||
if isinstance(opt.resume, str):
|
|
||||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
if opt.global_rank not in [-1, 0]: # For resuming DDP runs
|
|
||||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
|
||||||
api = wandb.Api()
|
|
||||||
artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
|
|
||||||
modeldir = artifact.download()
|
|
||||||
opt.weights = str(Path(modeldir) / "last.pt")
|
|
||||||
return True
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def process_wandb_config_ddp_mode(opt):
|
|
||||||
with open(check_file(opt.data)) as f:
|
|
||||||
data_dict = yaml.safe_load(f) # data dict
|
|
||||||
train_dir, val_dir = None, None
|
|
||||||
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
api = wandb.Api()
|
|
||||||
train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
|
|
||||||
train_dir = train_artifact.download()
|
|
||||||
train_path = Path(train_dir) / 'data/images/'
|
|
||||||
data_dict['train'] = str(train_path)
|
|
||||||
|
|
||||||
if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
api = wandb.Api()
|
|
||||||
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
|
|
||||||
val_dir = val_artifact.download()
|
|
||||||
val_path = Path(val_dir) / 'data/images/'
|
|
||||||
data_dict['val'] = str(val_path)
|
|
||||||
if train_dir or val_dir:
|
|
||||||
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
|
|
||||||
with open(ddp_data_path, 'w') as f:
|
|
||||||
yaml.safe_dump(data_dict, f)
|
|
||||||
opt.data = ddp_data_path
|
|
||||||
|
|
||||||
|
|
||||||
class WandbLogger():
|
|
||||||
"""Log training runs, datasets, models, and predictions to Weights & Biases.
|
|
||||||
|
|
||||||
This logger sends information to W&B at wandb.ai. By default, this information
|
|
||||||
includes hyperparameters, system configuration and metrics, model metrics,
|
|
||||||
and basic data metrics and analyses.
|
|
||||||
|
|
||||||
By providing additional command line arguments to train.py, datasets,
|
|
||||||
models and predictions can also be logged.
|
|
||||||
|
|
||||||
For more on how this logger is used, see the Weights & Biases documentation:
|
|
||||||
https://docs.wandb.com/guides/integrations/yolov5
|
|
||||||
"""
|
|
||||||
def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
|
|
||||||
# Pre-training routine --
|
|
||||||
self.job_type = job_type
|
|
||||||
self.wandb, self.wandb_run, self.data_dict = wandb, None if not wandb else wandb.run, data_dict
|
|
||||||
# It's more elegant to stick to 1 wandb.init call, but useful config data is overwritten in the WandbLogger's wandb.init call
|
|
||||||
if isinstance(opt.resume, str): # checks resume from artifact
|
|
||||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
|
|
||||||
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
|
|
||||||
assert wandb, 'install wandb to resume wandb runs'
|
|
||||||
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
|
|
||||||
self.wandb_run = wandb.init(id=run_id, project=project, entity=entity, resume='allow')
|
|
||||||
opt.resume = model_artifact_name
|
|
||||||
elif self.wandb:
|
|
||||||
self.wandb_run = wandb.init(config=opt,
|
|
||||||
resume="allow",
|
|
||||||
project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem,
|
|
||||||
entity=opt.entity,
|
|
||||||
name=name,
|
|
||||||
job_type=job_type,
|
|
||||||
id=run_id) if not wandb.run else wandb.run
|
|
||||||
if self.wandb_run:
|
|
||||||
if self.job_type == 'Training':
|
|
||||||
if not opt.resume:
|
|
||||||
wandb_data_dict = self.check_and_upload_dataset(opt) if opt.upload_dataset else data_dict
|
|
||||||
# Info useful for resuming from artifacts
|
|
||||||
self.wandb_run.config.opt = vars(opt)
|
|
||||||
self.wandb_run.config.data_dict = wandb_data_dict
|
|
||||||
self.data_dict = self.setup_training(opt, data_dict)
|
|
||||||
if self.job_type == 'Dataset Creation':
|
|
||||||
self.data_dict = self.check_and_upload_dataset(opt)
|
|
||||||
else:
|
|
||||||
prefix = colorstr('wandb: ')
|
|
||||||
print(f"{prefix}Install Weights & Biases for YOLOv3 logging with 'pip install wandb' (recommended)")
|
|
||||||
|
|
||||||
def check_and_upload_dataset(self, opt):
|
|
||||||
assert wandb, 'Install wandb to upload dataset'
|
|
||||||
check_dataset(self.data_dict)
|
|
||||||
config_path = self.log_dataset_artifact(check_file(opt.data),
|
|
||||||
opt.single_cls,
|
|
||||||
'YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem)
|
|
||||||
print("Created dataset config file ", config_path)
|
|
||||||
with open(config_path) as f:
|
|
||||||
wandb_data_dict = yaml.safe_load(f)
|
|
||||||
return wandb_data_dict
|
|
||||||
|
|
||||||
def setup_training(self, opt, data_dict):
|
|
||||||
self.log_dict, self.current_epoch, self.log_imgs = {}, 0, 16 # Logging Constants
|
|
||||||
self.bbox_interval = opt.bbox_interval
|
|
||||||
if isinstance(opt.resume, str):
|
|
||||||
modeldir, _ = self.download_model_artifact(opt)
|
|
||||||
if modeldir:
|
|
||||||
self.weights = Path(modeldir) / "last.pt"
|
|
||||||
config = self.wandb_run.config
|
|
||||||
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str(
|
|
||||||
self.weights), config.save_period, config.total_batch_size, config.bbox_interval, config.epochs, \
|
|
||||||
config.opt['hyp']
|
|
||||||
data_dict = dict(self.wandb_run.config.data_dict) # eliminates the need for config file to resume
|
|
||||||
if 'val_artifact' not in self.__dict__: # If --upload_dataset is set, use the existing artifact, don't download
|
|
||||||
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),
|
|
||||||
opt.artifact_alias)
|
|
||||||
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
|
|
||||||
opt.artifact_alias)
|
|
||||||
self.result_artifact, self.result_table, self.val_table, self.weights = None, None, None, None
|
|
||||||
if self.train_artifact_path is not None:
|
|
||||||
train_path = Path(self.train_artifact_path) / 'data/images/'
|
|
||||||
data_dict['train'] = str(train_path)
|
|
||||||
if self.val_artifact_path is not None:
|
|
||||||
val_path = Path(self.val_artifact_path) / 'data/images/'
|
|
||||||
data_dict['val'] = str(val_path)
|
|
||||||
self.val_table = self.val_artifact.get("val")
|
|
||||||
self.map_val_table_path()
|
|
||||||
if self.val_artifact is not None:
|
|
||||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
|
||||||
self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
|
|
||||||
if opt.bbox_interval == -1:
|
|
||||||
self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
|
|
||||||
return data_dict
|
|
||||||
|
|
||||||
def download_dataset_artifact(self, path, alias):
|
|
||||||
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
|
|
||||||
dataset_artifact = wandb.use_artifact(artifact_path.as_posix())
|
|
||||||
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
|
|
||||||
datadir = dataset_artifact.download()
|
|
||||||
return datadir, dataset_artifact
|
|
||||||
return None, None
|
|
||||||
|
|
||||||
def download_model_artifact(self, opt):
|
|
||||||
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
|
|
||||||
model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
|
|
||||||
assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
|
|
||||||
modeldir = model_artifact.download()
|
|
||||||
epochs_trained = model_artifact.metadata.get('epochs_trained')
|
|
||||||
total_epochs = model_artifact.metadata.get('total_epochs')
|
|
||||||
is_finished = total_epochs is None
|
|
||||||
assert not is_finished, 'training is finished, can only resume incomplete runs.'
|
|
||||||
return modeldir, model_artifact
|
|
||||||
return None, None
|
|
||||||
|
|
||||||
def log_model(self, path, opt, epoch, fitness_score, best_model=False):
|
|
||||||
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
|
|
||||||
'original_url': str(path),
|
|
||||||
'epochs_trained': epoch + 1,
|
|
||||||
'save period': opt.save_period,
|
|
||||||
'project': opt.project,
|
|
||||||
'total_epochs': opt.epochs,
|
|
||||||
'fitness_score': fitness_score
|
|
||||||
})
|
|
||||||
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
|
|
||||||
wandb.log_artifact(model_artifact,
|
|
||||||
aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
|
|
||||||
print("Saving model artifact on epoch ", epoch + 1)
|
|
||||||
|
|
||||||
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
|
|
||||||
with open(data_file) as f:
|
|
||||||
data = yaml.safe_load(f) # data dict
|
|
||||||
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
|
|
||||||
names = {k: v for k, v in enumerate(names)} # to index dictionary
|
|
||||||
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
|
||||||
data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None
|
|
||||||
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
|
|
||||||
data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
|
|
||||||
if data.get('train'):
|
|
||||||
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
|
|
||||||
if data.get('val'):
|
|
||||||
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
|
|
||||||
path = data_file if overwrite_config else '_wandb.'.join(data_file.rsplit('.', 1)) # updated data.yaml path
|
|
||||||
data.pop('download', None)
|
|
||||||
with open(path, 'w') as f:
|
|
||||||
yaml.safe_dump(data, f)
|
|
||||||
|
|
||||||
if self.job_type == 'Training': # builds correct artifact pipeline graph
|
|
||||||
self.wandb_run.use_artifact(self.val_artifact)
|
|
||||||
self.wandb_run.use_artifact(self.train_artifact)
|
|
||||||
self.val_artifact.wait()
|
|
||||||
self.val_table = self.val_artifact.get('val')
|
|
||||||
self.map_val_table_path()
|
|
||||||
else:
|
|
||||||
self.wandb_run.log_artifact(self.train_artifact)
|
|
||||||
self.wandb_run.log_artifact(self.val_artifact)
|
|
||||||
return path
|
|
||||||
|
|
||||||
def map_val_table_path(self):
|
|
||||||
self.val_table_map = {}
|
|
||||||
print("Mapping dataset")
|
|
||||||
for i, data in enumerate(tqdm(self.val_table.data)):
|
|
||||||
self.val_table_map[data[3]] = data[0]
|
|
||||||
|
|
||||||
def create_dataset_table(self, dataset, class_to_id, name='dataset'):
|
|
||||||
# TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
|
|
||||||
artifact = wandb.Artifact(name=name, type="dataset")
|
|
||||||
img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
|
|
||||||
img_files = tqdm(dataset.img_files) if not img_files else img_files
|
|
||||||
for img_file in img_files:
|
|
||||||
if Path(img_file).is_dir():
|
|
||||||
artifact.add_dir(img_file, name='data/images')
|
|
||||||
labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
|
|
||||||
artifact.add_dir(labels_path, name='data/labels')
|
|
||||||
else:
|
|
||||||
artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
|
|
||||||
label_file = Path(img2label_paths([img_file])[0])
|
|
||||||
artifact.add_file(str(label_file),
|
|
||||||
name='data/labels/' + label_file.name) if label_file.exists() else None
|
|
||||||
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
|
|
||||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
|
|
||||||
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
|
|
||||||
box_data, img_classes = [], {}
|
|
||||||
for cls, *xywh in labels[:, 1:].tolist():
|
|
||||||
cls = int(cls)
|
|
||||||
box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]},
|
|
||||||
"class_id": cls,
|
|
||||||
"box_caption": "%s" % (class_to_id[cls])})
|
|
||||||
img_classes[cls] = class_to_id[cls]
|
|
||||||
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
|
|
||||||
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes),
|
|
||||||
Path(paths).name)
|
|
||||||
artifact.add(table, name)
|
|
||||||
return artifact
|
|
||||||
|
|
||||||
def log_training_progress(self, predn, path, names):
|
|
||||||
if self.val_table and self.result_table:
|
|
||||||
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
|
|
||||||
box_data = []
|
|
||||||
total_conf = 0
|
|
||||||
for *xyxy, conf, cls in predn.tolist():
|
|
||||||
if conf >= 0.25:
|
|
||||||
box_data.append(
|
|
||||||
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
|
||||||
"class_id": int(cls),
|
|
||||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
|
||||||
"scores": {"class_score": conf},
|
|
||||||
"domain": "pixel"})
|
|
||||||
total_conf = total_conf + conf
|
|
||||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
|
||||||
id = self.val_table_map[Path(path).name]
|
|
||||||
self.result_table.add_data(self.current_epoch,
|
|
||||||
id,
|
|
||||||
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
|
|
||||||
total_conf / max(1, len(box_data))
|
|
||||||
)
|
|
||||||
|
|
||||||
def log(self, log_dict):
|
|
||||||
if self.wandb_run:
|
|
||||||
for key, value in log_dict.items():
|
|
||||||
self.log_dict[key] = value
|
|
||||||
|
|
||||||
def end_epoch(self, best_result=False):
|
|
||||||
if self.wandb_run:
|
|
||||||
wandb.log(self.log_dict)
|
|
||||||
self.log_dict = {}
|
|
||||||
if self.result_artifact:
|
|
||||||
train_results = wandb.JoinedTable(self.val_table, self.result_table, "id")
|
|
||||||
self.result_artifact.add(train_results, 'result')
|
|
||||||
wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch),
|
|
||||||
('best' if best_result else '')])
|
|
||||||
self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
|
|
||||||
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
|
|
||||||
|
|
||||||
def finish_run(self):
|
|
||||||
if self.wandb_run:
|
|
||||||
if self.log_dict:
|
|
||||||
wandb.log(self.log_dict)
|
|
||||||
wandb.run.finish()
|
|
||||||
367
val.py
Normal file
367
val.py
Normal file
@ -0,0 +1,367 @@
|
|||||||
|
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
|
||||||
|
"""
|
||||||
|
Validate a trained model accuracy on a custom dataset
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
$ python path/to/val.py --data coco128.yaml --weights yolov3.pt --img 640
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
from threading import Thread
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
FILE = Path(__file__).resolve()
|
||||||
|
ROOT = FILE.parents[0] # root directory
|
||||||
|
if str(ROOT) not in sys.path:
|
||||||
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||||
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||||
|
|
||||||
|
from models.common import DetectMultiBackend
|
||||||
|
from utils.callbacks import Callbacks
|
||||||
|
from utils.datasets import create_dataloader
|
||||||
|
from utils.general import (LOGGER, NCOLS, box_iou, check_dataset, check_img_size, check_requirements, check_yaml,
|
||||||
|
coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args,
|
||||||
|
scale_coords, xywh2xyxy, xyxy2xywh)
|
||||||
|
from utils.metrics import ConfusionMatrix, ap_per_class
|
||||||
|
from utils.plots import output_to_target, plot_images, plot_val_study
|
||||||
|
from utils.torch_utils import select_device, time_sync
|
||||||
|
|
||||||
|
|
||||||
|
def save_one_txt(predn, save_conf, shape, file):
|
||||||
|
# Save one txt result
|
||||||
|
gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||||||
|
for *xyxy, conf, cls in predn.tolist():
|
||||||
|
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||||||
|
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
||||||
|
with open(file, 'a') as f:
|
||||||
|
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
||||||
|
|
||||||
|
|
||||||
|
def save_one_json(predn, jdict, path, class_map):
|
||||||
|
# Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
|
||||||
|
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
|
||||||
|
box = xyxy2xywh(predn[:, :4]) # xywh
|
||||||
|
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
||||||
|
for p, b in zip(predn.tolist(), box.tolist()):
|
||||||
|
jdict.append({'image_id': image_id,
|
||||||
|
'category_id': class_map[int(p[5])],
|
||||||
|
'bbox': [round(x, 3) for x in b],
|
||||||
|
'score': round(p[4], 5)})
|
||||||
|
|
||||||
|
|
||||||
|
def process_batch(detections, labels, iouv):
|
||||||
|
"""
|
||||||
|
Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
|
||||||
|
Arguments:
|
||||||
|
detections (Array[N, 6]), x1, y1, x2, y2, conf, class
|
||||||
|
labels (Array[M, 5]), class, x1, y1, x2, y2
|
||||||
|
Returns:
|
||||||
|
correct (Array[N, 10]), for 10 IoU levels
|
||||||
|
"""
|
||||||
|
correct = torch.zeros(detections.shape[0], iouv.shape[0], dtype=torch.bool, device=iouv.device)
|
||||||
|
iou = box_iou(labels[:, 1:], detections[:, :4])
|
||||||
|
x = torch.where((iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5])) # IoU above threshold and classes match
|
||||||
|
if x[0].shape[0]:
|
||||||
|
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detection, iou]
|
||||||
|
if x[0].shape[0] > 1:
|
||||||
|
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||||
|
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
||||||
|
# matches = matches[matches[:, 2].argsort()[::-1]]
|
||||||
|
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
||||||
|
matches = torch.Tensor(matches).to(iouv.device)
|
||||||
|
correct[matches[:, 1].long()] = matches[:, 2:3] >= iouv
|
||||||
|
return correct
|
||||||
|
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def run(data,
|
||||||
|
weights=None, # model.pt path(s)
|
||||||
|
batch_size=32, # batch size
|
||||||
|
imgsz=640, # inference size (pixels)
|
||||||
|
conf_thres=0.001, # confidence threshold
|
||||||
|
iou_thres=0.6, # NMS IoU threshold
|
||||||
|
task='val', # train, val, test, speed or study
|
||||||
|
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||||
|
single_cls=False, # treat as single-class dataset
|
||||||
|
augment=False, # augmented inference
|
||||||
|
verbose=False, # verbose output
|
||||||
|
save_txt=False, # save results to *.txt
|
||||||
|
save_hybrid=False, # save label+prediction hybrid results to *.txt
|
||||||
|
save_conf=False, # save confidences in --save-txt labels
|
||||||
|
save_json=False, # save a COCO-JSON results file
|
||||||
|
project=ROOT / 'runs/val', # save to project/name
|
||||||
|
name='exp', # save to project/name
|
||||||
|
exist_ok=False, # existing project/name ok, do not increment
|
||||||
|
half=True, # use FP16 half-precision inference
|
||||||
|
dnn=False, # use OpenCV DNN for ONNX inference
|
||||||
|
model=None,
|
||||||
|
dataloader=None,
|
||||||
|
save_dir=Path(''),
|
||||||
|
plots=True,
|
||||||
|
callbacks=Callbacks(),
|
||||||
|
compute_loss=None,
|
||||||
|
):
|
||||||
|
# Initialize/load model and set device
|
||||||
|
training = model is not None
|
||||||
|
if training: # called by train.py
|
||||||
|
device, pt = next(model.parameters()).device, True # get model device, PyTorch model
|
||||||
|
|
||||||
|
half &= device.type != 'cpu' # half precision only supported on CUDA
|
||||||
|
model.half() if half else model.float()
|
||||||
|
else: # called directly
|
||||||
|
device = select_device(device, batch_size=batch_size)
|
||||||
|
|
||||||
|
# Directories
|
||||||
|
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
||||||
|
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
||||||
|
|
||||||
|
# Load model
|
||||||
|
model = DetectMultiBackend(weights, device=device, dnn=dnn)
|
||||||
|
stride, pt = model.stride, model.pt
|
||||||
|
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||||
|
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||||
|
if pt:
|
||||||
|
model.model.half() if half else model.model.float()
|
||||||
|
else:
|
||||||
|
half = False
|
||||||
|
batch_size = 1 # export.py models default to batch-size 1
|
||||||
|
device = torch.device('cpu')
|
||||||
|
LOGGER.info(f'Forcing --batch-size 1 square inference shape(1,3,{imgsz},{imgsz}) for non-PyTorch backends')
|
||||||
|
|
||||||
|
# Data
|
||||||
|
data = check_dataset(data) # check
|
||||||
|
|
||||||
|
# Configure
|
||||||
|
model.eval()
|
||||||
|
is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt') # COCO dataset
|
||||||
|
nc = 1 if single_cls else int(data['nc']) # number of classes
|
||||||
|
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
||||||
|
niou = iouv.numel()
|
||||||
|
|
||||||
|
# Dataloader
|
||||||
|
if not training:
|
||||||
|
if pt and device.type != 'cpu':
|
||||||
|
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.model.parameters()))) # warmup
|
||||||
|
pad = 0.0 if task == 'speed' else 0.5
|
||||||
|
task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
|
||||||
|
dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=pt,
|
||||||
|
prefix=colorstr(f'{task}: '))[0]
|
||||||
|
|
||||||
|
seen = 0
|
||||||
|
confusion_matrix = ConfusionMatrix(nc=nc)
|
||||||
|
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
|
||||||
|
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
|
||||||
|
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
|
||||||
|
dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
|
||||||
|
loss = torch.zeros(3, device=device)
|
||||||
|
jdict, stats, ap, ap_class = [], [], [], []
|
||||||
|
pbar = tqdm(dataloader, desc=s, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
|
||||||
|
for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
|
||||||
|
t1 = time_sync()
|
||||||
|
if pt:
|
||||||
|
im = im.to(device, non_blocking=True)
|
||||||
|
targets = targets.to(device)
|
||||||
|
im = im.half() if half else im.float() # uint8 to fp16/32
|
||||||
|
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||||
|
nb, _, height, width = im.shape # batch size, channels, height, width
|
||||||
|
t2 = time_sync()
|
||||||
|
dt[0] += t2 - t1
|
||||||
|
|
||||||
|
# Inference
|
||||||
|
out, train_out = model(im) if training else model(im, augment=augment, val=True) # inference, loss outputs
|
||||||
|
dt[1] += time_sync() - t2
|
||||||
|
|
||||||
|
# Loss
|
||||||
|
if compute_loss:
|
||||||
|
loss += compute_loss([x.float() for x in train_out], targets)[1] # box, obj, cls
|
||||||
|
|
||||||
|
# NMS
|
||||||
|
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
||||||
|
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
||||||
|
t3 = time_sync()
|
||||||
|
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
|
||||||
|
dt[2] += time_sync() - t3
|
||||||
|
|
||||||
|
# Metrics
|
||||||
|
for si, pred in enumerate(out):
|
||||||
|
labels = targets[targets[:, 0] == si, 1:]
|
||||||
|
nl = len(labels)
|
||||||
|
tcls = labels[:, 0].tolist() if nl else [] # target class
|
||||||
|
path, shape = Path(paths[si]), shapes[si][0]
|
||||||
|
seen += 1
|
||||||
|
|
||||||
|
if len(pred) == 0:
|
||||||
|
if nl:
|
||||||
|
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Predictions
|
||||||
|
if single_cls:
|
||||||
|
pred[:, 5] = 0
|
||||||
|
predn = pred.clone()
|
||||||
|
scale_coords(im[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred
|
||||||
|
|
||||||
|
# Evaluate
|
||||||
|
if nl:
|
||||||
|
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
|
||||||
|
scale_coords(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
|
||||||
|
labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
|
||||||
|
correct = process_batch(predn, labelsn, iouv)
|
||||||
|
if plots:
|
||||||
|
confusion_matrix.process_batch(predn, labelsn)
|
||||||
|
else:
|
||||||
|
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
|
||||||
|
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) # (correct, conf, pcls, tcls)
|
||||||
|
|
||||||
|
# Save/log
|
||||||
|
if save_txt:
|
||||||
|
save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt'))
|
||||||
|
if save_json:
|
||||||
|
save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
|
||||||
|
callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
|
||||||
|
|
||||||
|
# Plot images
|
||||||
|
if plots and batch_i < 3:
|
||||||
|
f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
|
||||||
|
Thread(target=plot_images, args=(im, targets, paths, f, names), daemon=True).start()
|
||||||
|
f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
|
||||||
|
Thread(target=plot_images, args=(im, output_to_target(out), paths, f, names), daemon=True).start()
|
||||||
|
|
||||||
|
# Compute metrics
|
||||||
|
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
|
||||||
|
if len(stats) and stats[0].any():
|
||||||
|
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
|
||||||
|
ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
|
||||||
|
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
|
||||||
|
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
|
||||||
|
else:
|
||||||
|
nt = torch.zeros(1)
|
||||||
|
|
||||||
|
# Print results
|
||||||
|
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
|
||||||
|
LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
|
||||||
|
|
||||||
|
# Print results per class
|
||||||
|
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
|
||||||
|
for i, c in enumerate(ap_class):
|
||||||
|
LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
|
||||||
|
|
||||||
|
# Print speeds
|
||||||
|
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
|
||||||
|
if not training:
|
||||||
|
shape = (batch_size, 3, imgsz, imgsz)
|
||||||
|
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
|
||||||
|
|
||||||
|
# Plots
|
||||||
|
if plots:
|
||||||
|
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
|
||||||
|
callbacks.run('on_val_end')
|
||||||
|
|
||||||
|
# Save JSON
|
||||||
|
if save_json and len(jdict):
|
||||||
|
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
|
||||||
|
anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json
|
||||||
|
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
|
||||||
|
LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
|
||||||
|
with open(pred_json, 'w') as f:
|
||||||
|
json.dump(jdict, f)
|
||||||
|
|
||||||
|
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
|
||||||
|
check_requirements(['pycocotools'])
|
||||||
|
from pycocotools.coco import COCO
|
||||||
|
from pycocotools.cocoeval import COCOeval
|
||||||
|
|
||||||
|
anno = COCO(anno_json) # init annotations api
|
||||||
|
pred = anno.loadRes(pred_json) # init predictions api
|
||||||
|
eval = COCOeval(anno, pred, 'bbox')
|
||||||
|
if is_coco:
|
||||||
|
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
|
||||||
|
eval.evaluate()
|
||||||
|
eval.accumulate()
|
||||||
|
eval.summarize()
|
||||||
|
map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
|
||||||
|
except Exception as e:
|
||||||
|
LOGGER.info(f'pycocotools unable to run: {e}')
|
||||||
|
|
||||||
|
# Return results
|
||||||
|
model.float() # for training
|
||||||
|
if not training:
|
||||||
|
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
||||||
|
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
|
||||||
|
maps = np.zeros(nc) + map
|
||||||
|
for i, c in enumerate(ap_class):
|
||||||
|
maps[c] = ap[i]
|
||||||
|
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
|
||||||
|
|
||||||
|
|
||||||
|
def parse_opt():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
||||||
|
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov3.pt', help='model.pt path(s)')
|
||||||
|
parser.add_argument('--batch-size', type=int, default=32, help='batch size')
|
||||||
|
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
|
||||||
|
parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
|
||||||
|
parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
|
||||||
|
parser.add_argument('--task', default='val', help='train, val, test, speed or study')
|
||||||
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||||
|
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
|
||||||
|
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||||
|
parser.add_argument('--verbose', action='store_true', help='report mAP by class')
|
||||||
|
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
||||||
|
parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
|
||||||
|
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
||||||
|
parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
|
||||||
|
parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
|
||||||
|
parser.add_argument('--name', default='exp', help='save to project/name')
|
||||||
|
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
||||||
|
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
||||||
|
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
||||||
|
opt = parser.parse_args()
|
||||||
|
opt.data = check_yaml(opt.data) # check YAML
|
||||||
|
opt.save_json |= opt.data.endswith('coco.yaml')
|
||||||
|
opt.save_txt |= opt.save_hybrid
|
||||||
|
print_args(FILE.stem, opt)
|
||||||
|
return opt
|
||||||
|
|
||||||
|
|
||||||
|
def main(opt):
|
||||||
|
check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
|
||||||
|
|
||||||
|
if opt.task in ('train', 'val', 'test'): # run normally
|
||||||
|
if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
|
||||||
|
LOGGER.info(f'WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.')
|
||||||
|
run(**vars(opt))
|
||||||
|
|
||||||
|
else:
|
||||||
|
weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
|
||||||
|
opt.half = True # FP16 for fastest results
|
||||||
|
if opt.task == 'speed': # speed benchmarks
|
||||||
|
# python val.py --task speed --data coco.yaml --batch 1 --weights yolov3.pt yolov3-spp.pt...
|
||||||
|
opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
|
||||||
|
for opt.weights in weights:
|
||||||
|
run(**vars(opt), plots=False)
|
||||||
|
|
||||||
|
elif opt.task == 'study': # speed vs mAP benchmarks
|
||||||
|
# python val.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt...
|
||||||
|
for opt.weights in weights:
|
||||||
|
f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt' # filename to save to
|
||||||
|
x, y = list(range(256, 1536 + 128, 128)), [] # x axis (image sizes), y axis
|
||||||
|
for opt.imgsz in x: # img-size
|
||||||
|
LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
|
||||||
|
r, _, t = run(**vars(opt), plots=False)
|
||||||
|
y.append(r + t) # results and times
|
||||||
|
np.savetxt(f, y, fmt='%10.4g') # save
|
||||||
|
os.system('zip -r study.zip study_*.txt')
|
||||||
|
plot_val_study(x=x) # plot
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
opt = parse_opt()
|
||||||
|
main(opt)
|
||||||
@ -1,12 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# Download latest models from https://github.com/ultralytics/yolov3/releases
|
|
||||||
# Usage:
|
|
||||||
# $ bash weights/download_weights.sh
|
|
||||||
|
|
||||||
python - <<EOF
|
|
||||||
from utils.google_utils import attempt_download
|
|
||||||
|
|
||||||
for x in ['yolov3', 'yolov3-spp', 'yolov3-tiny']:
|
|
||||||
attempt_download(f'{x}.pt')
|
|
||||||
|
|
||||||
EOF
|
|
||||||
Loading…
x
Reference in New Issue
Block a user