├── .dockerignore ├── .gitattributes ├── .github ├── FUNDING.yml ├── ISSUE_TEMPLATE │ ├── bug-report.yml │ ├── config.yml │ ├── feature-request.yml │ └── question.yml ├── dependabot.yml └── workflows │ ├── ci-testing.yml │ ├── codeql-analysis.yml │ ├── greetings.yml │ ├── rebase.yml │ └── stale.yml ├── .gitignore ├── .pre-commit-config.yaml ├── Arial.ttf ├── CONTRIBUTING.md ├── Dockerfile ├── LICENSE ├── README.md ├── data ├── Argoverse.yaml ├── GlobalWheat2020.yaml ├── Objects365.yaml ├── SKU-110K.yaml ├── VOC.yaml ├── VisDrone.yaml ├── coco.yaml ├── coco128.yaml ├── hyps │ ├── hyp.finetune.yaml │ ├── hyp.finetune_objects365.yaml │ ├── hyp.scratch-high.yaml │ ├── hyp.scratch-low.yaml │ ├── hyp.scratch-med.yaml │ └── hyp.scratch.yaml ├── images │ ├── bus.jpg │ └── zidane.jpg ├── scripts │ ├── download_weights.sh │ ├── get_coco.sh │ └── get_coco128.sh └── xView.yaml ├── detect.py ├── export.py ├── hubconf.py ├── models ├── __init__.py ├── common.py ├── experimental.py ├── hub │ ├── anchors.yaml │ ├── yolov3-spp.yaml │ ├── yolov3-tiny.yaml │ ├── yolov3.yaml │ ├── yolov5-bifpn.yaml │ ├── yolov5-fpn.yaml │ ├── yolov5-p2.yaml │ ├── yolov5-p6.yaml │ ├── yolov5-p7.yaml │ ├── yolov5-panet.yaml │ ├── yolov5l6.yaml │ ├── yolov5m6.yaml │ ├── yolov5n6.yaml │ ├── yolov5s-ghost.yaml │ ├── yolov5s-transformer.yaml │ ├── yolov5s6.yaml │ └── yolov5x6.yaml ├── tf.py ├── yolo.py ├── yolov5l.yaml ├── yolov5m.yaml ├── yolov5n.yaml ├── yolov5s.yaml └── yolov5x.yaml ├── prepare_data ├── README.md ├── xml_to_yolo_rotation.html ├── xml_to_yolo_rotation.ipynb └── xml_to_yolo_rotation.py ├── requirements.txt ├── setup.cfg ├── train.py ├── tutorial.ipynb ├── utils ├── __init__.py ├── activations.py ├── augmentations.py ├── autoanchor.py ├── autobatch.py ├── aws │ ├── __init__.py │ ├── mime.sh │ ├── resume.py │ └── userdata.sh ├── callbacks.py ├── datasets.py ├── downloads.py ├── flask_rest_api │ ├── README.md │ ├── example_request.py │ └── restapi.py ├── general.py ├── google_app_engine │ ├── Dockerfile │ ├── additional_requirements.txt │ └── app.yaml ├── loggers │ ├── __init__.py │ └── wandb │ │ ├── README.md │ │ ├── __init__.py │ │ ├── log_dataset.py │ │ ├── sweep.py │ │ ├── sweep.yaml │ │ └── wandb_utils.py ├── loss.py ├── metrics.py ├── plots.py └── torch_utils.py └── val.py /.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*.csv 12 | *.jpg 13 | 14 | # Neural Network weights ----------------------------------------------------------------------------------------------- 15 | **/*.pt 16 | **/*.pth 17 | **/*.onnx 18 | **/*.mlmodel 19 | **/*.torchscript 20 | **/*.torchscript.pt 21 | **/*.tflite 22 | **/*.h5 23 | **/*.pb 24 | *_saved_model/ 25 | *_web_model/ 26 | 27 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 28 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 29 | 30 | 31 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 32 | # Byte-compiled / optimized / DLL files 33 | __pycache__/ 34 | *.py[cod] 35 | *$py.class 36 | 37 | # C extensions 38 | *.so 39 | 40 | # Distribution / packaging 41 | .Python 42 | env/ 43 | build/ 44 | develop-eggs/ 45 | dist/ 46 | downloads/ 47 | eggs/ 48 | .eggs/ 49 | lib/ 50 | lib64/ 51 | parts/ 52 | sdist/ 53 | var/ 54 | wheels/ 55 | *.egg-info/ 56 | wandb/ 57 | .installed.cfg 58 | *.egg 59 | 60 | # PyInstaller 61 | # Usually these files are written by a python script from a template 62 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 63 | *.manifest 64 | *.spec 65 | 66 | # Installer logs 67 | pip-log.txt 68 | pip-delete-this-directory.txt 69 | 70 | # Unit test / coverage reports 71 | htmlcov/ 72 | .tox/ 73 | .coverage 74 | .coverage.* 75 | .cache 76 | nosetests.xml 77 | coverage.xml 78 | *.cover 79 | .hypothesis/ 80 | 81 | # Translations 82 | *.mo 83 | *.pot 84 | 85 | # Django stuff: 86 | *.log 87 | local_settings.py 88 | 89 | # Flask stuff: 90 | instance/ 91 | .webassets-cache 92 | 93 | # Scrapy stuff: 94 | .scrapy 95 | 96 | # Sphinx documentation 97 | docs/_build/ 98 | 99 | # PyBuilder 100 | target/ 101 | 102 | # Jupyter Notebook 103 | .ipynb_checkpoints 104 | 105 | # pyenv 106 | .python-version 107 | 108 | # celery beat schedule file 109 | celerybeat-schedule 110 | 111 | # SageMath parsed files 112 | *.sage.py 113 | 114 | # dotenv 115 | .env 116 | 117 | # virtualenv 118 | .venv* 119 | venv*/ 120 | ENV*/ 121 | 122 | # Spyder project settings 123 | .spyderproject 124 | .spyproject 125 | 126 | # Rope project settings 127 | .ropeproject 128 | 129 | # mkdocs documentation 130 | /site 131 | 132 | # mypy 133 | .mypy_cache/ 134 | 135 | 136 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 137 | 138 | # General 139 | .DS_Store 140 | .AppleDouble 141 | .LSOverride 142 | 143 | # Icon must end with two \r 144 | Icon 145 | Icon? 146 | 147 | # Thumbnails 148 | ._* 149 | 150 | # Files that might appear in the root of a volume 151 | .DocumentRevisions-V100 152 | .fseventsd 153 | .Spotlight-V100 154 | .TemporaryItems 155 | .Trashes 156 | .VolumeIcon.icns 157 | .com.apple.timemachine.donotpresent 158 | 159 | # Directories potentially created on remote AFP share 160 | .AppleDB 161 | .AppleDesktop 162 | Network Trash Folder 163 | Temporary Items 164 | .apdisk 165 | 166 | 167 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 168 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 169 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 170 | 171 | # User-specific stuff: 172 | .idea/* 173 | .idea/**/workspace.xml 174 | .idea/**/tasks.xml 175 | .idea/dictionaries 176 | .html # Bokeh Plots 177 | .pg # TensorFlow Frozen Graphs 178 | .avi # videos 179 | 180 | # Sensitive or high-churn files: 181 | .idea/**/dataSources/ 182 | .idea/**/dataSources.ids 183 | .idea/**/dataSources.local.xml 184 | .idea/**/sqlDataSources.xml 185 | .idea/**/dynamic.xml 186 | .idea/**/uiDesigner.xml 187 | 188 | # Gradle: 189 | .idea/**/gradle.xml 190 | .idea/**/libraries 191 | 192 | # CMake 193 | cmake-build-debug/ 194 | cmake-build-release/ 195 | 196 | # Mongo Explorer plugin: 197 | .idea/**/mongoSettings.xml 198 | 199 | ## File-based project format: 200 | *.iws 201 | 202 | ## Plugin-specific files: 203 | 204 | # IntelliJ 205 | out/ 206 | 207 | # mpeltonen/sbt-idea plugin 208 | .idea_modules/ 209 | 210 | # JIRA plugin 211 | atlassian-ide-plugin.xml 212 | 213 | # Cursive Clojure plugin 214 | .idea/replstate.xml 215 | 216 | # Crashlytics plugin (for Android Studio and IntelliJ) 217 | com_crashlytics_export_strings.xml 218 | crashlytics.properties 219 | crashlytics-build.properties 220 | fabric.properties 221 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # this drop notebooks from GitHub language stats 2 | *.ipynb linguist-vendored 3 | -------------------------------------------------------------------------------- /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | # These are supported funding model platforms 2 | 3 | github: glenn-jocher 4 | patreon: ultralytics 5 | open_collective: ultralytics 6 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug-report.yml: -------------------------------------------------------------------------------- 1 | name: 🐛 Bug Report 2 | # title: " " 3 | description: Problems with YOLOv5 4 | labels: [bug, triage] 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Thank you for submitting a YOLOv5 🐛 Bug Report! 10 | 11 | - type: checkboxes 12 | attributes: 13 | label: Search before asking 14 | description: > 15 | Please search the [issues](https://github.com/ultralytics/yolov5/issues) to see if a similar bug report already exists. 16 | options: 17 | - label: > 18 | I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and found no similar bug report. 19 | required: true 20 | 21 | - type: dropdown 22 | attributes: 23 | label: YOLOv5 Component 24 | description: | 25 | Please select the part of YOLOv5 where you found the bug. 26 | multiple: true 27 | options: 28 | - "Training" 29 | - "Validation" 30 | - "Detection" 31 | - "Export" 32 | - "PyTorch Hub" 33 | - "Multi-GPU" 34 | - "Evolution" 35 | - "Integrations" 36 | - "Other" 37 | validations: 38 | required: false 39 | 40 | - type: textarea 41 | attributes: 42 | label: Bug 43 | description: Provide console output with error messages and/or screenshots of the bug. 44 | placeholder: | 45 | 💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response. 46 | validations: 47 | required: true 48 | 49 | - type: textarea 50 | attributes: 51 | label: Environment 52 | description: Please specify the software and hardware you used to produce the bug. 53 | placeholder: | 54 | - YOLO: YOLOv5 🚀 v6.0-67-g60e42e1 torch 1.9.0+cu111 CUDA:0 (A100-SXM4-40GB, 40536MiB) 55 | - OS: Ubuntu 20.04 56 | - Python: 3.9.0 57 | validations: 58 | required: false 59 | 60 | - type: textarea 61 | attributes: 62 | label: Minimal Reproducible Example 63 | description: > 64 | When asking a question, people will be better able to provide help if you provide code that they can easily understand and use to **reproduce** the problem. 65 | This is referred to by community members as creating a [minimal reproducible example](https://stackoverflow.com/help/minimal-reproducible-example). 66 | placeholder: | 67 | ``` 68 | # Code to reproduce your issue here 69 | ``` 70 | validations: 71 | required: false 72 | 73 | - type: textarea 74 | attributes: 75 | label: Additional 76 | description: Anything else you would like to share? 77 | 78 | - type: checkboxes 79 | attributes: 80 | label: Are you willing to submit a PR? 81 | description: > 82 | (Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov5/pulls) (PR) to help improve YOLOv5 for everyone, especially if you have a good understanding of how to implement a fix or feature. 83 | See the YOLOv5 [Contributing Guide](https://github.com/ultralytics/yolov5/blob/master/CONTRIBUTING.md) to get started. 84 | options: 85 | - label: Yes I'd like to help by submitting a PR! 86 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: true 2 | contact_links: 3 | - name: Slack 4 | url: https://join.slack.com/t/ultralytics/shared_invite/zt-w29ei8bp-jczz7QYUmDtgo6r6KcMIAg 5 | about: Ask on Ultralytics Slack Forum 6 | - name: Stack Overflow 7 | url: https://stackoverflow.com/search?q=YOLOv5 8 | about: Ask on Stack Overflow with 'YOLOv5' tag 9 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.yml: -------------------------------------------------------------------------------- 1 | name: 🚀 Feature Request 2 | description: Suggest a YOLOv5 idea 3 | # title: " " 4 | labels: [enhancement] 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Thank you for submitting a YOLOv5 🚀 Feature Request! 10 | 11 | - type: checkboxes 12 | attributes: 13 | label: Search before asking 14 | description: > 15 | Please search the [issues](https://github.com/ultralytics/yolov5/issues) to see if a similar feature request already exists. 16 | options: 17 | - label: > 18 | I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and found no similar feature requests. 19 | required: true 20 | 21 | - type: textarea 22 | attributes: 23 | label: Description 24 | description: A short description of your feature. 25 | placeholder: | 26 | What new feature would you like to see in YOLOv5? 27 | validations: 28 | required: true 29 | 30 | - type: textarea 31 | attributes: 32 | label: Use case 33 | description: | 34 | Describe the use case of your feature request. It will help us understand and prioritize the feature request. 35 | placeholder: | 36 | How would this feature be used, and who would use it? 37 | 38 | - type: textarea 39 | attributes: 40 | label: Additional 41 | description: Anything else you would like to share? 42 | 43 | - type: checkboxes 44 | attributes: 45 | label: Are you willing to submit a PR? 46 | description: > 47 | (Optional) We encourage you to submit a [Pull Request](https://github.com/ultralytics/yolov5/pulls) (PR) to help improve YOLOv5 for everyone, especially if you have a good understanding of how to implement a fix or feature. 48 | See the YOLOv5 [Contributing Guide](https://github.com/ultralytics/yolov5/blob/master/CONTRIBUTING.md) to get started. 49 | options: 50 | - label: Yes I'd like to help by submitting a PR! 51 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.yml: -------------------------------------------------------------------------------- 1 | name: ❓ Question 2 | description: Ask a YOLOv5 question 3 | # title: " " 4 | labels: [question] 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Thank you for asking a YOLOv5 ❓ Question! 10 | 11 | - type: checkboxes 12 | attributes: 13 | label: Search before asking 14 | description: > 15 | Please search the [issues](https://github.com/ultralytics/yolov5/issues) and [discussions](https://github.com/ultralytics/yolov5/discussions) to see if a similar question already exists. 16 | options: 17 | - label: > 18 | I have searched the YOLOv5 [issues](https://github.com/ultralytics/yolov5/issues) and [discussions](https://github.com/ultralytics/yolov5/discussions) and found no similar questions. 19 | required: true 20 | 21 | - type: textarea 22 | attributes: 23 | label: Question 24 | description: What is your question? 25 | placeholder: | 26 | 💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response. 27 | validations: 28 | required: true 29 | 30 | - type: textarea 31 | attributes: 32 | label: Additional 33 | description: Anything else you would like to share? 34 | -------------------------------------------------------------------------------- /.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 | 14 | - package-ecosystem: github-actions 15 | directory: "/" 16 | schedule: 17 | interval: weekly 18 | time: "04:00" 19 | open-pull-requests-limit: 5 20 | reviewers: 21 | - glenn-jocher 22 | labels: 23 | - dependencies 24 | -------------------------------------------------------------------------------- /.github/workflows/ci-testing.yml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | name: CI CPU testing 4 | 5 | on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows 6 | push: 7 | branches: [ master ] 8 | pull_request: 9 | # The branches below must be a subset of the branches above 10 | branches: [ master ] 11 | schedule: 12 | - cron: '0 0 * * *' # Runs at 00:00 UTC every day 13 | 14 | jobs: 15 | cpu-tests: 16 | 17 | runs-on: ${{ matrix.os }} 18 | strategy: 19 | fail-fast: false 20 | matrix: 21 | os: [ ubuntu-latest, macos-latest, windows-latest ] 22 | python-version: [ 3.9 ] 23 | model: [ 'yolov5n' ] # models to test 24 | 25 | # Timeout: https://stackoverflow.com/a/59076067/4521646 26 | timeout-minutes: 50 27 | steps: 28 | - uses: actions/checkout@v2 29 | - name: Set up Python ${{ matrix.python-version }} 30 | uses: actions/setup-python@v2 31 | with: 32 | python-version: ${{ matrix.python-version }} 33 | 34 | # Note: This uses an internal pip API and may not always work 35 | # https://github.com/actions/cache/blob/master/examples.md#multiple-oss-in-a-workflow 36 | - name: Get pip cache 37 | id: pip-cache 38 | run: | 39 | python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)" 40 | 41 | - name: Cache pip 42 | uses: actions/cache@v2.1.6 43 | with: 44 | path: ${{ steps.pip-cache.outputs.dir }} 45 | key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }} 46 | restore-keys: | 47 | ${{ runner.os }}-${{ matrix.python-version }}-pip- 48 | 49 | # Known Keras 2.7.0 issue: https://github.com/ultralytics/yolov5/pull/5486 50 | - name: Install dependencies 51 | run: | 52 | python -m pip install --upgrade pip 53 | pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html 54 | pip install -q onnx tensorflow-cpu keras==2.6.0 # wandb # extras 55 | python --version 56 | pip --version 57 | pip list 58 | shell: bash 59 | 60 | # - name: W&B login 61 | # run: wandb login 345011b3fb26dc8337fd9b20e53857c1d403f2aa 62 | 63 | - name: Download data 64 | run: | 65 | # curl -L -o tmp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 66 | # unzip -q tmp.zip -d ../ 67 | # rm tmp.zip 68 | 69 | - name: Tests workflow 70 | run: | 71 | # export PYTHONPATH="$PWD" # to run '$ python *.py' files in subdirectories 72 | di=cpu # device 73 | 74 | # Train 75 | python train.py --img 64 --batch 32 --weights ${{ matrix.model }}.pt --cfg ${{ matrix.model }}.yaml --epochs 1 --device $di 76 | # Val 77 | python val.py --img 64 --batch 32 --weights ${{ matrix.model }}.pt --device $di 78 | python val.py --img 64 --batch 32 --weights runs/train/exp/weights/last.pt --device $di 79 | # Detect 80 | python detect.py --weights ${{ matrix.model }}.pt --device $di 81 | python detect.py --weights runs/train/exp/weights/last.pt --device $di 82 | python hubconf.py # hub 83 | # Export 84 | python models/yolo.py --cfg ${{ matrix.model }}.yaml # build PyTorch model 85 | python models/tf.py --weights ${{ matrix.model }}.pt # build TensorFlow model 86 | python export.py --img 64 --batch 1 --weights ${{ matrix.model }}.pt --include torchscript onnx # export 87 | # Python 88 | python - <=3.6.0**](https://www.python.org/) with all [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) installed including [**PyTorch>=1.7**](https://pytorch.org/get-started/locally/). To get started: 39 | ```bash 40 | $ git clone https://github.com/ultralytics/yolov5 41 | $ cd yolov5 42 | $ pip install -r requirements.txt 43 | ``` 44 | 45 | ## Environments 46 | 47 | 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): 48 | 49 | - **Google Colab and Kaggle** notebooks with free GPU: Open In Colab Open In Kaggle 50 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) 51 | - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/AWS-Quickstart) 52 | - **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) Docker Pulls 53 | 54 | 55 | ## Status 56 | 57 | CI CPU testing 58 | 59 | 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)), validation ([val.py](https://github.com/ultralytics/yolov5/blob/master/val.py)), inference ([detect.py](https://github.com/ultralytics/yolov5/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov5/blob/master/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit. 60 | -------------------------------------------------------------------------------- /.github/workflows/rebase.yml: -------------------------------------------------------------------------------- 1 | # https://github.com/marketplace/actions/automatic-rebase 2 | 3 | name: Automatic Rebase 4 | on: 5 | issue_comment: 6 | types: [created] 7 | jobs: 8 | rebase: 9 | name: Rebase 10 | if: github.event.issue.pull_request != '' && contains(github.event.comment.body, '/rebase') 11 | runs-on: ubuntu-latest 12 | steps: 13 | - name: Checkout the latest code 14 | uses: actions/checkout@v2 15 | with: 16 | token: ${{ secrets.ACTIONS_TOKEN }} 17 | fetch-depth: 0 # otherwise, you will fail to push refs to dest repo 18 | - name: Automatic Rebase 19 | uses: cirrus-actions/rebase@1.5 20 | env: 21 | GITHUB_TOKEN: ${{ secrets.ACTIONS_TOKEN }} 22 | -------------------------------------------------------------------------------- /.github/workflows/stale.yml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | name: Close stale issues 4 | on: 5 | schedule: 6 | - cron: "0 0 * * *" 7 | 8 | jobs: 9 | stale: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/stale@v4 13 | with: 14 | repo-token: ${{ secrets.GITHUB_TOKEN }} 15 | stale-issue-message: | 16 | 👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs. 17 | 18 | Access additional [YOLOv5](https://ultralytics.com/yolov5) 🚀 resources: 19 | - **Wiki** – https://github.com/ultralytics/yolov5/wiki 20 | - **Tutorials** – https://github.com/ultralytics/yolov5#tutorials 21 | - **Docs** – https://docs.ultralytics.com 22 | 23 | Access additional [Ultralytics](https://ultralytics.com) ⚡ resources: 24 | - **Ultralytics HUB** – https://ultralytics.com/hub 25 | - **Vision API** – https://ultralytics.com/yolov5 26 | - **About Us** – https://ultralytics.com/about 27 | - **Join Our Team** – https://ultralytics.com/work 28 | - **Contact Us** – https://ultralytics.com/contact 29 | 30 | Feel free to inform us of any other **issues** you discover or **feature requests** that come to mind in the future. Pull Requests (PRs) are also always welcomed! 31 | 32 | Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐! 33 | 34 | stale-pr-message: 'This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions YOLOv5 🚀 and Vision AI ⭐.' 35 | days-before-stale: 30 36 | days-before-close: 5 37 | exempt-issue-labels: 'documentation,tutorial' 38 | operations-per-run: 100 # The maximum number of operations per run, used to control rate limiting. 39 | -------------------------------------------------------------------------------- /.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 | *.cfg 23 | !setup.cfg 24 | !cfg/yolov3*.cfg 25 | 26 | storage.googleapis.com 27 | runs/* 28 | data/* 29 | !data/hyps/* 30 | !data/images/zidane.jpg 31 | !data/images/bus.jpg 32 | !data/*.sh 33 | 34 | results*.csv 35 | 36 | # Datasets ------------------------------------------------------------------------------------------------------------- 37 | coco/ 38 | coco128/ 39 | VOC/ 40 | 41 | # MATLAB GitIgnore ----------------------------------------------------------------------------------------------------- 42 | *.m~ 43 | *.mat 44 | !targets*.mat 45 | 46 | # Neural Network weights ----------------------------------------------------------------------------------------------- 47 | *.weights 48 | *.pt 49 | *.pb 50 | *.onnx 51 | *.mlmodel 52 | *.torchscript 53 | *.tflite 54 | *.h5 55 | *_saved_model/ 56 | *_web_model/ 57 | darknet53.conv.74 58 | yolov3-tiny.conv.15 59 | 60 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 61 | # Byte-compiled / optimized / DLL files 62 | __pycache__/ 63 | *.py[cod] 64 | *$py.class 65 | 66 | # C extensions 67 | *.so 68 | 69 | # Distribution / packaging 70 | .Python 71 | env/ 72 | build/ 73 | develop-eggs/ 74 | dist/ 75 | downloads/ 76 | eggs/ 77 | .eggs/ 78 | lib/ 79 | lib64/ 80 | parts/ 81 | sdist/ 82 | var/ 83 | wheels/ 84 | *.egg-info/ 85 | /wandb/ 86 | .installed.cfg 87 | *.egg 88 | 89 | 90 | # PyInstaller 91 | # Usually these files are written by a python script from a template 92 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 93 | *.manifest 94 | *.spec 95 | 96 | # Installer logs 97 | pip-log.txt 98 | pip-delete-this-directory.txt 99 | 100 | # Unit test / coverage reports 101 | htmlcov/ 102 | .tox/ 103 | .coverage 104 | .coverage.* 105 | .cache 106 | nosetests.xml 107 | coverage.xml 108 | *.cover 109 | .hypothesis/ 110 | 111 | # Translations 112 | *.mo 113 | *.pot 114 | 115 | # Django stuff: 116 | *.log 117 | local_settings.py 118 | 119 | # Flask stuff: 120 | instance/ 121 | .webassets-cache 122 | 123 | # Scrapy stuff: 124 | .scrapy 125 | 126 | # Sphinx documentation 127 | docs/_build/ 128 | 129 | # PyBuilder 130 | target/ 131 | 132 | # Jupyter Notebook 133 | .ipynb_checkpoints 134 | 135 | # pyenv 136 | .python-version 137 | 138 | # celery beat schedule file 139 | celerybeat-schedule 140 | 141 | # SageMath parsed files 142 | *.sage.py 143 | 144 | # dotenv 145 | .env 146 | 147 | # virtualenv 148 | .venv* 149 | venv*/ 150 | ENV*/ 151 | 152 | # Spyder project settings 153 | .spyderproject 154 | .spyproject 155 | 156 | # Rope project settings 157 | .ropeproject 158 | 159 | # mkdocs documentation 160 | /site 161 | 162 | # mypy 163 | .mypy_cache/ 164 | 165 | 166 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 167 | 168 | # General 169 | .DS_Store 170 | .AppleDouble 171 | .LSOverride 172 | 173 | # Icon must end with two \r 174 | Icon 175 | Icon? 176 | 177 | # Thumbnails 178 | ._* 179 | 180 | # Files that might appear in the root of a volume 181 | .DocumentRevisions-V100 182 | .fseventsd 183 | .Spotlight-V100 184 | .TemporaryItems 185 | .Trashes 186 | .VolumeIcon.icns 187 | .com.apple.timemachine.donotpresent 188 | 189 | # Directories potentially created on remote AFP share 190 | .AppleDB 191 | .AppleDesktop 192 | Network Trash Folder 193 | Temporary Items 194 | .apdisk 195 | 196 | 197 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 198 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 199 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 200 | 201 | # User-specific stuff: 202 | .idea/* 203 | .idea/**/workspace.xml 204 | .idea/**/tasks.xml 205 | .idea/dictionaries 206 | .html # Bokeh Plots 207 | .pg # TensorFlow Frozen Graphs 208 | .avi # videos 209 | 210 | # Sensitive or high-churn files: 211 | .idea/**/dataSources/ 212 | .idea/**/dataSources.ids 213 | .idea/**/dataSources.local.xml 214 | .idea/**/sqlDataSources.xml 215 | .idea/**/dynamic.xml 216 | .idea/**/uiDesigner.xml 217 | 218 | # Gradle: 219 | .idea/**/gradle.xml 220 | .idea/**/libraries 221 | 222 | # CMake 223 | cmake-build-debug/ 224 | cmake-build-release/ 225 | 226 | # Mongo Explorer plugin: 227 | .idea/**/mongoSettings.xml 228 | 229 | ## File-based project format: 230 | *.iws 231 | 232 | ## Plugin-specific files: 233 | 234 | # IntelliJ 235 | out/ 236 | 237 | # mpeltonen/sbt-idea plugin 238 | .idea_modules/ 239 | 240 | # JIRA plugin 241 | atlassian-ide-plugin.xml 242 | 243 | # Cursive Clojure plugin 244 | .idea/replstate.xml 245 | 246 | # Crashlytics plugin (for Android Studio and IntelliJ) 247 | com_crashlytics_export_strings.xml 248 | crashlytics.properties 249 | crashlytics-build.properties 250 | fabric.properties 251 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # Define hooks for code formations 2 | # Will be applied on any updated commit files if a user has installed and linked commit hook 3 | 4 | default_language_version: 5 | python: python3.8 6 | 7 | # Define bot property if installed via https://github.com/marketplace/pre-commit-ci 8 | ci: 9 | autofix_prs: true 10 | autoupdate_commit_msg: '[pre-commit.ci] pre-commit suggestions' 11 | autoupdate_schedule: quarterly 12 | # submodules: true 13 | 14 | repos: 15 | - repo: https://github.com/pre-commit/pre-commit-hooks 16 | rev: v4.0.1 17 | hooks: 18 | - id: end-of-file-fixer 19 | - id: trailing-whitespace 20 | - id: check-case-conflict 21 | - id: check-yaml 22 | - id: check-toml 23 | - id: pretty-format-json 24 | - id: check-docstring-first 25 | 26 | - repo: https://github.com/asottile/pyupgrade 27 | rev: v2.23.1 28 | hooks: 29 | - id: pyupgrade 30 | args: [--py36-plus] 31 | name: Upgrade code 32 | 33 | - repo: https://github.com/PyCQA/isort 34 | rev: 5.9.3 35 | hooks: 36 | - id: isort 37 | name: Sort imports 38 | 39 | # TODO 40 | #- repo: https://github.com/pre-commit/mirrors-yapf 41 | # rev: v0.31.0 42 | # hooks: 43 | # - id: yapf 44 | # name: formatting 45 | 46 | # TODO 47 | #- repo: https://github.com/executablebooks/mdformat 48 | # rev: 0.7.7 49 | # hooks: 50 | # - id: mdformat 51 | # additional_dependencies: 52 | # - mdformat-gfm 53 | # - mdformat-black 54 | # - mdformat_frontmatter 55 | 56 | # TODO 57 | #- repo: https://github.com/asottile/yesqa 58 | # rev: v1.2.3 59 | # hooks: 60 | # - id: yesqa 61 | 62 | - repo: https://github.com/PyCQA/flake8 63 | rev: 3.9.2 64 | hooks: 65 | - id: flake8 66 | name: PEP8 67 | -------------------------------------------------------------------------------- /Arial.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/Arial.ttf -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | ## Contributing to YOLOv5 🚀 2 | 3 | We love your input! We want to make contributing to YOLOv5 as easy and transparent as possible, whether it's: 4 | 5 | - Reporting a bug 6 | - Discussing the current state of the code 7 | - Submitting a fix 8 | - Proposing a new feature 9 | - Becoming a maintainer 10 | 11 | YOLOv5 works so well due to our combined community effort, and for every small improvement you contribute you will be 12 | helping push the frontiers of what's possible in AI 😃! 13 | 14 | ## Submitting a Pull Request (PR) 🛠️ 15 | 16 | Submitting a PR is easy! This example shows how to submit a PR for updating `requirements.txt` in 4 steps: 17 | 18 | ### 1. Select File to Update 19 | 20 | Select `requirements.txt` to update by clicking on it in GitHub. 21 |

