├── .dockerignore ├── .gitattributes ├── .github ├── ISSUE_TEMPLATE │ ├── bug-report.md │ ├── feature-request.md │ └── question.md ├── dependabot.yml └── workflows │ ├── ci-testing.yml │ ├── codeql-analysis.yml │ ├── greetings.yml │ ├── rebase.yml │ └── stale.yml ├── .gitignore ├── Dockerfile ├── LICENSE ├── README.md ├── convert ├── config.yaml ├── dataset 1.txt └── rknn_convert.py ├── detect.py ├── hubconf.py ├── models ├── __init__.py ├── common.py ├── common_rk_plug_in.py ├── experimental.py ├── export.py ├── hub │ ├── anchors.yaml │ ├── yolov3-spp.yaml │ ├── yolov3-tiny.yaml │ ├── yolov3.yaml │ ├── yolov5-fpn.yaml │ ├── yolov5-p2.yaml │ ├── yolov5-p6.yaml │ ├── yolov5-p7.yaml │ ├── yolov5-panet.yaml │ ├── yolov5l6.yaml │ ├── yolov5m6.yaml │ ├── yolov5s6.yaml │ └── yolov5x6.yaml ├── yolo.py ├── yolov5l.yaml ├── yolov5m.yaml ├── yolov5s.yaml └── yolov5x.yaml ├── pre_compile └── rknn2precompile.py ├── requirements.txt ├── rknn_detect ├── models │ └── yolov5_rknn_640x640.yaml └── rknn_detect_for_yolov5_original.py ├── test.py ├── train.py ├── tutorial.ipynb ├── utils ├── __init__.py ├── activations.py ├── autoanchor.py ├── aws │ ├── __init__.py │ ├── mime.sh │ ├── resume.py │ └── userdata.sh ├── datasets.py ├── general.py ├── google_app_engine │ ├── Dockerfile │ ├── additional_requirements.txt │ └── app.yaml ├── google_utils.py ├── loss.py ├── metrics.py ├── plots.py ├── torch_utils.py └── wandb_logging │ ├── __init__.py │ ├── log_dataset.py │ └── wandb_utils.py └── weights └── download_weights.sh /.dockerignore: -------------------------------------------------------------------------------- 1 | # Repo-specific DockerIgnore ------------------------------------------------------------------------------------------- 2 | #.git 3 | .cache 4 | .idea 5 | runs 6 | output 7 | coco 8 | storage.googleapis.com 9 | 10 | data/samples/* 11 | **/results*.txt 12 | *.jpg 13 | 14 | # Neural Network weights ----------------------------------------------------------------------------------------------- 15 | **/*.weights 16 | **/*.pt 17 | **/*.pth 18 | **/*.onnx 19 | **/*.mlmodel 20 | **/*.torchscript 21 | 22 | 23 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 24 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 25 | 26 | 27 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 28 | # Byte-compiled / optimized / DLL files 29 | __pycache__/ 30 | *.py[cod] 31 | *$py.class 32 | 33 | # C extensions 34 | *.so 35 | 36 | # Distribution / packaging 37 | .Python 38 | env/ 39 | build/ 40 | develop-eggs/ 41 | dist/ 42 | downloads/ 43 | eggs/ 44 | .eggs/ 45 | lib/ 46 | lib64/ 47 | parts/ 48 | sdist/ 49 | var/ 50 | wheels/ 51 | *.egg-info/ 52 | wandb/ 53 | .installed.cfg 54 | *.egg 55 | 56 | # PyInstaller 57 | # Usually these files are written by a python script from a template 58 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 59 | *.manifest 60 | *.spec 61 | 62 | # Installer logs 63 | pip-log.txt 64 | pip-delete-this-directory.txt 65 | 66 | # Unit test / coverage reports 67 | htmlcov/ 68 | .tox/ 69 | .coverage 70 | .coverage.* 71 | .cache 72 | nosetests.xml 73 | coverage.xml 74 | *.cover 75 | .hypothesis/ 76 | 77 | # Translations 78 | *.mo 79 | *.pot 80 | 81 | # Django stuff: 82 | *.log 83 | local_settings.py 84 | 85 | # Flask stuff: 86 | instance/ 87 | .webassets-cache 88 | 89 | # Scrapy stuff: 90 | .scrapy 91 | 92 | # Sphinx documentation 93 | docs/_build/ 94 | 95 | # PyBuilder 96 | target/ 97 | 98 | # Jupyter Notebook 99 | .ipynb_checkpoints 100 | 101 | # pyenv 102 | .python-version 103 | 104 | # celery beat schedule file 105 | celerybeat-schedule 106 | 107 | # SageMath parsed files 108 | *.sage.py 109 | 110 | # dotenv 111 | .env 112 | 113 | # virtualenv 114 | .venv* 115 | venv*/ 116 | ENV*/ 117 | 118 | # Spyder project settings 119 | .spyderproject 120 | .spyproject 121 | 122 | # Rope project settings 123 | .ropeproject 124 | 125 | # mkdocs documentation 126 | /site 127 | 128 | # mypy 129 | .mypy_cache/ 130 | 131 | 132 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 133 | 134 | # General 135 | .DS_Store 136 | .AppleDouble 137 | .LSOverride 138 | 139 | # Icon must end with two \r 140 | Icon 141 | Icon? 142 | 143 | # Thumbnails 144 | ._* 145 | 146 | # Files that might appear in the root of a volume 147 | .DocumentRevisions-V100 148 | .fseventsd 149 | .Spotlight-V100 150 | .TemporaryItems 151 | .Trashes 152 | .VolumeIcon.icns 153 | .com.apple.timemachine.donotpresent 154 | 155 | # Directories potentially created on remote AFP share 156 | .AppleDB 157 | .AppleDesktop 158 | Network Trash Folder 159 | Temporary Items 160 | .apdisk 161 | 162 | 163 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 164 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 165 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 166 | 167 | # User-specific stuff: 168 | .idea/* 169 | .idea/**/workspace.xml 170 | .idea/**/tasks.xml 171 | .idea/dictionaries 172 | .html # Bokeh Plots 173 | .pg # TensorFlow Frozen Graphs 174 | .avi # videos 175 | 176 | # Sensitive or high-churn files: 177 | .idea/**/dataSources/ 178 | .idea/**/dataSources.ids 179 | .idea/**/dataSources.local.xml 180 | .idea/**/sqlDataSources.xml 181 | .idea/**/dynamic.xml 182 | .idea/**/uiDesigner.xml 183 | 184 | # Gradle: 185 | .idea/**/gradle.xml 186 | .idea/**/libraries 187 | 188 | # CMake 189 | cmake-build-debug/ 190 | cmake-build-release/ 191 | 192 | # Mongo Explorer plugin: 193 | .idea/**/mongoSettings.xml 194 | 195 | ## File-based project format: 196 | *.iws 197 | 198 | ## Plugin-specific files: 199 | 200 | # IntelliJ 201 | out/ 202 | 203 | # mpeltonen/sbt-idea plugin 204 | .idea_modules/ 205 | 206 | # JIRA plugin 207 | atlassian-ide-plugin.xml 208 | 209 | # Cursive Clojure plugin 210 | .idea/replstate.xml 211 | 212 | # Crashlytics plugin (for Android Studio and IntelliJ) 213 | com_crashlytics_export_strings.xml 214 | crashlytics.properties 215 | crashlytics-build.properties 216 | fabric.properties 217 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # this drop notebooks from GitHub language stats 2 | *.ipynb linguist-vendored 3 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug-report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "🐛 Bug report" 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | 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: 11 | - **Current repo**: run `git fetch && git status -uno` to check and `git pull` to update repo 12 | - **Common dataset**: coco.yaml or coco128.yaml 13 | - **Common environment**: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov5#environments 14 | 15 | 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()`. 16 | 17 | 18 | ## 🐛 Bug 19 | A clear and concise description of what the bug is. 20 | 21 | 22 | ## To Reproduce (REQUIRED) 23 | 24 | Input: 25 | ``` 26 | import torch 27 | 28 | a = torch.tensor([5]) 29 | c = a / 0 30 | ``` 31 | 32 | Output: 33 | ``` 34 | Traceback (most recent call last): 35 | File "/Users/glennjocher/opt/anaconda3/envs/env1/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code 36 | exec(code_obj, self.user_global_ns, self.user_ns) 37 | File "", line 5, in 38 | c = a / 0 39 | RuntimeError: ZeroDivisionError 40 | ``` 41 | 42 | 43 | ## Expected behavior 44 | A clear and concise description of what you expected to happen. 45 | 46 | 47 | ## Environment 48 | If applicable, add screenshots to help explain your problem. 49 | 50 | - OS: [e.g. Ubuntu] 51 | - GPU [e.g. 2080 Ti] 52 | 53 | 54 | ## Additional context 55 | Add any other context about the problem here. 56 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "🚀 Feature request" 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: enhancement 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## 🚀 Feature 11 | 12 | 13 | ## Motivation 14 | 15 | 16 | 17 | ## Pitch 18 | 19 | 20 | 21 | ## Alternatives 22 | 23 | 24 | 25 | ## Additional context 26 | 27 | 28 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "❓Question" 3 | about: Ask a general question 4 | title: '' 5 | labels: question 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## ❔Question 11 | 12 | 13 | ## Additional context 14 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: pip 4 | directory: "/" 5 | schedule: 6 | interval: weekly 7 | time: "04:00" 8 | open-pull-requests-limit: 10 9 | reviewers: 10 | - glenn-jocher 11 | labels: 12 | - dependencies 13 | -------------------------------------------------------------------------------- /.github/workflows/ci-testing.yml: -------------------------------------------------------------------------------- 1 | name: CI CPU testing 2 | 3 | on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | # The branches below must be a subset of the branches above 8 | branches: [ master ] 9 | schedule: 10 | - cron: '0 0 * * *' # Runs at 00:00 UTC every day 11 | 12 | jobs: 13 | cpu-tests: 14 | 15 | runs-on: ${{ matrix.os }} 16 | strategy: 17 | fail-fast: false 18 | matrix: 19 | os: [ubuntu-latest, macos-latest, windows-latest] 20 | python-version: [3.8] 21 | model: ['yolov5s'] # models to test 22 | 23 | # Timeout: https://stackoverflow.com/a/59076067/4521646 24 | timeout-minutes: 50 25 | steps: 26 | - uses: actions/checkout@v2 27 | - name: Set up Python ${{ matrix.python-version }} 28 | uses: actions/setup-python@v2 29 | with: 30 | python-version: ${{ matrix.python-version }} 31 | 32 | # Note: This uses an internal pip API and may not always work 33 | # https://github.com/actions/cache/blob/master/examples.md#multiple-oss-in-a-workflow 34 | - name: Get pip cache 35 | id: pip-cache 36 | run: | 37 | python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)" 38 | 39 | - name: Cache pip 40 | uses: actions/cache@v1 41 | with: 42 | path: ${{ steps.pip-cache.outputs.dir }} 43 | key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }} 44 | restore-keys: | 45 | ${{ runner.os }}-${{ matrix.python-version }}-pip- 46 | 47 | - name: Install dependencies 48 | run: | 49 | python -m pip install --upgrade pip 50 | pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html 51 | pip install -q onnx 52 | python --version 53 | pip --version 54 | pip list 55 | shell: bash 56 | 57 | - name: Download data 58 | run: | 59 | # curl -L -o tmp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 60 | # unzip -q tmp.zip -d ../ 61 | # rm tmp.zip 62 | 63 | - name: Tests workflow 64 | run: | 65 | # export PYTHONPATH="$PWD" # to run '$ python *.py' files in subdirectories 66 | di=cpu # inference devices # define device 67 | 68 | # train 69 | python train.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di 70 | # detect 71 | python detect.py --weights weights/${{ matrix.model }}.pt --device $di 72 | python detect.py --weights runs/train/exp/weights/last.pt --device $di 73 | # test 74 | python test.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --device $di 75 | python test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di 76 | 77 | python hubconf.py # hub 78 | python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect 79 | python models/export.py --img 128 --batch 1 --weights weights/${{ matrix.model }}.pt # export 80 | shell: bash 81 | -------------------------------------------------------------------------------- /.github/workflows/codeql-analysis.yml: -------------------------------------------------------------------------------- 1 | # This action runs GitHub's industry-leading static analysis engine, CodeQL, against a repository's source code to find security vulnerabilities. 2 | # https://github.com/github/codeql-action 3 | 4 | name: "CodeQL" 5 | 6 | on: 7 | schedule: 8 | - cron: '0 0 1 * *' # Runs at 00:00 UTC on the 1st of every month 9 | 10 | jobs: 11 | analyze: 12 | name: Analyze 13 | runs-on: ubuntu-latest 14 | 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | language: [ 'python' ] 19 | # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ] 20 | # Learn more: 21 | # 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 22 | 23 | steps: 24 | - name: Checkout repository 25 | uses: actions/checkout@v2 26 | 27 | # Initializes the CodeQL tools for scanning. 28 | - name: Initialize CodeQL 29 | uses: github/codeql-action/init@v1 30 | with: 31 | languages: ${{ matrix.language }} 32 | # If you wish to specify custom queries, you can do so here or in a config file. 33 | # By default, queries listed here will override any specified in a config file. 34 | # Prefix the list here with "+" to use these queries and those in the config file. 35 | # queries: ./path/to/local/query, your-org/your-repo/queries@main 36 | 37 | # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). 38 | # If this step fails, then you should remove it and run the build manually (see below) 39 | - name: Autobuild 40 | uses: github/codeql-action/autobuild@v1 41 | 42 | # ℹ️ Command-line programs to run using the OS shell. 43 | # 📚 https://git.io/JvXDl 44 | 45 | # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines 46 | # and modify them (or add more) to build your code if your project 47 | # uses a compiled language 48 | 49 | #- run: | 50 | # make bootstrap 51 | # make release 52 | 53 | - name: Perform CodeQL Analysis 54 | uses: github/codeql-action/analyze@v1 55 | -------------------------------------------------------------------------------- /.github/workflows/greetings.yml: -------------------------------------------------------------------------------- 1 | name: Greetings 2 | 3 | on: [pull_request_target, issues] 4 | 5 | jobs: 6 | greeting: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/first-interaction@v1 10 | with: 11 | repo-token: ${{ secrets.GITHUB_TOKEN }} 12 | pr-message: | 13 | 👋 Hello @${{ github.actor }}, thank you for submitting a 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to: 14 | - ✅ Verify your PR is **up-to-date with origin/master.** If your PR is behind origin/master an automatic [GitHub actions](https://github.com/ultralytics/yolov5/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: 15 | ```bash 16 | git remote add upstream https://github.com/ultralytics/yolov5.git 17 | git fetch upstream 18 | git checkout feature # <----- replace 'feature' with local branch name 19 | git rebase upstream/master 20 | git push -u origin -f 21 | ``` 22 | - ✅ Verify all Continuous Integration (CI) **checks are passing**. 23 | - ✅ 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 24 | 25 | issue-message: | 26 | 👋 Hello @${{ github.actor }}, thank you for your interest in 🚀 YOLOv5! Please visit our ⭐️ [Tutorials](https://github.com/ultralytics/yolov5/wiki#tutorials) to get started, where you can find quickstart guides for simple tasks like [Custom Data Training](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) all the way to advanced concepts like [Hyperparameter Evolution](https://github.com/ultralytics/yolov5/issues/607). 27 | 28 | If this is a 🐛 Bug Report, please provide screenshots and **minimum viable code to reproduce your issue**, otherwise we can not help you. 29 | 30 | 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/yolov5/wiki/Train-Custom-Data#visualize) if available. 31 | 32 | For business inquiries or professional support requests please visit https://www.ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com. 33 | 34 | ## Requirements 35 | 36 | Python 3.8 or later with all [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) dependencies installed, including `torch>=1.7`. To install run: 37 | ```bash 38 | $ pip install -r requirements.txt 39 | ``` 40 | 41 | ## Environments 42 | 43 | YOLOv5 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): 44 | 45 | - **Google Colab and Kaggle** notebooks with free GPU: Open In Colab Open In Kaggle 46 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) 47 | - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/AWS-Quickstart) 48 | - **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) Docker Pulls 49 | 50 | 51 | ## Status 52 | 53 | ![CI CPU testing](https://github.com/ultralytics/yolov5/workflows/CI%20CPU%20testing/badge.svg) 54 | 55 | If this badge is green, all [YOLOv5 GitHub Actions](https://github.com/ultralytics/yolov5/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training ([train.py](https://github.com/ultralytics/yolov5/blob/master/train.py)), testing ([test.py](https://github.com/ultralytics/yolov5/blob/master/test.py)), inference ([detect.py](https://github.com/ultralytics/yolov5/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov5/blob/master/models/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit. 56 | 57 | -------------------------------------------------------------------------------- /.github/workflows/rebase.yml: -------------------------------------------------------------------------------- 1 | name: Automatic Rebase 2 | # https://github.com/marketplace/actions/automatic-rebase 3 | 4 | on: 5 | issue_comment: 6 | types: [created] 7 | 8 | jobs: 9 | rebase: 10 | name: Rebase 11 | if: github.event.issue.pull_request != '' && contains(github.event.comment.body, '/rebase') 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout the latest code 15 | uses: actions/checkout@v2 16 | with: 17 | fetch-depth: 0 18 | - name: Automatic Rebase 19 | uses: cirrus-actions/rebase@1.3.1 20 | env: 21 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 22 | -------------------------------------------------------------------------------- /.github/workflows/stale.yml: -------------------------------------------------------------------------------- 1 | name: Close stale issues 2 | on: 3 | schedule: 4 | - cron: "0 0 * * *" 5 | 6 | jobs: 7 | stale: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/stale@v3 11 | with: 12 | repo-token: ${{ secrets.GITHUB_TOKEN }} 13 | stale-issue-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.' 14 | 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.' 15 | days-before-stale: 30 16 | days-before-close: 5 17 | exempt-issue-labels: 'documentation,tutorial' 18 | operations-per-run: 100 # The maximum number of operations per run, used to control rate limiting. 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Repo-specific GitIgnore ---------------------------------------------------------------------------------------------- 2 | *.jpg 3 | *.jpeg 4 | *.png 5 | *.bmp 6 | *.tif 7 | *.tiff 8 | *.heic 9 | *.JPG 10 | *.JPEG 11 | *.PNG 12 | *.BMP 13 | *.TIF 14 | *.TIFF 15 | *.HEIC 16 | *.mp4 17 | *.mov 18 | *.MOV 19 | *.avi 20 | *.data 21 | *.json 22 | 23 | *.cfg 24 | !cfg/yolov3*.cfg 25 | 26 | storage.googleapis.com 27 | runs/* 28 | data/* 29 | !data/images/zidane.jpg 30 | !data/images/bus.jpg 31 | !data/coco.names 32 | !data/coco_paper.names 33 | !data/coco.data 34 | !data/coco_*.data 35 | !data/coco_*.txt 36 | !data/trainvalno5k.shapes 37 | !data/*.sh 38 | 39 | pycocotools/* 40 | results*.txt 41 | gcp_test*.sh 42 | 43 | # Datasets ------------------------------------------------------------------------------------------------------------- 44 | coco/ 45 | coco128/ 46 | VOC/ 47 | 48 | # MATLAB GitIgnore ----------------------------------------------------------------------------------------------------- 49 | *.m~ 50 | *.mat 51 | !targets*.mat 52 | 53 | # Neural Network weights ----------------------------------------------------------------------------------------------- 54 | *.weights 55 | *.pt 56 | *.onnx 57 | *.mlmodel 58 | *.torchscript 59 | darknet53.conv.74 60 | yolov3-tiny.conv.15 61 | 62 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 63 | # Byte-compiled / optimized / DLL files 64 | __pycache__/ 65 | *.py[cod] 66 | *$py.class 67 | 68 | # C extensions 69 | *.so 70 | 71 | # Distribution / packaging 72 | .Python 73 | env/ 74 | build/ 75 | develop-eggs/ 76 | dist/ 77 | downloads/ 78 | eggs/ 79 | .eggs/ 80 | lib/ 81 | lib64/ 82 | parts/ 83 | sdist/ 84 | var/ 85 | wheels/ 86 | *.egg-info/ 87 | wandb/ 88 | .installed.cfg 89 | *.egg 90 | 91 | 92 | # PyInstaller 93 | # Usually these files are written by a python script from a template 94 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 95 | *.manifest 96 | *.spec 97 | 98 | # Installer logs 99 | pip-log.txt 100 | pip-delete-this-directory.txt 101 | 102 | # Unit test / coverage reports 103 | htmlcov/ 104 | .tox/ 105 | .coverage 106 | .coverage.* 107 | .cache 108 | nosetests.xml 109 | coverage.xml 110 | *.cover 111 | .hypothesis/ 112 | 113 | # Translations 114 | *.mo 115 | *.pot 116 | 117 | # Django stuff: 118 | *.log 119 | local_settings.py 120 | 121 | # Flask stuff: 122 | instance/ 123 | .webassets-cache 124 | 125 | # Scrapy stuff: 126 | .scrapy 127 | 128 | # Sphinx documentation 129 | docs/_build/ 130 | 131 | # PyBuilder 132 | target/ 133 | 134 | # Jupyter Notebook 135 | .ipynb_checkpoints 136 | 137 | # pyenv 138 | .python-version 139 | 140 | # celery beat schedule file 141 | celerybeat-schedule 142 | 143 | # SageMath parsed files 144 | *.sage.py 145 | 146 | # dotenv 147 | .env 148 | 149 | # virtualenv 150 | .venv* 151 | venv*/ 152 | ENV*/ 153 | 154 | # Spyder project settings 155 | .spyderproject 156 | .spyproject 157 | 158 | # Rope project settings 159 | .ropeproject 160 | 161 | # mkdocs documentation 162 | /site 163 | 164 | # mypy 165 | .mypy_cache/ 166 | 167 | 168 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 169 | 170 | # General 171 | .DS_Store 172 | .AppleDouble 173 | .LSOverride 174 | 175 | # Icon must end with two \r 176 | Icon 177 | Icon? 178 | 179 | # Thumbnails 180 | ._* 181 | 182 | # Files that might appear in the root of a volume 183 | .DocumentRevisions-V100 184 | .fseventsd 185 | .Spotlight-V100 186 | .TemporaryItems 187 | .Trashes 188 | .VolumeIcon.icns 189 | .com.apple.timemachine.donotpresent 190 | 191 | # Directories potentially created on remote AFP share 192 | .AppleDB 193 | .AppleDesktop 194 | Network Trash Folder 195 | Temporary Items 196 | .apdisk 197 | 198 | 199 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 200 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 201 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 202 | 203 | # User-specific stuff: 204 | .idea/* 205 | .idea/**/workspace.xml 206 | .idea/**/tasks.xml 207 | .idea/dictionaries 208 | .html # Bokeh Plots 209 | .pg # TensorFlow Frozen Graphs 210 | .avi # videos 211 | 212 | # Sensitive or high-churn files: 213 | .idea/**/dataSources/ 214 | .idea/**/dataSources.ids 215 | .idea/**/dataSources.local.xml 216 | .idea/**/sqlDataSources.xml 217 | .idea/**/dynamic.xml 218 | .idea/**/uiDesigner.xml 219 | 220 | # Gradle: 221 | .idea/**/gradle.xml 222 | .idea/**/libraries 223 | 224 | # CMake 225 | cmake-build-debug/ 226 | cmake-build-release/ 227 | 228 | # Mongo Explorer plugin: 229 | .idea/**/mongoSettings.xml 230 | 231 | ## File-based project format: 232 | *.iws 233 | 234 | ## Plugin-specific files: 235 | 236 | # IntelliJ 237 | out/ 238 | 239 | # mpeltonen/sbt-idea plugin 240 | .idea_modules/ 241 | 242 | # JIRA plugin 243 | atlassian-ide-plugin.xml 244 | 245 | # Cursive Clojure plugin 246 | .idea/replstate.xml 247 | 248 | # Crashlytics plugin (for Android Studio and IntelliJ) 249 | com_crashlytics_export_strings.xml 250 | crashlytics.properties 251 | crashlytics-build.properties 252 | fabric.properties 253 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch 2 | FROM nvcr.io/nvidia/pytorch:21.02-py3 3 | 4 | # Install linux packages 5 | RUN apt update && apt install -y zip htop screen libgl1-mesa-glx 6 | 7 | # Install python dependencies 8 | COPY requirements.txt . 9 | RUN python -m pip install --upgrade pip 10 | RUN pip install --no-cache -r requirements.txt gsutil notebook 11 | 12 | # Create working directory 13 | RUN mkdir -p /usr/src/app 14 | WORKDIR /usr/src/app 15 | 16 | # Copy contents 17 | COPY . /usr/src/app 18 | 19 | # Set environment variables 20 | ENV HOME=/usr/src/app 21 | 22 | 23 | # --------------------------------------------------- Extras Below --------------------------------------------------- 24 | 25 | # Build and Push 26 | # t=ultralytics/yolov5:latest && sudo docker build -t $t . && sudo docker push $t 27 | # for v in {300..303}; do t=ultralytics/coco:v$v && sudo docker build -t $t . && sudo docker push $t; done 28 | 29 | # Pull and Run 30 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t 31 | 32 | # Pull and Run with local directory access 33 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/coco:/usr/src/coco $t 34 | 35 | # Kill all 36 | # sudo docker kill $(sudo docker ps -q) 37 | 38 | # Kill all image-based 39 | # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest) 40 | 41 | # Bash into running container 42 | # sudo docker exec -it 5a9b5863d93d bash 43 | 44 | # Bash into stopped container 45 | # id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash 46 | 47 | # Send weights to GCP 48 | # python -c "from utils.general import *; strip_optimizer('runs/train/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt 49 | 50 | # Clean up 51 | # docker system prune -a --volumes 52 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Description： 2 | 3 | Original repo：https://github.com/ultralytics/yolov5 4 | 5 | https://github.com/EASY-EAI/yolov5 6 | 7 | Requirement：python version >= 3.6 8 | 9 | Training：python3 train.py 10 | 11 | Model exporting：python3 models/export.py --rknn_mode 12 | 13 | Model inference：python3 detect.py --rknn_mode 14 | 15 | 16 | 17 | ## Explanation： 18 | 19 | The activation layer in the common file is changed to ReLU, and the model structure, training, testing and other operations are the same as the original version of Yolov5(4.0release). Add rknn_mode mode for model testing and exporting to export rknn-friendly models. (Based on opset_version=10, rknn_toolkit_1.6.0 test passed) 20 | 21 | Detail： 22 | 23 | - onnx.slice reports errors when loading models in rknn_toolkit_1.6.0. Add equivalent replacement operation. (convolution) 24 | 25 | 26 | 27 | ***4.29 update***： 28 | 29 | - Export models can optionally remove the trailing permute layer to be compatible with **rknn_yolov5_demo's c++ deployment code** 30 | 31 | `python3 models/export.py --rknn_mode --ignore_output_permute` 32 | 33 | - The export model can optionally add a **image preprocessing layer**, which can **effectively reduce the time consumption of the deployment segment rknn_input_set**. See the description of preprocess_conv_layer in models/common_rk_plug_in.py for details on how to use it. 34 | 35 | `python3 models/export.py --rknn_mode --add_image_preprocess_layer` 36 | 37 | （rknn_mode、ignore_output_permute、add_image_preprocess_layer . The three are not mutually exclusive and can be used simultaneously） 38 | 39 | - Add onnx->rknn model export tool, see rknn_convert_tools folder for details. 40 | 41 | - Add pre-compiled code to reduce the loading time of the model 42 | 43 | - Add model inference code in RK3399Pro 44 | 45 | ##### *5.12 update*： 46 | 47 | - When exporting the model using `--rknn_mode`, the large `maxpool` is equivalently replaced by multiple smaller `maxpools` by default, which has no effect on the computational results, but can significantly improve the inference speed on rknpu. 48 | 49 | 50 | 51 | ## Description of known problems (have no affect on current use)： 52 | 53 | - onnx.opset_version=12 does not support SiLU activation layer, add equivalent alternative model to solve it. (x* sigmoid(x)) But rknn_toolkit_1_6_0 works fine in simulations, deploying to the board side will cause an exception. Default is not used for now, waiting for **rockchip** to fix. 54 | - onnx.upsample.opset_version=12 Implementation in rknn_toolkit_1.6.0 Temporarily problematic, add equivalence replacement model. (deconvolution). rknn_toolkit_1_6_0 works fine in simulation, deploying to the board side results in an exception. Default is not used for now, waiting for **rockchip.inc** to fix. 55 | 56 | 57 | 58 | ------ 59 | 60 | ### rk_npu speed test[4](#脚注4) (ms)： 61 | 62 | | Model(416x416 input) | rknn_toolkit_1.6.0 Simulators (800MHZ)_rv1109 | rv1109[3](#脚注3) | rv1126 | rv1126(Model pre-compiling) | rknn_toolkit_1.6.0 Simulators (800MHZ)_rk1808 | rk1808 | rk1808(Model pre-compiling) | 63 | | :---------------------- | :-----------------------------------------: | :-------: | :-----: | :----------------: | :-----------------------------------------: | :----: | :----------------: | 64 | | yolov5s_int8[1](#脚注1) | 92 | 113 | 80 | 77 | 89 | 83 | 81 | 65 | | yolov5s_int8_optimize[2](#脚注2) | **18** | **45** | **36** | **33** | **15** | **30** | **29** | 66 | | yolov5s_int16 | 149 | 160 | 110 | 108 | 106 | 178 | 174 | 67 | | yolov5s_int16_optimize | 76 | 90 | 67 | 64 | 32 | 126 | 122 | 68 | | yolov5m_int8 | 158 | 192 | 132 | 120 | 144 | 132 | 123 | 69 | | yolov5m_int8_optimize | **47** | **88** | **66** | **55** | **33** | **54** | **45** | 70 | | yolov5m_int16 | 312 | 302 | 212 | 202 | 187 | 432 | 418 | 71 | | yolov5m_int16_optimize | 202 | 198 | 147 | 137 | 76 | 354 | 344 | 72 | | yolov5l_int8 | 246 | 293 | 199 | | 214 | 192 | | 73 | | yolov5l_int8_optimize | **98** | **155** | **110** | | **66** | **88** | | 74 | | yolov5l_int16 | 577 | 522 | 362 | | 301 | 697 | | 75 | | yolov5l_int16_optimize | 432 | 384 | 275 | | 154 | 592 | | 76 | 77 | 1: is based on the original yaml configuration, with the activation layer modified to relu. 78 | 79 | 2: optimize means to optimize the large size maxpool when exporting the model, now open source, used by default when exporting the parameter --rknn_mode. It does not affect the accuracy. 80 | 81 | 3: Statistical time includes **rknn_inputs_set**, **rknn_run**, **rknn_outputs_get** three parts of time, excluding post-processing time on the cpu side. This principle is followed for the tests on other platforms in this table except for the simulator evaluation. 82 | 83 | 4: This test is for reference only, the test is a single-threaded loop execution timing, only test npu efficiency. The actual use should consider the post-processing time. 84 | 85 | 86 | 87 | 88 | 89 | 90 | ## Reference： 91 | 92 | https://github.com/soloIife/yolov5_for_rknn 93 | 94 | https://github.com/ultralytics/yolov5 95 | 96 | https://github.com/EASY-EAI/yolov5 97 | 98 | RKNN QQ group: 1025468710 99 | 100 | 101 | 102 | 103 | 104 | -------------------------------------------------------------------------------- /convert/config.yaml: -------------------------------------------------------------------------------- 1 | running: 2 | model_type: onnx 3 | export: True 4 | inference: True 5 | eval_perf: True 6 | 7 | 8 | parameters: 9 | caffe: 10 | model: './mobilenet_v2.prototxt' 11 | proto: 'caffe' #lstm_caffe 12 | blobs: './mobilenet_v2.caffemodel' 13 | 14 | tensorflow: 15 | tf_pb: './ssd_mobilenet_v1_coco_2017_11_17.pb' 16 | inputs: ['FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_0/BatchNorm/batchnorm/mul_1'] 17 | outputs: ['concat', 'concat_1'] 18 | input_size_list: [[300, 300, 3]] 19 | 20 | tflite: 21 | model: './sample/tflite/mobilenet_v1/mobilenet_v1.tflite' 22 | 23 | onnx: 24 | model: './best_noop.onnx' #best_op.onnx #best_noop.onnx 25 | 26 | #C:\Users\HP\Desktop\CODE\yolov5_for_rknn-master\weights\best.onnx 27 | 28 | darknet: 29 | model: './yolov3-tiny.cfg' 30 | weight: './yolov3.weights' 31 | 32 | pytorch: 33 | model: './yolov5.pt' 34 | input_size_list: [[3, 512, 512]] 35 | 36 | mxnet: 37 | symbol: 'resnext50_32x4d-symbol.json' 38 | params: 'resnext50_32x4d-4ecf62e2.params' 39 | input_size_list: [[3, 224, 224]] 40 | 41 | rknn: 42 | path: './bestrk.rknn' 43 | 44 | config: 45 | #mean_value: [[0,0,0]] 46 | #std_value: [[58.82,58.82,58.82]] 47 | channel_mean_value: '0 0 0 255' # 123.675 116.28 103.53 58.395 # 0 0 0 255 48 | reorder_channel: '0 1 2' # '2 1 0' 49 | need_horizontal_merge: False 50 | batch_size: 1 51 | epochs: -1 52 | target_platform: ['rk3399pro'] 53 | quantized_dtype: 'asymmetric_quantized-u8' 54 | #asymmetric_quantized-u8,dynamic_fixed_point-8,dynamic_fixed_point-16 55 | optimization_level: 1 56 | 57 | build: 58 | do_quantization: True 59 | dataset: './dataset.txt' # '/home/zen/rknn_convert/quant_data/hand_dataset/pic_path_less.txt' 60 | pre_compile: False 61 | 62 | export_rknn: 63 | export_path: './best_noop1.rknn' 64 | 65 | init_runtime: 66 | target: rk3399pro 67 | device_id: null 68 | perf_debug: False 69 | eval_mem: False 70 | async_mode: False 71 | 72 | img: &img 73 | path: './test2.jpg' 74 | 75 | inference: 76 | inputs: *img 77 | data_type: 'uint8' 78 | data_format: 'nhwc' # 'nchw', 'nhwc' 79 | inputs_pass_through: None 80 | 81 | eval_perf: 82 | inputs: *img 83 | data_type: 'uint8' 84 | data_format: 'nhwc' 85 | is_print: True 86 | 87 | 88 | 89 | 90 | -------------------------------------------------------------------------------- /convert/dataset 1.txt: -------------------------------------------------------------------------------- 1 | 201.jpg -------------------------------------------------------------------------------- /convert/rknn_convert.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from rknn.api import RKNN 3 | import cv2 4 | 5 | _model_load_dict = { 6 | 'caffe': 'load_caffe', 7 | 'tensorflow': 'load_tensorflow', 8 | 'tflite': 'load_tflite', 9 | 'onnx': 'load_onnx', 10 | 'darknet': 'load_darknet', 11 | 'pytorch': 'load_pytorch', 12 | 'mxnet': 'load_mxnet', 13 | 'rknn': 'load_rknn', 14 | } 15 | 16 | yaml_file = './config.yaml' 17 | 18 | 19 | def main(): 20 | with open(yaml_file, 'r') as F: 21 | config = yaml.load(F) 22 | # print('config is:') 23 | # print(config) 24 | 25 | model_type = config['running']['model_type'] 26 | print('model_type is {}'.format(model_type))#检查模型的类型 27 | 28 | rknn = RKNN(verbose=True) 29 | 30 | 31 | 32 | #配置文件 33 | print('--> config model') 34 | rknn.config(**config['config']) 35 | print('done') 36 | 37 | 38 | print('--> Loading model') 39 | load_function = getattr(rknn, _model_load_dict[model_type]) 40 | ret = load_function(**config['parameters'][model_type]) 41 | if ret != 0: 42 | print('Load yolo failed! Ret = {}'.format(ret)) 43 | exit(ret) 44 | print('done') 45 | 46 | #### 47 | #print('hybrid_quantization') 48 | #ret = rknn.hybrid_quantization_step1(dataset=config['build']['dataset']) 49 | 50 | 51 | if model_type != 'rknn': 52 | print('--> Building model') 53 | ret = rknn.build(**config['build']) 54 | print('acc_eval') 55 | rknn.accuracy_analysis(inputs='./dataset1.txt', target='rk3399pro') 56 | print('acc_eval done!') 57 | 58 | if ret != 0: 59 | print('Build yolo failed!') 60 | exit(ret) 61 | else: 62 | print('--> skip Building model step, cause the model is already rknn') 63 | 64 | 65 | #导出RKNN模型 66 | if config['running']['export'] is True: 67 | print('--> Export RKNN model') 68 | ret = rknn.export_rknn(**config['export_rknn']) 69 | if ret != 0: 70 | print('Init runtime environment failed') 71 | exit(ret) 72 | else: 73 | print('--> skip Export model') 74 | 75 | 76 | #初始化 77 | print('--> Init runtime environment') 78 | ret = rknn.init_runtime(**config['init_runtime']) 79 | if ret != 0: 80 | print('Init runtime environment failed') 81 | exit(ret) 82 | print('done') 83 | 84 | 85 | print('--> load img') 86 | img = cv2.imread(config['img']['path']) 87 | print('img shape is {}'.format(img.shape)) 88 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 89 | inputs = [img] 90 | print(inputs[0][0:10,0,0]) 91 | #推理 92 | if config['running']['inference'] is True: 93 | print('--> Running model') 94 | config['inference']['inputs'] = inputs 95 | #print(config['inference']) 96 | outputs = rknn.inference(inputs) 97 | #outputs = rknn.inference(config['inference']) 98 | print('len of output {}'.format(len(outputs))) 99 | print('outputs[0] shape is {}'.format(outputs[0].shape)) 100 | print(outputs[0][0][0:2]) 101 | else: 102 | print('--> skip inference') 103 | #评价 104 | if config['running']['eval_perf'] is True: 105 | print('--> Begin evaluate model performance') 106 | config['inference']['inputs'] = inputs 107 | perf_results = rknn.eval_perf(inputs=[img]) 108 | else: 109 | print('--> skip eval_perf') 110 | 111 | 112 | if __name__ == '__main__': 113 | main() 114 | -------------------------------------------------------------------------------- /detect.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import time 3 | from pathlib import Path 4 | 5 | import cv2 6 | import torch 7 | import torch.nn as nn 8 | import torch.backends.cudnn as cudnn 9 | from numpy import random 10 | 11 | from models.experimental import attempt_load 12 | from utils.datasets import LoadStreams, LoadImages 13 | from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \ 14 | scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path 15 | from utils.plots import plot_one_box 16 | from utils.torch_utils import select_device, load_classifier, time_synchronized 17 | 18 | 19 | def detect(save_img=False): 20 | source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size 21 | webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( 22 | ('rtsp://', 'rtmp://', 'http://')) 23 | 24 | # Directories 25 | save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run 26 | (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir 27 | 28 | # Initialize 29 | set_logging() 30 | device = select_device(opt.device) 31 | half = device.type != 'cpu' # half precision only supported on CUDA 32 | 33 | # Load model 34 | model = attempt_load(weights, map_location=device) # load FP32 model 35 | stride = int(model.stride.max()) # model stride 36 | imgsz = check_img_size(imgsz, s=stride) # check img_size 37 | 38 | # Second-stage classifier 39 | classify = False 40 | if classify: 41 | modelc = load_classifier(name='resnet101', n=2) # initialize 42 | modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval() 43 | 44 | if opt.rknn_mode == True: 45 | print('model convert to rknn_mode') 46 | from models.common_rk_plug_in import surrogate_silu, surrogate_hardswish 47 | from models import common 48 | for k, m in model.named_modules(): 49 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 50 | if isinstance(m, common.Conv): # assign export-friendly activations 51 | if isinstance(m.act, torch.nn.Hardswish): 52 | m.act = torch.nn.Hardswish() 53 | elif isinstance(m.act, torch.nn.SiLU): 54 | # m.act = torch.nn.SiLU() 55 | m.act = surrogate_silu() 56 | # elif isinstance(m, models.yolo.Detect): 57 | # m.forward = m.forward_export # assign forward (optional) 58 | 59 | if isinstance(m, common.SPP): # assign export-friendly activations 60 | ### best 61 | # tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(2)]) 62 | # m.m[0] = tmp 63 | # m.m[1] = tmp 64 | # m.m[2] = tmp 65 | ### friendly to origin config 66 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(2)]) 67 | m.m[0] = tmp 68 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(4)]) 69 | m.m[1] = tmp 70 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(6)]) 71 | m.m[2] = tmp 72 | 73 | ### use deconv2d to surrogate upsample layer. 74 | # replace_one = torch.nn.ConvTranspose2d(model.model[10].conv.weight.shape[0], 75 | # model.model[10].conv.weight.shape[0], 76 | # (2, 2), 77 | # groups=model.model[10].conv.weight.shape[0], 78 | # bias=False, 79 | # stride=(2, 2)) 80 | # replace_one.weight.data.fill_(1) 81 | # replace_one.eval().to(device) 82 | # temp_i = model.model[11].i 83 | # temp_f = model.model[11].f 84 | # model.model[11] = replace_one 85 | # model.model[11].i = temp_i 86 | # model.model[11].f = temp_f 87 | 88 | # replace_one = torch.nn.ConvTranspose2d(model.model[14].conv.weight.shape[0], 89 | # model.model[14].conv.weight.shape[0], 90 | # (2, 2), 91 | # groups=model.model[14].conv.weight.shape[0], 92 | # bias=False, 93 | # stride=(2, 2)) 94 | # replace_one.weight.data.fill_(1) 95 | # replace_one.eval().to(device) 96 | # temp_i = model.model[11].i 97 | # temp_f = model.model[11].f 98 | # model.model[15] = replace_one 99 | # model.model[15].i = temp_i 100 | # model.model[15].f = temp_f 101 | 102 | ### use conv to surrogate slice operator 103 | from models.common_rk_plug_in import surrogate_focus 104 | surrogate_focous = surrogate_focus(int(model.model[0].conv.conv.weight.shape[1]/4), 105 | model.model[0].conv.conv.weight.shape[0], 106 | k=tuple(model.model[0].conv.conv.weight.shape[2:4]), 107 | s=model.model[0].conv.conv.stride, 108 | p=model.model[0].conv.conv.padding, 109 | g=model.model[0].conv.conv.groups, 110 | act=True) 111 | surrogate_focous.conv.conv.weight = model.model[0].conv.conv.weight 112 | surrogate_focous.conv.conv.bias = model.model[0].conv.conv.bias 113 | surrogate_focous.conv.act = model.model[0].conv.act 114 | temp_i = model.model[0].i 115 | temp_f = model.model[0].f 116 | 117 | model.model[0] = surrogate_focous 118 | model.model[0].i = temp_i 119 | model.model[0].f = temp_f 120 | model.model[0].eval().to(device) 121 | 122 | if half: 123 | model.half() # to FP16 124 | 125 | # Set Dataloader 126 | vid_path, vid_writer = None, None 127 | if webcam: 128 | view_img = check_imshow() 129 | cudnn.benchmark = True # set True to speed up constant image size inference 130 | dataset = LoadStreams(source, img_size=imgsz, stride=stride) 131 | else: 132 | save_img = True 133 | dataset = LoadImages(source, img_size=imgsz, stride=stride) 134 | 135 | # Get names and colors 136 | names = model.module.names if hasattr(model, 'module') else model.names 137 | print('names', names) 138 | colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] 139 | 140 | # Run inference 141 | if device.type != 'cpu': 142 | model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once 143 | t0 = time.time() 144 | for path, img, im0s, vid_cap in dataset: 145 | img = torch.from_numpy(img).to(device) 146 | img = img.half() if half else img.float() # uint8 to fp16/32 147 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 148 | if img.ndimension() == 3: 149 | img = img.unsqueeze(0) 150 | 151 | # Inference 152 | t1 = time_synchronized() 153 | pred = model(img, augment=opt.augment)[0] 154 | 155 | # Apply NMS 156 | pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) 157 | t2 = time_synchronized() 158 | 159 | # Apply Classifier 160 | if classify: 161 | pred = apply_classifier(pred, modelc, img, im0s) 162 | 163 | # Process detections 164 | for i, det in enumerate(pred): # detections per image 165 | if webcam: # batch_size >= 1 166 | p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count 167 | else: 168 | p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) 169 | 170 | p = Path(p) # to Path 171 | save_path = str(save_dir / p.name) # img.jpg 172 | txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt 173 | s += '%gx%g ' % img.shape[2:] # print string 174 | gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 175 | if len(det): 176 | # Rescale boxes from img_size to im0 size 177 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() 178 | 179 | # Print results 180 | for c in det[:, -1].unique(): 181 | n = (det[:, -1] == c).sum() # detections per class 182 | s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string 183 | 184 | # Write results 185 | for *xyxy, conf, cls in reversed(det): 186 | if save_txt: # Write to file 187 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 188 | line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format 189 | with open(txt_path + '.txt', 'a') as f: 190 | f.write(('%g ' * len(line)).rstrip() % line + '\n') 191 | 192 | if save_img or view_img: # Add bbox to image 193 | label = f'{names[int(cls)]} {conf:.2f}' 194 | plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=2) 195 | 196 | # Print time (inference + NMS) 197 | print(f'{s}Done. ({t2 - t1:.3f}s)') 198 | 199 | # Stream results 200 | if view_img: 201 | cv2.imshow(str(p), im0) 202 | cv2.waitKey(1) # 1 millisecond 203 | 204 | # Save results (image with detections) 205 | if save_img: 206 | if dataset.mode == 'image': 207 | cv2.imwrite(save_path, im0) 208 | else: # 'video' 209 | if vid_path != save_path: # new video 210 | vid_path = save_path 211 | if isinstance(vid_writer, cv2.VideoWriter): 212 | vid_writer.release() # release previous video writer 213 | 214 | fourcc = 'mp4v' # output video codec 215 | fps = vid_cap.get(cv2.CAP_PROP_FPS) 216 | w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 217 | h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 218 | vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) 219 | vid_writer.write(im0) 220 | 221 | if save_txt or save_img: 222 | s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' 223 | print(f"Results saved to {save_dir}{s}") 224 | 225 | print(f'Done. ({time.time() - t0:.3f}s)') 226 | 227 | 228 | if __name__ == '__main__': 229 | parser = argparse.ArgumentParser() 230 | parser.add_argument('--weights', nargs='+', type=str, default='./runs/train/exp/weights/best.pt', help='model.pt path(s)') 231 | parser.add_argument('--source', type=str, default=r'D:\data\liftcar_tuyang\convert_416x416\val_data\liftcar_tuyang_1.jpg', help='source') # file/folder, 0 for webcam 232 | parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)') 233 | parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold') 234 | parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS') 235 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 236 | parser.add_argument('--view-img', action='store_true', help='display results') 237 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 238 | parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') 239 | parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 240 | parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 241 | parser.add_argument('--augment', action='store_true', help='augmented inference') 242 | parser.add_argument('--update', action='store_true', help='update all models') 243 | parser.add_argument('--project', default='runs/detect', help='save results to project/name') 244 | parser.add_argument('--name', default='exp', help='save results to project/name') 245 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 246 | parser.add_argument('--rknn_mode', action='store_true', help='export rknn-friendly onnx model') 247 | opt = parser.parse_args() 248 | print(opt) 249 | # check_requirements(exclude=('pycocotools', 'thop')) 250 | 251 | with torch.no_grad(): 252 | if opt.update: # update all models (to fix SourceChangeWarning) 253 | for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 254 | detect() 255 | strip_optimizer(opt.weights) 256 | else: 257 | detect() 258 | -------------------------------------------------------------------------------- /hubconf.py: -------------------------------------------------------------------------------- 1 | """File for accessing YOLOv5 models via PyTorch Hub https://pytorch.org/hub/ultralytics_yolov5/ 2 | 3 | Usage: 4 | import torch 5 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 6 | """ 7 | 8 | from pathlib import Path 9 | 10 | import torch 11 | 12 | from models.yolo import Model 13 | from utils.general import check_requirements, set_logging 14 | from utils.google_utils import attempt_download 15 | from utils.torch_utils import select_device 16 | 17 | dependencies = ['torch', 'yaml'] 18 | check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop')) 19 | set_logging() 20 | 21 | 22 | def create(name, pretrained, channels, classes, autoshape): 23 | """Creates a specified YOLOv5 model 24 | 25 | Arguments: 26 | name (str): name of model, i.e. 'yolov5s' 27 | pretrained (bool): load pretrained weights into the model 28 | channels (int): number of input channels 29 | classes (int): number of model classes 30 | 31 | Returns: 32 | pytorch model 33 | """ 34 | config = Path(__file__).parent / 'models' / f'{name}.yaml' # model.yaml path 35 | try: 36 | model = Model(config, channels, classes) 37 | if pretrained: 38 | fname = f'{name}.pt' # checkpoint filename 39 | attempt_download(fname) # download if not found locally 40 | ckpt = torch.load(fname, map_location=torch.device('cpu')) # load 41 | state_dict = ckpt['model'].float().state_dict() # to FP32 42 | state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter 43 | model.load_state_dict(state_dict, strict=False) # load 44 | if len(ckpt['model'].names) == classes: 45 | model.names = ckpt['model'].names # set class names attribute 46 | if autoshape: 47 | model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS 48 | device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available 49 | return model.to(device) 50 | 51 | except Exception as e: 52 | help_url = 'https://github.com/ultralytics/yolov5/issues/36' 53 | s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url 54 | raise Exception(s) from e 55 | 56 | 57 | def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True): 58 | """YOLOv5-small model from https://github.com/ultralytics/yolov5 59 | 60 | Arguments: 61 | pretrained (bool): load pretrained weights into the model, default=False 62 | channels (int): number of input channels, default=3 63 | classes (int): number of model classes, default=80 64 | 65 | Returns: 66 | pytorch model 67 | """ 68 | return create('yolov5s', pretrained, channels, classes, autoshape) 69 | 70 | 71 | def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True): 72 | """YOLOv5-medium model from https://github.com/ultralytics/yolov5 73 | 74 | Arguments: 75 | pretrained (bool): load pretrained weights into the model, default=False 76 | channels (int): number of input channels, default=3 77 | classes (int): number of model classes, default=80 78 | 79 | Returns: 80 | pytorch model 81 | """ 82 | return create('yolov5m', pretrained, channels, classes, autoshape) 83 | 84 | 85 | def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True): 86 | """YOLOv5-large model from https://github.com/ultralytics/yolov5 87 | 88 | Arguments: 89 | pretrained (bool): load pretrained weights into the model, default=False 90 | channels (int): number of input channels, default=3 91 | classes (int): number of model classes, default=80 92 | 93 | Returns: 94 | pytorch model 95 | """ 96 | return create('yolov5l', pretrained, channels, classes, autoshape) 97 | 98 | 99 | def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True): 100 | """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5 101 | 102 | Arguments: 103 | pretrained (bool): load pretrained weights into the model, default=False 104 | channels (int): number of input channels, default=3 105 | classes (int): number of model classes, default=80 106 | 107 | Returns: 108 | pytorch model 109 | """ 110 | return create('yolov5x', pretrained, channels, classes, autoshape) 111 | 112 | 113 | def custom(path_or_model='path/to/model.pt', autoshape=True): 114 | """YOLOv5-custom model from https://github.com/ultralytics/yolov5 115 | 116 | Arguments (3 options): 117 | path_or_model (str): 'path/to/model.pt' 118 | path_or_model (dict): torch.load('path/to/model.pt') 119 | path_or_model (nn.Module): torch.load('path/to/model.pt')['model'] 120 | 121 | Returns: 122 | pytorch model 123 | """ 124 | model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint 125 | if isinstance(model, dict): 126 | model = model['ema' if model.get('ema') else 'model'] # load model 127 | 128 | hub_model = Model(model.yaml).to(next(model.parameters()).device) # create 129 | hub_model.load_state_dict(model.float().state_dict()) # load state_dict 130 | hub_model.names = model.names # class names 131 | return hub_model.autoshape() if autoshape else hub_model 132 | 133 | 134 | if __name__ == '__main__': 135 | model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example 136 | # model = custom(path_or_model='path/to/model.pt') # custom example 137 | 138 | # Verify inference 139 | import numpy as np 140 | from PIL import Image 141 | 142 | imgs = [Image.open('data/images/bus.jpg'), # PIL 143 | 'data/images/zidane.jpg', # filename 144 | 'https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', # URI 145 | np.zeros((640, 480, 3))] # numpy 146 | 147 | results = model(imgs) # batched inference 148 | results.print() 149 | results.save() 150 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/littledeep/YOLOv5-RK3399Pro/37a81a703ba7f82261c09d4316a55182ae84a248/models/__init__.py -------------------------------------------------------------------------------- /models/common.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | 3 | import math 4 | from pathlib import Path 5 | 6 | import numpy as np 7 | import requests 8 | import torch 9 | import torch.nn as nn 10 | from PIL import Image 11 | 12 | from utils.datasets import letterbox 13 | from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh 14 | from utils.plots import color_list, plot_one_box 15 | from utils.torch_utils import time_synchronized 16 | 17 | 18 | def autopad(k, p=None): # kernel, padding 19 | # Pad to 'same' 20 | if p is None: 21 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 22 | return p 23 | 24 | 25 | def DWConv(c1, c2, k=1, s=1, act=True): 26 | # Depthwise convolution 27 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) 28 | 29 | 30 | class Conv(nn.Module): 31 | # Standard convolution 32 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 33 | super(Conv, self).__init__() 34 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) 35 | self.bn = nn.BatchNorm2d(c2) 36 | # self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity()) 37 | self.act = nn.ReLU() 38 | 39 | def forward(self, x): 40 | return self.act(self.bn(self.conv(x))) 41 | 42 | def fuseforward(self, x): 43 | return self.act(self.conv(x)) 44 | 45 | 46 | class Bottleneck(nn.Module): 47 | # Standard bottleneck 48 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion 49 | super(Bottleneck, self).__init__() 50 | c_ = int(c2 * e) # hidden channels 51 | self.cv1 = Conv(c1, c_, 1, 1) 52 | self.cv2 = Conv(c_, c2, 3, 1, g=g) 53 | self.add = shortcut and c1 == c2 54 | 55 | def forward(self, x): 56 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 57 | 58 | 59 | class BottleneckCSP(nn.Module): 60 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks 61 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 62 | super(BottleneckCSP, self).__init__() 63 | c_ = int(c2 * e) # hidden channels 64 | self.cv1 = Conv(c1, c_, 1, 1) 65 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 66 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 67 | self.cv4 = Conv(2 * c_, c2, 1, 1) 68 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 69 | # self.act = nn.LeakyReLU(0.1, inplace=True) 70 | self.act = nn.ReLU() 71 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 72 | 73 | def forward(self, x): 74 | y1 = self.cv3(self.m(self.cv1(x))) 75 | y2 = self.cv2(x) 76 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 77 | 78 | 79 | class C3(nn.Module): 80 | # CSP Bottleneck with 3 convolutions 81 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 82 | super(C3, self).__init__() 83 | c_ = int(c2 * e) # hidden channels 84 | self.cv1 = Conv(c1, c_, 1, 1) 85 | self.cv2 = Conv(c1, c_, 1, 1) 86 | self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2) 87 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 88 | # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) 89 | 90 | def forward(self, x): 91 | return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) 92 | 93 | 94 | class SPP(nn.Module): 95 | # Spatial pyramid pooling layer used in YOLOv3-SPP 96 | def __init__(self, c1, c2, k=(5, 9, 13)): 97 | super(SPP, self).__init__() 98 | c_ = c1 // 2 # hidden channels 99 | self.cv1 = Conv(c1, c_, 1, 1) 100 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) 101 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) 102 | 103 | def forward(self, x): 104 | x = self.cv1(x) 105 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) 106 | 107 | 108 | class Focus(nn.Module): 109 | # Focus wh information into c-space 110 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 111 | super(Focus, self).__init__() 112 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 113 | # self.contract = Contract(gain=2) 114 | 115 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 116 | return self.conv(torch.cat([x[:, :, ::2, ::2], x[:, :, 1::2, ::2], x[:, :, ::2, 1::2], x[:, :, 1::2, 1::2]], 1)) 117 | 118 | 119 | class Contract(nn.Module): 120 | # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40) 121 | def __init__(self, gain=2): 122 | super().__init__() 123 | self.gain = gain 124 | 125 | def forward(self, x): 126 | N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain' 127 | s = self.gain 128 | x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2) 129 | x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40) 130 | return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40) 131 | 132 | 133 | class Expand(nn.Module): 134 | # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160) 135 | def __init__(self, gain=2): 136 | super().__init__() 137 | self.gain = gain 138 | 139 | def forward(self, x): 140 | N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain' 141 | s = self.gain 142 | x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80) 143 | x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2) 144 | return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160) 145 | 146 | 147 | class Concat(nn.Module): 148 | # Concatenate a list of tensors along dimension 149 | def __init__(self, dimension=1): 150 | super(Concat, self).__init__() 151 | self.d = dimension 152 | 153 | def forward(self, x): 154 | return torch.cat(x, self.d) 155 | 156 | 157 | class NMS(nn.Module): 158 | # Non-Maximum Suppression (NMS) module 159 | conf = 0.25 # confidence threshold 160 | iou = 0.45 # IoU threshold 161 | classes = None # (optional list) filter by class 162 | 163 | def __init__(self): 164 | super(NMS, self).__init__() 165 | 166 | def forward(self, x): 167 | return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) 168 | 169 | 170 | class autoShape(nn.Module): 171 | # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS 172 | conf = 0.25 # NMS confidence threshold 173 | iou = 0.45 # NMS IoU threshold 174 | classes = None # (optional list) filter by class 175 | 176 | def __init__(self, model): 177 | super(autoShape, self).__init__() 178 | self.model = model.eval() 179 | 180 | def autoshape(self): 181 | print('autoShape already enabled, skipping... ') # model already converted to model.autoshape() 182 | return self 183 | 184 | def forward(self, imgs, size=640, augment=False, profile=False): 185 | # Inference from various sources. For height=720, width=1280, RGB images example inputs are: 186 | # filename: imgs = 'data/samples/zidane.jpg' 187 | # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg' 188 | # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3) 189 | # PIL: = Image.open('image.jpg') # HWC x(720,1280,3) 190 | # numpy: = np.zeros((720,1280,3)) # HWC 191 | # torch: = torch.zeros(16,3,720,1280) # BCHW 192 | # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images 193 | 194 | t = [time_synchronized()] 195 | p = next(self.model.parameters()) # for device and type 196 | if isinstance(imgs, torch.Tensor): # torch 197 | return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference 198 | 199 | # Pre-process 200 | n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images 201 | shape0, shape1, files = [], [], [] # image and inference shapes, filenames 202 | for i, im in enumerate(imgs): 203 | if isinstance(im, str): # filename or uri 204 | im, f = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im), im # open 205 | im.filename = f # for uri 206 | files.append(Path(im.filename).with_suffix('.jpg').name if isinstance(im, Image.Image) else f'image{i}.jpg') 207 | im = np.array(im) # to numpy 208 | if im.shape[0] < 5: # image in CHW 209 | im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) 210 | im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input 211 | s = im.shape[:2] # HWC 212 | shape0.append(s) # image shape 213 | g = (size / max(s)) # gain 214 | shape1.append([y * g for y in s]) 215 | imgs[i] = im # update 216 | shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape 217 | x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad 218 | x = np.stack(x, 0) if n > 1 else x[0][None] # stack 219 | x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW 220 | x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32 221 | t.append(time_synchronized()) 222 | 223 | # Inference 224 | with torch.no_grad(): 225 | y = self.model(x, augment, profile)[0] # forward 226 | t.append(time_synchronized()) 227 | 228 | # Post-process 229 | y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS 230 | for i in range(n): 231 | scale_coords(shape1, y[i][:, :4], shape0[i]) 232 | t.append(time_synchronized()) 233 | 234 | return Detections(imgs, y, files, t, self.names, x.shape) 235 | 236 | 237 | class Detections: 238 | # detections class for YOLOv5 inference results 239 | def __init__(self, imgs, pred, files, times=None, names=None, shape=None): 240 | super(Detections, self).__init__() 241 | d = pred[0].device # device 242 | gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations 243 | self.imgs = imgs # list of images as numpy arrays 244 | self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) 245 | self.names = names # class names 246 | self.files = files # image filenames 247 | self.xyxy = pred # xyxy pixels 248 | self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels 249 | self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized 250 | self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized 251 | self.n = len(self.pred) 252 | self.t = ((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms) 253 | self.s = shape # inference BCHW shape 254 | 255 | def display(self, pprint=False, show=False, save=False, render=False, save_dir=''): 256 | colors = color_list() 257 | for i, (img, pred) in enumerate(zip(self.imgs, self.pred)): 258 | str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} ' 259 | if pred is not None: 260 | for c in pred[:, -1].unique(): 261 | n = (pred[:, -1] == c).sum() # detections per class 262 | str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string 263 | if show or save or render: 264 | for *box, conf, cls in pred: # xyxy, confidence, class 265 | label = f'{self.names[int(cls)]} {conf:.2f}' 266 | plot_one_box(box, img, label=label, color=colors[int(cls) % 10]) 267 | img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np 268 | if pprint: 269 | print(str.rstrip(', ')) 270 | if show: 271 | img.show(self.files[i]) # show 272 | if save: 273 | f = Path(save_dir) / self.files[i] 274 | img.save(f) # save 275 | print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') 276 | if render: 277 | self.imgs[i] = np.asarray(img) 278 | 279 | def print(self): 280 | self.display(pprint=True) # print results 281 | print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % 282 | tuple(self.t)) 283 | 284 | def show(self): 285 | self.display(show=True) # show results 286 | 287 | def save(self, save_dir='results/'): 288 | Path(save_dir).mkdir(exist_ok=True) 289 | self.display(save=True, save_dir=save_dir) # save results 290 | 291 | def render(self): 292 | self.display(render=True) # render results 293 | return self.imgs 294 | 295 | def __len__(self): 296 | return self.n 297 | 298 | def tolist(self): 299 | # return a list of Detections objects, i.e. 'for result in results.tolist():' 300 | x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)] 301 | for d in x: 302 | for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: 303 | setattr(d, k, getattr(d, k)[0]) # pop out of list 304 | return x 305 | 306 | 307 | class Classify(nn.Module): 308 | # Classification head, i.e. x(b,c1,20,20) to x(b,c2) 309 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups 310 | super(Classify, self).__init__() 311 | self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) 312 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1) 313 | self.flat = nn.Flatten() 314 | 315 | def forward(self, x): 316 | z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list 317 | return self.flat(self.conv(z)) # flatten to x(b,c2) 318 | -------------------------------------------------------------------------------- /models/common_rk_plug_in.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | 3 | import torch 4 | import torch.nn as nn 5 | from models.common import Conv 6 | 7 | 8 | class surrogate_silu(nn.Module): 9 | """docstring for surrogate_silu""" 10 | def __init__(self): 11 | super(surrogate_silu, self).__init__() 12 | self.act = nn.Sigmoid() 13 | 14 | def forward(self, x): 15 | return x*self.act(x) 16 | 17 | 18 | class surrogate_hardswish(nn.Module): 19 | """docstring for surrogate_hardswish""" 20 | def __init__(self): 21 | super(surrogate_hardswish, self).__init__() 22 | self.relu6 = nn.ReLU() 23 | 24 | def forward(self, x): 25 | return x *(self.relu6(torch.add(x, 3))/6) 26 | 27 | 28 | class surrogate_focus(nn.Module): 29 | # surrogate_focus wh information into c-space 30 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 31 | super(surrogate_focus, self).__init__() 32 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 33 | 34 | with torch.no_grad(): 35 | self.conv1 = nn.Conv2d(3, 3, (2, 2), groups=3, bias=False, stride=(2, 2)) 36 | self.conv1.weight[:, :, 0, 0] = 1 37 | self.conv1.weight[:, :, 0, 1] = 0 38 | self.conv1.weight[:, :, 1, 0] = 0 39 | self.conv1.weight[:, :, 1, 1] = 0 40 | 41 | self.conv2 = nn.Conv2d(3, 3, (2, 2), groups=3, bias=False, stride=(2, 2)) 42 | self.conv2.weight[:, :, 0, 0] = 0 43 | self.conv2.weight[:, :, 0, 1] = 0 44 | self.conv2.weight[:, :, 1, 0] = 1 45 | self.conv2.weight[:, :, 1, 1] = 0 46 | 47 | self.conv3 = nn.Conv2d(3, 3, (2, 2), groups=3, bias=False, stride=(2, 2)) 48 | self.conv3.weight[:, :, 0, 0] = 0 49 | self.conv3.weight[:, :, 0, 1] = 1 50 | self.conv3.weight[:, :, 1, 0] = 0 51 | self.conv3.weight[:, :, 1, 1] = 0 52 | 53 | self.conv4 = nn.Conv2d(3, 3, (2, 2), groups=3, bias=False, stride=(2, 2)) 54 | self.conv4.weight[:, :, 0, 0] = 0 55 | self.conv4.weight[:, :, 0, 1] = 0 56 | self.conv4.weight[:, :, 1, 0] = 0 57 | self.conv4.weight[:, :, 1, 1] = 1 58 | 59 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 60 | return self.conv(torch.cat([self.conv1(x), self.conv2(x), self.conv3(x), self.conv4(x)], 1)) 61 | 62 | 63 | class preprocess_conv_layer(nn.Module): 64 | """docstring for preprocess_conv_layer""" 65 | # input_module 为输入模型，即为想要导出模型 66 | # mean_value 的值可以是 [m1, m2, m3] 或 常数m 67 | # std_value 的值可以是 [s1, s2, s3] 或 常数s 68 | # BGR2RGB的操作默认为首先执行，既替代的原有操作顺序为 69 | # BGR2RGB -> minus mean -> minus std (与rknn config 设置保持一致) -> nhwc2nchw 70 | # 71 | # 使用示例-伪代码： 72 | # from add_preprocess_conv_layer import preprocess_conv_layer 73 | # model_A = create_model() 74 | # model_output = preprocess_co_nv_layer(model_A, mean_value, std_value, BGR2RGB) 75 | # onnx_export(model_output) 76 | # 77 | # 量化时： 78 | # rknn.config的中 channel_mean_value 、reorder_channel 均不赋值。 79 | # 80 | # 部署代码： 81 | # rknn_input 的属性 82 | # pass_through = 1 83 | # 84 | # 另外： 85 | # 由于加入permute操作，c端输入为opencv mat(hwc格式)即可，无需在外部将hwc改成chw格式。 86 | # 87 | 88 | def __init__(self, input_module, mean_value, std_value, BGR2RGB=False): 89 | super(preprocess_conv_layer, self).__init__() 90 | if isinstance(mean_value, int): 91 | mean_value = [mean_value for i in range(3)] 92 | if isinstance(std_value, int): 93 | std_value = [std_value for i in range(3)] 94 | 95 | assert len(mean_value) <= 3, 'mean_value should be int, or list with 3 element' 96 | assert len(std_value) <= 3, 'std_value should be int, or list with 3 element' 97 | 98 | self.input_module = input_module 99 | 100 | with torch.no_grad(): 101 | self.conv1 = nn.Conv2d(3, 3, (1, 1), groups=1, bias=True, stride=(1, 1)) 102 | 103 | if BGR2RGB is False: 104 | self.conv1.weight[:, :, :, :] = 0 105 | self.conv1.weight[0, 0, :, :] = 1/std_value[0] 106 | self.conv1.weight[1, 1, :, :] = 1/std_value[1] 107 | self.conv1.weight[2, 2, :, :] = 1/std_value[2] 108 | elif BGR2RGB is True: 109 | self.conv1.weight[:, :, :, :] = 0 110 | self.conv1.weight[0, 2, :, :] = 1/std_value[0] 111 | self.conv1.weight[1, 1, :, :] = 1/std_value[1] 112 | self.conv1.weight[2, 0, :, :] = 1/std_value[2] 113 | 114 | self.conv1.bias[0] = -mean_value[0]/std_value[0] 115 | self.conv1.bias[1] = -mean_value[1]/std_value[1] 116 | self.conv1.bias[2] = -mean_value[2]/std_value[2] 117 | 118 | self.conv1.eval() 119 | 120 | def forward(self, x): 121 | x = x.permute(0, 3, 1, 2) # NHWC -> NCHW, apply for rknn_pass_through 122 | x = self.conv1(x) 123 | return self.input_module(x) -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # This file contains experimental modules 2 | 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | 7 | from models.common import Conv, DWConv 8 | from utils.google_utils import attempt_download 9 | 10 | 11 | class CrossConv(nn.Module): 12 | # Cross Convolution Downsample 13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 15 | super(CrossConv, self).__init__() 16 | c_ = int(c2 * e) # hidden channels 17 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 19 | self.add = shortcut and c1 == c2 20 | 21 | def forward(self, x): 22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 23 | 24 | 25 | class Sum(nn.Module): 26 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 27 | def __init__(self, n, weight=False): # n: number of inputs 28 | super(Sum, self).__init__() 29 | self.weight = weight # apply weights boolean 30 | self.iter = range(n - 1) # iter object 31 | if weight: 32 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 33 | 34 | def forward(self, x): 35 | y = x[0] # no weight 36 | if self.weight: 37 | w = torch.sigmoid(self.w) * 2 38 | for i in self.iter: 39 | y = y + x[i + 1] * w[i] 40 | else: 41 | for i in self.iter: 42 | y = y + x[i + 1] 43 | return y 44 | 45 | 46 | class GhostConv(nn.Module): 47 | # Ghost Convolution https://github.com/huawei-noah/ghostnet 48 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups 49 | super(GhostConv, self).__init__() 50 | c_ = c2 // 2 # hidden channels 51 | self.cv1 = Conv(c1, c_, k, s, None, g, act) 52 | self.cv2 = Conv(c_, c_, 5, 1, None, c_, act) 53 | 54 | def forward(self, x): 55 | y = self.cv1(x) 56 | return torch.cat([y, self.cv2(y)], 1) 57 | 58 | 59 | class GhostBottleneck(nn.Module): 60 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet 61 | def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride 62 | super(GhostBottleneck, self).__init__() 63 | c_ = c2 // 2 64 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw 65 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw 66 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear 67 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), 68 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 69 | 70 | def forward(self, x): 71 | return self.conv(x) + self.shortcut(x) 72 | 73 | 74 | class MixConv2d(nn.Module): 75 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 76 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 77 | super(MixConv2d, self).__init__() 78 | groups = len(k) 79 | if equal_ch: # equal c_ per group 80 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 81 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 82 | else: # equal weight.numel() per group 83 | b = [c2] + [0] * groups 84 | a = np.eye(groups + 1, groups, k=-1) 85 | a -= np.roll(a, 1, axis=1) 86 | a *= np.array(k) ** 2 87 | a[0] = 1 88 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 89 | 90 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 91 | self.bn = nn.BatchNorm2d(c2) 92 | self.act = nn.LeakyReLU(0.1, inplace=True) 93 | 94 | def forward(self, x): 95 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 96 | 97 | 98 | class Ensemble(nn.ModuleList): 99 | # Ensemble of models 100 | def __init__(self): 101 | super(Ensemble, self).__init__() 102 | 103 | def forward(self, x, augment=False): 104 | y = [] 105 | for module in self: 106 | y.append(module(x, augment)[0]) 107 | # y = torch.stack(y).max(0)[0] # max ensemble 108 | # y = torch.stack(y).mean(0) # mean ensemble 109 | y = torch.cat(y, 1) # nms ensemble 110 | return y, None # inference, train output 111 | 112 | 113 | def attempt_load(weights, map_location=None): 114 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 115 | model = Ensemble() 116 | for w in weights if isinstance(weights, list) else [weights]: 117 | attempt_download(w) 118 | ckpt = torch.load(w, map_location=map_location) # load 119 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model 120 | 121 | # Compatibility updates 122 | for m in model.modules(): 123 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]: 124 | m.inplace = True # pytorch 1.7.0 compatibility 125 | elif type(m) is Conv: 126 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 127 | 128 | if len(model) == 1: 129 | return model[-1] # return model 130 | else: 131 | print('Ensemble created with %s\n' % weights) 132 | for k in ['names', 'stride']: 133 | setattr(model, k, getattr(model[-1], k)) 134 | return model # return ensemble 135 | -------------------------------------------------------------------------------- /models/export.py: -------------------------------------------------------------------------------- 1 | """Exports a YOLOv5 *.pt model to ONNX and TorchScript formats 2 | 3 | Usage: 4 | $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 5 | """ 6 | 7 | import argparse 8 | import sys 9 | import time 10 | 11 | sys.path.append('./') # to run '$ python *.py' files in subdirectories 12 | 13 | import torch 14 | import torch.nn as nn 15 | 16 | import models 17 | from models.experimental import attempt_load 18 | from utils.activations import Hardswish, SiLU 19 | from utils.general import set_logging, check_img_size 20 | from utils.torch_utils import select_device 21 | 22 | if __name__ == '__main__': 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--weights', type=str, default=r'D:\workspace\LMO\opengit\yolov5\runs\train\exp\weights\best.pt', help='weights path') # from yolov5/models/ 25 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width 26 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 27 | parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') 28 | parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') 29 | parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 30 | parser.add_argument('--rknn_mode', action='store_true', help='export rknn-friendly onnx model') 31 | parser.add_argument('--ignore_output_permute', action='store_true', help='export model without permute layer,which can be used for rknn_yolov5_demo c++ code') 32 | parser.add_argument('--add_image_preprocess_layer', action='store_true', help='add image preprocess layer, benefit for decreasing rknn_input_set time-cost') 33 | opt = parser.parse_args() 34 | opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand 35 | print(opt) 36 | set_logging() 37 | t = time.time() 38 | 39 | # Load PyTorch model 40 | device = select_device(opt.device) 41 | model = attempt_load(opt.weights, map_location=device) # load FP32 model 42 | labels = model.names 43 | 44 | # Checks 45 | gs = int(max(model.stride)) # grid size (max stride) 46 | opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples 47 | 48 | # Update model 49 | if opt.rknn_mode != True: 50 | for k, m in model.named_modules(): 51 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 52 | if isinstance(m, models.common.Conv): # assign export-friendly activations 53 | if isinstance(m.act, nn.Hardswish): 54 | m.act = Hardswish() 55 | elif isinstance(m.act, nn.SiLU): 56 | m.act = SiLU() 57 | # elif isinstance(m, models.yolo.Detect): 58 | # m.forward = m.forward_export # assign forward (optional) 59 | 60 | else: 61 | # Update model 62 | from models.common_rk_plug_in import surrogate_silu 63 | from models import common 64 | for k, m in model.named_modules(): 65 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 66 | if isinstance(m, common.Conv): # assign export-friendly activations 67 | if isinstance(m.act, torch.nn.Hardswish): 68 | m.act = torch.nn.Hardswish() 69 | elif isinstance(m.act, torch.nn.SiLU): 70 | # m.act = SiLU() 71 | m.act = surrogate_silu() 72 | # elif isinstance(m, models.yolo.Detect): 73 | # m.forward = m.forward_export # assign forward (optional) 74 | 75 | if isinstance(m, models.common.SPP): # assign export-friendly activations 76 | ### best 77 | # tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(2)]) 78 | # m.m[0] = tmp 79 | # m.m[1] = tmp 80 | # m.m[2] = tmp 81 | ### friendly to origin config 82 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(2)]) 83 | m.m[0] = tmp 84 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(4)]) 85 | m.m[1] = tmp 86 | tmp = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1) for i in range(6)]) 87 | m.m[2] = tmp 88 | 89 | ### use deconv2d to surrogate upsample layer. 90 | # replace_one = torch.nn.ConvTranspose2d(model.model[10].conv.weight.shape[0], 91 | # model.model[10].conv.weight.shape[0], 92 | # (2, 2), 93 | # groups=model.model[10].conv.weight.shape[0], 94 | # bias=False, 95 | # stride=(2, 2)) 96 | # replace_one.weight.data.fill_(1) 97 | # replace_one.eval() 98 | # temp_i = model.model[11].i 99 | # temp_f = model.model[11].f 100 | # model.model[11] = replace_one 101 | # model.model[11].i = temp_i 102 | # model.model[11].f = temp_f 103 | 104 | # replace_one = torch.nn.ConvTranspose2d(model.model[14].conv.weight.shape[0], 105 | # model.model[14].conv.weight.shape[0], 106 | # (2, 2), 107 | # groups=model.model[14].conv.weight.shape[0], 108 | # bias=False, 109 | # stride=(2, 2)) 110 | # replace_one.weight.data.fill_(1) 111 | # replace_one.eval() 112 | # temp_i = model.model[11].i 113 | # temp_f = model.model[11].f 114 | # model.model[15] = replace_one 115 | # model.model[15].i = temp_i 116 | # model.model[15].f = temp_f 117 | 118 | ### use conv to surrogate slice operator 119 | from models.common_rk_plug_in import surrogate_focus 120 | surrogate_focous = surrogate_focus(int(model.model[0].conv.conv.weight.shape[1]/4), 121 | model.model[0].conv.conv.weight.shape[0], 122 | k=tuple(model.model[0].conv.conv.weight.shape[2:4]), 123 | s=model.model[0].conv.conv.stride, 124 | p=model.model[0].conv.conv.padding, 125 | g=model.model[0].conv.conv.groups, 126 | act=True) 127 | surrogate_focous.conv.conv.weight = model.model[0].conv.conv.weight 128 | surrogate_focous.conv.conv.bias = model.model[0].conv.conv.bias 129 | surrogate_focous.conv.act = model.model[0].conv.act 130 | temp_i = model.model[0].i 131 | temp_f = model.model[0].f 132 | 133 | model.model[0] = surrogate_focous 134 | model.model[0].i = temp_i 135 | model.model[0].f = temp_f 136 | model.model[0].eval() 137 | 138 | model.model[-1].export = not opt.grid # set Detect() layer grid export 139 | 140 | if opt.ignore_output_permute is True: 141 | model.model[-1].ignore_permute_layer = True 142 | 143 | if opt.add_image_preprocess_layer is True: 144 | from models.common_rk_plug_in import preprocess_conv_layer 145 | model = preprocess_conv_layer(model, 0, 255, True) 146 | img = torch.zeros(opt.batch_size, *opt.img_size, 3).to(device) 147 | else: 148 | img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) 149 | 150 | y = model(img) # dry run 151 | # ONNX export 152 | try: 153 | import onnx 154 | 155 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__) 156 | f = opt.weights.replace('.pt', '.onnx') # filename 157 | torch.onnx.export(model, img, f, verbose=False, opset_version=10, input_names=['images'], 158 | output_names=['classes', 'boxes'] if y is None else ['output'], 159 | dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640) 160 | 'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None) 161 | 162 | # Checks 163 | onnx_model = onnx.load(f) # load onnx model 164 | onnx.checker.check_model(onnx_model) # check onnx model 165 | # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model 166 | print('ONNX export success, saved as %s' % f) 167 | except Exception as e: 168 | print('ONNX export failure: %s' % e) 169 | -------------------------------------------------------------------------------- /models/hub/anchors.yaml: -------------------------------------------------------------------------------- 1 | # Default YOLOv5 anchors for COCO data 2 | 3 | 4 | # P5 ------------------------------------------------------------------------------------------------------------------- 5 | # P5-640: 6 | anchors_p5_640: 7 | - [ 10,13, 16,30, 33,23 ] # P3/8 8 | - [ 30,61, 62,45, 59,119 ] # P4/16 9 | - [ 116,90, 156,198, 373,326 ] # P5/32 10 | 11 | 12 | # P6 ------------------------------------------------------------------------------------------------------------------- 13 | # P6-640: thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11, 21,19, 17,41, 43,32, 39,70, 86,64, 65,131, 134,130, 120,265, 282,180, 247,354, 512,387 14 | anchors_p6_640: 15 | - [ 9,11, 21,19, 17,41 ] # P3/8 16 | - [ 43,32, 39,70, 86,64 ] # P4/16 17 | - [ 65,131, 134,130, 120,265 ] # P5/32 18 | - [ 282,180, 247,354, 512,387 ] # P6/64 19 | 20 | # P6-1280: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1280, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 19,27, 44,40, 38,94, 96,68, 86,152, 180,137, 140,301, 303,264, 238,542, 436,615, 739,380, 925,792 21 | anchors_p6_1280: 22 | - [ 19,27, 44,40, 38,94 ] # P3/8 23 | - [ 96,68, 86,152, 180,137 ] # P4/16 24 | - [ 140,301, 303,264, 238,542 ] # P5/32 25 | - [ 436,615, 739,380, 925,792 ] # P6/64 26 | 27 | # P6-1920: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1920, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 28,41, 67,59, 57,141, 144,103, 129,227, 270,205, 209,452, 455,396, 358,812, 653,922, 1109,570, 1387,1187 28 | anchors_p6_1920: 29 | - [ 28,41, 67,59, 57,141 ] # P3/8 30 | - [ 144,103, 129,227, 270,205 ] # P4/16 31 | - [ 209,452, 455,396, 358,812 ] # P5/32 32 | - [ 653,922, 1109,570, 1387,1187 ] # P6/64 33 | 34 | 35 | # P7 ------------------------------------------------------------------------------------------------------------------- 36 | # P7-640: thr=0.25: 0.9962 BPR, 6.76 anchors past thr, n=15, img_size=640, metric_all=0.275/0.733-mean/best, past_thr=0.466-mean: 11,11, 13,30, 29,20, 30,46, 61,38, 39,92, 78,80, 146,66, 79,163, 149,150, 321,143, 157,303, 257,402, 359,290, 524,372 37 | anchors_p7_640: 38 | - [ 11,11, 13,30, 29,20 ] # P3/8 39 | - [ 30,46, 61,38, 39,92 ] # P4/16 40 | - [ 78,80, 146,66, 79,163 ] # P5/32 41 | - [ 149,150, 321,143, 157,303 ] # P6/64 42 | - [ 257,402, 359,290, 524,372 ] # P7/128 43 | 44 | # P7-1280: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1280, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 19,22, 54,36, 32,77, 70,83, 138,71, 75,173, 165,159, 148,334, 375,151, 334,317, 251,626, 499,474, 750,326, 534,814, 1079,818 45 | anchors_p7_1280: 46 | - [ 19,22, 54,36, 32,77 ] # P3/8 47 | - [ 70,83, 138,71, 75,173 ] # P4/16 48 | - [ 165,159, 148,334, 375,151 ] # P5/32 49 | - [ 334,317, 251,626, 499,474 ] # P6/64 50 | - [ 750,326, 534,814, 1079,818 ] # P7/128 51 | 52 | # P7-1920: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1920, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 29,34, 81,55, 47,115, 105,124, 207,107, 113,259, 247,238, 222,500, 563,227, 501,476, 376,939, 749,711, 1126,489, 801,1222, 1618,1227 53 | anchors_p7_1920: 54 | - [ 29,34, 81,55, 47,115 ] # P3/8 55 | - [ 105,124, 207,107, 113,259 ] # P4/16 56 | - [ 247,238, 222,500, 563,227 ] # P5/32 57 | - [ 501,476, 376,939, 749,711 ] # P6/64 58 | - [ 1126,489, 801,1222, 1618,1227 ] # P7/128 59 | -------------------------------------------------------------------------------- /models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3-SPP head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, SPP, [512, [5, 9, 13]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /models/hub/yolov3-tiny.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,14, 23,27, 37,58] # P4/16 9 | - [81,82, 135,169, 344,319] # P5/32 10 | 11 | # YOLOv3-tiny backbone 12 | backbone: 13 | # [from, number, module, args] 14 | [[-1, 1, Conv, [16, 3, 1]], # 0 15 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 1-P1/2 16 | [-1, 1, Conv, [32, 3, 1]], 17 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 3-P2/4 18 | [-1, 1, Conv, [64, 3, 1]], 19 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 5-P3/8 20 | [-1, 1, Conv, [128, 3, 1]], 21 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 7-P4/16 22 | [-1, 1, Conv, [256, 3, 1]], 23 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 9-P5/32 24 | [-1, 1, Conv, [512, 3, 1]], 25 | [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]], # 11 26 | [-1, 1, nn.MaxPool2d, [2, 1, 0]], # 12 27 | ] 28 | 29 | # YOLOv3-tiny head 30 | head: 31 | [[-1, 1, Conv, [1024, 3, 1]], 32 | [-1, 1, Conv, [256, 1, 1]], 33 | [-1, 1, Conv, [512, 3, 1]], # 15 (P5/32-large) 34 | 35 | [-2, 1, Conv, [128, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 38 | [-1, 1, Conv, [256, 3, 1]], # 19 (P4/16-medium) 39 | 40 | [[19, 15], 1, Detect, [nc, anchors]], # Detect(P4, P5) 41 | ] 42 | -------------------------------------------------------------------------------- /models/hub/yolov3.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3 head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, Conv, [512, [1, 1]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /models/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, Bottleneck, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 6, BottleneckCSP, [1024]], # 9 25 | ] 26 | 27 | # YOLOv5 FPN head 28 | head: 29 | [[-1, 3, BottleneckCSP, [1024, False]], # 10 (P5/32-large) 30 | 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 3, BottleneckCSP, [512, False]], # 14 (P4/16-medium) 35 | 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 3, BottleneckCSP, [256, False]], # 18 (P3/8-small) 40 | 41 | [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 42 | ] 43 | -------------------------------------------------------------------------------- /models/hub/yolov5-p2.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 3 8 | 9 | # YOLOv5 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 13 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 14 | [ -1, 3, C3, [ 128 ] ], 15 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 16 | [ -1, 9, C3, [ 256 ] ], 17 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 18 | [ -1, 9, C3, [ 512 ] ], 19 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32 20 | [ -1, 1, SPP, [ 1024, [ 5, 9, 13 ] ] ], 21 | [ -1, 3, C3, [ 1024, False ] ], # 9 22 | ] 23 | 24 | # YOLOv5 head 25 | head: 26 | [ [ -1, 1, Conv, [ 512, 1, 1 ] ], 27 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 28 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 29 | [ -1, 3, C3, [ 512, False ] ], # 13 30 | 31 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 32 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 33 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 34 | [ -1, 3, C3, [ 256, False ] ], # 17 (P3/8-small) 35 | 36 | [ -1, 1, Conv, [ 128, 1, 1 ] ], 37 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 38 | [ [ -1, 2 ], 1, Concat, [ 1 ] ], # cat backbone P2 39 | [ -1, 1, C3, [ 128, False ] ], # 21 (P2/4-xsmall) 40 | 41 | [ -1, 1, Conv, [ 128, 3, 2 ] ], 42 | [ [ -1, 18 ], 1, Concat, [ 1 ] ], # cat head P3 43 | [ -1, 3, C3, [ 256, False ] ], # 24 (P3/8-small) 44 | 45 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 46 | [ [ -1, 14 ], 1, Concat, [ 1 ] ], # cat head P4 47 | [ -1, 3, C3, [ 512, False ] ], # 27 (P4/16-medium) 48 | 49 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 50 | [ [ -1, 10 ], 1, Concat, [ 1 ] ], # cat head P5 51 | [ -1, 3, C3, [ 1024, False ] ], # 30 (P5/32-large) 52 | 53 | [ [ 24, 27, 30 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5) 54 | ] 55 | -------------------------------------------------------------------------------- /models/hub/yolov5-p6.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 3 8 | 9 | # YOLOv5 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 13 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 14 | [ -1, 3, C3, [ 128 ] ], 15 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 16 | [ -1, 9, C3, [ 256 ] ], 17 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 18 | [ -1, 9, C3, [ 512 ] ], 19 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 20 | [ -1, 3, C3, [ 768 ] ], 21 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 22 | [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ], 23 | [ -1, 3, C3, [ 1024, False ] ], # 11 24 | ] 25 | 26 | # YOLOv5 head 27 | head: 28 | [ [ -1, 1, Conv, [ 768, 1, 1 ] ], 29 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 30 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 31 | [ -1, 3, C3, [ 768, False ] ], # 15 32 | 33 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 34 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 35 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 36 | [ -1, 3, C3, [ 512, False ] ], # 19 37 | 38 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 39 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 40 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 41 | [ -1, 3, C3, [ 256, False ] ], # 23 (P3/8-small) 42 | 43 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 44 | [ [ -1, 20 ], 1, Concat, [ 1 ] ], # cat head P4 45 | [ -1, 3, C3, [ 512, False ] ], # 26 (P4/16-medium) 46 | 47 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 48 | [ [ -1, 16 ], 1, Concat, [ 1 ] ], # cat head P5 49 | [ -1, 3, C3, [ 768, False ] ], # 29 (P5/32-large) 50 | 51 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 52 | [ [ -1, 12 ], 1, Concat, [ 1 ] ], # cat head P6 53 | [ -1, 3, C3, [ 1024, False ] ], # 32 (P5/64-xlarge) 54 | 55 | [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6) 56 | ] 57 | -------------------------------------------------------------------------------- /models/hub/yolov5-p7.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 3 8 | 9 | # YOLOv5 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 13 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 14 | [ -1, 3, C3, [ 128 ] ], 15 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 16 | [ -1, 9, C3, [ 256 ] ], 17 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 18 | [ -1, 9, C3, [ 512 ] ], 19 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 20 | [ -1, 3, C3, [ 768 ] ], 21 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 22 | [ -1, 3, C3, [ 1024 ] ], 23 | [ -1, 1, Conv, [ 1280, 3, 2 ] ], # 11-P7/128 24 | [ -1, 1, SPP, [ 1280, [ 3, 5 ] ] ], 25 | [ -1, 3, C3, [ 1280, False ] ], # 13 26 | ] 27 | 28 | # YOLOv5 head 29 | head: 30 | [ [ -1, 1, Conv, [ 1024, 1, 1 ] ], 31 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 32 | [ [ -1, 10 ], 1, Concat, [ 1 ] ], # cat backbone P6 33 | [ -1, 3, C3, [ 1024, False ] ], # 17 34 | 35 | [ -1, 1, Conv, [ 768, 1, 1 ] ], 36 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 37 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 38 | [ -1, 3, C3, [ 768, False ] ], # 21 39 | 40 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 41 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 42 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 43 | [ -1, 3, C3, [ 512, False ] ], # 25 44 | 45 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 46 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 47 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 48 | [ -1, 3, C3, [ 256, False ] ], # 29 (P3/8-small) 49 | 50 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 51 | [ [ -1, 26 ], 1, Concat, [ 1 ] ], # cat head P4 52 | [ -1, 3, C3, [ 512, False ] ], # 32 (P4/16-medium) 53 | 54 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 55 | [ [ -1, 22 ], 1, Concat, [ 1 ] ], # cat head P5 56 | [ -1, 3, C3, [ 768, False ] ], # 35 (P5/32-large) 57 | 58 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 59 | [ [ -1, 18 ], 1, Concat, [ 1 ] ], # cat head P6 60 | [ -1, 3, C3, [ 1024, False ] ], # 38 (P6/64-xlarge) 61 | 62 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], 63 | [ [ -1, 14 ], 1, Concat, [ 1 ] ], # cat head P7 64 | [ -1, 3, C3, [ 1280, False ] ], # 41 (P7/128-xxlarge) 65 | 66 | [ [ 29, 32, 35, 38, 41 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6, P7) 67 | ] 68 | -------------------------------------------------------------------------------- /models/hub/yolov5-panet.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 PANet head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5l6.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [ 19,27, 44,40, 38,94 ] # P3/8 9 | - [ 96,68, 86,152, 180,137 ] # P4/16 10 | - [ 140,301, 303,264, 238,542 ] # P5/32 11 | - [ 436,615, 739,380, 925,792 ] # P6/64 12 | 13 | # YOLOv5 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 17 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 18 | [ -1, 3, C3, [ 128 ] ], 19 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 20 | [ -1, 9, C3, [ 256 ] ], 21 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 22 | [ -1, 9, C3, [ 512 ] ], 23 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 24 | [ -1, 3, C3, [ 768 ] ], 25 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 26 | [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ], 27 | [ -1, 3, C3, [ 1024, False ] ], # 11 28 | ] 29 | 30 | # YOLOv5 head 31 | head: 32 | [ [ -1, 1, Conv, [ 768, 1, 1 ] ], 33 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 34 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 35 | [ -1, 3, C3, [ 768, False ] ], # 15 36 | 37 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 38 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 39 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 40 | [ -1, 3, C3, [ 512, False ] ], # 19 41 | 42 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 43 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 44 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 45 | [ -1, 3, C3, [ 256, False ] ], # 23 (P3/8-small) 46 | 47 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 48 | [ [ -1, 20 ], 1, Concat, [ 1 ] ], # cat head P4 49 | [ -1, 3, C3, [ 512, False ] ], # 26 (P4/16-medium) 50 | 51 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 52 | [ [ -1, 16 ], 1, Concat, [ 1 ] ], # cat head P5 53 | [ -1, 3, C3, [ 768, False ] ], # 29 (P5/32-large) 54 | 55 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 56 | [ [ -1, 12 ], 1, Concat, [ 1 ] ], # cat head P6 57 | [ -1, 3, C3, [ 1024, False ] ], # 32 (P6/64-xlarge) 58 | 59 | [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5m6.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.67 # model depth multiple 4 | width_multiple: 0.75 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [ 19,27, 44,40, 38,94 ] # P3/8 9 | - [ 96,68, 86,152, 180,137 ] # P4/16 10 | - [ 140,301, 303,264, 238,542 ] # P5/32 11 | - [ 436,615, 739,380, 925,792 ] # P6/64 12 | 13 | # YOLOv5 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 17 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 18 | [ -1, 3, C3, [ 128 ] ], 19 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 20 | [ -1, 9, C3, [ 256 ] ], 21 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 22 | [ -1, 9, C3, [ 512 ] ], 23 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 24 | [ -1, 3, C3, [ 768 ] ], 25 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 26 | [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ], 27 | [ -1, 3, C3, [ 1024, False ] ], # 11 28 | ] 29 | 30 | # YOLOv5 head 31 | head: 32 | [ [ -1, 1, Conv, [ 768, 1, 1 ] ], 33 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 34 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 35 | [ -1, 3, C3, [ 768, False ] ], # 15 36 | 37 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 38 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 39 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 40 | [ -1, 3, C3, [ 512, False ] ], # 19 41 | 42 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 43 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 44 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 45 | [ -1, 3, C3, [ 256, False ] ], # 23 (P3/8-small) 46 | 47 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 48 | [ [ -1, 20 ], 1, Concat, [ 1 ] ], # cat head P4 49 | [ -1, 3, C3, [ 512, False ] ], # 26 (P4/16-medium) 50 | 51 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 52 | [ [ -1, 16 ], 1, Concat, [ 1 ] ], # cat head P5 53 | [ -1, 3, C3, [ 768, False ] ], # 29 (P5/32-large) 54 | 55 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 56 | [ [ -1, 12 ], 1, Concat, [ 1 ] ], # cat head P6 57 | [ -1, 3, C3, [ 1024, False ] ], # 32 (P6/64-xlarge) 58 | 59 | [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5s6.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.33 # model depth multiple 4 | width_multiple: 0.50 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [ 19,27, 44,40, 38,94 ] # P3/8 9 | - [ 96,68, 86,152, 180,137 ] # P4/16 10 | - [ 140,301, 303,264, 238,542 ] # P5/32 11 | - [ 436,615, 739,380, 925,792 ] # P6/64 12 | 13 | # YOLOv5 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 17 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 18 | [ -1, 3, C3, [ 128 ] ], 19 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 20 | [ -1, 9, C3, [ 256 ] ], 21 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 22 | [ -1, 9, C3, [ 512 ] ], 23 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 24 | [ -1, 3, C3, [ 768 ] ], 25 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 26 | [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ], 27 | [ -1, 3, C3, [ 1024, False ] ], # 11 28 | ] 29 | 30 | # YOLOv5 head 31 | head: 32 | [ [ -1, 1, Conv, [ 768, 1, 1 ] ], 33 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 34 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 35 | [ -1, 3, C3, [ 768, False ] ], # 15 36 | 37 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 38 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 39 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 40 | [ -1, 3, C3, [ 512, False ] ], # 19 41 | 42 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 43 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 44 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 45 | [ -1, 3, C3, [ 256, False ] ], # 23 (P3/8-small) 46 | 47 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 48 | [ [ -1, 20 ], 1, Concat, [ 1 ] ], # cat head P4 49 | [ -1, 3, C3, [ 512, False ] ], # 26 (P4/16-medium) 50 | 51 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 52 | [ [ -1, 16 ], 1, Concat, [ 1 ] ], # cat head P5 53 | [ -1, 3, C3, [ 768, False ] ], # 29 (P5/32-large) 54 | 55 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 56 | [ [ -1, 12 ], 1, Concat, [ 1 ] ], # cat head P6 57 | [ -1, 3, C3, [ 1024, False ] ], # 32 (P6/64-xlarge) 58 | 59 | [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5x6.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.33 # model depth multiple 4 | width_multiple: 1.25 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [ 19,27, 44,40, 38,94 ] # P3/8 9 | - [ 96,68, 86,152, 180,137 ] # P4/16 10 | - [ 140,301, 303,264, 238,542 ] # P5/32 11 | - [ 436,615, 739,380, 925,792 ] # P6/64 12 | 13 | # YOLOv5 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 17 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 18 | [ -1, 3, C3, [ 128 ] ], 19 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 20 | [ -1, 9, C3, [ 256 ] ], 21 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 22 | [ -1, 9, C3, [ 512 ] ], 23 | [ -1, 1, Conv, [ 768, 3, 2 ] ], # 7-P5/32 24 | [ -1, 3, C3, [ 768 ] ], 25 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 9-P6/64 26 | [ -1, 1, SPP, [ 1024, [ 3, 5, 7 ] ] ], 27 | [ -1, 3, C3, [ 1024, False ] ], # 11 28 | ] 29 | 30 | # YOLOv5 head 31 | head: 32 | [ [ -1, 1, Conv, [ 768, 1, 1 ] ], 33 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 34 | [ [ -1, 8 ], 1, Concat, [ 1 ] ], # cat backbone P5 35 | [ -1, 3, C3, [ 768, False ] ], # 15 36 | 37 | [ -1, 1, Conv, [ 512, 1, 1 ] ], 38 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 39 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 40 | [ -1, 3, C3, [ 512, False ] ], # 19 41 | 42 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 43 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 44 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 45 | [ -1, 3, C3, [ 256, False ] ], # 23 (P3/8-small) 46 | 47 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 48 | [ [ -1, 20 ], 1, Concat, [ 1 ] ], # cat head P4 49 | [ -1, 3, C3, [ 512, False ] ], # 26 (P4/16-medium) 50 | 51 | [ -1, 1, Conv, [ 512, 3, 2 ] ], 52 | [ [ -1, 16 ], 1, Concat, [ 1 ] ], # cat head P5 53 | [ -1, 3, C3, [ 768, False ] ], # 29 (P5/32-large) 54 | 55 | [ -1, 1, Conv, [ 768, 3, 2 ] ], 56 | [ [ -1, 12 ], 1, Concat, [ 1 ] ], # cat head P6 57 | [ -1, 3, C3, [ 1024, False ] ], # 32 (P6/64-xlarge) 58 | 59 | [ [ 23, 26, 29, 32 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/yolo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import sys 4 | from copy import deepcopy 5 | 6 | sys.path.append('./') # to run '$ python *.py' files in subdirectories 7 | logger = logging.getLogger(__name__) 8 | 9 | from models.common import * 10 | from models.experimental import * 11 | from utils.autoanchor import check_anchor_order 12 | from utils.general import make_divisible, check_file, set_logging 13 | from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \ 14 | select_device, copy_attr 15 | 16 | try: 17 | import thop # for FLOPS computation 18 | except ImportError: 19 | thop = None 20 | 21 | 22 | class Detect(nn.Module): 23 | stride = None # strides computed during build 24 | export = False # onnx export 25 | ignore_permute_layer = False 26 | 27 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 28 | super(Detect, self).__init__() 29 | self.nc = nc # number of classes 30 | self.no = nc + 5 # number of outputs per anchor 31 | self.nl = len(anchors) # number of detection layers 32 | self.na = len(anchors[0]) // 2 # number of anchors 33 | self.grid = [torch.zeros(1)] * self.nl # init grid 34 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 35 | self.register_buffer('anchors', a) # shape(nl,na,2) 36 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 37 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 38 | 39 | def forward(self, x): 40 | # x = x.copy() # for profiling 41 | z = [] # inference output 42 | self.training |= self.export 43 | for i in range(self.nl): 44 | x[i] = self.m[i](x[i]) # conv 45 | 46 | if self.ignore_permute_layer is True: 47 | continue 48 | 49 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 50 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 51 | 52 | if not self.training: # inference 53 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 54 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 55 | 56 | y = x[i].sigmoid() 57 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy 58 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 59 | z.append(y.view(bs, -1, self.no)) 60 | 61 | return x if self.training else (torch.cat(z, 1), x) 62 | 63 | @staticmethod 64 | def _make_grid(nx=20, ny=20): 65 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 66 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 67 | 68 | 69 | class Model(nn.Module): 70 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes 71 | super(Model, self).__init__() 72 | if isinstance(cfg, dict): 73 | self.yaml = cfg # model dict 74 | else: # is *.yaml 75 | import yaml # for torch hub 76 | self.yaml_file = Path(cfg).name 77 | with open(cfg) as f: 78 | self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict 79 | 80 | # Define model 81 | ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels 82 | if nc and nc != self.yaml['nc']: 83 | logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") 84 | self.yaml['nc'] = nc # override yaml value 85 | if anchors: 86 | logger.info(f'Overriding model.yaml anchors with anchors={anchors}') 87 | self.yaml['anchors'] = round(anchors) # override yaml value 88 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist 89 | self.names = [str(i) for i in range(self.yaml['nc'])] # default names 90 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 91 | 92 | # Build strides, anchors 93 | m = self.model[-1] # Detect() 94 | if isinstance(m, Detect): 95 | s = 256 # 2x min stride 96 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 97 | m.anchors /= m.stride.view(-1, 1, 1) 98 | check_anchor_order(m) 99 | self.stride = m.stride 100 | self._initialize_biases() # only run once 101 | # print('Strides: %s' % m.stride.tolist()) 102 | 103 | # Init weights, biases 104 | initialize_weights(self) 105 | self.info() 106 | logger.info('') 107 | 108 | def forward(self, x, augment=False, profile=False): 109 | if augment: 110 | img_size = x.shape[-2:] # height, width 111 | s = [1, 0.83, 0.67] # scales 112 | f = [None, 3, None] # flips (2-ud, 3-lr) 113 | y = [] # outputs 114 | for si, fi in zip(s, f): 115 | xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) 116 | yi = self.forward_once(xi)[0] # forward 117 | # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 118 | yi[..., :4] /= si # de-scale 119 | if fi == 2: 120 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 121 | elif fi == 3: 122 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 123 | y.append(yi) 124 | return torch.cat(y, 1), None # augmented inference, train 125 | else: 126 | return self.forward_once(x, profile) # single-scale inference, train 127 | 128 | def forward_once(self, x, profile=False): 129 | y, dt = [], [] # outputs 130 | for m in self.model: 131 | if m.f != -1: # if not from previous layer 132 | 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 133 | 134 | if profile: 135 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPS 136 | t = time_synchronized() 137 | for _ in range(10): 138 | _ = m(x) 139 | dt.append((time_synchronized() - t) * 100) 140 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 141 | 142 | x = m(x) # run 143 | y.append(x if m.i in self.save else None) # save output 144 | 145 | if profile: 146 | print('%.1fms total' % sum(dt)) 147 | return x 148 | 149 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 150 | # https://arxiv.org/abs/1708.02002 section 3.3 151 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 152 | m = self.model[-1] # Detect() module 153 | for mi, s in zip(m.m, m.stride): # from 154 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 155 | b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 156 | b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 157 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 158 | 159 | def _print_biases(self): 160 | m = self.model[-1] # Detect() module 161 | for mi in m.m: # from 162 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 163 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 164 | 165 | # def _print_weights(self): 166 | # for m in self.model.modules(): 167 | # if type(m) is Bottleneck: 168 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 169 | 170 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 171 | print('Fusing layers... ') 172 | for m in self.model.modules(): 173 | if type(m) is Conv and hasattr(m, 'bn'): 174 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 175 | delattr(m, 'bn') # remove batchnorm 176 | m.forward = m.fuseforward # update forward 177 | self.info() 178 | return self 179 | 180 | def nms(self, mode=True): # add or remove NMS module 181 | present = type(self.model[-1]) is NMS # last layer is NMS 182 | if mode and not present: 183 | print('Adding NMS... ') 184 | m = NMS() # module 185 | m.f = -1 # from 186 | m.i = self.model[-1].i + 1 # index 187 | self.model.add_module(name='%s' % m.i, module=m) # add 188 | self.eval() 189 | elif not mode and present: 190 | print('Removing NMS... ') 191 | self.model = self.model[:-1] # remove 192 | return self 193 | 194 | def autoshape(self): # add autoShape module 195 | print('Adding autoShape... ') 196 | m = autoShape(self) # wrap model 197 | copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes 198 | return m 199 | 200 | def info(self, verbose=False, img_size=640): # print model information 201 | model_info(self, verbose, img_size) 202 | 203 | 204 | def parse_model(d, ch): # model_dict, input_channels(3) 205 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 206 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 207 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 208 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 209 | 210 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 211 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 212 | m = eval(m) if isinstance(m, str) else m # eval strings 213 | for j, a in enumerate(args): 214 | try: 215 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 216 | except: 217 | pass 218 | 219 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 220 | if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, 221 | C3]: 222 | c1, c2 = ch[f], args[0] 223 | if c2 != no: # if not output 224 | c2 = make_divisible(c2 * gw, 8) 225 | 226 | args = [c1, c2, *args[1:]] 227 | if m in [BottleneckCSP, C3]: 228 | args.insert(2, n) # number of repeats 229 | n = 1 230 | elif m is nn.BatchNorm2d: 231 | args = [ch[f]] 232 | elif m is Concat: 233 | c2 = sum([ch[x] for x in f]) 234 | elif m is Detect: 235 | args.append([ch[x] for x in f]) 236 | if isinstance(args[1], int): # number of anchors 237 | args[1] = [list(range(args[1] * 2))] * len(f) 238 | elif m is Contract: 239 | c2 = ch[f] * args[0] ** 2 240 | elif m is Expand: 241 | c2 = ch[f] // args[0] ** 2 242 | else: 243 | c2 = ch[f] 244 | 245 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 246 | t = str(m)[8:-2].replace('__main__.', '') # module type 247 | np = sum([x.numel() for x in m_.parameters()]) # number params 248 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 249 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 250 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 251 | layers.append(m_) 252 | if i == 0: 253 | ch = [] 254 | ch.append(c2) 255 | return nn.Sequential(*layers), sorted(save) 256 | 257 | 258 | if __name__ == '__main__': 259 | parser = argparse.ArgumentParser() 260 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 261 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 262 | opt = parser.parse_args() 263 | opt.cfg = check_file(opt.cfg) # check file 264 | set_logging() 265 | device = select_device(opt.device) 266 | 267 | # Create model 268 | model = Model(opt.cfg).to(device) 269 | model.train() 270 | 271 | # Profile 272 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 273 | # y = model(img, profile=True) 274 | 275 | # Tensorboard 276 | # from torch.utils.tensorboard import SummaryWriter 277 | # tb_writer = SummaryWriter() 278 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 279 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 280 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 281 | -------------------------------------------------------------------------------- /models/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, C3, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.67 # model depth multiple 4 | width_multiple: 0.75 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, C3, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5s.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.33 # model depth multiple 4 | width_multiple: 0.50 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, C3, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5x.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.33 # model depth multiple 4 | width_multiple: 1.25 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, C3, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /pre_compile/rknn2precompile.py: -------------------------------------------------------------------------------- 1 | from rknn.api import RKNN 2 | if __name__ == '__main__': 3 | # Create RKNN object 4 | rknn = RKNN() 5 | 6 | 7 | # Load rknn model 8 | ret = rknn.load_rknn('./best_as_200.rknn') 9 | if ret != 0: 10 | print('Load RKNN model failed.') 11 | exit(ret) 12 | # init runtime 13 | ret = rknn.init_runtime(target='rk3399pro', rknn2precompile=True) 14 | if ret != 0: 15 | print('Init runtime failed.') 16 | exit(ret) 17 | # Note: the rknn2precompile must be set True when call init_runtime 18 | ret = rknn.export_rknn_precompile_model('./best_pre_compile.rknn') 19 | if ret != 0: 20 | print('export pre-compile model failed.') 21 | exit(ret) 22 | rknn.release() -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # base ---------------------------------------- 4 | Cython 5 | matplotlib>=3.2.2 6 | numpy>=1.18.5 7 | opencv-python>=4.1.2 8 | Pillow 9 | PyYAML>=5.3.1 10 | scipy>=1.4.1 11 | torch>=1.7.0 12 | torchvision>=0.8.1 13 | tqdm>=4.41.0 14 | 15 | # logging ------------------------------------- 16 | tensorboard>=2.4.1 17 | # wandb 18 | 19 | # plotting ------------------------------------ 20 | seaborn>=0.11.0 21 | pandas 22 | 23 | # export -------------------------------------- 24 | # coremltools>=4.1 25 | # onnx>=1.8.1 26 | # scikit-learn==0.19.2 # for coreml quantization 27 | 28 | # extras -------------------------------------- 29 | thop # FLOPS computation 30 | pycocotools>=2.0 # COCO mAP 31 | -------------------------------------------------------------------------------- /rknn_detect/models/yolov5_rknn_640x640.yaml: -------------------------------------------------------------------------------- 1 | class: rknn_detect_yolov5.Detector 2 | opt: 3 | model: "weights/yolov5s.rknn" 4 | size: [ 640, 640 ] 5 | masks: [ [ 0, 1, 2 ], [ 3, 4, 5 ], [ 6, 7, 8 ] ] 6 | anchors: [ [ 10,13 ], [ 16,30 ], [ 33,23 ], [ 30,61 ], [ 62,45 ], [ 59,119 ], [ 116,90 ], [ 156,198 ], [ 373,326 ] ] 7 | names: [ 'person', 'rider', 'car', 'bus', 'truck', 'bike', 'motor' ] 8 | conf_thres: 0.3 9 | iou_thres: 0.5 10 | platform: 0 11 | -------------------------------------------------------------------------------- /rknn_detect/rknn_detect_for_yolov5_original.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import time 3 | import random 4 | import numpy as np 5 | from rknn.api import RKNN 6 | 7 | """ 8 | yolov5 官方原版 预测脚本 for rknn 9 | """ 10 | 11 | 12 | def get_max_scale(img, max_w, max_h): 13 | h, w = img.shape[:2] 14 | scale = min(max_w / w, max_h / h, 1) 15 | return scale 16 | 17 | 18 | def get_new_size(img, scale): 19 | return tuple(map(int, np.array(img.shape[:2][::-1]) * scale)) 20 | 21 | 22 | def sigmoid(x): 23 | return 1 / (1 + np.exp(-x)) 24 | 25 | 26 | def filter_boxes(boxes, box_confidences, box_class_probs, conf_thres): 27 | box_scores = box_confidences * box_class_probs # 条件概率， 在该cell存在物体的概率的基础上是某个类别的概率 28 | box_classes = np.argmax(box_scores, axis=-1) # 找出概率最大的类别索引 29 | box_class_scores = np.max(box_scores, axis=-1) # 最大类别对应的概率值 30 | pos = np.where(box_class_scores >= conf_thres) # 找出概率大于阈值的item 31 | # pos = box_class_scores >= OBJ_THRESH # 找出概率大于阈值的item 32 | boxes = boxes[pos] 33 | classes = box_classes[pos] 34 | scores = box_class_scores[pos] 35 | return boxes, classes, scores 36 | 37 | 38 | def nms_boxes(boxes, scores, iou_thres): 39 | x = boxes[:, 0] 40 | y = boxes[:, 1] 41 | w = boxes[:, 2] 42 | h = boxes[:, 3] 43 | 44 | areas = w * h 45 | order = scores.argsort()[::-1] 46 | 47 | keep = [] 48 | while order.size > 0: 49 | i = order[0] 50 | keep.append(i) 51 | 52 | xx1 = np.maximum(x[i], x[order[1:]]) 53 | yy1 = np.maximum(y[i], y[order[1:]]) 54 | xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) 55 | yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) 56 | 57 | w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) 58 | h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) 59 | inter = w1 * h1 60 | 61 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 62 | inds = np.where(ovr <= iou_thres)[0] 63 | order = order[inds + 1] 64 | keep = np.array(keep) 65 | return keep 66 | 67 | 68 | def plot_one_box(x, img, color=None, label=None, line_thickness=None): 69 | tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness 70 | color = color or [random.randint(0, 255) for _ in range(3)] 71 | c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) 72 | cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) 73 | if label: 74 | tf = max(tl - 1, 1) # font thickness 75 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] 76 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 77 | cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled 78 | cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) 79 | 80 | 81 | def auto_resize(img, max_w, max_h): 82 | h, w = img.shape[:2] 83 | scale = min(max_w / w, max_h / h, 1) 84 | new_size = tuple(map(int, np.array(img.shape[:2][::-1]) * scale)) 85 | return cv2.resize(img, new_size), scale 86 | 87 | 88 | def letterbox(img, new_wh=(416, 416), color=(114, 114, 114)): 89 | new_img, scale = auto_resize(img, *new_wh) 90 | shape = new_img.shape 91 | new_img = cv2.copyMakeBorder(new_img, 0, new_wh[1] - shape[0], 0, new_wh[0] - shape[1], cv2.BORDER_CONSTANT, 92 | value=color) 93 | return new_img, (new_wh[0] / scale, new_wh[1] / scale) 94 | 95 | 96 | def load_model(model_path, npu_id): 97 | rknn = RKNN() 98 | devs = rknn.list_devices() 99 | device_id_dict = {} 100 | for index, dev_id in enumerate(devs[-1]): 101 | if dev_id[:2] != 'TS': 102 | device_id_dict[0] = dev_id 103 | if dev_id[:2] == 'TS': 104 | device_id_dict[1] = dev_id 105 | print('-->loading model : ' + model_path) 106 | rknn.load_rknn(model_path) 107 | print('--> Init runtime environment on: ' + device_id_dict[npu_id]) 108 | ret = rknn.init_runtime(device_id=device_id_dict[npu_id]) 109 | if ret != 0: 110 | print('Init runtime environment failed') 111 | exit(ret) 112 | print('done') 113 | return rknn 114 | 115 | 116 | class Detector: 117 | def __init__(self, opt): 118 | self.opt = opt 119 | 120 | model = opt['model'] 121 | wh = opt['size'] 122 | masks = opt['masks'] 123 | anchors = opt['anchors'] 124 | names = opt['names'] 125 | conf_thres = opt['conf_thres'] 126 | iou_thres = opt['iou_thres'] 127 | platform = opt['platform'] 128 | 129 | self.wh = wh 130 | self.size = wh 131 | self._masks = masks 132 | self._anchors = anchors 133 | self.names = list( 134 | filter(lambda a: len(a) > 0, map(lambda x: x.strip(), open(names, "r").read().split()))) if isinstance( 135 | names, str) else names 136 | self.conf_thres = conf_thres 137 | self.iou_thres = iou_thres 138 | if isinstance(model, str): 139 | model = load_model(model, platform) 140 | self._rknn = model 141 | self.draw_box = False 142 | 143 | def _predict(self, img_src, img, gain): 144 | src_h, src_w = img_src.shape[:2] 145 | # _img = cv2.cvtColor(_img, cv2.COLOR_BGR2RGB) 146 | img = img[..., ::-1] # 147 | img = np.concatenate([img[::2, ::2], img[1::2, ::2], img[::2, 1::2], img[1::2, 1::2]], 2) 148 | t0 = time.time() 149 | pred_onx = self._rknn.inference(inputs=[img]) 150 | print("inference time:\t", time.time() - t0) 151 | boxes, classes, scores = [], [], [] 152 | for t in range(3): 153 | input0_data = sigmoid(pred_onx[t][0]) 154 | input0_data = np.transpose(input0_data, (1, 2, 0, 3)) 155 | grid_h, grid_w, channel_n, predict_n = input0_data.shape 156 | anchors = [self._anchors[i] for i in self._masks[t]] 157 | box_confidence = input0_data[..., 4] 158 | box_confidence = np.expand_dims(box_confidence, axis=-1) 159 | box_class_probs = input0_data[..., 5:] 160 | box_xy = input0_data[..., :2] 161 | box_wh = input0_data[..., 2:4] 162 | col = np.tile(np.arange(0, grid_w), grid_h).reshape(-1, grid_w) 163 | row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_w) 164 | col = col.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2) 165 | row = row.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2) 166 | grid = np.concatenate((col, row), axis=-1) 167 | box_xy = box_xy * 2 - 0.5 + grid 168 | box_wh = (box_wh * 2) ** 2 * anchors 169 | box_xy /= (grid_w, grid_h) # 计算原尺寸的中心 170 | box_wh /= self.wh # 计算原尺寸的宽高 171 | box_xy -= (box_wh / 2.) # 计算原尺寸的中心 172 | box = np.concatenate((box_xy, box_wh), axis=-1) 173 | res = filter_boxes(box, box_confidence, box_class_probs, self.conf_thres) 174 | boxes.append(res[0]) 175 | classes.append(res[1]) 176 | scores.append(res[2]) 177 | boxes, classes, scores = np.concatenate(boxes), np.concatenate(classes), np.concatenate(scores) 178 | nboxes, nclasses, nscores = [], [], [] 179 | for c in set(classes): 180 | inds = np.where(classes == c) 181 | b = boxes[inds] 182 | c = classes[inds] 183 | s = scores[inds] 184 | keep = nms_boxes(b, s, self.iou_thres) 185 | nboxes.append(b[keep]) 186 | nclasses.append(c[keep]) 187 | nscores.append(s[keep]) 188 | if len(nboxes) < 1: 189 | return [], [] 190 | boxes = np.concatenate(nboxes) 191 | classes = np.concatenate(nclasses) 192 | scores = np.concatenate(nscores) 193 | label_list = [] 194 | box_list = [] 195 | for (x, y, w, h), score, cl in zip(boxes, scores, classes): 196 | x *= gain[0] 197 | y *= gain[1] 198 | w *= gain[0] 199 | h *= gain[1] 200 | x1 = max(0, np.floor(x).astype(int)) 201 | y1 = max(0, np.floor(y).astype(int)) 202 | x2 = min(src_w, np.floor(x + w + 0.5).astype(int)) 203 | y2 = min(src_h, np.floor(y + h + 0.5).astype(int)) 204 | # label_list.append(self.names[cl]) 205 | label_list.append(cl) 206 | box_list.append((x1, y1, x2, y2)) 207 | if self.draw_box: 208 | plot_one_box((x1, y1, x2, y2), img_src, label=self.names[cl]) 209 | return label_list, np.array(box_list) 210 | 211 | def detect_resize(self, img_src): 212 | """ 213 | 预测一张图片，预处理使用resize 214 | return: labels,boxes 215 | """ 216 | _img = cv2.resize(img_src, self.wh) 217 | gain = img_src.shape[:2][::-1] 218 | return self._predict(img_src, _img, gain) 219 | 220 | def detect(self, img_src): 221 | """ 222 | 预测一张图片，预处理保持宽高比 223 | return: labels,boxes 224 | """ 225 | _img, gain = letterbox(img_src, self.wh) 226 | return self._predict(img_src, _img, gain) 227 | 228 | def close(self): 229 | self._rknn.release() 230 | 231 | def __enter__(self): 232 | return self 233 | 234 | def __exit__(self, exc_type, exc_val, exc_tb): 235 | self.close() 236 | 237 | def __del__(self): 238 | self.close() 239 | 240 | 241 | def test_video(det, video_path): 242 | reader = cv2.VideoCapture() 243 | reader.open(video_path) 244 | while True: 245 | ret, frame = reader.read() 246 | if not ret: 247 | break 248 | t0 = time.time() 249 | det.detect(frame) 250 | print("total time", time.time() - t0) 251 | cv2.imshow("res", auto_resize(frame, 1200, 600)[0]) 252 | cv2.waitKey(1) 253 | 254 | 255 | if __name__ == '__main__': 256 | import yaml 257 | import cv2 258 | 259 | image = cv2.imread("data/images/bus.jpg") 260 | with open("models/yolov5_rknn_640x640.yaml", "rb") as f: 261 | cfg = yaml.load(f, yaml.FullLoader) 262 | d = Detector(cfg["opt"]) 263 | d.draw_box = True 264 | d.detect(image) 265 | cv2.imshow("res", image) 266 | cv2.waitKey() 267 | cv2.destroyAllWindows() 268 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/littledeep/YOLOv5-RK3399Pro/37a81a703ba7f82261c09d4316a55182ae84a248/utils/__init__.py -------------------------------------------------------------------------------- /utils/activations.py: -------------------------------------------------------------------------------- 1 | # Activation functions 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | # SiLU https://arxiv.org/pdf/1606.08415.pdf ---------------------------------------------------------------------------- 9 | class SiLU(nn.Module): # export-friendly version of nn.SiLU() 10 | @staticmethod 11 | def forward(x): 12 | return x * torch.sigmoid(x) 13 | 14 | 15 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 16 | @staticmethod 17 | def forward(x): 18 | # return x * F.hardsigmoid(x) # for torchscript and CoreML 19 | return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX 20 | 21 | 22 | class MemoryEfficientSwish(nn.Module): 23 | class F(torch.autograd.Function): 24 | @staticmethod 25 | def forward(ctx, x): 26 | ctx.save_for_backward(x) 27 | return x * torch.sigmoid(x) 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | x = ctx.saved_tensors[0] 32 | sx = torch.sigmoid(x) 33 | return grad_output * (sx * (1 + x * (1 - sx))) 34 | 35 | def forward(self, x): 36 | return self.F.apply(x) 37 | 38 | 39 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 40 | class Mish(nn.Module): 41 | @staticmethod 42 | def forward(x): 43 | return x * F.softplus(x).tanh() 44 | 45 | 46 | class MemoryEfficientMish(nn.Module): 47 | class F(torch.autograd.Function): 48 | @staticmethod 49 | def forward(ctx, x): 50 | ctx.save_for_backward(x) 51 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 52 | 53 | @staticmethod 54 | def backward(ctx, grad_output): 55 | x = ctx.saved_tensors[0] 56 | sx = torch.sigmoid(x) 57 | fx = F.softplus(x).tanh() 58 | return grad_output * (fx + x * sx * (1 - fx * fx)) 59 | 60 | def forward(self, x): 61 | return self.F.apply(x) 62 | 63 | 64 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 65 | class FReLU(nn.Module): 66 | def __init__(self, c1, k=3): # ch_in, kernel 67 | super().__init__() 68 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False) 69 | self.bn = nn.BatchNorm2d(c1) 70 | 71 | def forward(self, x): 72 | return torch.max(x, self.bn(self.conv(x))) 73 | -------------------------------------------------------------------------------- /utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # Auto-anchor utils 2 | 3 | import numpy as np 4 | import torch 5 | import yaml 6 | from scipy.cluster.vq import kmeans 7 | from tqdm import tqdm 8 | 9 | from utils.general import colorstr 10 | 11 | 12 | def check_anchor_order(m): 13 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 14 | a = m.anchor_grid.prod(-1).view(-1) # anchor area 15 | da = a[-1] - a[0] # delta a 16 | ds = m.stride[-1] - m.stride[0] # delta s 17 | if da.sign() != ds.sign(): # same order 18 | print('Reversing anchor order') 19 | m.anchors[:] = m.anchors.flip(0) 20 | m.anchor_grid[:] = m.anchor_grid.flip(0) 21 | 22 | 23 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 24 | # Check anchor fit to data, recompute if necessary 25 | prefix = colorstr('autoanchor: ') 26 | print(f'\n{prefix}Analyzing anchors... ', end='') 27 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 28 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 29 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 30 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 31 | 32 | def metric(k): # compute metric 33 | r = wh[:, None] / k[None] 34 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 35 | best = x.max(1)[0] # best_x 36 | aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold 37 | bpr = (best > 1. / thr).float().mean() # best possible recall 38 | return bpr, aat 39 | 40 | anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors 41 | bpr, aat = metric(anchors) 42 | print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='') 43 | if bpr < 0.98: # threshold to recompute 44 | print('. Attempting to improve anchors, please wait...') 45 | na = m.anchor_grid.numel() // 2 # number of anchors 46 | try: 47 | anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 48 | except Exception as e: 49 | print(f'{prefix}ERROR: {e}') 50 | new_bpr = metric(anchors)[0] 51 | if new_bpr > bpr: # replace anchors 52 | anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors) 53 | m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference 54 | m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss 55 | check_anchor_order(m) 56 | print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.') 57 | else: 58 | print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.') 59 | print('') # newline 60 | 61 | 62 | def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 63 | """ Creates kmeans-evolved anchors from training dataset 64 | 65 | Arguments: 66 | path: path to dataset *.yaml, or a loaded dataset 67 | n: number of anchors 68 | img_size: image size used for training 69 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 70 | gen: generations to evolve anchors using genetic algorithm 71 | verbose: print all results 72 | 73 | Return: 74 | k: kmeans evolved anchors 75 | 76 | Usage: 77 | from utils.autoanchor import *; _ = kmean_anchors() 78 | """ 79 | thr = 1. / thr 80 | prefix = colorstr('autoanchor: ') 81 | 82 | def metric(k, wh): # compute metrics 83 | r = wh[:, None] / k[None] 84 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 85 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 86 | return x, x.max(1)[0] # x, best_x 87 | 88 | def anchor_fitness(k): # mutation fitness 89 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 90 | return (best * (best > thr).float()).mean() # fitness 91 | 92 | def print_results(k): 93 | k = k[np.argsort(k.prod(1))] # sort small to large 94 | x, best = metric(k, wh0) 95 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 96 | print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr') 97 | print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' 98 | f'past_thr={x[x > thr].mean():.3f}-mean: ', end='') 99 | for i, x in enumerate(k): 100 | print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg 101 | return k 102 | 103 | if isinstance(path, str): # *.yaml file 104 | with open(path) as f: 105 | data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict 106 | from utils.datasets import LoadImagesAndLabels 107 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 108 | else: 109 | dataset = path # dataset 110 | 111 | # Get label wh 112 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 113 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 114 | 115 | # Filter 116 | i = (wh0 < 3.0).any(1).sum() 117 | if i: 118 | print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.') 119 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels 120 | # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 121 | 122 | # Kmeans calculation 123 | print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...') 124 | s = wh.std(0) # sigmas for whitening 125 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance 126 | assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}') 127 | k *= s 128 | wh = torch.tensor(wh, dtype=torch.float32) # filtered 129 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered 130 | k = print_results(k) 131 | 132 | # Plot 133 | # k, d = [None] * 20, [None] * 20 134 | # for i in tqdm(range(1, 21)): 135 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 136 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 137 | # ax = ax.ravel() 138 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 139 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 140 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 141 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 142 | # fig.savefig('wh.png', dpi=200) 143 | 144 | # Evolve 145 | npr = np.random 146 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 147 | pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar 148 | for _ in pbar: 149 | v = np.ones(sh) 150 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 151 | v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 152 | kg = (k.copy() * v).clip(min=2.0) 153 | fg = anchor_fitness(kg) 154 | if fg > f: 155 | f, k = fg, kg.copy() 156 | pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' 157 | if verbose: 158 | print_results(k) 159 | 160 | return print_results(k) 161 | -------------------------------------------------------------------------------- /utils/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/littledeep/YOLOv5-RK3399Pro/37a81a703ba7f82261c09d4316a55182ae84a248/utils/aws/__init__.py -------------------------------------------------------------------------------- /utils/aws/mime.sh: -------------------------------------------------------------------------------- 1 | # AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/ 2 | # This script will run on every instance restart, not only on first start 3 | # --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA --- 4 | 5 | Content-Type: multipart/mixed; boundary="//" 6 | MIME-Version: 1.0 7 | 8 | --// 9 | Content-Type: text/cloud-config; charset="us-ascii" 10 | MIME-Version: 1.0 11 | Content-Transfer-Encoding: 7bit 12 | Content-Disposition: attachment; filename="cloud-config.txt" 13 | 14 | #cloud-config 15 | cloud_final_modules: 16 | - [scripts-user, always] 17 | 18 | --// 19 | Content-Type: text/x-shellscript; charset="us-ascii" 20 | MIME-Version: 1.0 21 | Content-Transfer-Encoding: 7bit 22 | Content-Disposition: attachment; filename="userdata.txt" 23 | 24 | #!/bin/bash 25 | # --- paste contents of userdata.sh here --- 26 | --// 27 | -------------------------------------------------------------------------------- /utils/aws/resume.py: -------------------------------------------------------------------------------- 1 | # Resume all interrupted trainings in yolov5/ dir including DPP trainings 2 | # Usage: $ python utils/aws/resume.py 3 | 4 | import os 5 | import sys 6 | from pathlib import Path 7 | 8 | import torch 9 | import yaml 10 | 11 | sys.path.append('./') # to run '$ python *.py' files in subdirectories 12 | 13 | port = 0 # --master_port 14 | path = Path('').resolve() 15 | for last in path.rglob('*/**/last.pt'): 16 | ckpt = torch.load(last) 17 | if ckpt['optimizer'] is None: 18 | continue 19 | 20 | # Load opt.yaml 21 | with open(last.parent.parent / 'opt.yaml') as f: 22 | opt = yaml.load(f, Loader=yaml.SafeLoader) 23 | 24 | # Get device count 25 | d = opt['device'].split(',') # devices 26 | nd = len(d) # number of devices 27 | ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel 28 | 29 | if ddp: # multi-GPU 30 | port += 1 31 | cmd = f'python -m torch.distributed.launch --nproc_per_node {nd} --master_port {port} train.py --resume {last}' 32 | else: # single-GPU 33 | cmd = f'python train.py --resume {last}' 34 | 35 | cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread 36 | print(cmd) 37 | os.system(cmd) 38 | -------------------------------------------------------------------------------- /utils/aws/userdata.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html 3 | # This script will run only once on first instance start (for a re-start script see mime.sh) 4 | # /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir 5 | # Use >300 GB SSD 6 | 7 | cd home/ubuntu 8 | if [ ! -d yolov5 ]; then 9 | echo "Running first-time script." # install dependencies, download COCO, pull Docker 10 | git clone https://github.com/ultralytics/yolov5 && sudo chmod -R 777 yolov5 11 | cd yolov5 12 | bash data/scripts/get_coco.sh && echo "Data done." & 13 | sudo docker pull ultralytics/yolov5:latest && echo "Docker done." & 14 | python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." & 15 | wait && echo "All tasks done." # finish background tasks 16 | else 17 | echo "Running re-start script." # resume interrupted runs 18 | i=0 19 | list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour' 20 | while IFS= read -r id; do 21 | ((i++)) 22 | echo "restarting container $i: $id" 23 | sudo docker start $id 24 | # sudo docker exec -it $id python train.py --resume # single-GPU 25 | sudo docker exec -d $id python utils/aws/resume.py # multi-scenario 26 | done <<<"$list" 27 | fi 28 | -------------------------------------------------------------------------------- /utils/google_app_engine/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM gcr.io/google-appengine/python 2 | 3 | # Create a virtualenv for dependencies. This isolates these packages from 4 | # system-level packages. 5 | # Use -p python3 or -p python3.7 to select python version. Default is version 2. 6 | RUN virtualenv /env -p python3 7 | 8 | # Setting these environment variables are the same as running 9 | # source /env/bin/activate. 10 | ENV VIRTUAL_ENV /env 11 | ENV PATH /env/bin:$PATH 12 | 13 | RUN apt-get update && apt-get install -y python-opencv 14 | 15 | # Copy the application's requirements.txt and run pip to install all 16 | # dependencies into the virtualenv. 17 | ADD requirements.txt /app/requirements.txt 18 | RUN pip install -r /app/requirements.txt 19 | 20 | # Add the application source code. 21 | ADD . /app 22 | 23 | # Run a WSGI server to serve the application. gunicorn must be declared as 24 | # a dependency in requirements.txt. 25 | CMD gunicorn -b :$PORT main:app 26 | -------------------------------------------------------------------------------- /utils/google_app_engine/additional_requirements.txt: -------------------------------------------------------------------------------- 1 | # add these requirements in your app on top of the existing ones 2 | pip==18.1 3 | Flask==1.0.2 4 | gunicorn==19.9.0 5 | -------------------------------------------------------------------------------- /utils/google_app_engine/app.yaml: -------------------------------------------------------------------------------- 1 | runtime: custom 2 | env: flex 3 | 4 | service: yolov5app 5 | 6 | liveness_check: 7 | initial_delay_sec: 600 8 | 9 | manual_scaling: 10 | instances: 1 11 | resources: 12 | cpu: 1 13 | memory_gb: 4 14 | disk_size_gb: 20 -------------------------------------------------------------------------------- /utils/google_utils.py: -------------------------------------------------------------------------------- 1 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries 2 | 3 | import os 4 | import platform 5 | import subprocess 6 | import time 7 | from pathlib import Path 8 | 9 | import requests 10 | import torch 11 | 12 | 13 | def gsutil_getsize(url=''): 14 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 15 | s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8') 16 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 17 | 18 | 19 | def attempt_download(file, repo='ultralytics/yolov5'): 20 | # Attempt file download if does not exist 21 | file = Path(str(file).strip().replace("'", '').lower()) 22 | 23 | if not file.exists(): 24 | try: 25 | response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api 26 | assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] 27 | tag = response['tag_name'] # i.e. 'v1.0' 28 | except: # fallback plan 29 | assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] 30 | tag = subprocess.check_output('git tag', shell=True).decode().split()[-1] 31 | 32 | name = file.name 33 | if name in assets: 34 | msg = f'{file} missing, try downloading from https://github.com/{repo}/releases/' 35 | redundant = False # second download option 36 | try: # GitHub 37 | url = f'https://github.com/{repo}/releases/download/{tag}/{name}' 38 | print(f'Downloading {url} to {file}...') 39 | torch.hub.download_url_to_file(url, file) 40 | assert file.exists() and file.stat().st_size > 1E6 # check 41 | except Exception as e: # GCP 42 | print(f'Download error: {e}') 43 | assert redundant, 'No secondary mirror' 44 | url = f'https://storage.googleapis.com/{repo}/ckpt/{name}' 45 | print(f'Downloading {url} to {file}...') 46 | os.system(f'curl -L {url} -o {file}') # torch.hub.download_url_to_file(url, weights) 47 | finally: 48 | if not file.exists() or file.stat().st_size < 1E6: # check 49 | file.unlink(missing_ok=True) # remove partial downloads 50 | print(f'ERROR: Download failure: {msg}') 51 | print('') 52 | return 53 | 54 | 55 | def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): 56 | # Downloads a file from Google Drive. from yolov5.utils.google_utils import *; gdrive_download() 57 | t = time.time() 58 | file = Path(file) 59 | cookie = Path('cookie') # gdrive cookie 60 | print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='') 61 | file.unlink(missing_ok=True) # remove existing file 62 | cookie.unlink(missing_ok=True) # remove existing cookie 63 | 64 | # Attempt file download 65 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 66 | os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}') 67 | if os.path.exists('cookie'): # large file 68 | s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}' 69 | else: # small file 70 | s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"' 71 | r = os.system(s) # execute, capture return 72 | cookie.unlink(missing_ok=True) # remove existing cookie 73 | 74 | # Error check 75 | if r != 0: 76 | file.unlink(missing_ok=True) # remove partial 77 | print('Download error ') # raise Exception('Download error') 78 | return r 79 | 80 | # Unzip if archive 81 | if file.suffix == '.zip': 82 | print('unzipping... ', end='') 83 | os.system(f'unzip -q {file}') # unzip 84 | file.unlink() # remove zip to free space 85 | 86 | print(f'Done ({time.time() - t:.1f}s)') 87 | return r 88 | 89 | 90 | def get_token(cookie="./cookie"): 91 | with open(cookie) as f: 92 | for line in f: 93 | if "download" in line: 94 | return line.split()[-1] 95 | return "" 96 | 97 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 98 | # # Uploads a file to a bucket 99 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 100 | # 101 | # storage_client = storage.Client() 102 | # bucket = storage_client.get_bucket(bucket_name) 103 | # blob = bucket.blob(destination_blob_name) 104 | # 105 | # blob.upload_from_filename(source_file_name) 106 | # 107 | # print('File {} uploaded to {}.'.format( 108 | # source_file_name, 109 | # destination_blob_name)) 110 | # 111 | # 112 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 113 | # # Uploads a blob from a bucket 114 | # storage_client = storage.Client() 115 | # bucket = storage_client.get_bucket(bucket_name) 116 | # blob = bucket.blob(source_blob_name) 117 | # 118 | # blob.download_to_filename(destination_file_name) 119 | # 120 | # print('Blob {} downloaded to {}.'.format( 121 | # source_blob_name, 122 | # destination_file_name)) 123 | -------------------------------------------------------------------------------- /utils/loss.py: -------------------------------------------------------------------------------- 1 | # Loss functions 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from utils.general import bbox_iou 7 | from utils.torch_utils import is_parallel 8 | 9 | 10 | def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441 11 | # return positive, negative label smoothing BCE targets 12 | return 1.0 - 0.5 * eps, 0.5 * eps 13 | 14 | 15 | class BCEBlurWithLogitsLoss(nn.Module): 16 | # BCEwithLogitLoss() with reduced missing label effects. 17 | def __init__(self, alpha=0.05): 18 | super(BCEBlurWithLogitsLoss, self).__init__() 19 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss() 20 | self.alpha = alpha 21 | 22 | def forward(self, pred, true): 23 | loss = self.loss_fcn(pred, true) 24 | pred = torch.sigmoid(pred) # prob from logits 25 | dx = pred - true # reduce only missing label effects 26 | # dx = (pred - true).abs() # reduce missing label and false label effects 27 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4)) 28 | loss *= alpha_factor 29 | return loss.mean() 30 | 31 | 32 | class FocalLoss(nn.Module): 33 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 34 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 35 | super(FocalLoss, self).__init__() 36 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 37 | self.gamma = gamma 38 | self.alpha = alpha 39 | self.reduction = loss_fcn.reduction 40 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 41 | 42 | def forward(self, pred, true): 43 | loss = self.loss_fcn(pred, true) 44 | # p_t = torch.exp(-loss) 45 | # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability 46 | 47 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py 48 | pred_prob = torch.sigmoid(pred) # prob from logits 49 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob) 50 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 51 | modulating_factor = (1.0 - p_t) ** self.gamma 52 | loss *= alpha_factor * modulating_factor 53 | 54 | if self.reduction == 'mean': 55 | return loss.mean() 56 | elif self.reduction == 'sum': 57 | return loss.sum() 58 | else: # 'none' 59 | return loss 60 | 61 | 62 | class QFocalLoss(nn.Module): 63 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 64 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 65 | super(QFocalLoss, self).__init__() 66 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 67 | self.gamma = gamma 68 | self.alpha = alpha 69 | self.reduction = loss_fcn.reduction 70 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 71 | 72 | def forward(self, pred, true): 73 | loss = self.loss_fcn(pred, true) 74 | 75 | pred_prob = torch.sigmoid(pred) # prob from logits 76 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 77 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma 78 | loss *= alpha_factor * modulating_factor 79 | 80 | if self.reduction == 'mean': 81 | return loss.mean() 82 | elif self.reduction == 'sum': 83 | return loss.sum() 84 | else: # 'none' 85 | return loss 86 | 87 | 88 | class ComputeLoss: 89 | # Compute losses 90 | def __init__(self, model, autobalance=False): 91 | super(ComputeLoss, self).__init__() 92 | device = next(model.parameters()).device # get model device 93 | h = model.hyp # hyperparameters 94 | 95 | # Define criteria 96 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) 97 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device)) 98 | 99 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 100 | self.cp, self.cn = smooth_BCE(eps=0.0) 101 | 102 | # Focal loss 103 | g = h['fl_gamma'] # focal loss gamma 104 | if g > 0: 105 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) 106 | 107 | det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module 108 | self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7 109 | self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index 110 | self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance 111 | for k in 'na', 'nc', 'nl', 'anchors': 112 | setattr(self, k, getattr(det, k)) 113 | 114 | def __call__(self, p, targets): # predictions, targets, model 115 | device = targets.device 116 | lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) 117 | tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets 118 | 119 | # Losses 120 | for i, pi in enumerate(p): # layer index, layer predictions 121 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx 122 | tobj = torch.zeros_like(pi[..., 0], device=device) # target obj 123 | 124 | n = b.shape[0] # number of targets 125 | if n: 126 | ps = pi[b, a, gj, gi] # prediction subset corresponding to targets 127 | 128 | # Regression 129 | pxy = ps[:, :2].sigmoid() * 2. - 0.5 130 | pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i] 131 | pbox = torch.cat((pxy, pwh), 1) # predicted box 132 | iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target) 133 | lbox += (1.0 - iou).mean() # iou loss 134 | 135 | # Objectness 136 | tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio 137 | 138 | # Classification 139 | if self.nc > 1: # cls loss (only if multiple classes) 140 | t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets 141 | t[range(n), tcls[i]] = self.cp 142 | lcls += self.BCEcls(ps[:, 5:], t) # BCE 143 | 144 | # Append targets to text file 145 | # with open('targets.txt', 'a') as file: 146 | # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] 147 | 148 | obji = self.BCEobj(pi[..., 4], tobj) 149 | lobj += obji * self.balance[i] # obj loss 150 | if self.autobalance: 151 | self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item() 152 | 153 | if self.autobalance: 154 | self.balance = [x / self.balance[self.ssi] for x in self.balance] 155 | lbox *= self.hyp['box'] 156 | lobj *= self.hyp['obj'] 157 | lcls *= self.hyp['cls'] 158 | bs = tobj.shape[0] # batch size 159 | 160 | loss = lbox + lobj + lcls 161 | return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach() 162 | 163 | def build_targets(self, p, targets): 164 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h) 165 | na, nt = self.na, targets.shape[0] # number of anchors, targets 166 | tcls, tbox, indices, anch = [], [], [], [] 167 | gain = torch.ones(7, device=targets.device) # normalized to gridspace gain 168 | ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) 169 | targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices 170 | 171 | g = 0.5 # bias 172 | off = torch.tensor([[0, 0], 173 | [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m 174 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm 175 | ], device=targets.device).float() * g # offsets 176 | 177 | for i in range(self.nl): 178 | anchors = self.anchors[i] 179 | gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain 180 | 181 | # Match targets to anchors 182 | t = targets * gain 183 | if nt: 184 | # Matches 185 | r = t[:, :, 4:6] / anchors[:, None] # wh ratio 186 | j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare 187 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) 188 | t = t[j] # filter 189 | 190 | # Offsets 191 | gxy = t[:, 2:4] # grid xy 192 | gxi = gain[[2, 3]] - gxy # inverse 193 | j, k = ((gxy % 1. < g) & (gxy > 1.)).T 194 | l, m = ((gxi % 1. < g) & (gxi > 1.)).T 195 | j = torch.stack((torch.ones_like(j), j, k, l, m)) 196 | t = t.repeat((5, 1, 1))[j] 197 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j] 198 | else: 199 | t = targets[0] 200 | offsets = 0 201 | 202 | # Define 203 | b, c = t[:, :2].long().T # image, class 204 | gxy = t[:, 2:4] # grid xy 205 | gwh = t[:, 4:6] # grid wh 206 | gij = (gxy - offsets).long() 207 | gi, gj = gij.T # grid xy indices 208 | 209 | # Append 210 | a = t[:, 6].long() # anchor indices 211 | indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices 212 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box 213 | anch.append(anchors[a]) # anchors 214 | tcls.append(c) # class 215 | 216 | return tcls, tbox, indices, anch 217 | -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | # Model validation metrics 2 | 3 | from pathlib import Path 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import torch 8 | 9 | from . import general 10 | 11 | 12 | def fitness(x): 13 | # Model fitness as a weighted combination of metrics 14 | w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95] 15 | return (x[:, :4] * w).sum(1) 16 | 17 | 18 | def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()): 19 | """ Compute the average precision, given the recall and precision curves. 20 | Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. 21 | # Arguments 22 | tp: True positives (nparray, nx1 or nx10). 23 | conf: Objectness value from 0-1 (nparray). 24 | pred_cls: Predicted object classes (nparray). 25 | target_cls: True object classes (nparray). 26 | plot: Plot precision-recall curve at mAP@0.5 27 | save_dir: Plot save directory 28 | # Returns 29 | The average precision as computed in py-faster-rcnn. 30 | """ 31 | 32 | # Sort by objectness 33 | i = np.argsort(-conf) 34 | tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] 35 | 36 | # Find unique classes 37 | unique_classes = np.unique(target_cls) 38 | nc = unique_classes.shape[0] # number of classes, number of detections 39 | 40 | # Create Precision-Recall curve and compute AP for each class 41 | px, py = np.linspace(0, 1, 1000), [] # for plotting 42 | ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000)) 43 | for ci, c in enumerate(unique_classes): 44 | i = pred_cls == c 45 | n_l = (target_cls == c).sum() # number of labels 46 | n_p = i.sum() # number of predictions 47 | 48 | if n_p == 0 or n_l == 0: 49 | continue 50 | else: 51 | # Accumulate FPs and TPs 52 | fpc = (1 - tp[i]).cumsum(0) 53 | tpc = tp[i].cumsum(0) 54 | 55 | # Recall 56 | recall = tpc / (n_l + 1e-16) # recall curve 57 | r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases 58 | 59 | # Precision 60 | precision = tpc / (tpc + fpc) # precision curve 61 | p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score 62 | 63 | # AP from recall-precision curve 64 | for j in range(tp.shape[1]): 65 | ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j]) 66 | if plot and j == 0: 67 | py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5 68 | 69 | # Compute F1 (harmonic mean of precision and recall) 70 | f1 = 2 * p * r / (p + r + 1e-16) 71 | if plot: 72 | plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names) 73 | plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1') 74 | plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision') 75 | plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall') 76 | 77 | i = f1.mean(0).argmax() # max F1 index 78 | return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32') 79 | 80 | 81 | def compute_ap(recall, precision): 82 | """ Compute the average precision, given the recall and precision curves 83 | # Arguments 84 | recall: The recall curve (list) 85 | precision: The precision curve (list) 86 | # Returns 87 | Average precision, precision curve, recall curve 88 | """ 89 | 90 | # Append sentinel values to beginning and end 91 | mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01])) 92 | mpre = np.concatenate(([1.], precision, [0.])) 93 | 94 | # Compute the precision envelope 95 | mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) 96 | 97 | # Integrate area under curve 98 | method = 'interp' # methods: 'continuous', 'interp' 99 | if method == 'interp': 100 | x = np.linspace(0, 1, 101) # 101-point interp (COCO) 101 | ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate 102 | else: # 'continuous' 103 | i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes 104 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve 105 | 106 | return ap, mpre, mrec 107 | 108 | 109 | class ConfusionMatrix: 110 | # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix 111 | def __init__(self, nc, conf=0.25, iou_thres=0.45): 112 | self.matrix = np.zeros((nc + 1, nc + 1)) 113 | self.nc = nc # number of classes 114 | self.conf = conf 115 | self.iou_thres = iou_thres 116 | 117 | def process_batch(self, detections, labels): 118 | """ 119 | Return intersection-over-union (Jaccard index) of boxes. 120 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 121 | Arguments: 122 | detections (Array[N, 6]), x1, y1, x2, y2, conf, class 123 | labels (Array[M, 5]), class, x1, y1, x2, y2 124 | Returns: 125 | None, updates confusion matrix accordingly 126 | """ 127 | detections = detections[detections[:, 4] > self.conf] 128 | gt_classes = labels[:, 0].int() 129 | detection_classes = detections[:, 5].int() 130 | iou = general.box_iou(labels[:, 1:], detections[:, :4]) 131 | 132 | x = torch.where(iou > self.iou_thres) 133 | if x[0].shape[0]: 134 | matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() 135 | if x[0].shape[0] > 1: 136 | matches = matches[matches[:, 2].argsort()[::-1]] 137 | matches = matches[np.unique(matches[:, 1], return_index=True)[1]] 138 | matches = matches[matches[:, 2].argsort()[::-1]] 139 | matches = matches[np.unique(matches[:, 0], return_index=True)[1]] 140 | else: 141 | matches = np.zeros((0, 3)) 142 | 143 | n = matches.shape[0] > 0 144 | m0, m1, _ = matches.transpose().astype(np.int16) 145 | for i, gc in enumerate(gt_classes): 146 | j = m0 == i 147 | if n and sum(j) == 1: 148 | self.matrix[gc, detection_classes[m1[j]]] += 1 # correct 149 | else: 150 | self.matrix[self.nc, gc] += 1 # background FP 151 | 152 | if n: 153 | for i, dc in enumerate(detection_classes): 154 | if not any(m1 == i): 155 | self.matrix[dc, self.nc] += 1 # background FN 156 | 157 | def matrix(self): 158 | return self.matrix 159 | 160 | def plot(self, save_dir='', names=()): 161 | try: 162 | import seaborn as sn 163 | 164 | array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize 165 | array[array < 0.005] = np.nan # don't annotate (would appear as 0.00) 166 | 167 | fig = plt.figure(figsize=(12, 9), tight_layout=True) 168 | sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size 169 | labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels 170 | sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, 171 | xticklabels=names + ['background FP'] if labels else "auto", 172 | yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1)) 173 | fig.axes[0].set_xlabel('True') 174 | fig.axes[0].set_ylabel('Predicted') 175 | fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) 176 | except Exception as e: 177 | pass 178 | 179 | def print(self): 180 | for i in range(self.nc + 1): 181 | print(' '.join(map(str, self.matrix[i]))) 182 | 183 | 184 | # Plots ---------------------------------------------------------------------------------------------------------------- 185 | 186 | def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()): 187 | # Precision-recall curve 188 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 189 | py = np.stack(py, axis=1) 190 | 191 | if 0 < len(names) < 21: # display per-class legend if < 21 classes 192 | for i, y in enumerate(py.T): 193 | ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}') # plot(recall, precision) 194 | else: 195 | ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) 196 | 197 | ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean()) 198 | ax.set_xlabel('Recall') 199 | ax.set_ylabel('Precision') 200 | ax.set_xlim(0, 1) 201 | ax.set_ylim(0, 1) 202 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 203 | fig.savefig(Path(save_dir), dpi=250) 204 | 205 | 206 | def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'): 207 | # Metric-confidence curve 208 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 209 | 210 | if 0 < len(names) < 21: # display per-class legend if < 21 classes 211 | for i, y in enumerate(py): 212 | ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric) 213 | else: 214 | ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric) 215 | 216 | y = py.mean(0) 217 | ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}') 218 | ax.set_xlabel(xlabel) 219 | ax.set_ylabel(ylabel) 220 | ax.set_xlim(0, 1) 221 | ax.set_ylim(0, 1) 222 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 223 | fig.savefig(Path(save_dir), dpi=250) 224 | -------------------------------------------------------------------------------- /utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | # PyTorch utils 2 | import logging 3 | import math 4 | import os 5 | import platform 6 | import subprocess 7 | import time 8 | from contextlib import contextmanager 9 | from copy import deepcopy 10 | from pathlib import Path 11 | 12 | import torch 13 | import torch.backends.cudnn as cudnn 14 | import torch.nn as nn 15 | import torch.nn.functional as F 16 | import torchvision 17 | 18 | try: 19 | import thop # for FLOPS computation 20 | except ImportError: 21 | thop = None 22 | logger = logging.getLogger(__name__) 23 | 24 | 25 | @contextmanager 26 | def torch_distributed_zero_first(local_rank: int): 27 | """ 28 | Decorator to make all processes in distributed training wait for each local_master to do something. 29 | """ 30 | if local_rank not in [-1, 0]: 31 | torch.distributed.barrier() 32 | yield 33 | if local_rank == 0: 34 | torch.distributed.barrier() 35 | 36 | 37 | def init_torch_seeds(seed=0): 38 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html 39 | torch.manual_seed(seed) 40 | if seed == 0: # slower, more reproducible 41 | cudnn.benchmark, cudnn.deterministic = False, True 42 | else: # faster, less reproducible 43 | cudnn.benchmark, cudnn.deterministic = True, False 44 | 45 | 46 | def git_describe(): 47 | # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe 48 | if Path('.git').exists(): 49 | return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1] 50 | else: 51 | return '' 52 | 53 | 54 | def select_device(device='', batch_size=None): 55 | # device = 'cpu' or '0' or '0,1,2,3' 56 | s = f'YOLOv5 🚀 {git_describe()} torch {torch.__version__} ' # string 57 | cpu = device.lower() == 'cpu' 58 | if cpu: 59 | os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False 60 | elif device: # non-cpu device requested 61 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 62 | assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability 63 | 64 | cuda = not cpu and torch.cuda.is_available() 65 | if cuda: 66 | n = torch.cuda.device_count() 67 | if n > 1 and batch_size: # check that batch_size is compatible with device_count 68 | assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}' 69 | space = ' ' * len(s) 70 | for i, d in enumerate(device.split(',') if device else range(n)): 71 | p = torch.cuda.get_device_properties(i) 72 | s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB 73 | else: 74 | s += 'CPU\n' 75 | 76 | logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe 77 | return torch.device('cuda:0' if cuda else 'cpu') 78 | 79 | 80 | def time_synchronized(): 81 | # pytorch-accurate time 82 | if torch.cuda.is_available(): 83 | torch.cuda.synchronize() 84 | return time.time() 85 | 86 | 87 | def profile(x, ops, n=100, device=None): 88 | # profile a pytorch module or list of modules. Example usage: 89 | # x = torch.randn(16, 3, 640, 640) # input 90 | # m1 = lambda x: x * torch.sigmoid(x) 91 | # m2 = nn.SiLU() 92 | # profile(x, [m1, m2], n=100) # profile speed over 100 iterations 93 | 94 | device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 95 | x = x.to(device) 96 | x.requires_grad = True 97 | print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '') 98 | print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}") 99 | for m in ops if isinstance(ops, list) else [ops]: 100 | m = m.to(device) if hasattr(m, 'to') else m # device 101 | m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m # type 102 | dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward 103 | try: 104 | flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPS 105 | except: 106 | flops = 0 107 | 108 | for _ in range(n): 109 | t[0] = time_synchronized() 110 | y = m(x) 111 | t[1] = time_synchronized() 112 | try: 113 | _ = y.sum().backward() 114 | t[2] = time_synchronized() 115 | except: # no backward method 116 | t[2] = float('nan') 117 | dtf += (t[1] - t[0]) * 1000 / n # ms per op forward 118 | dtb += (t[2] - t[1]) * 1000 / n # ms per op backward 119 | 120 | s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' 121 | s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list' 122 | p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters 123 | print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}') 124 | 125 | 126 | def is_parallel(model): 127 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 128 | 129 | 130 | def intersect_dicts(da, db, exclude=()): 131 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values 132 | 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} 133 | 134 | 135 | def initialize_weights(model): 136 | for m in model.modules(): 137 | t = type(m) 138 | if t is nn.Conv2d: 139 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 140 | elif t is nn.BatchNorm2d: 141 | m.eps = 1e-3 142 | m.momentum = 0.03 143 | elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 144 | m.inplace = True 145 | 146 | 147 | def find_modules(model, mclass=nn.Conv2d): 148 | # Finds layer indices matching module class 'mclass' 149 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] 150 | 151 | 152 | def sparsity(model): 153 | # Return global model sparsity 154 | a, b = 0., 0. 155 | for p in model.parameters(): 156 | a += p.numel() 157 | b += (p == 0).sum() 158 | return b / a 159 | 160 | 161 | def prune(model, amount=0.3): 162 | # Prune model to requested global sparsity 163 | import torch.nn.utils.prune as prune 164 | print('Pruning model... ', end='') 165 | for name, m in model.named_modules(): 166 | if isinstance(m, nn.Conv2d): 167 | prune.l1_unstructured(m, name='weight', amount=amount) # prune 168 | prune.remove(m, 'weight') # make permanent 169 | print(' %.3g global sparsity' % sparsity(model)) 170 | 171 | 172 | def fuse_conv_and_bn(conv, bn): 173 | # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 174 | fusedconv = nn.Conv2d(conv.in_channels, 175 | conv.out_channels, 176 | kernel_size=conv.kernel_size, 177 | stride=conv.stride, 178 | padding=conv.padding, 179 | groups=conv.groups, 180 | bias=True).requires_grad_(False).to(conv.weight.device) 181 | 182 | # prepare filters 183 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 184 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 185 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 186 | 187 | # prepare spatial bias 188 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias 189 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 190 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 191 | 192 | return fusedconv 193 | 194 | 195 | def model_info(model, verbose=False, img_size=640): 196 | # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320] 197 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 198 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 199 | if verbose: 200 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 201 | for i, (name, p) in enumerate(model.named_parameters()): 202 | name = name.replace('module_list.', '') 203 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 204 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 205 | 206 | try: # FLOPS 207 | from thop import profile 208 | stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32 209 | img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input 210 | flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS 211 | img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float 212 | fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS 213 | except (ImportError, Exception): 214 | fs = '' 215 | 216 | logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}") 217 | 218 | 219 | def load_classifier(name='resnet101', n=2): 220 | # Loads a pretrained model reshaped to n-class output 221 | model = torchvision.models.__dict__[name](pretrained=True) 222 | 223 | # ResNet model properties 224 | # input_size = [3, 224, 224] 225 | # input_space = 'RGB' 226 | # input_range = [0, 1] 227 | # mean = [0.485, 0.456, 0.406] 228 | # std = [0.229, 0.224, 0.225] 229 | 230 | # Reshape output to n classes 231 | filters = model.fc.weight.shape[1] 232 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True) 233 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True) 234 | model.fc.out_features = n 235 | return model 236 | 237 | 238 | def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416) 239 | # scales img(bs,3,y,x) by ratio constrained to gs-multiple 240 | if ratio == 1.0: 241 | return img 242 | else: 243 | h, w = img.shape[2:] 244 | s = (int(h * ratio), int(w * ratio)) # new size 245 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 246 | if not same_shape: # pad/crop img 247 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 248 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 249 | 250 | 251 | def copy_attr(a, b, include=(), exclude=()): 252 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 253 | for k, v in b.__dict__.items(): 254 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 255 | continue 256 | else: 257 | setattr(a, k, v) 258 | 259 | 260 | class ModelEMA: 261 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 262 | Keep a moving average of everything in the model state_dict (parameters and buffers). 263 | This is intended to allow functionality like 264 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 265 | A smoothed version of the weights is necessary for some training schemes to perform well. 266 | This class is sensitive where it is initialized in the sequence of model init, 267 | GPU assignment and distributed training wrappers. 268 | """ 269 | 270 | def __init__(self, model, decay=0.9999, updates=0): 271 | # Create EMA 272 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 273 | # if next(model.parameters()).device.type != 'cpu': 274 | # self.ema.half() # FP16 EMA 275 | self.updates = updates # number of EMA updates 276 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) 277 | for p in self.ema.parameters(): 278 | p.requires_grad_(False) 279 | 280 | def update(self, model): 281 | # Update EMA parameters 282 | with torch.no_grad(): 283 | self.updates += 1 284 | d = self.decay(self.updates) 285 | 286 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 287 | for k, v in self.ema.state_dict().items(): 288 | if v.dtype.is_floating_point: 289 | v *= d 290 | v += (1. - d) * msd[k].detach() 291 | 292 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 293 | # Update EMA attributes 294 | copy_attr(self.ema, model, include, exclude) 295 | -------------------------------------------------------------------------------- /utils/wandb_logging/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/littledeep/YOLOv5-RK3399Pro/37a81a703ba7f82261c09d4316a55182ae84a248/utils/wandb_logging/__init__.py -------------------------------------------------------------------------------- /utils/wandb_logging/log_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from pathlib import Path 3 | 4 | import yaml 5 | 6 | from wandb_utils import WandbLogger 7 | from utils.datasets import LoadImagesAndLabels 8 | 9 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' 10 | 11 | 12 | def create_dataset_artifact(opt): 13 | with open(opt.data) as f: 14 | data = yaml.load(f, Loader=yaml.SafeLoader) # data dict 15 | logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation') 16 | 17 | 18 | if __name__ == '__main__': 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') 21 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') 22 | parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project') 23 | opt = parser.parse_args() 24 | 25 | create_dataset_artifact(opt) 26 | -------------------------------------------------------------------------------- /utils/wandb_logging/wandb_utils.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import shutil 5 | import sys 6 | import torch 7 | import yaml 8 | from datetime import datetime 9 | from pathlib import Path 10 | from tqdm import tqdm 11 | 12 | sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path 13 | from utils.datasets import LoadImagesAndLabels 14 | from utils.datasets import img2label_paths 15 | from utils.general import colorstr, xywh2xyxy, check_dataset 16 | 17 | try: 18 | import wandb 19 | from wandb import init, finish 20 | except ImportError: 21 | wandb = None 22 | 23 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' 24 | 25 | 26 | def remove_prefix(from_string, prefix): 27 | return from_string[len(prefix):] 28 | 29 | 30 | def check_wandb_config_file(data_config_file): 31 | wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path 32 | if Path(wandb_config).is_file(): 33 | return wandb_config 34 | return data_config_file 35 | 36 | 37 | def resume_and_get_id(opt): 38 | # It's more elegant to stick to 1 wandb.init call, but as useful config data is overwritten in the WandbLogger's wandb.init call 39 | if isinstance(opt.resume, str): 40 | if opt.resume.startswith(WANDB_ARTIFACT_PREFIX): 41 | run_path = Path(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX)) 42 | run_id = run_path.stem 43 | project = run_path.parent.stem 44 | model_artifact_name = WANDB_ARTIFACT_PREFIX + 'run_' + run_id + '_model' 45 | assert wandb, 'install wandb to resume wandb runs' 46 | # Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config 47 | run = wandb.init(id=run_id, project=project, resume='allow') 48 | opt.resume = model_artifact_name 49 | return run 50 | return None 51 | 52 | 53 | class WandbLogger(): 54 | def __init__(self, opt, name, run_id, data_dict, job_type='Training'): 55 | # Pre-training routine -- 56 | self.job_type = job_type 57 | self.wandb, self.wandb_run, self.data_dict = wandb, None if not wandb else wandb.run, data_dict 58 | if self.wandb: 59 | self.wandb_run = wandb.init(config=opt, 60 | resume="allow", 61 | project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem, 62 | name=name, 63 | job_type=job_type, 64 | id=run_id) if not wandb.run else wandb.run 65 | if self.job_type == 'Training': 66 | if not opt.resume: 67 | wandb_data_dict = self.check_and_upload_dataset(opt) if opt.upload_dataset else data_dict 68 | # Info useful for resuming from artifacts 69 | self.wandb_run.config.opt = vars(opt) 70 | self.wandb_run.config.data_dict = wandb_data_dict 71 | self.data_dict = self.setup_training(opt, data_dict) 72 | if self.job_type == 'Dataset Creation': 73 | self.data_dict = self.check_and_upload_dataset(opt) 74 | else: 75 | print(f"{colorstr('wandb: ')}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)") 76 | 77 | 78 | def check_and_upload_dataset(self, opt): 79 | assert wandb, 'Install wandb to upload dataset' 80 | check_dataset(self.data_dict) 81 | config_path = self.log_dataset_artifact(opt.data, 82 | opt.single_cls, 83 | 'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem) 84 | print("Created dataset config file ", config_path) 85 | with open(config_path) as f: 86 | wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader) 87 | return wandb_data_dict 88 | 89 | def setup_training(self, opt, data_dict): 90 | self.log_dict, self.current_epoch, self.log_imgs = {}, 0, 16 # Logging Constants 91 | self.bbox_interval = opt.bbox_interval 92 | if isinstance(opt.resume, str): 93 | modeldir, _ = self.download_model_artifact(opt) 94 | if modeldir: 95 | self.weights = Path(modeldir) / "last.pt" 96 | config = self.wandb_run.config 97 | opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str( 98 | self.weights), config.save_period, config.total_batch_size, config.bbox_interval, config.epochs, \ 99 | config.opt['hyp'] 100 | data_dict = dict(self.wandb_run.config.data_dict) # eliminates the need for config file to resume 101 | if 'val_artifact' not in self.__dict__: # If --upload_dataset is set, use the existing artifact, don't download 102 | self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'), 103 | opt.artifact_alias) 104 | self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'), 105 | opt.artifact_alias) 106 | self.result_artifact, self.result_table, self.val_table, self.weights = None, None, None, None 107 | if self.train_artifact_path is not None: 108 | train_path = Path(self.train_artifact_path) / 'data/images/' 109 | data_dict['train'] = str(train_path) 110 | if self.val_artifact_path is not None: 111 | val_path = Path(self.val_artifact_path) / 'data/images/' 112 | data_dict['val'] = str(val_path) 113 | self.val_table = self.val_artifact.get("val") 114 | self.map_val_table_path() 115 | if self.val_artifact is not None: 116 | self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation") 117 | self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"]) 118 | if opt.bbox_interval == -1: 119 | self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1 120 | return data_dict 121 | 122 | def download_dataset_artifact(self, path, alias): 123 | if path.startswith(WANDB_ARTIFACT_PREFIX): 124 | dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias) 125 | assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'" 126 | datadir = dataset_artifact.download() 127 | return datadir, dataset_artifact 128 | return None, None 129 | 130 | def download_model_artifact(self, opt): 131 | if opt.resume.startswith(WANDB_ARTIFACT_PREFIX): 132 | model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest") 133 | assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist' 134 | modeldir = model_artifact.download() 135 | epochs_trained = model_artifact.metadata.get('epochs_trained') 136 | total_epochs = model_artifact.metadata.get('total_epochs') 137 | assert epochs_trained < total_epochs, 'training to %g epochs is finished, nothing to resume.' % ( 138 | total_epochs) 139 | return modeldir, model_artifact 140 | return None, None 141 | 142 | def log_model(self, path, opt, epoch, fitness_score, best_model=False): 143 | model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={ 144 | 'original_url': str(path), 145 | 'epochs_trained': epoch + 1, 146 | 'save period': opt.save_period, 147 | 'project': opt.project, 148 | 'total_epochs': opt.epochs, 149 | 'fitness_score': fitness_score 150 | }) 151 | model_artifact.add_file(str(path / 'last.pt'), name='last.pt') 152 | wandb.log_artifact(model_artifact, 153 | aliases=['latest', 'epoch ' + str(self.current_epoch), 'best' if best_model else '']) 154 | print("Saving model artifact on epoch ", epoch + 1) 155 | 156 | def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False): 157 | with open(data_file) as f: 158 | data = yaml.load(f, Loader=yaml.SafeLoader) # data dict 159 | nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names']) 160 | names = {k: v for k, v in enumerate(names)} # to index dictionary 161 | self.train_artifact = self.create_dataset_table(LoadImagesAndLabels( 162 | data['train']), names, name='train') if data.get('train') else None 163 | self.val_artifact = self.create_dataset_table(LoadImagesAndLabels( 164 | data['val']), names, name='val') if data.get('val') else None 165 | if data.get('train'): 166 | data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train') 167 | if data.get('val'): 168 | data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val') 169 | path = data_file if overwrite_config else '_wandb.'.join(data_file.rsplit('.', 1)) # updated data.yaml path 170 | data.pop('download', None) 171 | with open(path, 'w') as f: 172 | yaml.dump(data, f) 173 | 174 | if self.job_type == 'Training': # builds correct artifact pipeline graph 175 | self.wandb_run.use_artifact(self.val_artifact) 176 | self.wandb_run.use_artifact(self.train_artifact) 177 | self.val_artifact.wait() 178 | self.val_table = self.val_artifact.get('val') 179 | self.map_val_table_path() 180 | else: 181 | self.wandb_run.log_artifact(self.train_artifact) 182 | self.wandb_run.log_artifact(self.val_artifact) 183 | return path 184 | 185 | def map_val_table_path(self): 186 | self.val_table_map = {} 187 | print("Mapping dataset") 188 | for i, data in enumerate(tqdm(self.val_table.data)): 189 | self.val_table_map[data[3]] = data[0] 190 | 191 | def create_dataset_table(self, dataset, class_to_id, name='dataset'): 192 | # TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging 193 | artifact = wandb.Artifact(name=name, type="dataset") 194 | for img_file in tqdm([dataset.path]) if Path(dataset.path).is_dir() else tqdm(dataset.img_files): 195 | if Path(img_file).is_dir(): 196 | artifact.add_dir(img_file, name='data/images') 197 | labels_path = 'labels'.join(dataset.path.rsplit('images', 1)) 198 | artifact.add_dir(labels_path, name='data/labels') 199 | else: 200 | artifact.add_file(img_file, name='data/images/' + Path(img_file).name) 201 | label_file = Path(img2label_paths([img_file])[0]) 202 | artifact.add_file(str(label_file), 203 | name='data/labels/' + label_file.name) if label_file.exists() else None 204 | table = wandb.Table(columns=["id", "train_image", "Classes", "name"]) 205 | class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()]) 206 | for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)): 207 | height, width = shapes[0] 208 | labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4))) * torch.Tensor([width, height, width, height]) 209 | box_data, img_classes = [], {} 210 | for cls, *xyxy in labels[:, 1:].tolist(): 211 | cls = int(cls) 212 | box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, 213 | "class_id": cls, 214 | "box_caption": "%s" % (class_to_id[cls]), 215 | "scores": {"acc": 1}, 216 | "domain": "pixel"}) 217 | img_classes[cls] = class_to_id[cls] 218 | boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space 219 | table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes), 220 | Path(paths).name) 221 | artifact.add(table, name) 222 | return artifact 223 | 224 | def log_training_progress(self, predn, path, names): 225 | if self.val_table and self.result_table: 226 | class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()]) 227 | box_data = [] 228 | total_conf = 0 229 | for *xyxy, conf, cls in predn.tolist(): 230 | if conf >= 0.25: 231 | box_data.append( 232 | {"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, 233 | "class_id": int(cls), 234 | "box_caption": "%s %.3f" % (names[cls], conf), 235 | "scores": {"class_score": conf}, 236 | "domain": "pixel"}) 237 | total_conf = total_conf + conf 238 | boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space 239 | id = self.val_table_map[Path(path).name] 240 | self.result_table.add_data(self.current_epoch, 241 | id, 242 | wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set), 243 | total_conf / max(1, len(box_data)) 244 | ) 245 | 246 | def log(self, log_dict): 247 | if self.wandb_run: 248 | for key, value in log_dict.items(): 249 | self.log_dict[key] = value 250 | 251 | def end_epoch(self, best_result=False): 252 | if self.wandb_run: 253 | wandb.log(self.log_dict) 254 | self.log_dict = {} 255 | if self.result_artifact: 256 | train_results = wandb.JoinedTable(self.val_table, self.result_table, "id") 257 | self.result_artifact.add(train_results, 'result') 258 | wandb.log_artifact(self.result_artifact, aliases=['latest', 'epoch ' + str(self.current_epoch), 259 | ('best' if best_result else '')]) 260 | self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"]) 261 | self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation") 262 | 263 | def finish_run(self): 264 | if self.wandb_run: 265 | if self.log_dict: 266 | wandb.log(self.log_dict) 267 | wandb.run.finish() 268 | -------------------------------------------------------------------------------- /weights/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Download latest models from https://github.com/ultralytics/yolov5/releases 3 | # Usage: 4 | # $ bash weights/download_weights.sh 5 | 6 | python - <