PR_step1

22 | 23 | ### 2. Click 'Edit this file' 24 | 25 | Button is in top-right corner. 26 |

PR_step2

27 | 28 | ### 3. Make Changes 29 | 30 | Change `matplotlib` version from `3.2.2` to `3.3`. 31 |

PR_step3

32 | 33 | ### 4. Preview Changes and Submit PR 34 | 35 | Click on the **Preview changes** tab to verify your updates. At the bottom of the screen select 'Create a **new branch** 36 | for this commit', assign your branch a descriptive name such as `fix/matplotlib_version` and click the green **Propose 37 | changes** button. All done, your PR is now submitted to YOLOv5 for review and approval 😃! 38 |

PR_step4

39 | 40 | ### PR recommendations 41 | 42 | To allow your work to be integrated as seamlessly as possible, we advise you to: 43 | 44 | - ✅ Verify your PR is **up-to-date with upstream/master.** If your PR is behind upstream/master an 45 | automatic [GitHub actions](https://github.com/ultralytics/yolov5/blob/master/.github/workflows/rebase.yml) rebase may 46 | be attempted by including the /rebase command in a comment body, or by running the following code, replacing 'feature' 47 | with the name of your local branch: 48 | 49 | ```bash 50 | git remote add upstream https://github.com/ultralytics/yolov5.git 51 | git fetch upstream 52 | git checkout feature # <----- replace 'feature' with local branch name 53 | git merge upstream/master 54 | git push -u origin -f 55 | ``` 56 | 57 | - ✅ Verify all Continuous Integration (CI) **checks are passing**. 58 | - ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase 59 | but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ — Bruce Lee 60 | 61 | ## Submitting a Bug Report 🐛 62 | 63 | If you spot a problem with YOLOv5 please submit a Bug Report! 64 | 65 | For us to start investigating a possible problem we need to be able to reproduce it ourselves first. We've created a few 66 | short guidelines below to help users provide what we need in order to get started. 67 | 68 | When asking a question, people will be better able to provide help if you provide **code** that they can easily 69 | understand and use to **reproduce** the problem. This is referred to by community members as creating 70 | a [minimum reproducible example](https://stackoverflow.com/help/minimal-reproducible-example). Your code that reproduces 71 | the problem should be: 72 | 73 | * ✅ **Minimal** – Use as little code as possible that still produces the same problem 74 | * ✅ **Complete** – Provide **all** parts someone else needs to reproduce your problem in the question itself 75 | * ✅ **Reproducible** – Test the code you're about to provide to make sure it reproduces the problem 76 | 77 | In addition to the above requirements, for [Ultralytics](https://ultralytics.com/) to provide assistance your code 78 | should be: 79 | 80 | * ✅ **Current** – Verify that your code is up-to-date with current 81 | GitHub [master](https://github.com/ultralytics/yolov5/tree/master), and if necessary `git pull` or `git clone` a new 82 | copy to ensure your problem has not already been resolved by previous commits. 83 | * ✅ **Unmodified** – Your problem must be reproducible without any modifications to the codebase in this 84 | repository. [Ultralytics](https://ultralytics.com/) does not provide support for custom code ⚠️. 85 | 86 | If you believe your problem meets all of the above criteria, please close this issue and raise a new one using the 🐛 ** 87 | Bug Report** [template](https://github.com/ultralytics/yolov5/issues/new/choose) and providing 88 | a [minimum reproducible example](https://stackoverflow.com/help/minimal-reproducible-example) to help us better 89 | understand and diagnose your problem. 90 | 91 | ## License 92 | 93 | By contributing, you agree that your contributions will be licensed under 94 | the [GPL-3.0 license](https://choosealicense.com/licenses/gpl-3.0/) 95 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch 4 | FROM nvcr.io/nvidia/pytorch:21.10-py3 5 | 6 | # Install linux packages 7 | RUN apt update && apt install -y zip htop screen libgl1-mesa-glx 8 | 9 | # Install python dependencies 10 | COPY requirements.txt . 11 | RUN python -m pip install --upgrade pip 12 | RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof 13 | RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook wandb>=0.12.2 14 | RUN pip install --no-cache -U torch torchvision numpy Pillow 15 | # RUN pip install --no-cache torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html 16 | 17 | # Create working directory 18 | RUN mkdir -p /usr/src/app 19 | WORKDIR /usr/src/app 20 | 21 | # Copy contents 22 | COPY . /usr/src/app 23 | 24 | # Downloads to user config dir 25 | ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/ 26 | 27 | # Set environment variables 28 | # ENV HOME=/usr/src/app 29 | 30 | 31 | # Usage Examples ------------------------------------------------------------------------------------------------------- 32 | 33 | # Build and Push 34 | # t=ultralytics/yolov5:latest && sudo docker build -t $t . && sudo docker push $t 35 | 36 | # Pull and Run 37 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t 38 | 39 | # Pull and Run with local directory access 40 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t 41 | 42 | # Kill all 43 | # sudo docker kill $(sudo docker ps -q) 44 | 45 | # Kill all image-based 46 | # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest) 47 | 48 | # Bash into running container 49 | # sudo docker exec -it 5a9b5863d93d bash 50 | 51 | # Bash into stopped container 52 | # id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash 53 | 54 | # Clean up 55 | # docker system prune -a --volumes 56 | 57 | # Update Ubuntu drivers 58 | # https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/ 59 | 60 | # DDP test 61 | # python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3 62 | -------------------------------------------------------------------------------- /data/Argoverse.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/ 3 | # Example usage: python train.py --data Argoverse.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── Argoverse ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/Argoverse # dataset root dir 12 | train: Argoverse-1.1/images/train/ # train images (relative to 'path') 39384 images 13 | val: Argoverse-1.1/images/val/ # val images (relative to 'path') 15062 images 14 | test: Argoverse-1.1/images/test/ # test images (optional) https://eval.ai/web/challenges/challenge-page/800/overview 15 | 16 | # Classes 17 | nc: 8 # number of classes 18 | names: ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign'] # class names 19 | 20 | 21 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 22 | download: | 23 | import json 24 | 25 | from tqdm import tqdm 26 | from utils.general import download, Path 27 | 28 | 29 | def argoverse2yolo(set): 30 | labels = {} 31 | a = json.load(open(set, "rb")) 32 | for annot in tqdm(a['annotations'], desc=f"Converting {set} to YOLOv5 format..."): 33 | img_id = annot['image_id'] 34 | img_name = a['images'][img_id]['name'] 35 | img_label_name = img_name[:-3] + "txt" 36 | 37 | cls = annot['category_id'] # instance class id 38 | x_center, y_center, width, height = annot['bbox'] 39 | x_center = (x_center + width / 2) / 1920.0 # offset and scale 40 | y_center = (y_center + height / 2) / 1200.0 # offset and scale 41 | width /= 1920.0 # scale 42 | height /= 1200.0 # scale 43 | 44 | img_dir = set.parents[2] / 'Argoverse-1.1' / 'labels' / a['seq_dirs'][a['images'][annot['image_id']]['sid']] 45 | if not img_dir.exists(): 46 | img_dir.mkdir(parents=True, exist_ok=True) 47 | 48 | k = str(img_dir / img_label_name) 49 | if k not in labels: 50 | labels[k] = [] 51 | labels[k].append(f"{cls} {x_center} {y_center} {width} {height}\n") 52 | 53 | for k in labels: 54 | with open(k, "w") as f: 55 | f.writelines(labels[k]) 56 | 57 | 58 | # Download 59 | dir = Path('../datasets/Argoverse') # dataset root dir 60 | urls = ['https://argoverse-hd.s3.us-east-2.amazonaws.com/Argoverse-HD-Full.zip'] 61 | download(urls, dir=dir, delete=False) 62 | 63 | # Convert 64 | annotations_dir = 'Argoverse-HD/annotations/' 65 | (dir / 'Argoverse-1.1' / 'tracking').rename(dir / 'Argoverse-1.1' / 'images') # rename 'tracking' to 'images' 66 | for d in "train.json", "val.json": 67 | argoverse2yolo(dir / annotations_dir / d) # convert VisDrone annotations to YOLO labels 68 | -------------------------------------------------------------------------------- /data/GlobalWheat2020.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Global Wheat 2020 dataset http://www.global-wheat.com/ 3 | # Example usage: python train.py --data GlobalWheat2020.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── GlobalWheat2020 ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/GlobalWheat2020 # dataset root dir 12 | train: # train images (relative to 'path') 3422 images 13 | - images/arvalis_1 14 | - images/arvalis_2 15 | - images/arvalis_3 16 | - images/ethz_1 17 | - images/rres_1 18 | - images/inrae_1 19 | - images/usask_1 20 | val: # val images (relative to 'path') 748 images (WARNING: train set contains ethz_1) 21 | - images/ethz_1 22 | test: # test images (optional) 1276 images 23 | - images/utokyo_1 24 | - images/utokyo_2 25 | - images/nau_1 26 | - images/uq_1 27 | 28 | # Classes 29 | nc: 1 # number of classes 30 | names: ['wheat_head'] # class names 31 | 32 | 33 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 34 | download: | 35 | from utils.general import download, Path 36 | 37 | # Download 38 | dir = Path(yaml['path']) # dataset root dir 39 | urls = ['https://zenodo.org/record/4298502/files/global-wheat-codalab-official.zip', 40 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/GlobalWheat2020_labels.zip'] 41 | download(urls, dir=dir) 42 | 43 | # Make Directories 44 | for p in 'annotations', 'images', 'labels': 45 | (dir / p).mkdir(parents=True, exist_ok=True) 46 | 47 | # Move 48 | for p in 'arvalis_1', 'arvalis_2', 'arvalis_3', 'ethz_1', 'rres_1', 'inrae_1', 'usask_1', \ 49 | 'utokyo_1', 'utokyo_2', 'nau_1', 'uq_1': 50 | (dir / p).rename(dir / 'images' / p) # move to /images 51 | f = (dir / p).with_suffix('.json') # json file 52 | if f.exists(): 53 | f.rename((dir / 'annotations' / p).with_suffix('.json')) # move to /annotations 54 | -------------------------------------------------------------------------------- /data/Objects365.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Objects365 dataset https://www.objects365.org/ 3 | # Example usage: python train.py --data Objects365.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── Objects365 ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/Objects365 # dataset root dir 12 | train: images/train # train images (relative to 'path') 1742289 images 13 | val: images/val # val images (relative to 'path') 80000 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | nc: 365 # number of classes 18 | names: ['Person', 'Sneakers', 'Chair', 'Other Shoes', 'Hat', 'Car', 'Lamp', 'Glasses', 'Bottle', 'Desk', 'Cup', 19 | 'Street Lights', 'Cabinet/shelf', 'Handbag/Satchel', 'Bracelet', 'Plate', 'Picture/Frame', 'Helmet', 'Book', 20 | 'Gloves', 'Storage box', 'Boat', 'Leather Shoes', 'Flower', 'Bench', 'Potted Plant', 'Bowl/Basin', 'Flag', 21 | 'Pillow', 'Boots', 'Vase', 'Microphone', 'Necklace', 'Ring', 'SUV', 'Wine Glass', 'Belt', 'Monitor/TV', 22 | 'Backpack', 'Umbrella', 'Traffic Light', 'Speaker', 'Watch', 'Tie', 'Trash bin Can', 'Slippers', 'Bicycle', 23 | 'Stool', 'Barrel/bucket', 'Van', 'Couch', 'Sandals', 'Basket', 'Drum', 'Pen/Pencil', 'Bus', 'Wild Bird', 24 | 'High Heels', 'Motorcycle', 'Guitar', 'Carpet', 'Cell Phone', 'Bread', 'Camera', 'Canned', 'Truck', 25 | 'Traffic cone', 'Cymbal', 'Lifesaver', 'Towel', 'Stuffed Toy', 'Candle', 'Sailboat', 'Laptop', 'Awning', 26 | 'Bed', 'Faucet', 'Tent', 'Horse', 'Mirror', 'Power outlet', 'Sink', 'Apple', 'Air Conditioner', 'Knife', 27 | 'Hockey Stick', 'Paddle', 'Pickup Truck', 'Fork', 'Traffic Sign', 'Balloon', 'Tripod', 'Dog', 'Spoon', 'Clock', 28 | 'Pot', 'Cow', 'Cake', 'Dinning Table', 'Sheep', 'Hanger', 'Blackboard/Whiteboard', 'Napkin', 'Other Fish', 29 | 'Orange/Tangerine', 'Toiletry', 'Keyboard', 'Tomato', 'Lantern', 'Machinery Vehicle', 'Fan', 30 | 'Green Vegetables', 'Banana', 'Baseball Glove', 'Airplane', 'Mouse', 'Train', 'Pumpkin', 'Soccer', 'Skiboard', 31 | 'Luggage', 'Nightstand', 'Tea pot', 'Telephone', 'Trolley', 'Head Phone', 'Sports Car', 'Stop Sign', 32 | 'Dessert', 'Scooter', 'Stroller', 'Crane', 'Remote', 'Refrigerator', 'Oven', 'Lemon', 'Duck', 'Baseball Bat', 33 | 'Surveillance Camera', 'Cat', 'Jug', 'Broccoli', 'Piano', 'Pizza', 'Elephant', 'Skateboard', 'Surfboard', 34 | 'Gun', 'Skating and Skiing shoes', 'Gas stove', 'Donut', 'Bow Tie', 'Carrot', 'Toilet', 'Kite', 'Strawberry', 35 | 'Other Balls', 'Shovel', 'Pepper', 'Computer Box', 'Toilet Paper', 'Cleaning Products', 'Chopsticks', 36 | 'Microwave', 'Pigeon', 'Baseball', 'Cutting/chopping Board', 'Coffee Table', 'Side Table', 'Scissors', 37 | 'Marker', 'Pie', 'Ladder', 'Snowboard', 'Cookies', 'Radiator', 'Fire Hydrant', 'Basketball', 'Zebra', 'Grape', 38 | 'Giraffe', 'Potato', 'Sausage', 'Tricycle', 'Violin', 'Egg', 'Fire Extinguisher', 'Candy', 'Fire Truck', 39 | 'Billiards', 'Converter', 'Bathtub', 'Wheelchair', 'Golf Club', 'Briefcase', 'Cucumber', 'Cigar/Cigarette', 40 | 'Paint Brush', 'Pear', 'Heavy Truck', 'Hamburger', 'Extractor', 'Extension Cord', 'Tong', 'Tennis Racket', 41 | 'Folder', 'American Football', 'earphone', 'Mask', 'Kettle', 'Tennis', 'Ship', 'Swing', 'Coffee Machine', 42 | 'Slide', 'Carriage', 'Onion', 'Green beans', 'Projector', 'Frisbee', 'Washing Machine/Drying Machine', 43 | 'Chicken', 'Printer', 'Watermelon', 'Saxophone', 'Tissue', 'Toothbrush', 'Ice cream', 'Hot-air balloon', 44 | 'Cello', 'French Fries', 'Scale', 'Trophy', 'Cabbage', 'Hot dog', 'Blender', 'Peach', 'Rice', 'Wallet/Purse', 45 | 'Volleyball', 'Deer', 'Goose', 'Tape', 'Tablet', 'Cosmetics', 'Trumpet', 'Pineapple', 'Golf Ball', 46 | 'Ambulance', 'Parking meter', 'Mango', 'Key', 'Hurdle', 'Fishing Rod', 'Medal', 'Flute', 'Brush', 'Penguin', 47 | 'Megaphone', 'Corn', 'Lettuce', 'Garlic', 'Swan', 'Helicopter', 'Green Onion', 'Sandwich', 'Nuts', 48 | 'Speed Limit Sign', 'Induction Cooker', 'Broom', 'Trombone', 'Plum', 'Rickshaw', 'Goldfish', 'Kiwi fruit', 49 | 'Router/modem', 'Poker Card', 'Toaster', 'Shrimp', 'Sushi', 'Cheese', 'Notepaper', 'Cherry', 'Pliers', 'CD', 50 | 'Pasta', 'Hammer', 'Cue', 'Avocado', 'Hamimelon', 'Flask', 'Mushroom', 'Screwdriver', 'Soap', 'Recorder', 51 | 'Bear', 'Eggplant', 'Board Eraser', 'Coconut', 'Tape Measure/Ruler', 'Pig', 'Showerhead', 'Globe', 'Chips', 52 | 'Steak', 'Crosswalk Sign', 'Stapler', 'Camel', 'Formula 1', 'Pomegranate', 'Dishwasher', 'Crab', 53 | 'Hoverboard', 'Meat ball', 'Rice Cooker', 'Tuba', 'Calculator', 'Papaya', 'Antelope', 'Parrot', 'Seal', 54 | 'Butterfly', 'Dumbbell', 'Donkey', 'Lion', 'Urinal', 'Dolphin', 'Electric Drill', 'Hair Dryer', 'Egg tart', 55 | 'Jellyfish', 'Treadmill', 'Lighter', 'Grapefruit', 'Game board', 'Mop', 'Radish', 'Baozi', 'Target', 'French', 56 | 'Spring Rolls', 'Monkey', 'Rabbit', 'Pencil Case', 'Yak', 'Red Cabbage', 'Binoculars', 'Asparagus', 'Barbell', 57 | 'Scallop', 'Noddles', 'Comb', 'Dumpling', 'Oyster', 'Table Tennis paddle', 'Cosmetics Brush/Eyeliner Pencil', 58 | 'Chainsaw', 'Eraser', 'Lobster', 'Durian', 'Okra', 'Lipstick', 'Cosmetics Mirror', 'Curling', 'Table Tennis'] 59 | 60 | 61 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 62 | download: | 63 | from pycocotools.coco import COCO 64 | from tqdm import tqdm 65 | 66 | from utils.general import Path, download, np, xyxy2xywhn 67 | 68 | # Make Directories 69 | dir = Path(yaml['path']) # dataset root dir 70 | for p in 'images', 'labels': 71 | (dir / p).mkdir(parents=True, exist_ok=True) 72 | for q in 'train', 'val': 73 | (dir / p / q).mkdir(parents=True, exist_ok=True) 74 | 75 | # Train, Val Splits 76 | for split, patches in [('train', 50 + 1), ('val', 43 + 1)]: 77 | print(f"Processing {split} in {patches} patches ...") 78 | images, labels = dir / 'images' / split, dir / 'labels' / split 79 | 80 | # Download 81 | url = f"https://dorc.ks3-cn-beijing.ksyun.com/data-set/2020Objects365%E6%95%B0%E6%8D%AE%E9%9B%86/{split}/" 82 | if split == 'train': 83 | download([f'{url}zhiyuan_objv2_{split}.tar.gz'], dir=dir, delete=False) # annotations json 84 | download([f'{url}patch{i}.tar.gz' for i in range(patches)], dir=images, curl=True, delete=False, threads=8) 85 | elif split == 'val': 86 | download([f'{url}zhiyuan_objv2_{split}.json'], dir=dir, delete=False) # annotations json 87 | download([f'{url}images/v1/patch{i}.tar.gz' for i in range(15 + 1)], dir=images, curl=True, delete=False, threads=8) 88 | download([f'{url}images/v2/patch{i}.tar.gz' for i in range(16, patches)], dir=images, curl=True, delete=False, threads=8) 89 | 90 | # Move 91 | for f in tqdm(images.rglob('*.jpg'), desc=f'Moving {split} images'): 92 | f.rename(images / f.name) # move to /images/{split} 93 | 94 | # Labels 95 | coco = COCO(dir / f'zhiyuan_objv2_{split}.json') 96 | names = [x["name"] for x in coco.loadCats(coco.getCatIds())] 97 | for cid, cat in enumerate(names): 98 | catIds = coco.getCatIds(catNms=[cat]) 99 | imgIds = coco.getImgIds(catIds=catIds) 100 | for im in tqdm(coco.loadImgs(imgIds), desc=f'Class {cid + 1}/{len(names)} {cat}'): 101 | width, height = im["width"], im["height"] 102 | path = Path(im["file_name"]) # image filename 103 | try: 104 | with open(labels / path.with_suffix('.txt').name, 'a') as file: 105 | annIds = coco.getAnnIds(imgIds=im["id"], catIds=catIds, iscrowd=None) 106 | for a in coco.loadAnns(annIds): 107 | x, y, w, h = a['bbox'] # bounding box in xywh (xy top-left corner) 108 | xyxy = np.array([x, y, x + w, y + h])[None] # pixels(1,4) 109 | x, y, w, h = xyxy2xywhn(xyxy, w=width, h=height, clip=True)[0] # normalized and clipped 110 | file.write(f"{cid} {x:.5f} {y:.5f} {w:.5f} {h:.5f}\n") 111 | except Exception as e: 112 | print(e) 113 | -------------------------------------------------------------------------------- /data/SKU-110K.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # SKU-110K retail items dataset https://github.com/eg4000/SKU110K_CVPR19 3 | # Example usage: python train.py --data SKU-110K.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── SKU-110K ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/SKU-110K # dataset root dir 12 | train: train.txt # train images (relative to 'path') 8219 images 13 | val: val.txt # val images (relative to 'path') 588 images 14 | test: test.txt # test images (optional) 2936 images 15 | 16 | # Classes 17 | nc: 1 # number of classes 18 | names: ['object'] # class names 19 | 20 | 21 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 22 | download: | 23 | import shutil 24 | from tqdm import tqdm 25 | from utils.general import np, pd, Path, download, xyxy2xywh 26 | 27 | # Download 28 | dir = Path(yaml['path']) # dataset root dir 29 | parent = Path(dir.parent) # download dir 30 | urls = ['http://trax-geometry.s3.amazonaws.com/cvpr_challenge/SKU110K_fixed.tar.gz'] 31 | download(urls, dir=parent, delete=False) 32 | 33 | # Rename directories 34 | if dir.exists(): 35 | shutil.rmtree(dir) 36 | (parent / 'SKU110K_fixed').rename(dir) # rename dir 37 | (dir / 'labels').mkdir(parents=True, exist_ok=True) # create labels dir 38 | 39 | # Convert labels 40 | names = 'image', 'x1', 'y1', 'x2', 'y2', 'class', 'image_width', 'image_height' # column names 41 | for d in 'annotations_train.csv', 'annotations_val.csv', 'annotations_test.csv': 42 | x = pd.read_csv(dir / 'annotations' / d, names=names).values # annotations 43 | images, unique_images = x[:, 0], np.unique(x[:, 0]) 44 | with open((dir / d).with_suffix('.txt').__str__().replace('annotations_', ''), 'w') as f: 45 | f.writelines(f'./images/{s}\n' for s in unique_images) 46 | for im in tqdm(unique_images, desc=f'Converting {dir / d}'): 47 | cls = 0 # single-class dataset 48 | with open((dir / 'labels' / im).with_suffix('.txt'), 'a') as f: 49 | for r in x[images == im]: 50 | w, h = r[6], r[7] # image width, height 51 | xywh = xyxy2xywh(np.array([[r[1] / w, r[2] / h, r[3] / w, r[4] / h]]))[0] # instance 52 | f.write(f"{cls} {xywh[0]:.5f} {xywh[1]:.5f} {xywh[2]:.5f} {xywh[3]:.5f}\n") # write label 53 | -------------------------------------------------------------------------------- /data/VOC.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC 3 | # Example usage: python train.py --data VOC.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── VOC ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/VOC 12 | train: # train images (relative to 'path') 16551 images 13 | - images/train2012 14 | - images/train2007 15 | - images/val2012 16 | - images/val2007 17 | val: # val images (relative to 'path') 4952 images 18 | - images/test2007 19 | test: # test images (optional) 20 | - images/test2007 21 | 22 | # Classes 23 | nc: 20 # number of classes 24 | names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 25 | 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] # class names 26 | 27 | 28 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 29 | download: | 30 | import xml.etree.ElementTree as ET 31 | 32 | from tqdm import tqdm 33 | from utils.general import download, Path 34 | 35 | 36 | def convert_label(path, lb_path, year, image_id): 37 | def convert_box(size, box): 38 | dw, dh = 1. / size[0], 1. / size[1] 39 | x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2] 40 | return x * dw, y * dh, w * dw, h * dh 41 | 42 | in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml') 43 | out_file = open(lb_path, 'w') 44 | tree = ET.parse(in_file) 45 | root = tree.getroot() 46 | size = root.find('size') 47 | w = int(size.find('width').text) 48 | h = int(size.find('height').text) 49 | 50 | for obj in root.iter('object'): 51 | cls = obj.find('name').text 52 | if cls in yaml['names'] and not int(obj.find('difficult').text) == 1: 53 | xmlbox = obj.find('bndbox') 54 | bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')]) 55 | cls_id = yaml['names'].index(cls) # class id 56 | out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n') 57 | 58 | 59 | # Download 60 | dir = Path(yaml['path']) # dataset root dir 61 | url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/' 62 | urls = [url + 'VOCtrainval_06-Nov-2007.zip', # 446MB, 5012 images 63 | url + 'VOCtest_06-Nov-2007.zip', # 438MB, 4953 images 64 | url + 'VOCtrainval_11-May-2012.zip'] # 1.95GB, 17126 images 65 | download(urls, dir=dir / 'images', delete=False) 66 | 67 | # Convert 68 | path = dir / f'images/VOCdevkit' 69 | for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'): 70 | imgs_path = dir / 'images' / f'{image_set}{year}' 71 | lbs_path = dir / 'labels' / f'{image_set}{year}' 72 | imgs_path.mkdir(exist_ok=True, parents=True) 73 | lbs_path.mkdir(exist_ok=True, parents=True) 74 | 75 | image_ids = open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt').read().strip().split() 76 | for id in tqdm(image_ids, desc=f'{image_set}{year}'): 77 | f = path / f'VOC{year}/JPEGImages/{id}.jpg' # old img path 78 | lb_path = (lbs_path / f.name).with_suffix('.txt') # new label path 79 | f.rename(imgs_path / f.name) # move image 80 | convert_label(path, lb_path, year, id) # convert labels to YOLO format 81 | -------------------------------------------------------------------------------- /data/VisDrone.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset 3 | # Example usage: python train.py --data VisDrone.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── VisDrone ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/VisDrone # dataset root dir 12 | train: VisDrone2019-DET-train/images # train images (relative to 'path') 6471 images 13 | val: VisDrone2019-DET-val/images # val images (relative to 'path') 548 images 14 | test: VisDrone2019-DET-test-dev/images # test images (optional) 1610 images 15 | 16 | # Classes 17 | nc: 10 # number of classes 18 | names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor'] 19 | 20 | 21 | # Download script/URL (optional) --------------------------------------------------------------------------------------- 22 | download: | 23 | from utils.general import download, os, Path 24 | 25 | def visdrone2yolo(dir): 26 | from PIL import Image 27 | from tqdm import tqdm 28 | 29 | def convert_box(size, box): 30 | # Convert VisDrone box to YOLO xywh box 31 | dw = 1. / size[0] 32 | dh = 1. / size[1] 33 | return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh 34 | 35 | (dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory 36 | pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}') 37 | for f in pbar: 38 | img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size 39 | lines = [] 40 | with open(f, 'r') as file: # read annotation.txt 41 | for row in [x.split(',') for x in file.read().strip().splitlines()]: 42 | if row[4] == '0': # VisDrone 'ignored regions' class 0 43 | continue 44 | cls = int(row[5]) - 1 45 | box = convert_box(img_size, tuple(map(int, row[:4]))) 46 | lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n") 47 | with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl: 48 | fl.writelines(lines) # write label.txt 49 | 50 | 51 | # Download 52 | dir = Path(yaml['path']) # dataset root dir 53 | urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip', 54 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip', 55 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip', 56 | 'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip'] 57 | download(urls, dir=dir) 58 | 59 | # Convert 60 | for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev': 61 | visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels 62 | -------------------------------------------------------------------------------- /data/coco.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # COCO 2017 dataset http://cocodataset.org 3 | # Example usage: python train.py --data coco.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── coco ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/coco # dataset root dir 12 | train: train2017.txt # train images (relative to 'path') 118287 images 13 | val: val2017.txt # train images (relative to 'path') 5000 images 14 | test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794 15 | 16 | # Classes 17 | nc: 80 # number of classes 18 | names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 19 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 20 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 21 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 22 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 23 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 24 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 25 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 26 | 'hair drier', 'toothbrush'] # class names 27 | 28 | 29 | # Download script/URL (optional) 30 | download: | 31 | from utils.general import download, Path 32 | 33 | # Download labels 34 | segments = False # segment or box labels 35 | dir = Path(yaml['path']) # dataset root dir 36 | url = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/' 37 | urls = [url + ('coco2017labels-segments.zip' if segments else 'coco2017labels.zip')] # labels 38 | download(urls, dir=dir.parent) 39 | 40 | # Download data 41 | urls = ['http://images.cocodataset.org/zips/train2017.zip', # 19G, 118k images 42 | 'http://images.cocodataset.org/zips/val2017.zip', # 1G, 5k images 43 | 'http://images.cocodataset.org/zips/test2017.zip'] # 7G, 41k images (optional) 44 | download(urls, dir=dir / 'images', threads=3) 45 | -------------------------------------------------------------------------------- /data/coco128.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) 3 | # Example usage: python train.py --data coco128.yaml 4 | # parent 5 | # ├── yolov5 6 | # └── datasets 7 | # └── coco128 ← downloads here 8 | 9 | 10 | # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..] 11 | path: ../datasets/coco128 # dataset root dir 12 | train: images/train2017 # train images (relative to 'path') 128 images 13 | val: images/train2017 # val images (relative to 'path') 128 images 14 | test: # test images (optional) 15 | 16 | # Classes 17 | nc: 80 # number of classes 18 | names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 19 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 20 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 21 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 22 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 23 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 24 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 25 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 26 | 'hair drier', 'toothbrush'] # class names 27 | 28 | 29 | # Download script/URL (optional) 30 | download: https://ultralytics.com/assets/coco128.zip 31 | -------------------------------------------------------------------------------- /data/hyps/hyp.finetune.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Hyperparameters for VOC finetuning 3 | # python train.py --batch 64 --weights yolov5m.pt --data VOC.yaml --img 512 --epochs 50 4 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 5 | 6 | # Hyperparameter Evolution Results 7 | # Generations: 306 8 | # P R mAP.5 mAP.5:.95 box obj cls 9 | # Metrics: 0.6 0.936 0.896 0.684 0.0115 0.00805 0.00146 10 | 11 | lr0: 0.0032 12 | lrf: 0.12 13 | momentum: 0.843 14 | weight_decay: 0.00036 15 | warmup_epochs: 2.0 16 | warmup_momentum: 0.5 17 | warmup_bias_lr: 0.05 18 | box: 0.0296 19 | cls: 0.243 20 | cls_pw: 0.631 21 | obj: 0.301 22 | obj_pw: 0.911 23 | angle: 0.03 24 | angle_pw: 0.05 # angle BCELoss positive_weight 25 | iou_t: 0.2 26 | anchor_t: 2.91 27 | # anchors: 3.63 28 | fl_gamma: 0.0 29 | hsv_h: 0.0138 30 | hsv_s: 0.664 31 | hsv_v: 0.464 32 | degrees: 0.373 33 | translate: 0.245 34 | scale: 0.898 35 | shear: 0.602 36 | perspective: 0.0 37 | flipud: 0.00856 38 | fliplr: 0.5 39 | mosaic: 1.0 40 | mixup: 0.243 41 | copy_paste: 0.0 42 | -------------------------------------------------------------------------------- /data/hyps/hyp.finetune_objects365.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | lr0: 0.00258 4 | lrf: 0.17 5 | momentum: 0.779 6 | weight_decay: 0.00058 7 | warmup_epochs: 1.33 8 | warmup_momentum: 0.86 9 | warmup_bias_lr: 0.0711 10 | box: 0.0539 11 | cls: 0.299 12 | cls_pw: 0.825 13 | obj: 0.632 14 | obj_pw: 1.0 15 | angle: 0.03 16 | angle_pw: 0.05 # angle BCELoss positive_weight 17 | iou_t: 0.2 18 | anchor_t: 3.44 19 | anchors: 3.2 20 | fl_gamma: 0.0 21 | hsv_h: 0.0188 22 | hsv_s: 0.704 23 | hsv_v: 0.36 24 | degrees: 0.0 25 | translate: 0.0902 26 | scale: 0.491 27 | shear: 0.0 28 | perspective: 0.0 29 | flipud: 0.0 30 | fliplr: 0.5 31 | mosaic: 1.0 32 | mixup: 0.0 33 | copy_paste: 0.0 34 | -------------------------------------------------------------------------------- /data/hyps/hyp.scratch-high.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Hyperparameters for high-augmentation COCO training from scratch 3 | # python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300 4 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf) 8 | momentum: 0.937 # SGD momentum/Adam beta1 9 | weight_decay: 0.0005 # optimizer weight decay 5e-4 10 | warmup_epochs: 3.0 # warmup epochs (fractions ok) 11 | warmup_momentum: 0.8 # warmup initial momentum 12 | warmup_bias_lr: 0.1 # warmup initial bias lr 13 | box: 0.05 # box loss gain 14 | cls: 0.3 # cls loss gain 15 | cls_pw: 1.0 # cls BCELoss positive_weight 16 | obj: 0.7 # obj loss gain (scale with pixels) 17 | obj_pw: 1.0 # obj BCELoss positive_weight 18 | angle: 0.03 19 | angle_pw: 0.05 # angle BCELoss positive_weight 20 | iou_t: 0.20 # IoU training threshold 21 | anchor_t: 4.0 # anchor-multiple threshold 22 | # anchors: 3 # anchors per output layer (0 to ignore) 23 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 24 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 25 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 26 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 27 | degrees: 0.0 # image rotation (+/- deg) 28 | translate: 0.1 # image translation (+/- fraction) 29 | scale: 0.9 # image scale (+/- gain) 30 | shear: 0.0 # image shear (+/- deg) 31 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 32 | flipud: 0.0 # image flip up-down (probability) 33 | fliplr: 0.5 # image flip left-right (probability) 34 | mosaic: 1.0 # image mosaic (probability) 35 | mixup: 0.1 # image mixup (probability) 36 | copy_paste: 0.1 # segment copy-paste (probability) 37 | -------------------------------------------------------------------------------- /data/hyps/hyp.scratch-low.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Hyperparameters for low-augmentation COCO training from scratch 3 | # python train.py --batch 64 --cfg yolov5n6.yaml --weights '' --data coco.yaml --img 640 --epochs 300 --linear 4 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf) 8 | momentum: 0.937 # SGD momentum/Adam beta1 9 | weight_decay: 0.0005 # optimizer weight decay 5e-4 10 | warmup_epochs: 3.0 # warmup epochs (fractions ok) 11 | warmup_momentum: 0.8 # warmup initial momentum 12 | warmup_bias_lr: 0.1 # warmup initial bias lr 13 | box: 0.05 # box loss gain 14 | cls: 0.5 # cls loss gain 15 | cls_pw: 1.0 # cls BCELoss positive_weight 16 | obj: 1.0 # obj loss gain (scale with pixels) 17 | obj_pw: 1.0 # obj BCELoss positive_weight 18 | angle: 0.03 19 | angle_pw: 0.05 # angle BCELoss positive_weight 20 | iou_t: 0.20 # IoU training threshold 21 | anchor_t: 4.0 # anchor-multiple threshold 22 | # anchors: 3 # anchors per output layer (0 to ignore) 23 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 24 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 25 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 26 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 27 | degrees: 0.0 # image rotation (+/- deg) 28 | translate: 0.1 # image translation (+/- fraction) 29 | scale: 0.5 # image scale (+/- gain) 30 | shear: 0.0 # image shear (+/- deg) 31 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 32 | flipud: 0.0 # image flip up-down (probability) 33 | fliplr: 0.5 # image flip left-right (probability) 34 | mosaic: 1.0 # image mosaic (probability) 35 | mixup: 0.0 # image mixup (probability) 36 | copy_paste: 0.0 # segment copy-paste (probability) 37 | -------------------------------------------------------------------------------- /data/hyps/hyp.scratch-med.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Hyperparameters for medium-augmentation COCO training from scratch 3 | # python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300 4 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf) 8 | momentum: 0.937 # SGD momentum/Adam beta1 9 | weight_decay: 0.0005 # optimizer weight decay 5e-4 10 | warmup_epochs: 3.0 # warmup epochs (fractions ok) 11 | warmup_momentum: 0.8 # warmup initial momentum 12 | warmup_bias_lr: 0.1 # warmup initial bias lr 13 | box: 0.05 # box loss gain 14 | cls: 0.3 # cls loss gain 15 | cls_pw: 1.0 # cls BCELoss positive_weight 16 | obj: 0.7 # obj loss gain (scale with pixels) 17 | obj_pw: 1.0 # obj BCELoss positive_weight 18 | angle: 0.03 19 | angle_pw: 0.05 # angle BCELoss positive_weight 20 | iou_t: 0.20 # IoU training threshold 21 | anchor_t: 4.0 # anchor-multiple threshold 22 | # anchors: 3 # anchors per output layer (0 to ignore) 23 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 24 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 25 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 26 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 27 | degrees: 0.0 # image rotation (+/- deg) 28 | translate: 0.1 # image translation (+/- fraction) 29 | scale: 0.9 # image scale (+/- gain) 30 | shear: 0.0 # image shear (+/- deg) 31 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 32 | flipud: 0.0 # image flip up-down (probability) 33 | fliplr: 0.5 # image flip left-right (probability) 34 | mosaic: 1.0 # image mosaic (probability) 35 | mixup: 0.1 # image mixup (probability) 36 | copy_paste: 0.0 # segment copy-paste (probability) 37 | -------------------------------------------------------------------------------- /data/hyps/hyp.scratch.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Hyperparameters for COCO training from scratch 3 | # python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300 4 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf) 8 | momentum: 0.937 # SGD momentum/Adam beta1 9 | weight_decay: 0.0005 # optimizer weight decay 5e-4 10 | warmup_epochs: 3.0 # warmup epochs (fractions ok) 11 | warmup_momentum: 0.8 # warmup initial momentum 12 | warmup_bias_lr: 0.1 # warmup initial bias lr 13 | box: 0.05 # box loss gain 14 | cls: 0.5 # cls loss gain 15 | cls_pw: 1.0 # cls BCELoss positive_weight 16 | obj: 1.0 # obj loss gain (scale with pixels) 17 | obj_pw: 1.0 # obj BCELoss positive_weight 18 | angle: 0.03 19 | angle_pw: 0.05 # angle BCELoss positive_weight 20 | iou_t: 0.20 # IoU training threshold 21 | anchor_t: 4.0 # anchor-multiple threshold 22 | # anchors: 3 # anchors per output layer (0 to ignore) 23 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 24 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 25 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 26 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 27 | degrees: 0.0 # image rotation (+/- deg) 28 | translate: 0.1 # image translation (+/- fraction) 29 | scale: 0.5 # image scale (+/- gain) 30 | shear: 0.0 # image shear (+/- deg) 31 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 32 | flipud: 0.0 # image flip up-down (probability) 33 | fliplr: 0.5 # image flip left-right (probability) 34 | mosaic: 1.0 # image mosaic (probability) 35 | mixup: 0.0 # image mixup (probability) 36 | copy_paste: 0.0 # segment copy-paste (probability) 37 | -------------------------------------------------------------------------------- /data/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/data/images/bus.jpg -------------------------------------------------------------------------------- /data/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/data/images/zidane.jpg -------------------------------------------------------------------------------- /data/scripts/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 3 | # Download latest models from https://github.com/ultralytics/yolov5/releases 4 | # Example usage: bash path/to/download_weights.sh 5 | # parent 6 | # └── yolov5 7 | # ├── yolov5s.pt ← downloads here 8 | # ├── yolov5m.pt 9 | # └── ... 10 | 11 | python - <= cls >= 0, f'incorrect class index {cls}' 74 | 75 | # Write YOLO label 76 | if id not in shapes: 77 | shapes[id] = Image.open(file).size 78 | box = xyxy2xywhn(box[None].astype(np.float), w=shapes[id][0], h=shapes[id][1], clip=True) 79 | with open((labels / id).with_suffix('.txt'), 'a') as f: 80 | f.write(f"{cls} {' '.join(f'{x:.6f}' for x in box[0])}\n") # write label.txt 81 | except Exception as e: 82 | print(f'WARNING: skipping one label for {file}: {e}') 83 | 84 | 85 | # Download manually from https://challenge.xviewdataset.org 86 | dir = Path(yaml['path']) # dataset root dir 87 | # urls = ['https://d307kc0mrhucc3.cloudfront.net/train_labels.zip', # train labels 88 | # 'https://d307kc0mrhucc3.cloudfront.net/train_images.zip', # 15G, 847 train images 89 | # 'https://d307kc0mrhucc3.cloudfront.net/val_images.zip'] # 5G, 282 val images (no labels) 90 | # download(urls, dir=dir, delete=False) 91 | 92 | # Convert labels 93 | convert_labels(dir / 'xView_train.geojson') 94 | 95 | # Move images 96 | images = Path(dir / 'images') 97 | images.mkdir(parents=True, exist_ok=True) 98 | Path(dir / 'train_images').rename(dir / 'images' / 'train') 99 | Path(dir / 'val_images').rename(dir / 'images' / 'val') 100 | 101 | # Split 102 | autosplit(dir / 'images' / 'train') 103 | -------------------------------------------------------------------------------- /hubconf.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5/ 4 | 5 | Usage: 6 | import torch 7 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 8 | """ 9 | 10 | import torch 11 | 12 | 13 | def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 14 | """Creates a specified YOLOv5 model 15 | 16 | Arguments: 17 | name (str): name of model, i.e. 'yolov5s' 18 | pretrained (bool): load pretrained weights into the model 19 | channels (int): number of input channels 20 | classes (int): number of model classes 21 | autoshape (bool): apply YOLOv5 .autoshape() wrapper to model 22 | verbose (bool): print all information to screen 23 | device (str, torch.device, None): device to use for model parameters 24 | 25 | Returns: 26 | YOLOv5 pytorch model 27 | """ 28 | from pathlib import Path 29 | 30 | from models.common import AutoShape 31 | from models.experimental import attempt_load 32 | from models.yolo import Model 33 | from utils.downloads import attempt_download 34 | from utils.general import check_requirements, intersect_dicts, set_logging 35 | from utils.torch_utils import select_device 36 | 37 | file = Path(__file__).resolve() 38 | check_requirements(exclude=('tensorboard', 'thop', 'opencv-python')) 39 | set_logging(verbose=verbose) 40 | 41 | save_dir = Path('') if str(name).endswith('.pt') else file.parent 42 | path = (save_dir / name).with_suffix('.pt') # checkpoint path 43 | try: 44 | device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) 45 | 46 | if pretrained and channels == 3 and classes == 80: 47 | model = attempt_load(path, map_location=device) # download/load FP32 model 48 | else: 49 | cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path 50 | model = Model(cfg, channels, classes) # create model 51 | if pretrained: 52 | ckpt = torch.load(attempt_download(path), map_location=device) # load 53 | csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 54 | csd = intersect_dicts(csd, model.state_dict(), exclude=['anchors']) # intersect 55 | model.load_state_dict(csd, strict=False) # load 56 | if len(ckpt['model'].names) == classes: 57 | model.names = ckpt['model'].names # set class names attribute 58 | if autoshape: 59 | model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS 60 | return model.to(device) 61 | 62 | except Exception as e: 63 | help_url = 'https://github.com/ultralytics/yolov5/issues/36' 64 | s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url 65 | raise Exception(s) from e 66 | 67 | 68 | def custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None): 69 | # YOLOv5 custom or local model 70 | return _create(path, autoshape=autoshape, verbose=verbose, device=device) 71 | 72 | 73 | def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 74 | # YOLOv5-nano model https://github.com/ultralytics/yolov5 75 | return _create('yolov5n', pretrained, channels, classes, autoshape, verbose, device) 76 | 77 | 78 | def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 79 | # YOLOv5-small model https://github.com/ultralytics/yolov5 80 | return _create('yolov5s', pretrained, channels, classes, autoshape, verbose, device) 81 | 82 | 83 | def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 84 | # YOLOv5-medium model https://github.com/ultralytics/yolov5 85 | return _create('yolov5m', pretrained, channels, classes, autoshape, verbose, device) 86 | 87 | 88 | def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 89 | # YOLOv5-large model https://github.com/ultralytics/yolov5 90 | return _create('yolov5l', pretrained, channels, classes, autoshape, verbose, device) 91 | 92 | 93 | def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 94 | # YOLOv5-xlarge model https://github.com/ultralytics/yolov5 95 | return _create('yolov5x', pretrained, channels, classes, autoshape, verbose, device) 96 | 97 | 98 | def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 99 | # YOLOv5-nano-P6 model https://github.com/ultralytics/yolov5 100 | return _create('yolov5n6', pretrained, channels, classes, autoshape, verbose, device) 101 | 102 | 103 | def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 104 | # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5 105 | return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose, device) 106 | 107 | 108 | def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 109 | # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5 110 | return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose, device) 111 | 112 | 113 | def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 114 | # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5 115 | return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose, device) 116 | 117 | 118 | def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): 119 | # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5 120 | return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose, device) 121 | 122 | 123 | if __name__ == '__main__': 124 | model = _create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained 125 | # model = custom(path='path/to/model.pt') # custom 126 | 127 | # Verify inference 128 | from pathlib import Path 129 | 130 | import cv2 131 | import numpy as np 132 | from PIL import Image 133 | 134 | imgs = ['data/images/zidane.jpg', # filename 135 | Path('data/images/zidane.jpg'), # Path 136 | 'https://ultralytics.com/images/zidane.jpg', # URI 137 | cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV 138 | Image.open('data/images/bus.jpg'), # PIL 139 | np.zeros((320, 640, 3))] # numpy 140 | 141 | results = model(imgs) # batched inference 142 | results.print() 143 | results.save() 144 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/models/__init__.py -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Experimental modules 4 | """ 5 | import math 6 | 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | 11 | from models.common import Conv 12 | from utils.downloads import attempt_download 13 | 14 | 15 | class CrossConv(nn.Module): 16 | # Cross Convolution Downsample 17 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 18 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 19 | super().__init__() 20 | c_ = int(c2 * e) # hidden channels 21 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 22 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 23 | self.add = shortcut and c1 == c2 24 | 25 | def forward(self, x): 26 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 27 | 28 | 29 | class Sum(nn.Module): 30 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 31 | def __init__(self, n, weight=False): # n: number of inputs 32 | super().__init__() 33 | self.weight = weight # apply weights boolean 34 | self.iter = range(n - 1) # iter object 35 | if weight: 36 | self.w = nn.Parameter(-torch.arange(1.0, n) / 2, requires_grad=True) # layer weights 37 | 38 | def forward(self, x): 39 | y = x[0] # no weight 40 | if self.weight: 41 | w = torch.sigmoid(self.w) * 2 42 | for i in self.iter: 43 | y = y + x[i + 1] * w[i] 44 | else: 45 | for i in self.iter: 46 | y = y + x[i + 1] 47 | return y 48 | 49 | 50 | class MixConv2d(nn.Module): 51 | # Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595 52 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kernel, stride, ch_strategy 53 | super().__init__() 54 | n = len(k) # number of convolutions 55 | if equal_ch: # equal c_ per group 56 | i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices 57 | c_ = [(i == g).sum() for g in range(n)] # intermediate channels 58 | else: # equal weight.numel() per group 59 | b = [c2] + [0] * n 60 | a = np.eye(n + 1, n, k=-1) 61 | a -= np.roll(a, 1, axis=1) 62 | a *= np.array(k) ** 2 63 | a[0] = 1 64 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 65 | 66 | self.m = nn.ModuleList( 67 | [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]) 68 | self.bn = nn.BatchNorm2d(c2) 69 | self.act = nn.SiLU() 70 | 71 | def forward(self, x): 72 | return self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 73 | 74 | 75 | class Ensemble(nn.ModuleList): 76 | # Ensemble of models 77 | def __init__(self): 78 | super().__init__() 79 | 80 | def forward(self, x, augment=False, profile=False, visualize=False): 81 | y = [] 82 | for module in self: 83 | y.append(module(x, augment, profile, visualize)[0]) 84 | # y = torch.stack(y).max(0)[0] # max ensemble 85 | # y = torch.stack(y).mean(0) # mean ensemble 86 | y = torch.cat(y, 1) # nms ensemble 87 | return y, None # inference, train output 88 | 89 | 90 | def attempt_load(weights, map_location=None, inplace=True, fuse=True): 91 | from models.yolo import Detect, Model 92 | 93 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 94 | model = Ensemble() 95 | for w in weights if isinstance(weights, list) else [weights]: 96 | ckpt = torch.load(attempt_download(w), map_location=map_location) # load 97 | if fuse: 98 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model 99 | else: 100 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval()) # without layer fuse 101 | 102 | # Compatibility updates 103 | for m in model.modules(): 104 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model]: 105 | m.inplace = inplace # pytorch 1.7.0 compatibility 106 | if type(m) is Detect: 107 | if not isinstance(m.anchor_grid, list): # new Detect Layer compatibility 108 | delattr(m, 'anchor_grid') 109 | setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) 110 | elif type(m) is Conv: 111 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility 112 | 113 | if len(model) == 1: 114 | return model[-1] # return model 115 | else: 116 | print(f'Ensemble created with {weights}\n') 117 | for k in ['names']: 118 | setattr(model, k, getattr(model[-1], k)) 119 | model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride 120 | return model # return ensemble 121 | -------------------------------------------------------------------------------- /models/hub/anchors.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Default anchors for COCO data 3 | 4 | 5 | # P5 ------------------------------------------------------------------------------------------------------------------- 6 | # P5-640: 7 | anchors_p5_640: 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 | 13 | # P6 ------------------------------------------------------------------------------------------------------------------- 14 | # 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 15 | anchors_p6_640: 16 | - [9,11, 21,19, 17,41] # P3/8 17 | - [43,32, 39,70, 86,64] # P4/16 18 | - [65,131, 134,130, 120,265] # P5/32 19 | - [282,180, 247,354, 512,387] # P6/64 20 | 21 | # 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 22 | anchors_p6_1280: 23 | - [19,27, 44,40, 38,94] # P3/8 24 | - [96,68, 86,152, 180,137] # P4/16 25 | - [140,301, 303,264, 238,542] # P5/32 26 | - [436,615, 739,380, 925,792] # P6/64 27 | 28 | # 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 29 | anchors_p6_1920: 30 | - [28,41, 67,59, 57,141] # P3/8 31 | - [144,103, 129,227, 270,205] # P4/16 32 | - [209,452, 455,396, 358,812] # P5/32 33 | - [653,922, 1109,570, 1387,1187] # P6/64 34 | 35 | 36 | # P7 ------------------------------------------------------------------------------------------------------------------- 37 | # 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 38 | anchors_p7_640: 39 | - [11,11, 13,30, 29,20] # P3/8 40 | - [30,46, 61,38, 39,92] # P4/16 41 | - [78,80, 146,66, 79,163] # P5/32 42 | - [149,150, 321,143, 157,303] # P6/64 43 | - [257,402, 359,290, 524,372] # P7/128 44 | 45 | # 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 46 | anchors_p7_1280: 47 | - [19,22, 54,36, 32,77] # P3/8 48 | - [70,83, 138,71, 75,173] # P4/16 49 | - [165,159, 148,334, 375,151] # P5/32 50 | - [334,317, 251,626, 499,474] # P6/64 51 | - [750,326, 534,814, 1079,818] # P7/128 52 | 53 | # 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 54 | anchors_p7_1920: 55 | - [29,34, 81,55, 47,115] # P3/8 56 | - [105,124, 207,107, 113,259] # P4/16 57 | - [247,238, 222,500, 563,227] # P5/32 58 | - [501,476, 376,939, 749,711] # P6/64 59 | - [1126,489, 801,1222, 1618,1227] # P7/128 60 | -------------------------------------------------------------------------------- /models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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-bifpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 BiFPN 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, 6], 1, Concat, [1]], # cat P4 <--- BiFPN change 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/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 FPN head 28 | head: 29 | [[-1, 3, C3, [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, C3, [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, C3, [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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 21 | [-1, 1, SPPF, [1024, 5]], # 9 22 | ] 23 | 24 | # YOLOv5 v6.0 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 | [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5) 54 | ] 55 | -------------------------------------------------------------------------------- /models/hub/yolov5-p6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, SPPF, [1024, 5]], # 11 24 | ] 25 | 26 | # YOLOv5 v6.0 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 (P6/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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # auto-anchor 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1280]], 25 | [-1, 1, SPPF, [1280, 5]], # 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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, 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/hub/yolov5l6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 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 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 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/yolov5n6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 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 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 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/yolov5s-ghost.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, GhostConv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3Ghost, [128]], 18 | [-1, 1, GhostConv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3Ghost, [256]], 20 | [-1, 1, GhostConv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3Ghost, [512]], 22 | [-1, 1, GhostConv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3Ghost, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, GhostConv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3Ghost, [512, False]], # 13 33 | 34 | [-1, 1, GhostConv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3Ghost, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, GhostConv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3Ghost, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, GhostConv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3Ghost, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5s-transformer.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3TR, [1024]], # 9 <--- C3TR() Transformer module 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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/hub/yolov5s6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 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 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 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 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 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/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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/yolov5n.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 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 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, 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, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 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 | -------------------------------------------------------------------------------- /prepare_data/README.md: -------------------------------------------------------------------------------- 1 | # convert xml labels to yolov5-rotation format. 2 | 3 | - you can label images with following tools: 4 | 1. [labelimg2](https://github.com/chinakook/labelImg2) 5 | 2. [rolabelimg](https://github.com/cgvict/roLabelImg) 6 | 7 | - converted label format: 8 | 9 | ``` 10 | classid center_x center_y longside shortside angle 11 | ``` 12 | 13 | center_x, longside are divided by image width. 14 | 15 | center_y, shortside are divided by image height. 16 | 17 | angle is the angle between the long side and the x-axis, the range is [0, 180). The x-axis rotate clockwise by angle will parallel to the long side 18 | 19 | 20 | # 转换数据标注格式 21 | 22 | - 旋转框标注可以使用如下标注软件: 23 | 1. [labelimg2](https://github.com/chinakook/labelImg2) 24 | 2. [rolabelimg](https://github.com/cgvict/roLabelImg) 25 | 26 | - 转换后标签格式: 27 | 28 | ``` 29 | classid center_x center_y longside shortside angle 30 | ``` 31 | 32 | center_x, longside 是相对于图像宽度的比例 33 | 34 | center_y, shortside 是相对于图像高度的比例 35 | 36 | angle 是 长边 与 x 轴之间的夹角，范围为[0, 180)，x轴顺时针方向旋转angle度即与长边平行。 37 | 38 | -------------------------------------------------------------------------------- /prepare_data/xml_to_yolo_rotation.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # coding: utf-8 3 | 4 | # ## convert xml to [yolov5-rotation](https://github.com/acai66/yolov5_rotation) format. [旋转版yolov5](https://github.com/acai66/yolov5_rotation)标签格式转换 5 | 6 | # In[1]: 7 | 8 | 9 | import xml.etree.ElementTree as ET 10 | from tqdm import tqdm # pip install tqdm 11 | import os 12 | import math 13 | 14 | 15 | # In[2]: 16 | 17 | 18 | workdir = './data_wire3/' # datasets root path. 数据集路径 19 | images_dir = os.path.join(workdir, 'all_images') # images path. 图像路径 20 | labels_dir = os.path.join(workdir, 'all_labels') # xml labels path. xml标签路径 21 | # yolov5_all_images = os.path.join(workdir, 'yolov5_all_images') # all images for yolov5 rotation. 转换后旋转版yolov5可用的图像路径 22 | yolov5_all_labels = os.path.join(workdir, 'yolov5_all_labels') # all labels for yolov5 rotation. 转换后旋转版yolov5可用的txt标签路径 23 | for d in [yolov5_all_labels]: 24 | if not os.path.exists(d): 25 | os.mkdir(d) 26 | 27 | 28 | # In[3]: 29 | 30 | 31 | all_files = [i for i in os.listdir(labels_dir) if i[-4:] == '.xml'] 32 | 33 | 34 | # In[4]: 35 | 36 | 37 | print('labels count: ', len(all_files)) 38 | 39 | 40 | # In[5]: 41 | 42 | 43 | keep_class_names = [ ] # auto scan if blank. 如果留空，会自动扫描类别 44 | 45 | 46 | # In[6]: 47 | 48 | 49 | class_names = dict(zip(keep_class_names, range(len(keep_class_names)))) 50 | 51 | 52 | # ## convert. 转换 53 | 54 | # In[7]: 55 | 56 | 57 | auto_scan = False 58 | class_index = 0 59 | if len(class_names) == 0: 60 | auto_scan = True 61 | print('Auto scan classnames enabled.') 62 | 63 | for file in tqdm(all_files): 64 | file_path = os.path.join(labels_dir, file) 65 | tree = ET.parse(file_path) 66 | root = tree.getroot() 67 | 68 | img_size = root.find('size') 69 | width = float(img_size.find('width').text) 70 | height = float(img_size.find('height').text) 71 | 72 | objs = root.findall('object') 73 | 74 | with open(os.path.join(yolov5_all_labels, file[:-4] + '.txt'), 'w+') as f: 75 | for obj in objs: 76 | name = obj.find('name').text.strip() 77 | if name not in class_names.keys(): 78 | if auto_scan: 79 | class_names[name] = class_index 80 | class_index = class_index + 1 81 | else: 82 | continue 83 | rbb = obj.find('robndbox') 84 | if not rbb: 85 | print('no robndbox in %s' % (file_path)) 86 | cx = float(rbb.find('cx').text) 87 | cy = float(rbb.find('cy').text) 88 | w = float(rbb.find('w').text) 89 | h = float(rbb.find('h').text) 90 | angle = float(rbb.find('angle').text) 91 | if angle > math.pi: 92 | angle = angle - math.pi 93 | 94 | degree = round(angle / math.pi * 180) 95 | if h > w: # swap w,h if w is not longside. 宽不是长边时，交换宽高 96 | w, h = h, w 97 | if degree < 90: 98 | degree = degree + 90 99 | else: 100 | degree = degree - 90 101 | cv_degree = degree # 180 - degree # opencv angle format. opencv格式角度 102 | if cv_degree == 180: 103 | cv_degree = 0 104 | assert cv_degree >= 0 and cv_degree < 180 105 | 106 | f.write('{} {} {} {} {} {}\n'.format(class_names[name], cx/width, cy/height, w/width, h/height, cv_degree)) 107 | # break 108 | 109 | 110 | # In[8]: 111 | 112 | 113 | print('class_names: ', class_names) 114 | 115 | 116 | # In[9]: 117 | 118 | 119 | sorted_keys = [i[0] for i in sorted(class_names.items(), key = lambda kv:(kv[1], kv[0]))] 120 | 121 | 122 | # In[10]: 123 | 124 | 125 | print('names: [ "{}" ]'.format('", "'.join(sorted_keys))) 126 | 127 | 128 | # ## split datasets for yolov5. 将转换格式后的数据集按yolov5方式组织 129 | 130 | # In[11]: 131 | 132 | 133 | import random 134 | import shutil 135 | 136 | 137 | # In[12]: 138 | 139 | 140 | imgages = os.path.join(workdir, 'images') 141 | labels = os.path.join(workdir, 'labels') 142 | train_img = os.path.join(imgages, 'train') 143 | val_img = os.path.join(imgages, 'val') 144 | train_label = os.path.join(labels, 'train') 145 | val_label = os.path.join(labels, 'val') 146 | for d in [imgages, labels, train_img, val_img, train_label, val_label]: 147 | if not os.path.exists(d): 148 | os.mkdir(d) 149 | 150 | 151 | # In[13]: 152 | 153 | 154 | all_txts = [i for i in os.listdir(yolov5_all_labels) if i[-4:] == '.txt'] 155 | 156 | 157 | # In[14]: 158 | 159 | 160 | random.shuffle(all_txts) # shuffle data. 随机打乱顺序 161 | 162 | 163 | # ## split (train, test) rate. 随机划分数据集，比例 164 | 165 | # In[15]: 166 | 167 | 168 | train_factor = 0.8 169 | 170 | 171 | # In[16]: 172 | 173 | 174 | inds = int(train_factor * len(all_txts)) 175 | 176 | 177 | # In[17]: 178 | 179 | 180 | print('all count: ', len(all_txts)) 181 | print('train count: ', inds) 182 | print('test count: ', len(all_txts) - inds) 183 | 184 | 185 | # In[18]: 186 | 187 | 188 | train_txts = all_txts[:inds] 189 | val_txts = all_txts[inds:] 190 | 191 | 192 | # ## copy images and labels to train/val dirs, only run once. 只执行一次 193 | 194 | # In[19]: 195 | 196 | 197 | run = True 198 | 199 | 200 | # In[20]: 201 | 202 | 203 | if run: 204 | for txt in train_txts: 205 | shutil.copyfile(os.path.join(yolov5_all_labels, txt), os.path.join(train_label, txt)) 206 | src = os.path.join(images_dir, txt[:-4] + '.jpg') 207 | if os.path.exists(src): 208 | shutil.copyfile(src, os.path.join(train_img, txt[:-4] + '.jpg')) 209 | else: 210 | shutil.copyfile(os.path.join(images_dir, txt[:-4] + '.png'), os.path.join(train_img, txt[:-4] + '.png')) 211 | 212 | 213 | # In[21]: 214 | 215 | 216 | if run: 217 | for txt in val_txts: 218 | shutil.copyfile(os.path.join(yolov5_all_labels, txt), os.path.join(val_label, txt)) 219 | src = os.path.join(images_dir, txt[:-4] + '.jpg') 220 | if os.path.exists(src): 221 | shutil.copyfile(src, os.path.join(val_img, txt[:-4] + '.jpg')) 222 | else: 223 | shutil.copyfile(os.path.join(images_dir, txt[:-4] + '.png'), os.path.join(val_img, txt[:-4] + '.png')) 224 | 225 | 226 | # In[ ]: 227 | 228 | 229 | 230 | 231 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # Base ---------------------------------------- 4 | matplotlib>=3.2.2 5 | numpy>=1.18.5 6 | opencv-python>=4.1.2 7 | Pillow>=7.1.2 8 | PyYAML>=5.3.1 9 | requests>=2.23.0 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 | pandas>=1.1.4 21 | seaborn>=0.11.0 22 | 23 | # Export -------------------------------------- 24 | # coremltools>=4.1 # CoreML export 25 | # onnx>=1.9.0 # ONNX export 26 | # onnx-simplifier>=0.3.6 # ONNX simplifier 27 | # scikit-learn==0.19.2 # CoreML quantization 28 | # tensorflow>=2.4.1 # TFLite export 29 | # tensorflowjs>=3.9.0 # TF.js export 30 | 31 | # Extras -------------------------------------- 32 | # albumentations>=1.0.3 33 | # Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172 34 | # pycocotools>=2.0 # COCO mAP 35 | # roboflow 36 | thop # FLOPs computation 37 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | # Project-wide configuration file, can be used for package metadata and other toll configurations 2 | # Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments 3 | 4 | [metadata] 5 | license_file = LICENSE 6 | description-file = README.md 7 | 8 | 9 | [tool:pytest] 10 | norecursedirs = 11 | .git 12 | dist 13 | build 14 | addopts = 15 | --doctest-modules 16 | --durations=25 17 | --color=yes 18 | 19 | 20 | [flake8] 21 | max-line-length = 120 22 | exclude = .tox,*.egg,build,temp 23 | select = E,W,F 24 | doctests = True 25 | verbose = 2 26 | # https://pep8.readthedocs.io/en/latest/intro.html#error-codes 27 | format = pylint 28 | # see: https://www.flake8rules.com/ 29 | ignore = 30 | E731 # Do not assign a lambda expression, use a def 31 | F405 32 | E402 33 | F841 34 | E741 35 | F821 36 | E722 37 | F401 38 | W504 39 | E127 40 | W504 41 | E231 42 | E501 43 | F403 44 | E302 45 | F541 46 | 47 | 48 | [isort] 49 | # https://pycqa.github.io/isort/docs/configuration/options.html 50 | line_length = 120 51 | multi_line_output = 0 52 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | utils/initialization 4 | """ 5 | 6 | 7 | def notebook_init(): 8 | # For YOLOv5 notebooks 9 | print('Checking setup...') 10 | from IPython import display # to display images and clear console output 11 | 12 | from utils.general import emojis 13 | from utils.torch_utils import select_device # YOLOv5 imports 14 | 15 | display.clear_output() 16 | select_device(newline=False) 17 | print(emojis('Setup complete ✅')) 18 | return display 19 | -------------------------------------------------------------------------------- /utils/activations.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Activation functions 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | # SiLU https://arxiv.org/pdf/1606.08415.pdf ---------------------------------------------------------------------------- 12 | class SiLU(nn.Module): # export-friendly version of nn.SiLU() 13 | @staticmethod 14 | def forward(x): 15 | return x * torch.sigmoid(x) 16 | 17 | 18 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 19 | @staticmethod 20 | def forward(x): 21 | # return x * F.hardsigmoid(x) # for torchscript and CoreML 22 | return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for torchscript, CoreML and ONNX 23 | 24 | 25 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 26 | class Mish(nn.Module): 27 | @staticmethod 28 | def forward(x): 29 | return x * F.softplus(x).tanh() 30 | 31 | 32 | class MemoryEfficientMish(nn.Module): 33 | class F(torch.autograd.Function): 34 | @staticmethod 35 | def forward(ctx, x): 36 | ctx.save_for_backward(x) 37 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 38 | 39 | @staticmethod 40 | def backward(ctx, grad_output): 41 | x = ctx.saved_tensors[0] 42 | sx = torch.sigmoid(x) 43 | fx = F.softplus(x).tanh() 44 | return grad_output * (fx + x * sx * (1 - fx * fx)) 45 | 46 | def forward(self, x): 47 | return self.F.apply(x) 48 | 49 | 50 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 51 | class FReLU(nn.Module): 52 | def __init__(self, c1, k=3): # ch_in, kernel 53 | super().__init__() 54 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False) 55 | self.bn = nn.BatchNorm2d(c1) 56 | 57 | def forward(self, x): 58 | return torch.max(x, self.bn(self.conv(x))) 59 | 60 | 61 | # ACON https://arxiv.org/pdf/2009.04759.pdf ---------------------------------------------------------------------------- 62 | class AconC(nn.Module): 63 | r""" ACON activation (activate or not). 64 | AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter 65 | according to "Activate or Not: Learning Customized Activation" . 66 | """ 67 | 68 | def __init__(self, c1): 69 | super().__init__() 70 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 71 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 72 | self.beta = nn.Parameter(torch.ones(1, c1, 1, 1)) 73 | 74 | def forward(self, x): 75 | dpx = (self.p1 - self.p2) * x 76 | return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x 77 | 78 | 79 | class MetaAconC(nn.Module): 80 | r""" ACON activation (activate or not). 81 | MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network 82 | according to "Activate or Not: Learning Customized Activation" . 83 | """ 84 | 85 | def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r 86 | super().__init__() 87 | c2 = max(r, c1 // r) 88 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 89 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 90 | self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True) 91 | self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True) 92 | # self.bn1 = nn.BatchNorm2d(c2) 93 | # self.bn2 = nn.BatchNorm2d(c1) 94 | 95 | def forward(self, x): 96 | y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True) 97 | # batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891 98 | # beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y))))) # bug/unstable 99 | beta = torch.sigmoid(self.fc2(self.fc1(y))) # bug patch BN layers removed 100 | dpx = (self.p1 - self.p2) * x 101 | return dpx * torch.sigmoid(beta * dpx) + self.p2 * x 102 | -------------------------------------------------------------------------------- /utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Auto-anchor utils 4 | """ 5 | 6 | import random 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from utils.general import LOGGER, colorstr, emojis 14 | 15 | PREFIX = colorstr('AutoAnchor: ') 16 | 17 | 18 | def check_anchor_order(m): 19 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 20 | a = m.anchors.prod(-1).view(-1) # anchor area 21 | da = a[-1] - a[0] # delta a 22 | ds = m.stride[-1] - m.stride[0] # delta s 23 | if da.sign() != ds.sign(): # same order 24 | LOGGER.info(f'{PREFIX}Reversing anchor order') 25 | m.anchors[:] = m.anchors.flip(0) 26 | 27 | 28 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 29 | # Check anchor fit to data, recompute if necessary 30 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 31 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 32 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 33 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 34 | 35 | def metric(k): # compute metric 36 | r = wh[:, None] / k[None] 37 | x = torch.min(r, 1 / r).min(2)[0] # ratio metric 38 | best = x.max(1)[0] # best_x 39 | aat = (x > 1 / thr).float().sum(1).mean() # anchors above threshold 40 | bpr = (best > 1 / thr).float().mean() # best possible recall 41 | return bpr, aat 42 | 43 | anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors 44 | bpr, aat = metric(anchors.cpu().view(-1, 2)) 45 | s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). ' 46 | if bpr > 0.98: # threshold to recompute 47 | LOGGER.info(emojis(f'{s}Current anchors are a good fit to dataset ✅')) 48 | else: 49 | LOGGER.info(emojis(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...')) 50 | na = m.anchors.numel() // 2 # number of anchors 51 | try: 52 | anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 53 | except Exception as e: 54 | LOGGER.info(f'{PREFIX}ERROR: {e}') 55 | new_bpr = metric(anchors)[0] 56 | if new_bpr > bpr: # replace anchors 57 | anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors) 58 | m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss 59 | check_anchor_order(m) 60 | LOGGER.info(f'{PREFIX}New anchors saved to model. Update model *.yaml to use these anchors in the future.') 61 | else: 62 | LOGGER.info(f'{PREFIX}Original anchors better than new anchors. Proceeding with original anchors.') 63 | 64 | 65 | def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 66 | """ Creates kmeans-evolved anchors from training dataset 67 | 68 | Arguments: 69 | dataset: path to data.yaml, or a loaded dataset 70 | n: number of anchors 71 | img_size: image size used for training 72 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 73 | gen: generations to evolve anchors using genetic algorithm 74 | verbose: print all results 75 | 76 | Return: 77 | k: kmeans evolved anchors 78 | 79 | Usage: 80 | from utils.autoanchor import *; _ = kmean_anchors() 81 | """ 82 | from scipy.cluster.vq import kmeans 83 | 84 | thr = 1 / thr 85 | 86 | def metric(k, wh): # compute metrics 87 | r = wh[:, None] / k[None] 88 | x = torch.min(r, 1 / r).min(2)[0] # ratio metric 89 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 90 | return x, x.max(1)[0] # x, best_x 91 | 92 | def anchor_fitness(k): # mutation fitness 93 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 94 | return (best * (best > thr).float()).mean() # fitness 95 | 96 | def print_results(k, verbose=True): 97 | k = k[np.argsort(k.prod(1))] # sort small to large 98 | x, best = metric(k, wh0) 99 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 100 | s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \ 101 | f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \ 102 | f'past_thr={x[x > thr].mean():.3f}-mean: ' 103 | for i, x in enumerate(k): 104 | s += '%i,%i, ' % (round(x[0]), round(x[1])) 105 | if verbose: 106 | LOGGER.info(s[:-2]) 107 | return k 108 | 109 | if isinstance(dataset, str): # *.yaml file 110 | with open(dataset, errors='ignore') as f: 111 | data_dict = yaml.safe_load(f) # model dict 112 | from utils.datasets import LoadImagesAndLabels 113 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 114 | 115 | # Get label wh 116 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 117 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 118 | 119 | # Filter 120 | i = (wh0 < 3.0).any(1).sum() 121 | if i: 122 | LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.') 123 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels 124 | # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 125 | 126 | # Kmeans calculation 127 | LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...') 128 | s = wh.std(0) # sigmas for whitening 129 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance 130 | assert len(k) == n, f'{PREFIX}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}' 131 | k *= s 132 | wh = torch.tensor(wh, dtype=torch.float32) # filtered 133 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered 134 | k = print_results(k, verbose=False) 135 | 136 | # Plot 137 | # k, d = [None] * 20, [None] * 20 138 | # for i in tqdm(range(1, 21)): 139 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 140 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 141 | # ax = ax.ravel() 142 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 143 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 144 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 145 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 146 | # fig.savefig('wh.png', dpi=200) 147 | 148 | # Evolve 149 | npr = np.random 150 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 151 | pbar = tqdm(range(gen), desc=f'{PREFIX}Evolving anchors with Genetic Algorithm:') # progress bar 152 | for _ in pbar: 153 | v = np.ones(sh) 154 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 155 | v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 156 | kg = (k.copy() * v).clip(min=2.0) 157 | fg = anchor_fitness(kg) 158 | if fg > f: 159 | f, k = fg, kg.copy() 160 | pbar.desc = f'{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' 161 | if verbose: 162 | print_results(k, verbose) 163 | 164 | return print_results(k) 165 | -------------------------------------------------------------------------------- /utils/autobatch.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Auto-batch utils 4 | """ 5 | 6 | from copy import deepcopy 7 | 8 | import numpy as np 9 | import torch 10 | from torch.cuda import amp 11 | 12 | from utils.general import LOGGER, colorstr 13 | from utils.torch_utils import profile 14 | 15 | 16 | def check_train_batch_size(model, imgsz=640): 17 | # Check YOLOv5 training batch size 18 | with amp.autocast(): 19 | return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size 20 | 21 | 22 | def autobatch(model, imgsz=640, fraction=0.9, batch_size=16): 23 | # Automatically estimate best batch size to use `fraction` of available CUDA memory 24 | # Usage: 25 | # import torch 26 | # from utils.autobatch import autobatch 27 | # model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False) 28 | # print(autobatch(model)) 29 | 30 | prefix = colorstr('AutoBatch: ') 31 | LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') 32 | device = next(model.parameters()).device # get model device 33 | if device.type == 'cpu': 34 | LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') 35 | return batch_size 36 | 37 | d = str(device).upper() # 'CUDA:0' 38 | properties = torch.cuda.get_device_properties(device) # device properties 39 | t = properties.total_memory / 1024 ** 3 # (GiB) 40 | r = torch.cuda.memory_reserved(device) / 1024 ** 3 # (GiB) 41 | a = torch.cuda.memory_allocated(device) / 1024 ** 3 # (GiB) 42 | f = t - (r + a) # free inside reserved 43 | LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free') 44 | 45 | batch_sizes = [1, 2, 4, 8, 16] 46 | try: 47 | img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes] 48 | y = profile(img, model, n=3, device=device) 49 | except Exception as e: 50 | LOGGER.warning(f'{prefix}{e}') 51 | 52 | y = [x[2] for x in y if x] # memory [2] 53 | batch_sizes = batch_sizes[:len(y)] 54 | p = np.polyfit(batch_sizes, y, deg=1) # first degree polynomial fit 55 | b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) 56 | LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)') 57 | return b 58 | -------------------------------------------------------------------------------- /utils/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/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 DDP 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 | FILE = Path(__file__).resolve() 12 | ROOT = FILE.parents[2] # YOLOv5 root directory 13 | if str(ROOT) not in sys.path: 14 | sys.path.append(str(ROOT)) # add ROOT to PATH 15 | 16 | port = 0 # --master_port 17 | path = Path('').resolve() 18 | for last in path.rglob('*/**/last.pt'): 19 | ckpt = torch.load(last) 20 | if ckpt['optimizer'] is None: 21 | continue 22 | 23 | # Load opt.yaml 24 | with open(last.parent.parent / 'opt.yaml', errors='ignore') as f: 25 | opt = yaml.safe_load(f) 26 | 27 | # Get device count 28 | d = opt['device'].split(',') # devices 29 | nd = len(d) # number of devices 30 | ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel 31 | 32 | if ddp: # multi-GPU 33 | port += 1 34 | cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}' 35 | else: # single-GPU 36 | cmd = f'python train.py --resume {last}' 37 | 38 | cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread 39 | print(cmd) 40 | os.system(cmd) 41 | -------------------------------------------------------------------------------- /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 -b master && sudo chmod -R 777 yolov5 11 | cd yolov5 12 | bash data/scripts/get_coco.sh && echo "COCO 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/callbacks.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Callback utils 4 | """ 5 | 6 | 7 | class Callbacks: 8 | """" 9 | Handles all registered callbacks for YOLOv5 Hooks 10 | """ 11 | 12 | # Define the available callbacks 13 | _callbacks = { 14 | 'on_pretrain_routine_start': [], 15 | 'on_pretrain_routine_end': [], 16 | 17 | 'on_train_start': [], 18 | 'on_train_epoch_start': [], 19 | 'on_train_batch_start': [], 20 | 'optimizer_step': [], 21 | 'on_before_zero_grad': [], 22 | 'on_train_batch_end': [], 23 | 'on_train_epoch_end': [], 24 | 25 | 'on_val_start': [], 26 | 'on_val_batch_start': [], 27 | 'on_val_image_end': [], 28 | 'on_val_batch_end': [], 29 | 'on_val_end': [], 30 | 31 | 'on_fit_epoch_end': [], # fit = train + val 32 | 'on_model_save': [], 33 | 'on_train_end': [], 34 | 35 | 'teardown': [], 36 | } 37 | 38 | def register_action(self, hook, name='', callback=None): 39 | """ 40 | Register a new action to a callback hook 41 | 42 | Args: 43 | hook The callback hook name to register the action to 44 | name The name of the action for later reference 45 | callback The callback to fire 46 | """ 47 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 48 | assert callable(callback), f"callback '{callback}' is not callable" 49 | self._callbacks[hook].append({'name': name, 'callback': callback}) 50 | 51 | def get_registered_actions(self, hook=None): 52 | """" 53 | Returns all the registered actions by callback hook 54 | 55 | Args: 56 | hook The name of the hook to check, defaults to all 57 | """ 58 | if hook: 59 | return self._callbacks[hook] 60 | else: 61 | return self._callbacks 62 | 63 | def run(self, hook, *args, **kwargs): 64 | """ 65 | Loop through the registered actions and fire all callbacks 66 | 67 | Args: 68 | hook The name of the hook to check, defaults to all 69 | args Arguments to receive from YOLOv5 70 | kwargs Keyword Arguments to receive from YOLOv5 71 | """ 72 | 73 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 74 | 75 | for logger in self._callbacks[hook]: 76 | logger['callback'](*args, **kwargs) 77 | -------------------------------------------------------------------------------- /utils/downloads.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Download utils 4 | """ 5 | 6 | import os 7 | import platform 8 | import subprocess 9 | import time 10 | import urllib 11 | from pathlib import Path 12 | from zipfile import ZipFile 13 | 14 | import requests 15 | import torch 16 | 17 | 18 | def gsutil_getsize(url=''): 19 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 20 | s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8') 21 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 22 | 23 | 24 | def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''): 25 | # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes 26 | file = Path(file) 27 | assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}" 28 | try: # url1 29 | print(f'Downloading {url} to {file}...') 30 | torch.hub.download_url_to_file(url, str(file)) 31 | assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check 32 | except Exception as e: # url2 33 | file.unlink(missing_ok=True) # remove partial downloads 34 | print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...') 35 | os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail 36 | finally: 37 | if not file.exists() or file.stat().st_size < min_bytes: # check 38 | file.unlink(missing_ok=True) # remove partial downloads 39 | print(f"ERROR: {assert_msg}\n{error_msg}") 40 | print('') 41 | 42 | 43 | def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads import *; attempt_download() 44 | # Attempt file download if does not exist 45 | file = Path(str(file).strip().replace("'", '')) 46 | 47 | if not file.exists(): 48 | # URL specified 49 | name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc. 50 | if str(file).startswith(('http:/', 'https:/')): # download 51 | url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ 52 | name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... 53 | safe_download(file=name, url=url, min_bytes=1E5) 54 | return name 55 | 56 | # GitHub assets 57 | file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required) 58 | try: 59 | response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api 60 | assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] 61 | tag = response['tag_name'] # i.e. 'v1.0' 62 | except: # fallback plan 63 | assets = ['yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 64 | 'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt'] 65 | try: 66 | tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] 67 | except: 68 | tag = 'v6.0' # current release 69 | 70 | if name in assets: 71 | safe_download(file, 72 | url=f'https://github.com/{repo}/releases/download/{tag}/{name}', 73 | # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) 74 | min_bytes=1E5, 75 | error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/') 76 | 77 | return str(file) 78 | 79 | 80 | def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): 81 | # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download() 82 | t = time.time() 83 | file = Path(file) 84 | cookie = Path('cookie') # gdrive cookie 85 | print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='') 86 | file.unlink(missing_ok=True) # remove existing file 87 | cookie.unlink(missing_ok=True) # remove existing cookie 88 | 89 | # Attempt file download 90 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 91 | os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}') 92 | if os.path.exists('cookie'): # large file 93 | s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}' 94 | else: # small file 95 | s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"' 96 | r = os.system(s) # execute, capture return 97 | cookie.unlink(missing_ok=True) # remove existing cookie 98 | 99 | # Error check 100 | if r != 0: 101 | file.unlink(missing_ok=True) # remove partial 102 | print('Download error ') # raise Exception('Download error') 103 | return r 104 | 105 | # Unzip if archive 106 | if file.suffix == '.zip': 107 | print('unzipping... ', end='') 108 | ZipFile(file).extractall(path=file.parent) # unzip 109 | file.unlink() # remove zip 110 | 111 | print(f'Done ({time.time() - t:.1f}s)') 112 | return r 113 | 114 | 115 | def get_token(cookie="./cookie"): 116 | with open(cookie) as f: 117 | for line in f: 118 | if "download" in line: 119 | return line.split()[-1] 120 | return "" 121 | 122 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries ---------------------------------------------- 123 | # 124 | # 125 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 126 | # # Uploads a file to a bucket 127 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 128 | # 129 | # storage_client = storage.Client() 130 | # bucket = storage_client.get_bucket(bucket_name) 131 | # blob = bucket.blob(destination_blob_name) 132 | # 133 | # blob.upload_from_filename(source_file_name) 134 | # 135 | # print('File {} uploaded to {}.'.format( 136 | # source_file_name, 137 | # destination_blob_name)) 138 | # 139 | # 140 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 141 | # # Uploads a blob from a bucket 142 | # storage_client = storage.Client() 143 | # bucket = storage_client.get_bucket(bucket_name) 144 | # blob = bucket.blob(source_blob_name) 145 | # 146 | # blob.download_to_filename(destination_file_name) 147 | # 148 | # print('Blob {} downloaded to {}.'.format( 149 | # source_blob_name, 150 | # destination_file_name)) 151 | -------------------------------------------------------------------------------- /utils/flask_rest_api/README.md: -------------------------------------------------------------------------------- 1 | # Flask REST API 2 | 3 | [REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are 4 | commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API 5 | created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/). 6 | 7 | ## Requirements 8 | 9 | [Flask](https://palletsprojects.com/p/flask/) is required. Install with: 10 | 11 | ```shell 12 | $ pip install Flask 13 | ``` 14 | 15 | ## Run 16 | 17 | After Flask installation run: 18 | 19 | ```shell 20 | $ python3 restapi.py --port 5000 21 | ``` 22 | 23 | Then use [curl](https://curl.se/) to perform a request: 24 | 25 | ```shell 26 | $ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s' 27 | ``` 28 | 29 | The model inference results are returned as a JSON response: 30 | 31 | ```json 32 | [ 33 | { 34 | "class": 0, 35 | "confidence": 0.8900438547, 36 | "height": 0.9318675399, 37 | "name": "person", 38 | "width": 0.3264600933, 39 | "xcenter": 0.7438579798, 40 | "ycenter": 0.5207948685 41 | }, 42 | { 43 | "class": 0, 44 | "confidence": 0.8440024257, 45 | "height": 0.7155083418, 46 | "name": "person", 47 | "width": 0.6546785235, 48 | "xcenter": 0.427829951, 49 | "ycenter": 0.6334488392 50 | }, 51 | { 52 | "class": 27, 53 | "confidence": 0.3771208823, 54 | "height": 0.3902671337, 55 | "name": "tie", 56 | "width": 0.0696444362, 57 | "xcenter": 0.3675483763, 58 | "ycenter": 0.7991207838 59 | }, 60 | { 61 | "class": 27, 62 | "confidence": 0.3527112305, 63 | "height": 0.1540903747, 64 | "name": "tie", 65 | "width": 0.0336618312, 66 | "xcenter": 0.7814827561, 67 | "ycenter": 0.5065554976 68 | } 69 | ] 70 | ``` 71 | 72 | An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given 73 | in `example_request.py` 74 | -------------------------------------------------------------------------------- /utils/flask_rest_api/example_request.py: -------------------------------------------------------------------------------- 1 | """Perform test request""" 2 | import pprint 3 | 4 | import requests 5 | 6 | DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s" 7 | TEST_IMAGE = "zidane.jpg" 8 | 9 | image_data = open(TEST_IMAGE, "rb").read() 10 | 11 | response = requests.post(DETECTION_URL, files={"image": image_data}).json() 12 | 13 | pprint.pprint(response) 14 | -------------------------------------------------------------------------------- /utils/flask_rest_api/restapi.py: -------------------------------------------------------------------------------- 1 | """ 2 | Run a rest API exposing the yolov5s object detection model 3 | """ 4 | import argparse 5 | import io 6 | 7 | import torch 8 | from flask import Flask, request 9 | from PIL import Image 10 | 11 | app = Flask(__name__) 12 | 13 | DETECTION_URL = "/v1/object-detection/yolov5s" 14 | 15 | 16 | @app.route(DETECTION_URL, methods=["POST"]) 17 | def predict(): 18 | if not request.method == "POST": 19 | return 20 | 21 | if request.files.get("image"): 22 | image_file = request.files["image"] 23 | image_bytes = image_file.read() 24 | 25 | img = Image.open(io.BytesIO(image_bytes)) 26 | 27 | results = model(img, size=640) # reduce size=320 for faster inference 28 | return results.pandas().xyxy[0].to_json(orient="records") 29 | 30 | 31 | if __name__ == "__main__": 32 | parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model") 33 | parser.add_argument("--port", default=5000, type=int, help="port number") 34 | args = parser.parse_args() 35 | 36 | model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache 37 | app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat 38 | -------------------------------------------------------------------------------- /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==21.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 15 | -------------------------------------------------------------------------------- /utils/loggers/__init__.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Logging utils 4 | """ 5 | 6 | import os 7 | import warnings 8 | from threading import Thread 9 | 10 | import pkg_resources as pkg 11 | import torch 12 | from torch.utils.tensorboard import SummaryWriter 13 | 14 | from utils.general import colorstr, emojis 15 | from utils.loggers.wandb.wandb_utils import WandbLogger 16 | from utils.plots import plot_images, plot_results 17 | from utils.torch_utils import de_parallel 18 | 19 | LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases 20 | RANK = int(os.getenv('RANK', -1)) 21 | 22 | try: 23 | import wandb 24 | 25 | assert hasattr(wandb, '__version__') # verify package import not local dir 26 | if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]: 27 | wandb_login_success = wandb.login(timeout=30) 28 | if not wandb_login_success: 29 | wandb = None 30 | except (ImportError, AssertionError): 31 | wandb = None 32 | 33 | 34 | class Loggers(): 35 | # YOLOv5 Loggers class 36 | def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS): 37 | self.save_dir = save_dir 38 | self.weights = weights 39 | self.opt = opt 40 | self.hyp = hyp 41 | self.logger = logger # for printing results to console 42 | self.include = include 43 | self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', 'train/angle_loss', # train loss 44 | 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics 45 | 'val/box_loss', 'val/obj_loss', 'val/cls_loss', 'val/angle_loss', # val loss 46 | 'x/lr0', 'x/lr1', 'x/lr2'] # params 47 | for k in LOGGERS: 48 | setattr(self, k, None) # init empty logger dictionary 49 | self.csv = True # always log to csv 50 | 51 | # Message 52 | if not wandb: 53 | prefix = colorstr('Weights & Biases: ') 54 | s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)" 55 | print(emojis(s)) 56 | 57 | # TensorBoard 58 | s = self.save_dir 59 | if 'tb' in self.include and not self.opt.evolve: 60 | prefix = colorstr('TensorBoard: ') 61 | self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/") 62 | self.tb = SummaryWriter(str(s)) 63 | 64 | # W&B 65 | if wandb and 'wandb' in self.include: 66 | wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://') 67 | run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None 68 | self.opt.hyp = self.hyp # add hyperparameters 69 | self.wandb = WandbLogger(self.opt, run_id) 70 | else: 71 | self.wandb = None 72 | 73 | def on_pretrain_routine_end(self): 74 | # Callback runs on pre-train routine end 75 | paths = self.save_dir.glob('*labels*.jpg') # training labels 76 | if self.wandb: 77 | self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]}) 78 | 79 | def on_train_batch_end(self, ni, model, imgs, targets, paths, plots, sync_bn): 80 | # Callback runs on train batch end 81 | if plots: 82 | if ni == 0: 83 | if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754 84 | with warnings.catch_warnings(): 85 | warnings.simplefilter('ignore') # suppress jit trace warning 86 | self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), []) 87 | if ni < 3: 88 | f = self.save_dir / f'train_batch{ni}.jpg' # filename 89 | Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() 90 | if self.wandb and ni == 10: 91 | files = sorted(self.save_dir.glob('train*.jpg')) 92 | self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]}) 93 | 94 | def on_train_epoch_end(self, epoch): 95 | # Callback runs on train epoch end 96 | if self.wandb: 97 | self.wandb.current_epoch = epoch + 1 98 | 99 | def on_val_image_end(self, pred, predn, path, names, im): 100 | # Callback runs on val image end 101 | if self.wandb: 102 | self.wandb.val_one_image(pred, predn, path, names, im) 103 | 104 | def on_val_end(self): 105 | # Callback runs on val end 106 | if self.wandb: 107 | files = sorted(self.save_dir.glob('val*.jpg')) 108 | self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]}) 109 | 110 | def on_fit_epoch_end(self, vals, epoch, best_fitness, fi): 111 | # Callback runs at the end of each fit (train+val) epoch 112 | x = {k: v for k, v in zip(self.keys, vals)} # dict 113 | if self.csv: 114 | file = self.save_dir / 'results.csv' 115 | n = len(x) + 1 # number of cols 116 | s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n') # add header 117 | with open(file, 'a') as f: 118 | f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n') 119 | 120 | if self.tb: 121 | for k, v in x.items(): 122 | self.tb.add_scalar(k, v, epoch) 123 | 124 | if self.wandb: 125 | self.wandb.log(x) 126 | self.wandb.end_epoch(best_result=best_fitness == fi) 127 | 128 | def on_model_save(self, last, epoch, final_epoch, best_fitness, fi): 129 | # Callback runs on model save event 130 | if self.wandb: 131 | if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1: 132 | self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi) 133 | 134 | def on_train_end(self, last, best, plots, epoch, results): 135 | # Callback runs on training end 136 | if plots: 137 | plot_results(file=self.save_dir / 'results.csv') # save results.png 138 | files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))] 139 | files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter 140 | 141 | if self.tb: 142 | import cv2 143 | for f in files: 144 | self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC') 145 | 146 | if self.wandb: 147 | self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]}) 148 | # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model 149 | if not self.opt.evolve: 150 | wandb.log_artifact(str(best if best.exists() else last), type='model', 151 | name='run_' + self.wandb.wandb_run.id + '_model', 152 | aliases=['latest', 'best', 'stripped']) 153 | self.wandb.finish_run() 154 | else: 155 | self.wandb.finish_run() 156 | self.wandb = WandbLogger(self.opt) 157 | -------------------------------------------------------------------------------- /utils/loggers/wandb/README.md: -------------------------------------------------------------------------------- 1 | 📚 This guide explains how to use **Weights & Biases** (W&B) with YOLOv5 🚀. UPDATED 29 September 2021. 2 | * [About Weights & Biases](#about-weights-&-biases) 3 | * [First-Time Setup](#first-time-setup) 4 | * [Viewing runs](#viewing-runs) 5 | * [Advanced Usage: Dataset Versioning and Evaluation](#advanced-usage) 6 | * [Reports: Share your work with the world!](#reports) 7 | 8 | ## About Weights & Biases 9 | Think of [W&B](https://wandb.ai/site?utm_campaign=repo_yolo_wandbtutorial) like GitHub for machine learning models. With a few lines of code, save everything you need to debug, compare and reproduce your models — architecture, hyperparameters, git commits, model weights, GPU usage, and even datasets and predictions. 10 | 11 | Used by top researchers including teams at OpenAI, Lyft, Github, and MILA, W&B is part of the new standard of best practices for machine learning. How W&B can help you optimize your machine learning workflows: 12 | 13 | * [Debug](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Free-2) model performance in real time 14 | * [GPU usage](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#System-4) visualized automatically 15 | * [Custom charts](https://wandb.ai/wandb/customizable-charts/reports/Powerful-Custom-Charts-To-Debug-Model-Peformance--VmlldzoyNzY4ODI) for powerful, extensible visualization 16 | * [Share insights](https://wandb.ai/wandb/getting-started/reports/Visualize-Debug-Machine-Learning-Models--VmlldzoyNzY5MDk#Share-8) interactively with collaborators 17 | * [Optimize hyperparameters](https://docs.wandb.com/sweeps) efficiently 18 | * [Track](https://docs.wandb.com/artifacts) datasets, pipelines, and production models 19 | 20 | ## First-Time Setup 21 |
22 | Toggle Details 23 | When you first train, W&B will prompt you to create a new account and will generate an **API key** for you. If you are an existing user you can retrieve your key from https://wandb.ai/authorize. This key is used to tell W&B where to log your data. You only need to supply your key once, and then it is remembered on the same device. 24 | 25 | W&B will create a cloud **project** (default is 'YOLOv5') for your training runs, and each new training run will be provided a unique run **name** within that project as project/name. You can also manually set your project and run name as: 26 | 27 | ```shell 28 | $ python train.py --project ... --name ... 29 | ``` 30 | 31 | YOLOv5 notebook example: Open In Colab Open In Kaggle 32 | Screen Shot 2021-09-29 at 10 23 13 PM 33 | 34 | 35 |
36 | 37 | ## Viewing Runs 38 |
39 | Toggle Details 40 | Run information streams from your environment to the W&B cloud console as you train. This allows you to monitor and even cancel runs in realtime . All important information is logged: 41 | 42 | * Training & Validation losses 43 | * Metrics: Precision, Recall, mAP@0.5, mAP@0.5:0.95 44 | * Learning Rate over time 45 | * A bounding box debugging panel, showing the training progress over time 46 | * GPU: Type, **GPU Utilization**, power, temperature, **CUDA memory usage** 47 | * System: Disk I/0, CPU utilization, RAM memory usage 48 | * Your trained model as W&B Artifact 49 | * Environment: OS and Python types, Git repository and state, **training command** 50 | 51 |

Weights & Biases dashboard

52 | 53 | 54 |
55 | 56 | ## Advanced Usage 57 | You can leverage W&B artifacts and Tables integration to easily visualize and manage your datasets, models and training evaluations. Here are some quick examples to get you started. 58 |
59 |

1. Visualize and Version Datasets

60 | Log, visualize, dynamically query, and understand your data with W&B Tables. You can use the following command to log your dataset as a W&B Table. This will generate a {dataset}_wandb.yaml file which can be used to train from dataset artifact. 61 |
62 | Usage 63 | Code $ python utils/logger/wandb/log_dataset.py --project ... --name ... --data .. 64 | 65 | ![Screenshot (64)](https://user-images.githubusercontent.com/15766192/128486078-d8433890-98a3-4d12-8986-b6c0e3fc64b9.png) 66 |
67 | 68 |

2: Train and Log Evaluation simultaneousy

69 | This is an extension of the previous section, but it'll also training after uploading the dataset. This also evaluation Table 70 | Evaluation table compares your predictions and ground truths across the validation set for each epoch. It uses the references to the already uploaded datasets, 71 | so no images will be uploaded from your system more than once. 72 |
73 | Usage 74 | Code $ python utils/logger/wandb/log_dataset.py --data .. --upload_data 75 | 76 | ![Screenshot (72)](https://user-images.githubusercontent.com/15766192/128979739-4cf63aeb-a76f-483f-8861-1c0100b938a5.png) 77 |
78 | 79 |

3: Train using dataset artifact

80 | When you upload a dataset as described in the first section, you get a new config file with an added `_wandb` to its name. This file contains the information that 81 | can be used to train a model directly from the dataset artifact. This also logs evaluation 82 |
83 | Usage 84 | Code $ python utils/logger/wandb/log_dataset.py --data {data}_wandb.yaml 85 | 86 | ![Screenshot (72)](https://user-images.githubusercontent.com/15766192/128979739-4cf63aeb-a76f-483f-8861-1c0100b938a5.png) 87 |
88 | 89 |

4: Save model checkpoints as artifacts

90 | To enable saving and versioning checkpoints of your experiment, pass `--save_period n` with the base cammand, where `n` represents checkpoint interval. 91 | You can also log both the dataset and model checkpoints simultaneously. If not passed, only the final model will be logged 92 | 93 |
94 | Usage 95 | Code $ python train.py --save_period 1 96 | 97 | ![Screenshot (68)](https://user-images.githubusercontent.com/15766192/128726138-ec6c1f60-639d-437d-b4ee-3acd9de47ef3.png) 98 |
99 | 100 |
101 | 102 |

5: Resume runs from checkpoint artifacts.

103 | Any run can be resumed using artifacts if the --resume argument starts with wandb-artifact:// prefix followed by the run path, i.e, wandb-artifact://username/project/runid . This doesn't require the model checkpoint to be present on the local system. 104 | 105 |
106 | Usage 107 | Code $ python train.py --resume wandb-artifact://{run_path} 108 | 109 | ![Screenshot (70)](https://user-images.githubusercontent.com/15766192/128728988-4e84b355-6c87-41ae-a591-14aecf45343e.png) 110 |
111 | 112 |

6: Resume runs from dataset artifact & checkpoint artifacts.

113 | Local dataset or model checkpoints are not required. This can be used to resume runs directly on a different device 114 | The syntax is same as the previous section, but you'll need to lof both the dataset and model checkpoints as artifacts, i.e, set bot --upload_dataset or 115 | train from _wandb.yaml file and set --save_period 116 | 117 |
118 | Usage 119 | Code $ python train.py --resume wandb-artifact://{run_path} 120 | 121 | ![Screenshot (70)](https://user-images.githubusercontent.com/15766192/128728988-4e84b355-6c87-41ae-a591-14aecf45343e.png) 122 |
123 | 124 | 125 | 126 | 127 |

Reports

128 | W&B Reports can be created from your saved runs for sharing online. Once a report is created you will receive a link you can use to publically share your results. Here is an example report created from the COCO128 tutorial trainings of all four YOLOv5 models ([link](https://wandb.ai/glenn-jocher/yolov5_tutorial/reports/YOLOv5-COCO128-Tutorial-Results--VmlldzozMDI5OTY)). 129 | 130 | Weights & Biases Reports 131 | 132 | 133 | ## Environments 134 | 135 | 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): 136 | 137 | - **Google Colab and Kaggle** notebooks with free GPU: Open In Colab Open In Kaggle 138 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) 139 | - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/AWS-Quickstart) 140 | - **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) Docker Pulls 141 | 142 | 143 | ## Status 144 | 145 | ![CI CPU testing](https://github.com/ultralytics/yolov5/workflows/CI%20CPU%20testing/badge.svg) 146 | 147 | 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)), validation ([val.py](https://github.com/ultralytics/yolov5/blob/master/val.py)), inference ([detect.py](https://github.com/ultralytics/yolov5/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov5/blob/master/export.py)) on MacOS, Windows, and Ubuntu every 24 hours and on every commit. 148 | -------------------------------------------------------------------------------- /utils/loggers/wandb/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/acai66/yolov5_rotation/883d0cc566db09b9f9e356cbb53eb5c5a5efc10d/utils/loggers/wandb/__init__.py -------------------------------------------------------------------------------- /utils/loggers/wandb/log_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from wandb_utils import WandbLogger 4 | 5 | from utils.general import LOGGER 6 | 7 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' 8 | 9 | 10 | def create_dataset_artifact(opt): 11 | logger = WandbLogger(opt, None, job_type='Dataset Creation') # TODO: return value unused 12 | if not logger.wandb: 13 | LOGGER.info("install wandb using `pip install wandb` to log the dataset") 14 | 15 | 16 | if __name__ == '__main__': 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') 19 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') 20 | parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project') 21 | parser.add_argument('--entity', default=None, help='W&B entity') 22 | parser.add_argument('--name', type=str, default='log dataset', help='name of W&B run') 23 | 24 | opt = parser.parse_args() 25 | opt.resume = False # Explicitly disallow resume check for dataset upload job 26 | 27 | create_dataset_artifact(opt) 28 | -------------------------------------------------------------------------------- /utils/loggers/wandb/sweep.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | 4 | import wandb 5 | 6 | FILE = Path(__file__).resolve() 7 | ROOT = FILE.parents[3] # YOLOv5 root directory 8 | if str(ROOT) not in sys.path: 9 | sys.path.append(str(ROOT)) # add ROOT to PATH 10 | 11 | from train import parse_opt, train 12 | from utils.callbacks import Callbacks 13 | from utils.general import increment_path 14 | from utils.torch_utils import select_device 15 | 16 | 17 | def sweep(): 18 | wandb.init() 19 | # Get hyp dict from sweep agent 20 | hyp_dict = vars(wandb.config).get("_items") 21 | 22 | # Workaround: get necessary opt args 23 | opt = parse_opt(known=True) 24 | opt.batch_size = hyp_dict.get("batch_size") 25 | opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve)) 26 | opt.epochs = hyp_dict.get("epochs") 27 | opt.nosave = True 28 | opt.data = hyp_dict.get("data") 29 | opt.weights = str(opt.weights) 30 | opt.cfg = str(opt.cfg) 31 | opt.data = str(opt.data) 32 | opt.hyp = str(opt.hyp) 33 | opt.project = str(opt.project) 34 | device = select_device(opt.device, batch_size=opt.batch_size) 35 | 36 | # train 37 | train(hyp_dict, opt, device, callbacks=Callbacks()) 38 | 39 | 40 | if __name__ == "__main__": 41 | sweep() 42 | -------------------------------------------------------------------------------- /utils/loggers/wandb/sweep.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for training 2 | # To set range- 3 | # Provide min and max values as: 4 | # parameter: 5 | # 6 | # min: scalar 7 | # max: scalar 8 | # OR 9 | # 10 | # Set a specific list of search space- 11 | # parameter: 12 | # values: [scalar1, scalar2, scalar3...] 13 | # 14 | # You can use grid, bayesian and hyperopt search strategy 15 | # For more info on configuring sweeps visit - https://docs.wandb.ai/guides/sweeps/configuration 16 | 17 | program: utils/loggers/wandb/sweep.py 18 | method: random 19 | metric: 20 | name: metrics/mAP_0.5 21 | goal: maximize 22 | 23 | parameters: 24 | # hyperparameters: set either min, max range or values list 25 | data: 26 | value: "data/coco128.yaml" 27 | batch_size: 28 | values: [64] 29 | epochs: 30 | values: [10] 31 | 32 | lr0: 33 | distribution: uniform 34 | min: 1e-5 35 | max: 1e-1 36 | lrf: 37 | distribution: uniform 38 | min: 0.01 39 | max: 1.0 40 | momentum: 41 | distribution: uniform 42 | min: 0.6 43 | max: 0.98 44 | weight_decay: 45 | distribution: uniform 46 | min: 0.0 47 | max: 0.001 48 | warmup_epochs: 49 | distribution: uniform 50 | min: 0.0 51 | max: 5.0 52 | warmup_momentum: 53 | distribution: uniform 54 | min: 0.0 55 | max: 0.95 56 | warmup_bias_lr: 57 | distribution: uniform 58 | min: 0.0 59 | max: 0.2 60 | box: 61 | distribution: uniform 62 | min: 0.02 63 | max: 0.2 64 | cls: 65 | distribution: uniform 66 | min: 0.2 67 | max: 4.0 68 | cls_pw: 69 | distribution: uniform 70 | min: 0.5 71 | max: 2.0 72 | obj: 73 | distribution: uniform 74 | min: 0.2 75 | max: 4.0 76 | obj_pw: 77 | distribution: uniform 78 | min: 0.5 79 | max: 2.0 80 | iou_t: 81 | distribution: uniform 82 | min: 0.1 83 | max: 0.7 84 | anchor_t: 85 | distribution: uniform 86 | min: 2.0 87 | max: 8.0 88 | fl_gamma: 89 | distribution: uniform 90 | min: 0.0 91 | max: 0.1 92 | hsv_h: 93 | distribution: uniform 94 | min: 0.0 95 | max: 0.1 96 | hsv_s: 97 | distribution: uniform 98 | min: 0.0 99 | max: 0.9 100 | hsv_v: 101 | distribution: uniform 102 | min: 0.0 103 | max: 0.9 104 | degrees: 105 | distribution: uniform 106 | min: 0.0 107 | max: 45.0 108 | translate: 109 | distribution: uniform 110 | min: 0.0 111 | max: 0.9 112 | scale: 113 | distribution: uniform 114 | min: 0.0 115 | max: 0.9 116 | shear: 117 | distribution: uniform 118 | min: 0.0 119 | max: 10.0 120 | perspective: 121 | distribution: uniform 122 | min: 0.0 123 | max: 0.001 124 | flipud: 125 | distribution: uniform 126 | min: 0.0 127 | max: 1.0 128 | fliplr: 129 | distribution: uniform 130 | min: 0.0 131 | max: 1.0 132 | mosaic: 133 | distribution: uniform 134 | min: 0.0 135 | max: 1.0 136 | mixup: 137 | distribution: uniform 138 | min: 0.0 139 | max: 1.0 140 | copy_paste: 141 | distribution: uniform 142 | min: 0.0 143 | max: 1.0 144 | -------------------------------------------------------------------------------- /utils/loss.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Loss functions 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import numpy as np 9 | import math 10 | 11 | from utils.metrics import bbox_iou 12 | from utils.torch_utils import is_parallel 13 | 14 | 15 | def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441 16 | # return positive, negative label smoothing BCE targets 17 | return 1.0 - 0.5 * eps, 0.5 * eps 18 | 19 | 20 | class BCEBlurWithLogitsLoss(nn.Module): 21 | # BCEwithLogitLoss() with reduced missing label effects. 22 | def __init__(self, alpha=0.05): 23 | super().__init__() 24 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss() 25 | self.alpha = alpha 26 | 27 | def forward(self, pred, true): 28 | loss = self.loss_fcn(pred, true) 29 | pred = torch.sigmoid(pred) # prob from logits 30 | dx = pred - true # reduce only missing label effects 31 | # dx = (pred - true).abs() # reduce missing label and false label effects 32 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4)) 33 | loss *= alpha_factor 34 | return loss.mean() 35 | 36 | 37 | class FocalLoss(nn.Module): 38 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 39 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 40 | super().__init__() 41 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 42 | self.gamma = gamma 43 | self.alpha = alpha 44 | self.reduction = loss_fcn.reduction 45 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 46 | 47 | def forward(self, pred, true): 48 | loss = self.loss_fcn(pred, true) 49 | # p_t = torch.exp(-loss) 50 | # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability 51 | 52 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py 53 | pred_prob = torch.sigmoid(pred) # prob from logits 54 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob) 55 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 56 | modulating_factor = (1.0 - p_t) ** self.gamma 57 | loss *= alpha_factor * modulating_factor 58 | 59 | if self.reduction == 'mean': 60 | return loss.mean() 61 | elif self.reduction == 'sum': 62 | return loss.sum() 63 | else: # 'none' 64 | return loss 65 | 66 | 67 | class QFocalLoss(nn.Module): 68 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 69 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 70 | super().__init__() 71 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 72 | self.gamma = gamma 73 | self.alpha = alpha 74 | self.reduction = loss_fcn.reduction 75 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 76 | 77 | def forward(self, pred, true): 78 | loss = self.loss_fcn(pred, true) 79 | 80 | pred_prob = torch.sigmoid(pred) # prob from logits 81 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 82 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma 83 | loss *= alpha_factor * modulating_factor 84 | 85 | if self.reduction == 'mean': 86 | return loss.mean() 87 | elif self.reduction == 'sum': 88 | return loss.sum() 89 | else: # 'none' 90 | return loss 91 | 92 | 93 | def gaussian_label(label, num_class, u=0, sig=4.0): 94 | ''' 95 | 转换成CSL Labels： 96 | 用高斯窗口函数根据角度θ的周期性赋予gt labels同样的周期性，使得损失函数在计算边界处时可以做到“差值很大但loss很小”； 97 | 并且使得其labels具有环形特征，能够反映各个θ之间的角度距离 98 | @param label: 当前box的θ类别 shape(1) 99 | @param num_class: θ类别数量=180 100 | @param u: 高斯函数中的μ 101 | @param sig: 高斯函数中的σ 102 | @return: 高斯离散数组:将高斯函数的最高值设置在θ所在的位置，例如label为45，则将高斯分布数列向右移动直至x轴为45时，取值为1 shape(180) 103 | ''' 104 | # floor()返回数字的下舍整数 ceil() 函数返回数字的上入整数 range(-90,90) 105 | # 以num_class=180为例，生成从-90到89的数字整形list shape(180) 106 | x = np.array(range(math.floor(-num_class / 2), math.ceil(num_class / 2), 1)) 107 | y_sig = np.exp(-(x - u) ** 2 / (2 * sig ** 2)) # shape(180) 为-90到89的经高斯公式计算后的数值 108 | # 将高斯函数的最高值设置在θ所在的位置，例如label为45，则将高斯分布数列向右移动直至x轴为45时，取值为1 109 | return np.concatenate([y_sig[math.ceil(num_class / 2) - int(label.item()):], 110 | y_sig[:math.ceil(num_class / 2) - int(label.item())]], axis=0) 111 | 112 | 113 | class ComputeLoss: 114 | # Compute losses 115 | def __init__(self, model, autobalance=False): 116 | self.sort_obj_iou = False 117 | device = next(model.parameters()).device # get model device 118 | h = model.hyp # hyperparameters 119 | self.class_index = 5 + model.nc 120 | 121 | # Define criteria 122 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) 123 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device)) 124 | BCEangle = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['angle_pw']])).to(device) 125 | 126 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 127 | self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets 128 | 129 | # Focal loss 130 | g = h['fl_gamma'] # focal loss gamma 131 | if g > 0: 132 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) 133 | BCEangle = FocalLoss(BCEangle, g) 134 | 135 | det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module 136 | self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7 137 | self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index 138 | self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance 139 | self.BCEangle = BCEangle 140 | for k in 'na', 'nc', 'nl', 'anchors': 141 | setattr(self, k, getattr(det, k)) 142 | 143 | def __call__(self, p, targets): # predictions, targets, model 144 | device = targets.device 145 | lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) 146 | langle = torch.zeros(1, device=device) 147 | tcls, tangle, tbox, indices, anchors = self.build_targets(p, targets) # targets 148 | 149 | # Losses 150 | for i, pi in enumerate(p): # layer index, layer predictions 151 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx 152 | tobj = torch.zeros_like(pi[..., 0], device=device) # target obj 153 | 154 | n = b.shape[0] # number of targets 155 | if n: 156 | ps = pi[b, a, gj, gi] # prediction subset corresponding to targets 157 | 158 | # Regression 159 | pxy = ps[:, :2].sigmoid() * 2 - 0.5 160 | pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i] 161 | pbox = torch.cat((pxy, pwh), 1) # predicted box 162 | iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target) 163 | lbox += (1.0 - iou).mean() # iou loss 164 | 165 | # Objectness 166 | score_iou = iou.detach().clamp(0).type(tobj.dtype) 167 | if self.sort_obj_iou: 168 | sort_id = torch.argsort(score_iou) 169 | b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id] 170 | tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio 171 | 172 | # Classification 173 | if self.nc > 1: # cls loss (only if multiple classes) 174 | t = torch.full_like(ps[:, 5:self.class_index], self.cn, device=device) # targets 175 | t[range(n), tcls[i]] = self.cp 176 | lcls += self.BCEcls(ps[:, 5:self.class_index], t) # BCE 177 | 178 | # Θ类别损失 179 | ttheta = torch.zeros_like(ps[:, self.class_index:]) # size(num, 180) 180 | for idx in range(len(ps)): # idx start from 0 to len(ps)-1 181 | # 3个tensor组成的list (tensor_angle_list[i]) 对每个步长网络生成对应的class tensor tangle[i].shape=(num_i, 1) 182 | theta = tangle[i][idx] # 取出第i个layer中的第idx个目标的角度数值 例如取值θ=90 183 | # CSL论文中窗口半径为6效果最佳，过小无法学到角度信息，过大则角度预测偏差加大 184 | csl_label = gaussian_label(theta, 180, u=0, sig=6) # 用长度为1的θ值构建长度为180的csl_label 185 | ttheta[idx] = torch.from_numpy(csl_label) # 将cls_label放入对应的目标中 186 | langle += self.BCEangle(ps[:, self.class_index:], ttheta) 187 | 188 | # Append targets to text file 189 | # with open('targets.txt', 'a') as file: 190 | # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] 191 | 192 | obji = self.BCEobj(pi[..., 4], tobj) 193 | lobj += obji * self.balance[i] # obj loss 194 | if self.autobalance: 195 | self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item() 196 | 197 | if self.autobalance: 198 | self.balance = [x / self.balance[self.ssi] for x in self.balance] 199 | lbox *= self.hyp['box'] 200 | lobj *= self.hyp['obj'] 201 | lcls *= self.hyp['cls'] 202 | langle *= self.hyp['angle'] 203 | bs = tobj.shape[0] # batch size 204 | 205 | return (lbox + lobj + lcls + langle) * bs, torch.cat((lbox, lobj, lcls, langle)).detach() 206 | 207 | def build_targets(self, p, targets): 208 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h) 209 | na, nt = self.na, targets.shape[0] # number of anchors, targets 210 | tcls, tbox, indices, anch = [], [], [], [] 211 | tangle = [] 212 | gain = torch.ones(8, device=targets.device) # normalized to gridspace gain 213 | ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) 214 | targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices 215 | 216 | g = 0.5 # bias 217 | off = torch.tensor([[0, 0], 218 | [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m 219 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm 220 | ], device=targets.device).float() * g # offsets 221 | 222 | for i in range(self.nl): 223 | anchors = self.anchors[i] 224 | gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain 225 | 226 | # Match targets to anchors 227 | t = targets * gain 228 | if nt: 229 | # Matches 230 | r = t[:, :, 4:6] / anchors[:, None] # wh ratio 231 | j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare 232 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) 233 | t = t[j] # filter 234 | 235 | # Offsets 236 | gxy = t[:, 2:4] # grid xy 237 | gxi = gain[[2, 3]] - gxy # inverse 238 | j, k = ((gxy % 1 < g) & (gxy > 1)).T 239 | l, m = ((gxi % 1 < g) & (gxi > 1)).T 240 | j = torch.stack((torch.ones_like(j), j, k, l, m)) 241 | t = t.repeat((5, 1, 1))[j] 242 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j] 243 | else: 244 | t = targets[0] 245 | offsets = 0 246 | 247 | # Define 248 | b, c = t[:, :2].long().T # image, class 249 | angle = t[:, 6].long() # angle 250 | gxy = t[:, 2:4] # grid xy 251 | gwh = t[:, 4:6] # grid wh 252 | gij = (gxy - offsets).long() 253 | gi, gj = gij.T # grid xy indices 254 | 255 | # Append 256 | a = t[:, 7].long() # anchor indices 257 | indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices 258 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box 259 | anch.append(anchors[a]) # anchors 260 | tcls.append(c) # class 261 | tangle.append(angle) # angle 262 | 263 | return tcls, tangle, tbox, indices, anch 264 | --------------------------------------------------------------------------------