apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

YOLOv5 v5.0 release compatibility update for YOLOv3 (#1737)

Browse Source 
* YOLOv5 v5.0 release compatibility update

* Update README

* Update README

* Conv act LeakyReLU(0.1)

* update plots_study()

* update speeds
This commit is contained in:
Glenn Jocher 2021-04-12 18:00:47 +02:00 committed by GitHub
parent 5d8f03020c
commit 8eb4cde090
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
37 changed files with 1791 additions and 908 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitattributes vendored
View File

@ -1,2 +1,2 @@
# remove notebooks from GitHub language stats # this drop notebooks from GitHub language stats
*.ipynb linguist-vendored *.ipynb linguist-vendored

9
.github/workflows/ci-testing.yml vendored
View File

@ -66,14 +66,15 @@ jobs:
di=cpu # inference devices # define device di=cpu # inference devices # define device
# train # train
python train.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di python train.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di
# detect # detect
python detect.py --weights weights/${{ matrix.model }}.pt --device $di python detect.py --weights weights/${{ matrix.model }}.pt --device $di
python detect.py --weights runs/train/exp/weights/last.pt --device $di python detect.py --weights runs/train/exp/weights/last.pt --device $di
# test # test
python test.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --device $di python test.py --img 128 --batch 16 --weights weights/${{ matrix.model }}.pt --device $di
python test.py --img 256 --batch 8 --weights runs/train/exp/weights/last.pt --device $di python test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
python hubconf.py # hub
python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect
python models/export.py --img 256 --batch 1 --weights weights/${{ matrix.model }}.pt # export python models/export.py --img 128 --batch 1 --weights weights/${{ matrix.model }}.pt # export
shell: bash shell: bash

9
.github/workflows/greetings.yml vendored
View File

@ -23,7 +23,7 @@ jobs:
- ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ -Bruce Lee - ✅ Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ -Bruce Lee
issue-message: | issue-message: |
👋 Hello @${{ github.actor }}, thank you for your interest in 🚀 YOLOv3! Please visit our ⭐️ [Tutorials](https://github.com/ultralytics/yolov3/wiki#tutorials) to get started, where you can find quickstart guides for simple tasks like [Custom Data Training](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data) all the way to advanced concepts like [Hyperparameter Evolution](https://github.com/ultralytics/yolov5/issues/607). 👋 Hello @${{ github.actor }}, thank you for your interest in YOLOv3 🚀! Please visit our ⭐️ [Tutorials](https://github.com/ultralytics/yolov3/wiki#tutorials) to get started, where you can find quickstart guides for simple tasks like [Custom Data Training](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data) all the way to advanced concepts like [Hyperparameter Evolution](https://github.com/ultralytics/yolov5/issues/607).
If this is a 🐛 Bug Report, please provide screenshots and **minimum viable code to reproduce your issue**, otherwise we can not help you. If this is a 🐛 Bug Report, please provide screenshots and **minimum viable code to reproduce your issue**, otherwise we can not help you.
@ -42,10 +42,11 @@ jobs:
YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled): YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
- **Google Colab Notebook** with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> - **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
- **Kaggle Notebook** with free GPU: [https://www.kaggle.com/ultralytics/yolov3](https://www.kaggle.com/ultralytics/yolov3)
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart) - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov3. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) ![Docker Pulls](https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker) - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
## Status ## Status

19
Dockerfile
View File

@ -1,14 +1,14 @@
# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch # Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
FROM nvcr.io/nvidia/pytorch:20.12-py3 FROM nvcr.io/nvidia/pytorch:21.03-py3
# Install linux packages # Install linux packages
RUN apt update && apt install -y screen libgl1-mesa-glx RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
# Install python dependencies # Install python dependencies
RUN pip install --upgrade pip
COPY requirements.txt . COPY requirements.txt .
RUN pip install -r requirements.txt RUN python -m pip install --upgrade pip
RUN pip install gsutil RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof
RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook
# Create working directory # Create working directory
RUN mkdir -p /usr/src/app RUN mkdir -p /usr/src/app
@ -17,6 +17,9 @@ WORKDIR /usr/src/app
# Copy contents # Copy contents
COPY . /usr/src/app COPY . /usr/src/app
# Set environment variables
ENV HOME=/usr/src/app
# --------------------------------------------------- Extras Below --------------------------------------------------- # --------------------------------------------------- Extras Below ---------------------------------------------------
@ -34,13 +37,13 @@ COPY . /usr/src/app
# sudo docker kill $(sudo docker ps -q) # sudo docker kill $(sudo docker ps -q)
# Kill all image-based # Kill all image-based
# sudo docker kill $(sudo docker ps -a -q --filter ancestor=ultralytics/yolov3:latest) # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest)
# Bash into running container # Bash into running container
# sudo docker container exec -it ba65811811ab bash # sudo docker exec -it 5a9b5863d93d bash
# Bash into stopped container # Bash into stopped container
# sudo docker commit 092b16b25c5b usr/resume && sudo docker run -it --gpus all --ipc=host -v "$(pwd)"/coco:/usr/src/coco --entrypoint=sh usr/resume # id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash
# Send weights to GCP # Send weights to GCP
# python -c "from utils.general import *; strip_optimizer('runs/train/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt # python -c "from utils.general import *; strip_optimizer('runs/train/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt

108
README.md
View File

@ -1,35 +1,61 @@
<a href="https://apps.apple.com/app/id1452689527" target="_blank"> <a align="left" href="https://apps.apple.com/app/id1452689527" target="_blank">
<img src="https://user-images.githubusercontent.com/26833433/99805965-8f2ca800-2b3d-11eb-8fad-13a96b222a23.jpg" width="1000"></a> <img width="800" src="https://user-images.githubusercontent.com/26833433/99805965-8f2ca800-2b3d-11eb-8fad-13a96b222a23.jpg"></a>
&nbsp &nbsp
<a href="https://github.com/ultralytics/yolov3/actions"><img src="https://github.com/ultralytics/yolov3/workflows/CI%20CPU%20testing/badge.svg" alt="CI CPU testing"></a> <a href="https://github.com/ultralytics/yolov3/actions"><img src="https://github.com/ultralytics/yolov3/workflows/CI%20CPU%20testing/badge.svg" alt="CI CPU testing"></a>
BRANCH NOTICE: The [ultralytics/yolov3](https://github.com/ultralytics/yolov3) repository is now divided into two branches: This repository represents Ultralytics open-source research into future object detection methods, and incorporates lessons learned and best practices evolved over thousands of hours of training and evolution on anonymized client datasets. **All code and models are under active development, and are subject to modification or deletion without notice.** Use at your own risk.
* [Master branch](https://github.com/ultralytics/yolov3/tree/master): Forward-compatible with all [YOLOv5](https://github.com/ultralytics/yolov5) models and methods (**recommended**).
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/114424655-a0dc1e00-9bb8-11eb-9a2e-cbe21803f05c.png"></p>
<details>
<summary>YOLOv5-P5 640 Figure (click to expand)</summary>
<p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/114313219-f1d70e00-9af5-11eb-9973-52b1f98d321a.png"></p>
</details>
<details>
<summary>Figure Notes (click to expand)</summary>
* GPU Speed measures end-to-end time per image averaged over 5000 COCO val2017 images using a V100 GPU with batch size 32, and includes image preprocessing, PyTorch FP16 inference, postprocessing and NMS.
* EfficientDet data from [google/automl](https://github.com/google/automl) at batch size 8.
* **Reproduce** by `python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt yolov5l.pt`
</details>
## Branch Notice
The [ultralytics/yolov3](https://github.com/ultralytics/yolov3) repository is now divided into two branches:
* [Master branch](https://github.com/ultralytics/yolov3/tree/master): Forward-compatible with all [YOLOv5](https://github.com/ultralytics/yolov5) models and methods (**recommended** ✅).
```bash ```bash
$ git clone https://github.com/ultralytics/yolov3 # master branch (default) $ git clone https://github.com/ultralytics/yolov3 # master branch (default)
``` ```
* [Archive branch](https://github.com/ultralytics/yolov3/tree/archive): Backwards-compatible with original [darknet](https://pjreddie.com/darknet/) *.cfg models (⚠️ no longer maintained). * [Archive branch](https://github.com/ultralytics/yolov3/tree/archive): Backwards-compatible with original [darknet](https://pjreddie.com/darknet/) *.cfg models (**no longer maintained** ⚠️).
```bash ```bash
$ git clone https://github.com/ultralytics/yolov3 -b archive # archive branch $ git clone https://github.com/ultralytics/yolov3 -b archive # archive branch
``` ```
<img src="https://user-images.githubusercontent.com/26833433/100382066-c8bc5200-301a-11eb-907b-799a0301595e.png" width="1000">** GPU Speed measures end-to-end time per image averaged over 5000 COCO val2017 images using a V100 GPU with batch size 32, and includes image preprocessing, PyTorch FP16 inference, postprocessing and NMS. EfficientDet data from [google/automl](https://github.com/google/automl) at batch size 8.
## Pretrained Checkpoints ## Pretrained Checkpoints
| Model | AP<sup>val</sup> | AP<sup>test</sup> | AP<sub>50</sub> | Speed<sub>GPU</sub> | FPS<sub>GPU</sub> || params | FLOPS | [assets3]: https://github.com/ultralytics/yolov3/releases
|---------- |------ |------ |------ | -------- | ------| ------ |------ | :------: | [assets5]: https://github.com/ultralytics/yolov5/releases
| [YOLOv3](https://github.com/ultralytics/yolov3/releases) | 43.3 | 43.3 | 63.0 | 4.8ms | 208 || 61.9M | 156.4B
| [YOLOv3-SPP](https://github.com/ultralytics/yolov3/releases) | **44.3** | **44.3** | **64.6** | 4.9ms | 204 || 63.0M | 157.0B Model |size<br><sup>(pixels) |mAP<sup>val<br>0.5:0.95 |mAP<sup>test<br>0.5:0.95 |mAP<sup>val<br>0.5 |Speed<br><sup>V100 (ms) | |params<br><sup>(M) |FLOPS<br><sup>640 (B)
| [YOLOv3-tiny](https://github.com/ultralytics/yolov3/releases) | 17.6 | 34.9 | 34.9 | **1.7ms** | **588** || 8.9M | 13.3B --- |--- |--- |--- |--- |--- |---|--- |---
[YOLOv3-tiny][assets3] |640 |17.6 |17.6 |34.8 |**1.2** | |8.8 |13.2
[YOLOv3][assets3] |640 |43.3 |43.3 |63.0 |4.1 | |61.9 |156.3
[YOLOv3-SPP][assets3] |640 |44.3 |44.3 |64.6 |4.1 | |63.0 |157.1
| | | | | | || |
[YOLOv5l][assets5] |640 |**48.2** |**48.2** |**66.9** |3.7 | |47.0 |115.4
<details>
<summary>Table Notes (click to expand)</summary>
* AP<sup>test</sup> denotes COCO [test-dev2017](http://cocodataset.org/#upload) server results, all other AP results denote val2017 accuracy.
* AP values are for single-model single-scale unless otherwise noted. **Reproduce mAP** by `python test.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
* Speed<sub>GPU</sub> averaged over 5000 COCO val2017 images using a GCP [n1-standard-16](https://cloud.google.com/compute/docs/machine-types#n1_standard_machine_types) V100 instance, and includes FP16 inference, postprocessing and NMS. **Reproduce speed** by `python test.py --data coco.yaml --img 640 --conf 0.25 --iou 0.45`
* All checkpoints are trained to 300 epochs with default settings and hyperparameters (no autoaugmentation).
</details>
** AP<sup>test</sup> denotes COCO [test-dev2017](http://cocodataset.org/#upload) server results, all other AP results denote val2017 accuracy.
** All AP numbers are for single-model single-scale without ensemble or TTA. **Reproduce mAP** by `python test.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
** Speed<sub>GPU</sub> averaged over 5000 COCO val2017 images using a GCP [n1-standard-16](https://cloud.google.com/compute/docs/machine-types#n1_standard_machine_types) V100 instance, and includes image preprocessing, FP16 inference, postprocessing and NMS. NMS is 1-2ms/img. **Reproduce speed** by `python test.py --data coco.yaml --img 640 --conf 0.25 --iou 0.45`
** All checkpoints are trained to 300 epochs with default settings and hyperparameters (no autoaugmentation).
** Test Time Augmentation ([TTA](https://github.com/ultralytics/yolov5/issues/303)) runs at 3 image sizes. **Reproduce TTA** by `python test.py --data coco.yaml --img 832 --iou 0.65 --augment`
## Requirements ## Requirements
@ -42,7 +68,9 @@ $ pip install -r requirements.txt
## Tutorials ## Tutorials
* [Train Custom Data](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data)&nbsp; 🚀 RECOMMENDED * [Train Custom Data](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data)&nbsp; 🚀 RECOMMENDED
* [Tips for Best Training Results](https://github.com/ultralytics/yolov5/wiki/Tips-for-Best-Training-Results)&nbsp; ☘️ RECOMMENDED
* [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289)&nbsp; 🌟 NEW * [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289)&nbsp; 🌟 NEW
* [Supervisely Ecosystem](https://github.com/ultralytics/yolov5/issues/2518)&nbsp; 🌟 NEW
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475) * [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36)&nbsp; ⭐ NEW * [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36)&nbsp; ⭐ NEW
* [ONNX and TorchScript Export](https://github.com/ultralytics/yolov5/issues/251) * [ONNX and TorchScript Export](https://github.com/ultralytics/yolov5/issues/251)
@ -58,73 +86,59 @@ $ pip install -r requirements.txt
YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled): YOLOv3 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
- **Google Colab Notebook** with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> - **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
- **Kaggle Notebook** with free GPU: [https://www.kaggle.com/ultralytics/yolov3](https://www.kaggle.com/ultralytics/yolov3)
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart) - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov3. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) ![Docker Pulls](https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker) - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
## Inference ## Inference
detect.py runs inference on a variety of sources, downloading models automatically from the [latest YOLOv3 release](https://github.com/ultralytics/yolov3/releases) and saving results to `runs/detect`. `detect.py` runs inference on a variety of sources, downloading models automatically from the [latest YOLOv3 release](https://github.com/ultralytics/yolov3/releases) and saving results to `runs/detect`.
```bash ```bash
$ python detect.py --source 0 # webcam $ python detect.py --source 0 # webcam
file.jpg # image file.jpg # image
file.mp4 # video file.mp4 # video
path/ # directory path/ # directory
path/*.jpg # glob path/*.jpg # glob
rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp stream 'https://youtu.be/NUsoVlDFqZg' # YouTube video
rtmp://192.168.1.105/live/test # rtmp stream 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8 # http stream
``` ```
To run inference on example images in `data/images`: To run inference on example images in `data/images`:
```bash ```bash
$ python detect.py --source data/images --weights yolov3.pt --conf 0.25 $ python detect.py --source data/images --weights yolov3.pt --conf 0.25
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, img_size=640, iou_thres=0.45, name='exp', project='runs/detect', save_conf=False, save_txt=False, source='data/images/', update=False, view_img=False, weights=['yolov3.pt'])
Using torch 1.7.0+cu101 CUDA:0 (Tesla V100-SXM2-16GB, 16130MB)
Downloading https://github.com/ultralytics/yolov3/releases/download/v1.0/yolov3.pt to yolov3.pt... 100% 118M/118M [00:05<00:00, 24.2MB/s]
Fusing layers...
Model Summary: 261 layers, 61922845 parameters, 0 gradients
image 1/2 /content/yolov3/data/images/bus.jpg: 640x480 4 persons, 1 buss, Done. (0.014s)
image 2/2 /content/yolov3/data/images/zidane.jpg: 384x640 2 persons, 3 ties, Done. (0.014s)
Results saved to runs/detect/exp
Done. (0.133s)
``` ```
<img src="https://user-images.githubusercontent.com/26833433/100375993-06b37900-300f-11eb-8d2d-5fc7b22fbfbd.jpg" width="500"> <img width="500" src="https://user-images.githubusercontent.com/26833433/100375993-06b37900-300f-11eb-8d2d-5fc7b22fbfbd.jpg">
### PyTorch Hub ### PyTorch Hub
To run **batched inference** with YOLO3 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36): To run **batched inference** with YOLOv5 and [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36):
```python ```python
import torch import torch
from PIL import Image
# Model # Model
model = torch.hub.load('ultralytics/yolov3', 'yolov3', pretrained=True).autoshape() # for PIL/cv2/np inputs and NMS model = torch.hub.load('ultralytics/yolov3', 'yolov3') # or 'yolov3_spp', 'yolov3_tiny'
# Images # Images
img1 = Image.open('zidane.jpg') dir = 'https://github.com/ultralytics/yolov3/raw/master/data/images/'
img2 = Image.open('bus.jpg') imgs = [dir + f for f in ('zidane.jpg', 'bus.jpg')] # batch of images
imgs = [img1, img2] # batched list of images
# Inference # Inference
prediction = model(imgs, size=640) # includes NMS results = model(imgs)
results.print() # or .show(), .save()
``` ```
## Training ## Training
Download [COCO](https://github.com/ultralytics/yolov3/blob/master/data/scripts/get_coco.sh) and run command below. Training times for YOLOv3/YOLOv3-SPP/YOLOv3-tiny are 6/6/2 days on a single V100 (multi-GPU times faster). Use the largest `--batch-size` your GPU allows (batch sizes shown for 16 GB devices). Run commands below to reproduce results on [COCO](https://github.com/ultralytics/yolov3/blob/master/data/scripts/get_coco.sh) dataset (dataset auto-downloads on first use). Training times for YOLOv3/YOLOv3-SPP/YOLOv3-tiny are 6/6/2 days on a single V100 (multi-GPU times faster). Use the largest `--batch-size` your GPU allows (batch sizes shown for 16 GB devices).
```bash ```bash
$ python train.py --data coco.yaml --cfg yolov3.yaml --weights '' --batch-size 24 $ python train.py --data coco.yaml --cfg yolov3.yaml --weights '' --batch-size 24
yolov3-spp.yaml 24 yolov3-spp.yaml 24
yolov3-tiny.yaml 64 yolov3-tiny.yaml 64
``` ```
<img src="https://user-images.githubusercontent.com/26833433/100378028-af170c80-3012-11eb-8521-f0d2a8d021bc.png" width="900"> <img width="800" src="https://user-images.githubusercontent.com/26833433/100378028-af170c80-3012-11eb-8521-f0d2a8d021bc.png">
## Citation ## Citation

21
data/argoverse_hd.yaml Normal file
View File

@ -0,0 +1,21 @@
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
# Train command: python train.py --data argoverse_hd.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /argoverse
# /yolov5
# download command/URL (optional)
download: bash data/scripts/get_argoverse_hd.sh
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../argoverse/Argoverse-1.1/images/train/ # 39384 images
val: ../argoverse/Argoverse-1.1/images/val/ # 15062 iamges
test: ../argoverse/Argoverse-1.1/images/test/ # Submit to: https://eval.ai/web/challenges/challenge-page/800/overview
# number of classes
nc: 8
# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic_light', 'stop_sign' ]

62
data/scripts/get_argoverse_hd.sh Normal file
View File

@ -0,0 +1,62 @@
#!/bin/bash
# Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
# Download command: bash data/scripts/get_argoverse_hd.sh
# Train command: python train.py --data argoverse_hd.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /argoverse
# /yolov5
# Download/unzip images
d='../argoverse/' # unzip directory
mkdir $d
url=https://argoverse-hd.s3.us-east-2.amazonaws.com/
f=Argoverse-HD-Full.zip
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f &# download, unzip, remove in background
wait # finish background tasks
cd ../argoverse/Argoverse-1.1/
ln -s tracking images
cd ../Argoverse-HD/annotations/
python3 - "$@" <<END
import json
from pathlib import Path
annotation_files = ["train.json", "val.json"]
print("Converting annotations to YOLOv5 format...")
for val in annotation_files:
a = json.load(open(val, "rb"))
label_dict = {}
for annot in a['annotations']:
img_id = annot['image_id']
img_name = a['images'][img_id]['name']
img_label_name = img_name[:-3] + "txt"
obj_class = annot['category_id']
x_center, y_center, width, height = annot['bbox']
x_center = (x_center + width / 2) / 1920. # offset and scale
y_center = (y_center + height / 2) / 1200. # offset and scale
width /= 1920. # scale
height /= 1200. # scale
img_dir = "./labels/" + a['seq_dirs'][a['images'][annot['image_id']]['sid']]
Path(img_dir).mkdir(parents=True, exist_ok=True)
if img_dir + "/" + img_label_name not in label_dict:
label_dict[img_dir + "/" + img_label_name] = []
label_dict[img_dir + "/" + img_label_name].append(f"{obj_class} {x_center} {y_center} {width} {height}\n")
for filename in label_dict:
with open(filename, "w") as file:
for string in label_dict[filename]:
file.write(string)
END
mv ./labels ../../Argoverse-1.1/

9
data/scripts/get_coco.sh
View File

@ -10,8 +10,9 @@
# Download/unzip labels # Download/unzip labels
d='../' # unzip directory d='../' # unzip directory
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/ url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
f='coco2017labels.zip' # 68 MB f='coco2017labels.zip' # or 'coco2017labels-segments.zip', 68 MB
echo 'Downloading' $url$f ' ...' && curl -L $url$f -o $f && unzip -q $f -d $d && rm $f # download, unzip, remove echo 'Downloading' $url$f ' ...'
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
# Download/unzip images # Download/unzip images
d='../coco/images' # unzip directory d='../coco/images' # unzip directory
@ -20,5 +21,7 @@ f1='train2017.zip' # 19G, 118k images
f2='val2017.zip' # 1G, 5k images f2='val2017.zip' # 1G, 5k images
f3='test2017.zip' # 7G, 41k images (optional) f3='test2017.zip' # 7G, 41k images (optional)
for f in $f1 $f2; do for f in $f1 $f2; do
echo 'Downloading' $url$f ' ...' && curl -L $url$f -o $f && unzip -q $f -d $d && rm $f # download, unzip, remove echo 'Downloading' $url$f '...'
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
done done
wait # finish background tasks

6
data/scripts/get_voc.sh
View File

@ -17,9 +17,11 @@ url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
f1=VOCtrainval_06-Nov-2007.zip # 446MB, 5012 images f1=VOCtrainval_06-Nov-2007.zip # 446MB, 5012 images
f2=VOCtest_06-Nov-2007.zip # 438MB, 4953 images f2=VOCtest_06-Nov-2007.zip # 438MB, 4953 images
f3=VOCtrainval_11-May-2012.zip # 1.95GB, 17126 images f3=VOCtrainval_11-May-2012.zip # 1.95GB, 17126 images
for f in $f1 $f2 $f3; do for f in $f3 $f2 $f1; do
echo 'Downloading' $url$f ' ...' && curl -L $url$f -o $f && unzip -q $f -d $d && rm $f # download, unzip, remove echo 'Downloading' $url$f '...'
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
done done
wait # finish background tasks
end=$(date +%s) end=$(date +%s)
runtime=$((end - start)) runtime=$((end - start))

37
detect.py
View File

@ -9,16 +9,17 @@ from numpy import random
from models.experimental import attempt_load from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, non_max_suppression, apply_classifier, scale_coords, \ from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
xyxy2xywh, strip_optimizer, set_logging, increment_path scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
from utils.plots import plot_one_box from utils.plots import plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized from utils.torch_utils import select_device, load_classifier, time_synchronized
def detect(save_img=False): def detect(save_img=False):
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://')) ('rtsp://', 'rtmp://', 'http://', 'https://'))
# Directories # Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
@ -31,7 +32,8 @@ def detect(save_img=False):
# Load model # Load model
model = attempt_load(weights, map_location=device) # load FP32 model model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half: if half:
model.half() # to FP16 model.half() # to FP16
@ -44,21 +46,20 @@ def detect(save_img=False):
# Set Dataloader # Set Dataloader
vid_path, vid_writer = None, None vid_path, vid_writer = None, None
if webcam: if webcam:
view_img = True view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz) dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else: else:
save_img = True dataset = LoadImages(source, img_size=imgsz, stride=stride)
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors # Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
# Run inference # Run inference
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time() t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset: for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device) img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32 img = img.half() if half else img.float() # uint8 to fp16/32
@ -97,7 +98,7 @@ def detect(save_img=False):
# Print results # Print results
for c in det[:, -1].unique(): for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class n = (det[:, -1] == c).sum() # detections per class
s += f'{n} {names[int(c)]}s, ' # add to string s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results # Write results
for *xyxy, conf, cls in reversed(det): for *xyxy, conf, cls in reversed(det):
@ -117,22 +118,25 @@ def detect(save_img=False):
# Stream results # Stream results
if view_img: if view_img:
cv2.imshow(str(p), im0) cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections) # Save results (image with detections)
if save_img: if save_img:
if dataset.mode == 'image': if dataset.mode == 'image':
cv2.imwrite(save_path, im0) cv2.imwrite(save_path, im0)
else: # 'video' else: # 'video' or 'stream'
if vid_path != save_path: # new video if vid_path != save_path: # new video
vid_path = save_path vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter): if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer vid_writer.release() # release previous video writer
if vid_cap: # video
fourcc = 'mp4v' # output video codec
fps = vid_cap.get(cv2.CAP_PROP_FPS) fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer.write(im0) vid_writer.write(im0)
if save_txt or save_img: if save_txt or save_img:
@ -153,6 +157,7 @@ if __name__ == '__main__':
parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--augment', action='store_true', help='augmented inference')
@ -162,7 +167,7 @@ if __name__ == '__main__':
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
check_requirements() check_requirements(exclude=('pycocotools', 'thop'))
with torch.no_grad(): with torch.no_grad():
if opt.update: # update all models (to fix SourceChangeWarning) if opt.update: # update all models (to fix SourceChangeWarning)

101
hubconf.py
View File

@ -1,8 +1,8 @@
"""File for accessing YOLOv3 via PyTorch Hub https://pytorch.org/hub/ """YOLOv3 PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov3/
Usage: Usage:
import torch import torch
model = torch.hub.load('ultralytics/yolov3', 'yolov3', pretrained=True, channels=3, classes=80) model = torch.hub.load('ultralytics/yolov3', 'yolov3tiny')
""" """
from pathlib import Path from pathlib import Path
@ -10,10 +10,12 @@ from pathlib import Path
import torch import torch
from models.yolo import Model from models.yolo import Model
from utils.general import set_logging from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download from utils.google_utils import attempt_download
from utils.torch_utils import select_device
dependencies = ['torch', 'yaml'] dependencies = ['torch', 'yaml']
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
set_logging() set_logging()
@ -21,7 +23,7 @@ def create(name, pretrained, channels, classes, autoshape):
"""Creates a specified YOLOv3 model """Creates a specified YOLOv3 model
Arguments: Arguments:
name (str): name of model, i.e. 'yolov3_spp' name (str): name of model, i.e. 'yolov3'
pretrained (bool): load pretrained weights into the model pretrained (bool): load pretrained weights into the model
channels (int): number of input channels channels (int): number of input channels
classes (int): number of model classes classes (int): number of model classes
@ -29,21 +31,23 @@ def create(name, pretrained, channels, classes, autoshape):
Returns: Returns:
pytorch model pytorch model
""" """
config = Path(__file__).parent / 'models' / f'{name}.yaml' # model.yaml path
try: try:
model = Model(config, channels, classes) cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
model = Model(cfg, channels, classes)
if pretrained: if pretrained:
fname = f'{name}.pt' # checkpoint filename fname = f'{name}.pt' # checkpoint filename
attempt_download(fname) # download if not found locally attempt_download(fname) # download if not found locally
ckpt = torch.load(fname, map_location=torch.device('cpu')) # load ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
state_dict = ckpt['model'].float().state_dict() # to FP32 msd = model.state_dict() # model state_dict
state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
model.load_state_dict(state_dict, strict=False) # load csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
model.load_state_dict(csd, strict=False) # load
if len(ckpt['model'].names) == classes: if len(ckpt['model'].names) == classes:
model.names = ckpt['model'].names # set class names attribute model.names = ckpt['model'].names # set class names attribute
if autoshape: if autoshape:
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
return model device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
return model.to(device)
except Exception as e: except Exception as e:
help_url = 'https://github.com/ultralytics/yolov5/issues/36' help_url = 'https://github.com/ultralytics/yolov5/issues/36'
@ -51,50 +55,8 @@ def create(name, pretrained, channels, classes, autoshape):
raise Exception(s) from e raise Exception(s) from e
def yolov3(pretrained=False, channels=3, classes=80, autoshape=True):
"""YOLOv3 model from https://github.com/ultralytics/yolov3
Arguments:
pretrained (bool): load pretrained weights into the model, default=False
channels (int): number of input channels, default=3
classes (int): number of model classes, default=80
Returns:
pytorch model
"""
return create('yolov3', pretrained, channels, classes, autoshape)
def yolov3_spp(pretrained=False, channels=3, classes=80, autoshape=True):
"""YOLOv3-SPP model from https://github.com/ultralytics/yolov3
Arguments:
pretrained (bool): load pretrained weights into the model, default=False
channels (int): number of input channels, default=3
classes (int): number of model classes, default=80
Returns:
pytorch model
"""
return create('yolov3-spp', pretrained, channels, classes, autoshape)
def yolov3_tiny(pretrained=False, channels=3, classes=80, autoshape=True):
"""YOLOv3-tiny model from https://github.com/ultralytics/yolov3
Arguments:
pretrained (bool): load pretrained weights into the model, default=False
channels (int): number of input channels, default=3
classes (int): number of model classes, default=80
Returns:
pytorch model
"""
return create('yolov3-tiny', pretrained, channels, classes, autoshape)
def custom(path_or_model='path/to/model.pt', autoshape=True): def custom(path_or_model='path/to/model.pt', autoshape=True):
"""YOLOv3-custom model from https://github.com/ultralytics/yolov3 """YOLOv3-custom model https://github.com/ultralytics/yolov3
Arguments (3 options): Arguments (3 options):
path_or_model (str): 'path/to/model.pt' path_or_model (str): 'path/to/model.pt'
@ -106,12 +68,30 @@ def custom(path_or_model='path/to/model.pt', autoshape=True):
""" """
model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
if isinstance(model, dict): if isinstance(model, dict):
model = model['model'] # load model model = model['ema' if model.get('ema') else 'model'] # load model
hub_model = Model(model.yaml).to(next(model.parameters()).device) # create hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
hub_model.load_state_dict(model.float().state_dict()) # load state_dict hub_model.load_state_dict(model.float().state_dict()) # load state_dict
hub_model.names = model.names # class names hub_model.names = model.names # class names
return hub_model.autoshape() if autoshape else hub_model if autoshape:
hub_model = hub_model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
return hub_model.to(device)
def yolov3(pretrained=True, channels=3, classes=80, autoshape=True):
# YOLOv3 model https://github.com/ultralytics/yolov3
return create('yolov3', pretrained, channels, classes, autoshape)
def yolov3_spp(pretrained=True, channels=3, classes=80, autoshape=True):
# YOLOv3-SPP model https://github.com/ultralytics/yolov3
return create('yolov3-spp', pretrained, channels, classes, autoshape)
def yolov3_tiny(pretrained=True, channels=3, classes=80, autoshape=True):
# YOLOv3-tiny model https://github.com/ultralytics/yolov3
return create('yolov3-tiny', pretrained, channels, classes, autoshape)
if __name__ == '__main__': if __name__ == '__main__':
@ -119,9 +99,14 @@ if __name__ == '__main__':
# model = custom(path_or_model='path/to/model.pt') # custom example # model = custom(path_or_model='path/to/model.pt') # custom example
# Verify inference # Verify inference
import numpy as np
from PIL import Image from PIL import Image
imgs = [Image.open(x) for x in Path('data/images').glob('*.jpg')] imgs = [Image.open('data/images/bus.jpg'), # PIL
results = model(imgs) 'data/images/zidane.jpg', # filename
results.show() 'https://github.com/ultralytics/yolov3/raw/master/data/images/bus.jpg', # URI
np.zeros((640, 480, 3))] # numpy
results = model(imgs) # batched inference
results.print() results.print()
results.save()

153
models/common.py
View File

@ -1,16 +1,21 @@
# This file contains modules common to various models # YOLOv3 common modules
import math import math
from copy import copy
from pathlib import Path
import numpy as np import numpy as np
import pandas as pd
import requests import requests
import torch import torch
import torch.nn as nn import torch.nn as nn
from PIL import Image, ImageDraw from PIL import Image
from torch.cuda import amp
from utils.datasets import letterbox from utils.datasets import letterbox
from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh
from utils.plots import color_list from utils.plots import color_list, plot_one_box
from utils.torch_utils import time_synchronized
def autopad(k, p=None): # kernel, padding def autopad(k, p=None): # kernel, padding
@ -40,6 +45,52 @@ class Conv(nn.Module):
return self.act(self.conv(x)) return self.act(self.conv(x))
class TransformerLayer(nn.Module):
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
def __init__(self, c, num_heads):
super().__init__()
self.q = nn.Linear(c, c, bias=False)
self.k = nn.Linear(c, c, bias=False)
self.v = nn.Linear(c, c, bias=False)
self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
self.fc1 = nn.Linear(c, c, bias=False)
self.fc2 = nn.Linear(c, c, bias=False)
def forward(self, x):
x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
x = self.fc2(self.fc1(x)) + x
return x
class TransformerBlock(nn.Module):
# Vision Transformer https://arxiv.org/abs/2010.11929
def __init__(self, c1, c2, num_heads, num_layers):
super().__init__()
self.conv = None
if c1 != c2:
self.conv = Conv(c1, c2)
self.linear = nn.Linear(c2, c2) # learnable position embedding
self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)])
self.c2 = c2
def forward(self, x):
if self.conv is not None:
x = self.conv(x)
b, _, w, h = x.shape
p = x.flatten(2)
p = p.unsqueeze(0)
p = p.transpose(0, 3)
p = p.squeeze(3)
e = self.linear(p)
x = p + e
x = self.tr(x)
x = x.unsqueeze(3)
x = x.transpose(0, 3)
x = x.reshape(b, self.c2, w, h)
return x
class Bottleneck(nn.Module): class Bottleneck(nn.Module):
# Standard bottleneck # Standard bottleneck
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
@ -87,6 +138,14 @@ class C3(nn.Module):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
class C3TR(C3):
# C3 module with TransformerBlock()
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e)
self.m = TransformerBlock(c_, c_, 4, n)
class SPP(nn.Module): class SPP(nn.Module):
# Spatial pyramid pooling layer used in YOLOv3-SPP # Spatial pyramid pooling layer used in YOLOv3-SPP
def __init__(self, c1, c2, k=(5, 9, 13)): def __init__(self, c1, c2, k=(5, 9, 13)):
@ -166,7 +225,6 @@ class NMS(nn.Module):
class autoShape(nn.Module): class autoShape(nn.Module):
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
img_size = 640 # inference size (pixels)
conf = 0.25 # NMS confidence threshold conf = 0.25 # NMS confidence threshold
iou = 0.45 # NMS IoU threshold iou = 0.45 # NMS IoU threshold
classes = None # (optional list) filter by class classes = None # (optional list) filter by class
@ -179,27 +237,33 @@ class autoShape(nn.Module):
print('autoShape already enabled, skipping... ') # model already converted to model.autoshape() print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()
return self return self
@torch.no_grad()
def forward(self, imgs, size=640, augment=False, profile=False): def forward(self, imgs, size=640, augment=False, profile=False):
# Inference from various sources. For height=720, width=1280, RGB images example inputs are: # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
# filename: imgs = 'data/samples/zidane.jpg' # filename: imgs = 'data/samples/zidane.jpg'
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg' # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3) # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
# PIL: = Image.open('image.jpg') # HWC x(720,1280,3) # PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
# numpy: = np.zeros((720,1280,3)) # HWC # numpy: = np.zeros((640,1280,3)) # HWC
# torch: = torch.zeros(16,3,720,1280) # BCHW # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
t = [time_synchronized()]
p = next(self.model.parameters()) # for device and type p = next(self.model.parameters()) # for device and type
if isinstance(imgs, torch.Tensor): # torch if isinstance(imgs, torch.Tensor): # torch
with amp.autocast(enabled=p.device.type != 'cpu'):
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
# Pre-process # Pre-process
n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
shape0, shape1 = [], [] # image and inference shapes shape0, shape1, files = [], [], [] # image and inference shapes, filenames
for i, im in enumerate(imgs): for i, im in enumerate(imgs):
f = f'image{i}' # filename
if isinstance(im, str): # filename or uri if isinstance(im, str): # filename or uri
im = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im) # open im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
im = np.array(im) # to numpy elif isinstance(im, Image.Image): # PIL Image
im, f = np.asarray(im), getattr(im, 'filename', f) or f
files.append(Path(f).with_suffix('.jpg').name)
if im.shape[0] < 5: # image in CHW if im.shape[0] < 5: # image in CHW
im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
@ -213,82 +277,101 @@ class autoShape(nn.Module):
x = np.stack(x, 0) if n > 1 else x[0][None] # stack x = np.stack(x, 0) if n > 1 else x[0][None] # stack
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32 x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
t.append(time_synchronized())
with amp.autocast(enabled=p.device.type != 'cpu'):
# Inference # Inference
with torch.no_grad():
y = self.model(x, augment, profile)[0] # forward y = self.model(x, augment, profile)[0] # forward
y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS t.append(time_synchronized())
# Post-process # Post-process
y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
for i in range(n): for i in range(n):
scale_coords(shape1, y[i][:, :4], shape0[i]) scale_coords(shape1, y[i][:, :4], shape0[i])
return Detections(imgs, y, self.names) t.append(time_synchronized())
return Detections(imgs, y, files, t, self.names, x.shape)
class Detections: class Detections:
# detections class for YOLOv5 inference results # detections class for YOLOv3 inference results
def __init__(self, imgs, pred, names=None): def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
super(Detections, self).__init__() super(Detections, self).__init__()
d = pred[0].device # device d = pred[0].device # device
gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
self.imgs = imgs # list of images as numpy arrays self.imgs = imgs # list of images as numpy arrays
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
self.names = names # class names self.names = names # class names
self.files = files # image filenames
self.xyxy = pred # xyxy pixels self.xyxy = pred # xyxy pixels
self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
self.n = len(self.pred) self.n = len(self.pred) # number of images (batch size)
self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
self.s = shape # inference BCHW shape
def display(self, pprint=False, show=False, save=False, render=False): def display(self, pprint=False, show=False, save=False, render=False, save_dir=''):
colors = color_list() colors = color_list()
for i, (img, pred) in enumerate(zip(self.imgs, self.pred)): for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} ' str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
if pred is not None: if pred is not None:
for c in pred[:, -1].unique(): for c in pred[:, -1].unique():
n = (pred[:, -1] == c).sum() # detections per class n = (pred[:, -1] == c).sum() # detections per class
str += f'{n} {self.names[int(c)]}s, ' # add to string str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
if show or save or render: if show or save or render:
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
for *box, conf, cls in pred: # xyxy, confidence, class for *box, conf, cls in pred: # xyxy, confidence, class
# str += '%s %.2f, ' % (names[int(cls)], conf) # label label = f'{self.names[int(cls)]} {conf:.2f}'
ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
if pprint: if pprint:
print(str) print(str.rstrip(', '))
if show: if show:
img.show(f'Image {i}') # show img.show(self.files[i]) # show
if save: if save:
f = f'results{i}.jpg' f = self.files[i]
str += f"saved to '{f}'" img.save(Path(save_dir) / f) # save
img.save(f) # save print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
if render: if render:
self.imgs[i] = np.asarray(img) self.imgs[i] = np.asarray(img)
def print(self): def print(self):
self.display(pprint=True) # print results self.display(pprint=True) # print results
print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)
def show(self): def show(self):
self.display(show=True) # show results self.display(show=True) # show results
def save(self): def save(self, save_dir='runs/hub/exp'):
self.display(save=True) # save results save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp') # increment save_dir
Path(save_dir).mkdir(parents=True, exist_ok=True)
self.display(save=True, save_dir=save_dir) # save results
def render(self): def render(self):
self.display(render=True) # render results self.display(render=True) # render results
return self.imgs return self.imgs
def __len__(self): def pandas(self):
return self.n # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
new = copy(self) # return copy
ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
return new
def tolist(self): def tolist(self):
# return a list of Detections objects, i.e. 'for result in results.tolist():' # return a list of Detections objects, i.e. 'for result in results.tolist():'
x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)] x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)]
for d in x: for d in x:
for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
setattr(d, k, getattr(d, k)[0]) # pop out of list setattr(d, k, getattr(d, k)[0]) # pop out of list
return x return x
def __len__(self):
return self.n
class Classify(nn.Module): class Classify(nn.Module):
# Classification head, i.e. x(b,c1,20,20) to x(b,c2) # Classification head, i.e. x(b,c1,20,20) to x(b,c2)

9
models/experimental.py
View File

@ -1,4 +1,4 @@
# This file contains experimental modules # YOLOv3 experimental modules
import numpy as np import numpy as np
import torch import torch
@ -58,7 +58,7 @@ class GhostConv(nn.Module):
class GhostBottleneck(nn.Module): class GhostBottleneck(nn.Module):
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
def __init__(self, c1, c2, k, s): def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
super(GhostBottleneck, self).__init__() super(GhostBottleneck, self).__init__()
c_ = c2 // 2 c_ = c2 // 2
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
@ -115,11 +115,12 @@ def attempt_load(weights, map_location=None):
model = Ensemble() model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]: for w in weights if isinstance(weights, list) else [weights]:
attempt_download(w) attempt_download(w)
model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model ckpt = torch.load(w, map_location=map_location) # load
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
# Compatibility updates # Compatibility updates
for m in model.modules(): for m in model.modules():
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
m.inplace = True # pytorch 1.7.0 compatibility m.inplace = True # pytorch 1.7.0 compatibility
elif type(m) is Conv: elif type(m) is Conv:
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility

19
models/export.py
View File

@ -1,4 +1,4 @@
"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats """Exports a YOLOv3 *.pt model to ONNX and TorchScript formats
Usage: Usage:
$ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov3.pt --img 640 --batch 1 $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov3.pt --img 640 --batch 1
@ -17,12 +17,16 @@ import models
from models.experimental import attempt_load from models.experimental import attempt_load
from utils.activations import Hardswish, SiLU from utils.activations import Hardswish, SiLU
from utils.general import set_logging, check_img_size from utils.general import set_logging, check_img_size
from utils.torch_utils import select_device
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path') # from yolov3/models/ parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path') # from yolov3/models/
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args() opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
print(opt) print(opt)
@ -30,7 +34,8 @@ if __name__ == '__main__':
t = time.time() t = time.time()
# Load PyTorch model # Load PyTorch model
model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model device = select_device(opt.device)
model = attempt_load(opt.weights, map_location=device) # load FP32 model
labels = model.names labels = model.names
# Checks # Checks
@ -38,7 +43,7 @@ if __name__ == '__main__':
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
# Input # Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size) # image size(1,3,320,192) iDetection img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
# Update model # Update model
for k, m in model.named_modules(): for k, m in model.named_modules():
@ -50,14 +55,14 @@ if __name__ == '__main__':
m.act = SiLU() m.act = SiLU()
# elif isinstance(m, models.yolo.Detect): # elif isinstance(m, models.yolo.Detect):
# m.forward = m.forward_export # assign forward (optional) # m.forward = m.forward_export # assign forward (optional)
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = not opt.grid # set Detect() layer grid export
y = model(img) # dry run y = model(img) # dry run
# TorchScript export # TorchScript export
try: try:
print('\nStarting TorchScript export with torch %s...' % torch.__version__) print('\nStarting TorchScript export with torch %s...' % torch.__version__)
f = opt.weights.replace('.pt', '.torchscript.pt') # filename f = opt.weights.replace('.pt', '.torchscript.pt') # filename
ts = torch.jit.trace(model, img) ts = torch.jit.trace(model, img, strict=False)
ts.save(f) ts.save(f)
print('TorchScript export success, saved as %s' % f) print('TorchScript export success, saved as %s' % f)
except Exception as e: except Exception as e:
@ -70,7 +75,9 @@ if __name__ == '__main__':
print('\nStarting ONNX export with onnx %s...' % onnx.__version__) print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
f = opt.weights.replace('.pt', '.onnx') # filename f = opt.weights.replace('.pt', '.onnx') # filename
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
output_names=['classes', 'boxes'] if y is None else ['output']) output_names=['classes', 'boxes'] if y is None else ['output'],
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
# Checks # Checks
onnx_model = onnx.load(f) # load onnx model onnx_model = onnx.load(f) # load onnx model

57
models/yolo.py
View File

@ -1,14 +1,15 @@
# YOLOv3 YOLO-specific modules
import argparse import argparse
import logging import logging
import sys import sys
from copy import deepcopy from copy import deepcopy
from pathlib import Path
sys.path.append('./') # to run '$ python *.py' files in subdirectories sys.path.append('./') # to run '$ python *.py' files in subdirectories
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
from models.common import * from models.common import *
from models.experimental import MixConv2d, CrossConv from models.experimental import *
from utils.autoanchor import check_anchor_order from utils.autoanchor import check_anchor_order
from utils.general import make_divisible, check_file, set_logging from utils.general import make_divisible, check_file, set_logging
from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \ from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
@ -50,7 +51,7 @@ class Detect(nn.Module):
self.grid[i] = self._make_grid(nx, ny).to(x[i].device) self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
y = x[i].sigmoid() y = x[i].sigmoid()
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
z.append(y.view(bs, -1, self.no)) z.append(y.view(bs, -1, self.no))
@ -63,7 +64,7 @@ class Detect(nn.Module):
class Model(nn.Module): class Model(nn.Module):
def __init__(self, cfg='yolov3.yaml', ch=3, nc=None): # model, input channels, number of classes def __init__(self, cfg='yolov3.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
super(Model, self).__init__() super(Model, self).__init__()
if isinstance(cfg, dict): if isinstance(cfg, dict):
self.yaml = cfg # model dict self.yaml = cfg # model dict
@ -71,13 +72,16 @@ class Model(nn.Module):
import yaml # for torch hub import yaml # for torch hub
self.yaml_file = Path(cfg).name self.yaml_file = Path(cfg).name
with open(cfg) as f: with open(cfg) as f:
self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict
# Define model # Define model
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
if nc and nc != self.yaml['nc']: if nc and nc != self.yaml['nc']:
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc)) logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
self.yaml['nc'] = nc # override yaml value self.yaml['nc'] = nc # override yaml value
if anchors:
logger.info(f'Overriding model.yaml anchors with anchors={anchors}')
self.yaml['anchors'] = round(anchors) # override yaml value
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
self.names = [str(i) for i in range(self.yaml['nc'])] # default names self.names = [str(i) for i in range(self.yaml['nc'])] # default names
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
@ -107,7 +111,7 @@ class Model(nn.Module):
for si, fi in zip(s, f): for si, fi in zip(s, f):
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
yi = self.forward_once(xi)[0] # forward yi = self.forward_once(xi)[0] # forward
# cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
yi[..., :4] /= si # de-scale yi[..., :4] /= si # de-scale
if fi == 2: if fi == 2:
yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud
@ -210,45 +214,30 @@ def parse_model(d, ch): # model_dict, input_channels(3)
pass pass
n = max(round(n * gd), 1) if n > 1 else n # depth gain n = max(round(n * gd), 1) if n > 1 else n # depth gain
if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
C3, C3TR]:
c1, c2 = ch[f], args[0] c1, c2 = ch[f], args[0]
if c2 != no: # if not output
# Normal c2 = make_divisible(c2 * gw, 8)
# if i > 0 and args[0] != no: # channel expansion factor
# ex = 1.75 # exponential (default 2.0)
# e = math.log(c2 / ch[1]) / math.log(2)
# c2 = int(ch[1] * ex ** e)
# if m != Focus:
c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
# Experimental
# if i > 0 and args[0] != no: # channel expansion factor
# ex = 1 + gw # exponential (default 2.0)
# ch1 = 32 # ch[1]
# e = math.log(c2 / ch1) / math.log(2) # level 1-n
# c2 = int(ch1 * ex ** e)
# if m != Focus:
# c2 = make_divisible(c2, 8) if c2 != no else c2
args = [c1, c2, *args[1:]] args = [c1, c2, *args[1:]]
if m in [BottleneckCSP, C3]: if m in [BottleneckCSP, C3, C3TR]:
args.insert(2, n) args.insert(2, n) # number of repeats
n = 1 n = 1
elif m is nn.BatchNorm2d: elif m is nn.BatchNorm2d:
args = [ch[f]] args = [ch[f]]
elif m is Concat: elif m is Concat:
c2 = sum([ch[x if x < 0 else x + 1] for x in f]) c2 = sum([ch[x] for x in f])
elif m is Detect: elif m is Detect:
args.append([ch[x + 1] for x in f]) args.append([ch[x] for x in f])
if isinstance(args[1], int): # number of anchors if isinstance(args[1], int): # number of anchors
args[1] = [list(range(args[1] * 2))] * len(f) args[1] = [list(range(args[1] * 2))] * len(f)
elif m is Contract: elif m is Contract:
c2 = ch[f if f < 0 else f + 1] * args[0] ** 2 c2 = ch[f] * args[0] ** 2
elif m is Expand: elif m is Expand:
c2 = ch[f if f < 0 else f + 1] // args[0] ** 2 c2 = ch[f] // args[0] ** 2
else: else:
c2 = ch[f if f < 0 else f + 1] c2 = ch[f]
m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module
t = str(m)[8:-2].replace('__main__.', '') # module type t = str(m)[8:-2].replace('__main__.', '') # module type
@ -257,6 +246,8 @@ def parse_model(d, ch): # model_dict, input_channels(3)
logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
layers.append(m_) layers.append(m_)
if i == 0:
ch = []
ch.append(c2) ch.append(c2)
return nn.Sequential(*layers), sorted(save) return nn.Sequential(*layers), sorted(save)

9
requirements.txt
View File

@ -1,19 +1,18 @@
# pip install -r requirements.txt # pip install -r requirements.txt
# base ---------------------------------------- # base ----------------------------------------
Cython
matplotlib>=3.2.2 matplotlib>=3.2.2
numpy>=1.18.5 numpy>=1.18.5
opencv-python>=4.1.2 opencv-python>=4.1.2
Pillow Pillow
PyYAML>=5.3 PyYAML>=5.3.1
scipy>=1.4.1 scipy>=1.4.1
tensorboard>=2.2
torch>=1.7.0 torch>=1.7.0
torchvision>=0.8.1 torchvision>=0.8.1
tqdm>=4.41.0 tqdm>=4.41.0
# logging ------------------------------------- # logging -------------------------------------
tensorboard>=2.4.1
# wandb # wandb
# plotting ------------------------------------ # plotting ------------------------------------
@ -21,8 +20,8 @@ seaborn>=0.11.0
pandas pandas
# export -------------------------------------- # export --------------------------------------
# coremltools==4.0 # coremltools>=4.1
# onnx>=1.8.0 # onnx>=1.8.1
# scikit-learn==0.19.2 # for coreml quantization # scikit-learn==0.19.2 # for coreml quantization
# extras -------------------------------------- # extras --------------------------------------

95
test.py
View File

@ -13,7 +13,6 @@ from models.experimental import attempt_load
from utils.datasets import create_dataloader from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \ from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
from utils.loss import compute_loss
from utils.metrics import ap_per_class, ConfusionMatrix from utils.metrics import ap_per_class, ConfusionMatrix
from utils.plots import plot_images, output_to_target, plot_study_txt from utils.plots import plot_images, output_to_target, plot_study_txt
from utils.torch_utils import select_device, time_synchronized from utils.torch_utils import select_device, time_synchronized
@ -36,8 +35,10 @@ def test(data,
save_hybrid=False, # for hybrid auto-labelling save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences save_conf=False, # save auto-label confidences
plots=True, plots=True,
log_imgs=0): # number of logged images wandb_logger=None,
compute_loss=None,
half_precision=True,
is_coco=False):
# Initialize/load model and set device # Initialize/load model and set device
training = model is not None training = model is not None
if training: # called by train.py if training: # called by train.py
@ -53,47 +54,46 @@ def test(data,
# Load model # Load model
model = attempt_load(weights, map_location=device) # load FP32 model model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(imgsz, s=gs) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99 # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1: # if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model) # model = nn.DataParallel(model)
# Half # Half
half = device.type != 'cpu' # half precision only supported on CUDA half = device.type != 'cpu' and half_precision # half precision only supported on CUDA
if half: if half:
model.half() model.half()
# Configure # Configure
model.eval() model.eval()
is_coco = data.endswith('coco.yaml') # is COCO dataset if isinstance(data, str):
is_coco = data.endswith('coco.yaml')
with open(data) as f: with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict data = yaml.load(f, Loader=yaml.SafeLoader)
check_dataset(data) # check check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
niou = iouv.numel() niou = iouv.numel()
# Logging # Logging
log_imgs, wandb = min(log_imgs, 100), None # ceil
try:
import wandb # Weights & Biases
except ImportError:
log_imgs = 0 log_imgs = 0
if wandb_logger and wandb_logger.wandb:
log_imgs = min(wandb_logger.log_imgs, 100)
# Dataloader # Dataloader
if not training: if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img if device.type != 'cpu':
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True, dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,
prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0] prefix=colorstr(f'{task}: '))[0]
seen = 0 seen = 0
confusion_matrix = ConfusionMatrix(nc=nc) confusion_matrix = ConfusionMatrix(nc=nc)
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)} names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
coco91class = coco80_to_coco91_class() coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95') s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0. p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device) loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], [] jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
@ -107,22 +107,22 @@ def test(data,
with torch.no_grad(): with torch.no_grad():
# Run model # Run model
t = time_synchronized() t = time_synchronized()
inf_out, train_out = model(img, augment=augment) # inference and training outputs out, train_out = model(img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t t0 += time_synchronized() - t
# Compute loss # Compute loss
if training: if compute_loss:
loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # box, obj, cls loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
# Run NMS # Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized() t = time_synchronized()
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb) out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)
t1 += time_synchronized() - t t1 += time_synchronized() - t
# Statistics per image # Statistics per image
for si, pred in enumerate(output): for si, pred in enumerate(out):
labels = targets[targets[:, 0] == si, 1:] labels = targets[targets[:, 0] == si, 1:]
nl = len(labels) nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class tcls = labels[:, 0].tolist() if nl else [] # target class
@ -147,15 +147,17 @@ def test(data,
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f: with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n') f.write(('%g ' * len(line)).rstrip() % line + '\n')
# W&B logging # W&B logging - Media Panel Plots
if plots and len(wandb_images) < log_imgs: if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0: # Check for test operation
if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": int(cls), "class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf), "box_caption": "%s %.3f" % (names[cls], conf),
"scores": {"class_score": conf}, "scores": {"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()] "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name)) wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))
wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None
# Append to pycocotools JSON dictionary # Append to pycocotools JSON dictionary
if save_json: if save_json:
@ -179,7 +181,7 @@ def test(data,
tbox = xywh2xyxy(labels[:, 1:5]) tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
if plots: if plots:
confusion_matrix.process_batch(pred, torch.cat((labels[:, 0:1], tbox), 1)) confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class # Per target class
for cls in torch.unique(tcls_tensor): for cls in torch.unique(tcls_tensor):
@ -210,24 +212,24 @@ def test(data,
f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start() Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images, args=(img, output_to_target(output), paths, f, names), daemon=True).start() Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start()
# Compute statistics # Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any(): if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names) p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95] ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
else: else:
nt = torch.zeros(1) nt = torch.zeros(1)
# Print results # Print results
pf = '%20s' + '%12.3g' * 6 # print format pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class # Print results per class
if (verbose or (nc <= 20 and not training)) and nc > 1 and len(stats): if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class): for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
@ -239,9 +241,11 @@ def test(data,
# Plots # Plots
if plots: if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values())) confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb and wandb.run: if wandb_logger and wandb_logger.wandb:
wandb.log({"Images": wandb_images}) val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]}) wandb_logger.log({"Validation": val_batches})
if wandb_images:
wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
# Save JSON # Save JSON
if save_json and len(jdict): if save_json and len(jdict):
@ -269,10 +273,10 @@ def test(data,
print(f'pycocotools unable to run: {e}') print(f'pycocotools unable to run: {e}')
# Return results # Return results
model.float() # for training
if not training: if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}") print(f"Results saved to {save_dir}{s}")
model.float() # for training
maps = np.zeros(nc) + map maps = np.zeros(nc) + map
for i, c in enumerate(ap_class): for i, c in enumerate(ap_class):
maps[c] = ap[i] maps[c] = ap[i]
@ -287,7 +291,7 @@ if __name__ == '__main__':
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS') parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
parser.add_argument('--task', default='val', help="'val', 'test', 'study'") parser.add_argument('--task', default='val', help='train, val, test, speed or study')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--augment', action='store_true', help='augmented inference')
@ -305,7 +309,7 @@ if __name__ == '__main__':
print(opt) print(opt)
check_requirements() check_requirements()
if opt.task in ['val', 'test']: # run normally if opt.task in ('train', 'val', 'test'): # run normally
test(opt.data, test(opt.data,
opt.weights, opt.weights,
opt.batch_size, opt.batch_size,
@ -321,16 +325,21 @@ if __name__ == '__main__':
save_conf=opt.save_conf, save_conf=opt.save_conf,
) )
elif opt.task == 'speed': # speed benchmarks
for w in opt.weights:
test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False)
elif opt.task == 'study': # run over a range of settings and save/plot elif opt.task == 'study': # run over a range of settings and save/plot
for weights in ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']: # python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt
f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to x = list(range(256, 1536 + 128, 128)) # x axis (image sizes)
x = list(range(320, 800, 64)) # x axis for w in opt.weights:
f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt' # filename to save to
y = [] # y axis y = [] # y axis
for i in x: # img-size for i in x: # img-size
print('\nRunning %s point %s...' % (f, i)) print(f'\nRunning {f} point {i}...')
r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json, r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
plots=False) plots=False)
y.append(r + t) # results and times y.append(r + t) # results and times
np.savetxt(f, y, fmt='%10.4g') # save np.savetxt(f, y, fmt='%10.4g') # save
os.system('zip -r study.zip study_*.txt') os.system('zip -r study.zip study_*.txt')
plot_study_txt(f, x) # plot plot_study_txt(x=x) # plot

217
train.py
View File

@ -4,6 +4,7 @@ import math
import os import os
import random import random
import time import time
from copy import deepcopy
from pathlib import Path from pathlib import Path
from threading import Thread from threading import Thread
@ -29,14 +30,15 @@ from utils.general import labels_to_class_weights, increment_path, labels_to_ima
fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \ fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
check_requirements, print_mutation, set_logging, one_cycle, colorstr check_requirements, print_mutation, set_logging, one_cycle, colorstr
from utils.google_utils import attempt_download from utils.google_utils import attempt_download
from utils.loss import compute_loss from utils.loss import ComputeLoss
from utils.plots import plot_images, plot_labels, plot_results, plot_evolution from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel
from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def train(hyp, opt, device, tb_writer=None, wandb=None): def train(hyp, opt, device, tb_writer=None):
logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
save_dir, epochs, batch_size, total_batch_size, weights, rank = \ save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
@ -60,10 +62,19 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
init_seeds(2 + rank) init_seeds(2 + rank)
with open(opt.data) as f: with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
with torch_distributed_zero_first(rank): is_coco = opt.data.endswith('coco.yaml')
check_dataset(data_dict) # check
train_path = data_dict['train'] # Logging- Doing this before checking the dataset. Might update data_dict
test_path = data_dict['val'] loggers = {'wandb': None} # loggers dict
if rank in [-1, 0]:
opt.hyp = hyp # add hyperparameters
run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
loggers['wandb'] = wandb_logger.wandb
data_dict = wandb_logger.data_dict
if wandb_logger.wandb:
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
@ -74,16 +85,18 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
with torch_distributed_zero_first(rank): with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'): model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32 state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load model.load_state_dict(state_dict, strict=False) # load
logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report
else: else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
# Freeze # Freeze
freeze = [] # parameter names to freeze (full or partial) freeze = [] # parameter names to freeze (full or partial)
@ -120,18 +133,15 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Scheduler https://arxiv.org/pdf/1812.01187.pdf # Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs) # plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging # EMA
if rank in [-1, 0] and wandb and wandb.run is None: ema = ModelEMA(model) if rank in [-1, 0] else None
opt.hyp = hyp # add hyperparameters
wandb_run = wandb.init(config=opt, resume="allow",
project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem,
name=save_dir.stem,
id=ckpt.get('wandb_id') if 'ckpt' in locals() else None)
loggers = {'wandb': wandb} # loggers dict
# Resume # Resume
start_epoch, best_fitness = 0, 0.0 start_epoch, best_fitness = 0, 0.0
@ -141,10 +151,14 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
optimizer.load_state_dict(ckpt['optimizer']) optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness'] best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Results # Results
if ckpt.get('training_results') is not None: if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file: results_file.write_text(ckpt['training_results']) # write results.txt
file.write(ckpt['training_results']) # write results.txt
# Epochs # Epochs
start_epoch = ckpt['epoch'] + 1 start_epoch = ckpt['epoch'] + 1
@ -158,7 +172,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
del ckpt, state_dict del ckpt, state_dict
# Image sizes # Image sizes
gs = int(model.stride.max()) # grid size (max stride) gs = max(int(model.stride.max()), 32) # grid size (max stride)
nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj']) nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj'])
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
@ -171,13 +185,6 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()') logger.info('Using SyncBatchNorm()')
# EMA
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)
# Trainloader # Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank, hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
@ -189,8 +196,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Process 0 # Process 0
if rank in [-1, 0]: if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader
testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, # testloader
hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1, hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
world_size=opt.world_size, workers=opt.workers, world_size=opt.world_size, workers=opt.workers,
pad=0.5, prefix=colorstr('val: '))[0] pad=0.5, prefix=colorstr('val: '))[0]
@ -201,18 +207,26 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device)) # model._initialize_biases(cf.to(device))
if plots: if plots:
plot_labels(labels, save_dir, loggers) plot_labels(labels, names, save_dir, loggers)
if tb_writer: if tb_writer:
tb_writer.add_histogram('classes', c, 0) tb_writer.add_histogram('classes', c, 0)
# Anchors # Anchors
if not opt.noautoanchor: if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
# DDP mode
if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
# Model parameters # Model parameters
hyp['box'] *= 3. / nl # scale to layers hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou) model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
@ -227,6 +241,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda) scaler = amp.GradScaler(enabled=cuda)
compute_loss = ComputeLoss(model) # init loss class
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n' logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
f'Using {dataloader.num_workers} dataloader workers\n' f'Using {dataloader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n' f'Logging results to {save_dir}\n'
@ -256,7 +271,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
if rank != -1: if rank != -1:
dataloader.sampler.set_epoch(epoch) dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader) pbar = enumerate(dataloader)
logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'targets', 'img_size')) logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
if rank in [-1, 0]: if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad() optimizer.zero_grad()
@ -286,7 +301,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Forward # Forward
with amp.autocast(enabled=cuda): with amp.autocast(enabled=cuda):
pred = model(imgs) # forward pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device), model) # loss scaled by batch_size loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if rank != -1: if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode loss *= opt.world_size # gradient averaged between devices in DDP mode
if opt.quad: if opt.quad:
@ -317,9 +332,10 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
# if tb_writer: # if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard # tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
elif plots and ni == 3 and wandb: elif plots and ni == 10 and wandb_logger.wandb:
wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')]}) wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
save_dir.glob('train*.jpg') if x.exists()]})
# end batch ------------------------------------------------------------------------------------------------ # end batch ------------------------------------------------------------------------------------------------
# end epoch ---------------------------------------------------------------------------------------------------- # end epoch ----------------------------------------------------------------------------------------------------
@ -331,23 +347,26 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# DDP process 0 or single-GPU # DDP process 0 or single-GPU
if rank in [-1, 0]: if rank in [-1, 0]:
# mAP # mAP
if ema:
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights']) ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
final_epoch = epoch + 1 == epochs final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(opt.data, wandb_logger.current_epoch = epoch + 1
batch_size=total_batch_size, results, maps, times = test.test(data_dict,
batch_size=batch_size * 2,
imgsz=imgsz_test, imgsz=imgsz_test,
model=ema.ema, model=ema.ema,
single_cls=opt.single_cls, single_cls=opt.single_cls,
dataloader=testloader, dataloader=testloader,
save_dir=save_dir, save_dir=save_dir,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch, plots=plots and final_epoch,
log_imgs=opt.log_imgs if wandb else 0) wandb_logger=wandb_logger,
compute_loss=compute_loss,
is_coco=is_coco)
# Write # Write
with open(results_file, 'a') as f: with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results + '\n') # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
if len(opt.name) and opt.bucket: if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
@ -359,72 +378,77 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags): for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer: if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb: if wandb_logger.wandb:
wandb.log({tag: x}) # W&B wandb_logger.log({tag: x}) # W&B
# Update best mAP # Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
if fi > best_fitness: if fi > best_fitness:
best_fitness = fi best_fitness = fi
wandb_logger.end_epoch(best_result=best_fitness == fi)
# Save model # Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve) if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {'epoch': epoch, ckpt = {'epoch': epoch,
'best_fitness': best_fitness, 'best_fitness': best_fitness,
'training_results': f.read(), 'training_results': results_file.read_text(),
'model': ema.ema, 'model': deepcopy(model.module if is_parallel(model) else model).half(),
'optimizer': None if final_epoch else optimizer.state_dict(), 'ema': deepcopy(ema.ema).half(),
'wandb_id': wandb_run.id if wandb else None} 'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None}
# Save last, best and delete # Save last, best and delete
torch.save(ckpt, last) torch.save(ckpt, last)
if best_fitness == fi: if best_fitness == fi:
torch.save(ckpt, best) torch.save(ckpt, best)
if wandb_logger.wandb:
if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
wandb_logger.log_model(
last.parent, opt, epoch, fi, best_model=best_fitness == fi)
del ckpt del ckpt
# end epoch ---------------------------------------------------------------------------------------------------- # end epoch ----------------------------------------------------------------------------------------------------
# end training # end training
if rank in [-1, 0]: if rank in [-1, 0]:
# Plots
if plots:
plot_results(save_dir=save_dir) # save as results.png
if wandb_logger.wandb:
files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
if (save_dir / f).exists()]})
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for m in (last, best) if best.exists() else (last): # speed, mAP tests
results, _, _ = test.test(opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
conf_thres=0.001,
iou_thres=0.7,
model=attempt_load(m, device).half(),
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=True,
plots=False,
is_coco=is_coco)
# Strip optimizers # Strip optimizers
final = best if best.exists() else last # final model final = best if best.exists() else last # final model
for f in [last, best]: for f in last, best:
if f.exists(): if f.exists():
strip_optimizer(f) # strip optimizers strip_optimizer(f) # strip optimizers
if opt.bucket: if opt.bucket:
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
if wandb_logger.wandb and not opt.evolve: # Log the stripped model
# Plots wandb_logger.wandb.log_artifact(str(final), type='model',
if plots: name='run_' + wandb_logger.wandb_run.id + '_model',
plot_results(save_dir=save_dir) # save as results.png aliases=['last', 'best', 'stripped'])
if wandb: wandb_logger.finish_run()
files = ['results.png', 'precision_recall_curve.png', 'confusion_matrix.png']
wandb.log({"Results": [wandb.Image(str(save_dir / f), caption=f) for f in files
if (save_dir / f).exists()]})
if opt.log_artifacts:
wandb.log_artifact(artifact_or_path=str(final), type='model', name=save_dir.stem)
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for conf, iou, save_json in ([0.25, 0.45, False], [0.001, 0.65, True]): # speed, mAP tests
results, _, _ = test.test(opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
conf_thres=conf,
iou_thres=iou,
model=attempt_load(final, device).half(),
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=save_json,
plots=False)
else: else:
dist.destroy_process_group() dist.destroy_process_group()
wandb.run.finish() if wandb and wandb.run else None
torch.cuda.empty_cache() torch.cuda.empty_cache()
return results return results
@ -453,13 +477,18 @@ if __name__ == '__main__':
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100')
parser.add_argument('--log-artifacts', action='store_true', help='log artifacts, i.e. final trained model')
parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers') parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
parser.add_argument('--project', default='runs/train', help='save to project/name') parser.add_argument('--project', default='runs/train', help='save to project/name')
parser.add_argument('--entity', default=None, help='W&B entity')
parser.add_argument('--name', default='exp', help='save to project/name') parser.add_argument('--name', default='exp', help='save to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--quad', action='store_true', help='quad dataloader') parser.add_argument('--quad', action='store_true', help='quad dataloader')
parser.add_argument('--linear-lr', action='store_true', help='linear LR')
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
opt = parser.parse_args() opt = parser.parse_args()
# Set DDP variables # Set DDP variables
@ -471,13 +500,14 @@ if __name__ == '__main__':
check_requirements() check_requirements()
# Resume # Resume
if opt.resume: # resume an interrupted run wandb_run = check_wandb_resume(opt)
if opt.resume and not wandb_run: # resume an interrupted run
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
apriori = opt.global_rank, opt.local_rank apriori = opt.global_rank, opt.local_rank
with open(Path(ckpt).parent.parent / 'opt.yaml') as f: with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader)) # replace
opt.cfg, opt.weights, opt.resume, opt.global_rank, opt.local_rank = '', ckpt, True, *apriori # reinstate opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori # reinstate
logger.info('Resuming training from %s' % ckpt) logger.info('Resuming training from %s' % ckpt)
else: else:
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml') # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
@ -504,18 +534,13 @@ if __name__ == '__main__':
# Train # Train
logger.info(opt) logger.info(opt)
try:
import wandb
except ImportError:
wandb = None
prefix = colorstr('wandb: ')
logger.info(f"{prefix}Install Weights & Biases for YOLOv3 logging with 'pip install wandb' (recommended)")
if not opt.evolve: if not opt.evolve:
tb_writer = None # init loggers tb_writer = None # init loggers
if opt.global_rank in [-1, 0]: if opt.global_rank in [-1, 0]:
logger.info(f'Start Tensorboard with "tensorboard --logdir {opt.project}", view at http://localhost:6006/') prefix = colorstr('tensorboard: ')
logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
tb_writer = SummaryWriter(opt.save_dir) # Tensorboard tb_writer = SummaryWriter(opt.save_dir) # Tensorboard
train(hyp, opt, device, tb_writer, wandb) train(hyp, opt, device, tb_writer)
# Evolve hyperparameters (optional) # Evolve hyperparameters (optional)
else: else:
@ -589,7 +614,7 @@ if __name__ == '__main__':
hyp[k] = round(hyp[k], 5) # significant digits hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation # Train mutation
results = train(hyp.copy(), opt, device, wandb=wandb) results = train(hyp.copy(), opt, device)
# Write mutation results # Write mutation results
print_mutation(hyp.copy(), results, yaml_file, opt.bucket) print_mutation(hyp.copy(), results, yaml_file, opt.bucket)

448
tutorial.ipynb vendored
View File

File diff suppressed because one or more lines are too long

19
utils/autoanchor.py
View File

@ -37,17 +37,21 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
bpr = (best > 1. / thr).float().mean() # best possible recall bpr = (best > 1. / thr).float().mean() # best possible recall
return bpr, aat return bpr, aat
bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2)) anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
bpr, aat = metric(anchors)
print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='') print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
if bpr < 0.98: # threshold to recompute if bpr < 0.98: # threshold to recompute
print('. Attempting to improve anchors, please wait...') print('. Attempting to improve anchors, please wait...')
na = m.anchor_grid.numel() // 2 # number of anchors na = m.anchor_grid.numel() // 2 # number of anchors
new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) try:
new_bpr = metric(new_anchors.reshape(-1, 2))[0] anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
except Exception as e:
print(f'{prefix}ERROR: {e}')
new_bpr = metric(anchors)[0]
if new_bpr > bpr: # replace anchors if new_bpr > bpr: # replace anchors
new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors) anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
check_anchor_order(m) check_anchor_order(m)
print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.') print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
else: else:
@ -98,7 +102,7 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
if isinstance(path, str): # *.yaml file if isinstance(path, str): # *.yaml file
with open(path) as f: with open(path) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict
from utils.datasets import LoadImagesAndLabels from utils.datasets import LoadImagesAndLabels
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
else: else:
@ -119,6 +123,7 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...') print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
s = wh.std(0) # sigmas for whitening s = wh.std(0) # sigmas for whitening
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}')
k *= s k *= s
wh = torch.tensor(wh, dtype=torch.float32) # filtered wh = torch.tensor(wh, dtype=torch.float32) # filtered
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered

0
utils/aws/__init__.py Normal file
View File

26
utils/aws/mime.sh Normal file
View File

@ -0,0 +1,26 @@
# AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/
# This script will run on every instance restart, not only on first start
# --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA ---
Content-Type: multipart/mixed; boundary="//"
MIME-Version: 1.0
--//
Content-Type: text/cloud-config; charset="us-ascii"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Content-Disposition: attachment; filename="cloud-config.txt"
#cloud-config
cloud_final_modules:
- [scripts-user, always]
--//
Content-Type: text/x-shellscript; charset="us-ascii"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Content-Disposition: attachment; filename="userdata.txt"
#!/bin/bash
# --- paste contents of userdata.sh here ---
--//

37
utils/aws/resume.py Normal file
View File

@ -0,0 +1,37 @@
# Resume all interrupted trainings in yolov5/ dir including DDP trainings
# Usage: $ python utils/aws/resume.py
import os
import sys
from pathlib import Path
import torch
import yaml
sys.path.append('./') # to run '$ python *.py' files in subdirectories
port = 0 # --master_port
path = Path('').resolve()
for last in path.rglob('*/**/last.pt'):
ckpt = torch.load(last)
if ckpt['optimizer'] is None:
continue
# Load opt.yaml
with open(last.parent.parent / 'opt.yaml') as f:
opt = yaml.load(f, Loader=yaml.SafeLoader)
# Get device count
d = opt['device'].split(',') # devices
nd = len(d) # number of devices
ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel
if ddp: # multi-GPU
port += 1
cmd = f'python -m torch.distributed.launch --nproc_per_node {nd} --master_port {port} train.py --resume {last}'
else: # single-GPU
cmd = f'python train.py --resume {last}'
cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread
print(cmd)
os.system(cmd)

27
utils/aws/userdata.sh Normal file
View File

@ -0,0 +1,27 @@
#!/bin/bash
# AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html
# This script will run only once on first instance start (for a re-start script see mime.sh)
# /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir
# Use >300 GB SSD
cd home/ubuntu
if [ ! -d yolov5 ]; then
echo "Running first-time script." # install dependencies, download COCO, pull Docker
git clone https://github.com/ultralytics/yolov5 && sudo chmod -R 777 yolov5
cd yolov5
bash data/scripts/get_coco.sh && echo "Data done." &
sudo docker pull ultralytics/yolov5:latest && echo "Docker done." &
python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." &
wait && echo "All tasks done." # finish background tasks
else
echo "Running re-start script." # resume interrupted runs
i=0
list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour'
while IFS= read -r id; do
((i++))
echo "restarting container $i: $id"
sudo docker start $id
# sudo docker exec -it $id python train.py --resume # single-GPU
sudo docker exec -d $id python utils/aws/resume.py # multi-scenario
done <<<"$list"
fi

226
utils/datasets.py
View File

@ -20,12 +20,13 @@ from PIL import Image, ExifTags
from torch.utils.data import Dataset from torch.utils.data import Dataset
from tqdm import tqdm from tqdm import tqdm
from utils.general import xyxy2xywh, xywh2xyxy, clean_str from utils.general import check_requirements, xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, \
resample_segments, clean_str
from utils.torch_utils import torch_distributed_zero_first from utils.torch_utils import torch_distributed_zero_first
# Parameters # Parameters
help_url = 'https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data' help_url = 'https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data'
img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng'] # acceptable image suffixes img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes
vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -119,9 +120,8 @@ class _RepeatSampler(object):
class LoadImages: # for inference class LoadImages: # for inference
def __init__(self, path, img_size=640): def __init__(self, path, img_size=640, stride=32):
p = str(Path(path)) # os-agnostic p = str(Path(path).absolute()) # os-agnostic absolute path
p = os.path.abspath(p) # absolute path
if '*' in p: if '*' in p:
files = sorted(glob.glob(p, recursive=True)) # glob files = sorted(glob.glob(p, recursive=True)) # glob
elif os.path.isdir(p): elif os.path.isdir(p):
@ -136,6 +136,7 @@ class LoadImages: # for inference
ni, nv = len(images), len(videos) ni, nv = len(images), len(videos)
self.img_size = img_size self.img_size = img_size
self.stride = stride
self.files = images + videos self.files = images + videos
self.nf = ni + nv # number of files self.nf = ni + nv # number of files
self.video_flag = [False] * ni + [True] * nv self.video_flag = [False] * ni + [True] * nv
@ -181,7 +182,7 @@ class LoadImages: # for inference
print(f'image {self.count}/{self.nf} {path}: ', end='') print(f'image {self.count}/{self.nf} {path}: ', end='')
# Padded resize # Padded resize
img = letterbox(img0, new_shape=self.img_size)[0] img = letterbox(img0, self.img_size, stride=self.stride)[0]
# Convert # Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
@ -199,8 +200,9 @@ class LoadImages: # for inference
class LoadWebcam: # for inference class LoadWebcam: # for inference
def __init__(self, pipe='0', img_size=640): def __init__(self, pipe='0', img_size=640, stride=32):
self.img_size = img_size self.img_size = img_size
self.stride = stride
if pipe.isnumeric(): if pipe.isnumeric():
pipe = eval(pipe) # local camera pipe = eval(pipe) # local camera
@ -243,7 +245,7 @@ class LoadWebcam: # for inference
print(f'webcam {self.count}: ', end='') print(f'webcam {self.count}: ', end='')
# Padded resize # Padded resize
img = letterbox(img0, new_shape=self.img_size)[0] img = letterbox(img0, self.img_size, stride=self.stride)[0]
# Convert # Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
@ -256,9 +258,10 @@ class LoadWebcam: # for inference
class LoadStreams: # multiple IP or RTSP cameras class LoadStreams: # multiple IP or RTSP cameras
def __init__(self, sources='streams.txt', img_size=640): def __init__(self, sources='streams.txt', img_size=640, stride=32):
self.mode = 'stream' self.mode = 'stream'
self.img_size = img_size self.img_size = img_size
self.stride = stride
if os.path.isfile(sources): if os.path.isfile(sources):
with open(sources, 'r') as f: with open(sources, 'r') as f:
@ -272,19 +275,25 @@ class LoadStreams: # multiple IP or RTSP cameras
for i, s in enumerate(sources): for i, s in enumerate(sources):
# Start the thread to read frames from the video stream # Start the thread to read frames from the video stream
print(f'{i + 1}/{n}: {s}... ', end='') print(f'{i + 1}/{n}: {s}... ', end='')
cap = cv2.VideoCapture(eval(s) if s.isnumeric() else s) url = eval(s) if s.isnumeric() else s
if 'youtube.com/' in url or 'youtu.be/' in url: # if source is YouTube video
check_requirements(('pafy', 'youtube_dl'))
import pafy
url = pafy.new(url).getbest(preftype="mp4").url
cap = cv2.VideoCapture(url)
assert cap.isOpened(), f'Failed to open {s}' assert cap.isOpened(), f'Failed to open {s}'
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS) % 100 self.fps = cap.get(cv2.CAP_PROP_FPS) % 100
_, self.imgs[i] = cap.read() # guarantee first frame _, self.imgs[i] = cap.read() # guarantee first frame
thread = Thread(target=self.update, args=([i, cap]), daemon=True) thread = Thread(target=self.update, args=([i, cap]), daemon=True)
print(f' success ({w}x{h} at {fps:.2f} FPS).') print(f' success ({w}x{h} at {self.fps:.2f} FPS).')
thread.start() thread.start()
print('') # newline print('') # newline
# check for common shapes # check for common shapes
s = np.stack([letterbox(x, new_shape=self.img_size)[0].shape for x in self.imgs], 0) # inference shapes s = np.stack([letterbox(x, self.img_size, stride=self.stride)[0].shape for x in self.imgs], 0) # shapes
self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
if not self.rect: if not self.rect:
print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.') print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
@ -297,9 +306,10 @@ class LoadStreams: # multiple IP or RTSP cameras
# _, self.imgs[index] = cap.read() # _, self.imgs[index] = cap.read()
cap.grab() cap.grab()
if n == 4: # read every 4th frame if n == 4: # read every 4th frame
_, self.imgs[index] = cap.retrieve() success, im = cap.retrieve()
self.imgs[index] = im if success else self.imgs[index] * 0
n = 0 n = 0
time.sleep(0.01) # wait time time.sleep(1 / self.fps) # wait time
def __iter__(self): def __iter__(self):
self.count = -1 self.count = -1
@ -313,7 +323,7 @@ class LoadStreams: # multiple IP or RTSP cameras
raise StopIteration raise StopIteration
# Letterbox # Letterbox
img = [letterbox(x, new_shape=self.img_size, auto=self.rect)[0] for x in img0] img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]
# Stack # Stack
img = np.stack(img, 0) img = np.stack(img, 0)
@ -331,7 +341,7 @@ class LoadStreams: # multiple IP or RTSP cameras
def img2label_paths(img_paths): def img2label_paths(img_paths):
# Define label paths as a function of image paths # Define label paths as a function of image paths
sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings
return [x.replace(sa, sb, 1).replace('.' + x.split('.')[-1], '.txt') for x in img_paths] return ['txt'.join(x.replace(sa, sb, 1).rsplit(x.split('.')[-1], 1)) for x in img_paths]
class LoadImagesAndLabels(Dataset): # for training/testing class LoadImagesAndLabels(Dataset): # for training/testing
@ -345,6 +355,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training) self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
self.mosaic_border = [-img_size // 2, -img_size // 2] self.mosaic_border = [-img_size // 2, -img_size // 2]
self.stride = stride self.stride = stride
self.path = path
try: try:
f = [] # image files f = [] # image files
@ -352,37 +363,42 @@ class LoadImagesAndLabels(Dataset): # for training/testing
p = Path(p) # os-agnostic p = Path(p) # os-agnostic
if p.is_dir(): # dir if p.is_dir(): # dir
f += glob.glob(str(p / '**' / '*.*'), recursive=True) f += glob.glob(str(p / '**' / '*.*'), recursive=True)
# f = list(p.rglob('**/*.*')) # pathlib
elif p.is_file(): # file elif p.is_file(): # file
with open(p, 'r') as t: with open(p, 'r') as t:
t = t.read().strip().splitlines() t = t.read().strip().splitlines()
parent = str(p.parent) + os.sep parent = str(p.parent) + os.sep
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
else: else:
raise Exception(f'{prefix}{p} does not exist') raise Exception(f'{prefix}{p} does not exist')
self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats]) self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats])
# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
assert self.img_files, f'{prefix}No images found' assert self.img_files, f'{prefix}No images found'
except Exception as e: except Exception as e:
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}') raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')
# Check cache # Check cache
self.label_files = img2label_paths(self.img_files) # labels self.label_files = img2label_paths(self.img_files) # labels
cache_path = Path(self.label_files[0]).parent.with_suffix('.cache') # cached labels cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') # cached labels
if cache_path.is_file(): if cache_path.is_file():
cache = torch.load(cache_path) # load cache, exists = torch.load(cache_path), True # load
if cache['hash'] != get_hash(self.label_files + self.img_files) or 'results' not in cache: # changed if cache['hash'] != get_hash(self.label_files + self.img_files) or 'version' not in cache: # changed
cache = self.cache_labels(cache_path, prefix) # re-cache cache, exists = self.cache_labels(cache_path, prefix), False # re-cache
else: else:
cache = self.cache_labels(cache_path, prefix) # cache cache, exists = self.cache_labels(cache_path, prefix), False # cache
# Display cache # Display cache
[nf, nm, ne, nc, n] = cache.pop('results') # found, missing, empty, corrupted, total nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
desc = f"Scanning '{cache_path}' for images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted" if exists:
tqdm(None, desc=prefix + desc, total=n, initial=n) d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {help_url}' assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {help_url}'
# Read cache # Read cache
cache.pop('hash') # remove hash cache.pop('hash') # remove hash
labels, shapes = zip(*cache.values()) cache.pop('version') # remove version
labels, shapes, self.segments = zip(*cache.values())
self.labels = list(labels) self.labels = list(labels)
self.shapes = np.array(shapes, dtype=np.float64) self.shapes = np.array(shapes, dtype=np.float64)
self.img_files = list(cache.keys()) # update self.img_files = list(cache.keys()) # update
@ -433,6 +449,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # img, hw_original, hw_resized = load_image(self, i) self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # img, hw_original, hw_resized = load_image(self, i)
gb += self.imgs[i].nbytes gb += self.imgs[i].nbytes
pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)' pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)'
pbar.close()
def cache_labels(self, path=Path('./labels.cache'), prefix=''): def cache_labels(self, path=Path('./labels.cache'), prefix=''):
# Cache dataset labels, check images and read shapes # Cache dataset labels, check images and read shapes
@ -445,13 +462,20 @@ class LoadImagesAndLabels(Dataset): # for training/testing
im = Image.open(im_file) im = Image.open(im_file)
im.verify() # PIL verify im.verify() # PIL verify
shape = exif_size(im) # image size shape = exif_size(im) # image size
assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels' segments = [] # instance segments
assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
assert im.format.lower() in img_formats, f'invalid image format {im.format}'
# verify labels # verify labels
if os.path.isfile(lb_file): if os.path.isfile(lb_file):
nf += 1 # label found nf += 1 # label found
with open(lb_file, 'r') as f: with open(lb_file, 'r') as f:
l = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels l = [x.split() for x in f.read().strip().splitlines()]
if any([len(x) > 8 for x in l]): # is segment
classes = np.array([x[0] for x in l], dtype=np.float32)
segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
l = np.array(l, dtype=np.float32)
if len(l): if len(l):
assert l.shape[1] == 5, 'labels require 5 columns each' assert l.shape[1] == 5, 'labels require 5 columns each'
assert (l >= 0).all(), 'negative labels' assert (l >= 0).all(), 'negative labels'
@ -463,19 +487,21 @@ class LoadImagesAndLabels(Dataset): # for training/testing
else: else:
nm += 1 # label missing nm += 1 # label missing
l = np.zeros((0, 5), dtype=np.float32) l = np.zeros((0, 5), dtype=np.float32)
x[im_file] = [l, shape] x[im_file] = [l, shape, segments]
except Exception as e: except Exception as e:
nc += 1 nc += 1
print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}') print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' for images and labels... " \ pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels... " \
f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted" f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
pbar.close()
if nf == 0: if nf == 0:
print(f'{prefix}WARNING: No labels found in {path}. See {help_url}') print(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
x['hash'] = get_hash(self.label_files + self.img_files) x['hash'] = get_hash(self.label_files + self.img_files)
x['results'] = [nf, nm, ne, nc, i + 1] x['results'] = nf, nm, ne, nc, i + 1
x['version'] = 0.1 # cache version
torch.save(x, path) # save for next time torch.save(x, path) # save for next time
logging.info(f'{prefix}New cache created: {path}') logging.info(f'{prefix}New cache created: {path}')
return x return x
@ -515,16 +541,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment) img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
# Load labels labels = self.labels[index].copy()
labels = [] if labels.size: # normalized xywh to pixel xyxy format
x = self.labels[index] labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
if x.size > 0:
# Normalized xywh to pixel xyxy format
labels = x.copy()
labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width
labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1] # pad height
labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]
labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]
if self.augment: if self.augment:
# Augment imagespace # Augment imagespace
@ -637,19 +656,25 @@ def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype) img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed
# Histogram equalization
# if random.random() < 0.2: def hist_equalize(img, clahe=True, bgr=False):
# for i in range(3): # Equalize histogram on BGR image 'img' with img.shape(n,m,3) and range 0-255
# img[:, :, i] = cv2.equalizeHist(img[:, :, i]) yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
if clahe:
c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
yuv[:, :, 0] = c.apply(yuv[:, :, 0])
else:
yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0]) # equalize Y channel histogram
return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB) # convert YUV image to RGB
def load_mosaic(self, index): def load_mosaic(self, index):
# loads images in a 4-mosaic # loads images in a 4-mosaic
labels4 = [] labels4, segments4 = [], []
s = self.img_size s = self.img_size
yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y
indices = [index] + [self.indices[random.randint(0, self.n - 1)] for _ in range(3)] # 3 additional image indices indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
for i, index in enumerate(indices): for i, index in enumerate(indices):
# Load image # Load image
img, _, (h, w) = load_image(self, index) img, _, (h, w) = load_image(self, index)
@ -674,23 +699,21 @@ def load_mosaic(self, index):
padh = y1a - y1b padh = y1a - y1b
# Labels # Labels
x = self.labels[index] labels, segments = self.labels[index].copy(), self.segments[index].copy()
labels = x.copy() if labels.size:
if x.size > 0: # Normalized xywh to pixel xyxy format labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh
labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw
labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh
labels4.append(labels) labels4.append(labels)
segments4.extend(segments)
# Concat/clip labels # Concat/clip labels
if len(labels4):
labels4 = np.concatenate(labels4, 0) labels4 = np.concatenate(labels4, 0)
np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_perspective for x in (labels4[:, 1:], *segments4):
np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
# img4, labels4 = replicate(img4, labels4) # replicate # img4, labels4 = replicate(img4, labels4) # replicate
# Augment # Augment
img4, labels4 = random_perspective(img4, labels4, img4, labels4 = random_perspective(img4, labels4, segments4,
degrees=self.hyp['degrees'], degrees=self.hyp['degrees'],
translate=self.hyp['translate'], translate=self.hyp['translate'],
scale=self.hyp['scale'], scale=self.hyp['scale'],
@ -704,9 +727,9 @@ def load_mosaic(self, index):
def load_mosaic9(self, index): def load_mosaic9(self, index):
# loads images in a 9-mosaic # loads images in a 9-mosaic
labels9 = [] labels9, segments9 = [], []
s = self.img_size s = self.img_size
indices = [index] + [self.indices[random.randint(0, self.n - 1)] for _ in range(8)] # 8 additional image indices indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
for i, index in enumerate(indices): for i, index in enumerate(indices):
# Load image # Load image
img, _, (h, w) = load_image(self, index) img, _, (h, w) = load_image(self, index)
@ -737,34 +760,34 @@ def load_mosaic9(self, index):
x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords
# Labels # Labels
x = self.labels[index] labels, segments = self.labels[index].copy(), self.segments[index].copy()
labels = x.copy() if labels.size:
if x.size > 0: # Normalized xywh to pixel xyxy format labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padx segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + pady
labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padx
labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + pady
labels9.append(labels) labels9.append(labels)
segments9.extend(segments)
# Image # Image
img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax] img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
hp, wp = h, w # height, width previous hp, wp = h, w # height, width previous
# Offset # Offset
yc, xc = [int(random.uniform(0, s)) for x in self.mosaic_border] # mosaic center x, y yc, xc = [int(random.uniform(0, s)) for _ in self.mosaic_border] # mosaic center x, y
img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s] img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]
# Concat/clip labels # Concat/clip labels
if len(labels9):
labels9 = np.concatenate(labels9, 0) labels9 = np.concatenate(labels9, 0)
labels9[:, [1, 3]] -= xc labels9[:, [1, 3]] -= xc
labels9[:, [2, 4]] -= yc labels9[:, [2, 4]] -= yc
c = np.array([xc, yc]) # centers
segments9 = [x - c for x in segments9]
np.clip(labels9[:, 1:], 0, 2 * s, out=labels9[:, 1:]) # use with random_perspective for x in (labels9[:, 1:], *segments9):
np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
# img9, labels9 = replicate(img9, labels9) # replicate # img9, labels9 = replicate(img9, labels9) # replicate
# Augment # Augment
img9, labels9 = random_perspective(img9, labels9, img9, labels9 = random_perspective(img9, labels9, segments9,
degrees=self.hyp['degrees'], degrees=self.hyp['degrees'],
translate=self.hyp['translate'], translate=self.hyp['translate'],
scale=self.hyp['scale'], scale=self.hyp['scale'],
@ -792,8 +815,8 @@ def replicate(img, labels):
return img, labels return img, labels
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
# Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 # Resize and pad image while meeting stride-multiple constraints
shape = img.shape[:2] # current shape [height, width] shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int): if isinstance(new_shape, int):
new_shape = (new_shape, new_shape) new_shape = (new_shape, new_shape)
@ -808,7 +831,7 @@ def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
if auto: # minimum rectangle if auto: # minimum rectangle
dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
elif scaleFill: # stretch elif scaleFill: # stretch
dw, dh = 0.0, 0.0 dw, dh = 0.0, 0.0
new_unpad = (new_shape[1], new_shape[0]) new_unpad = (new_shape[1], new_shape[0])
@ -825,7 +848,8 @@ def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale
return img, ratio, (dw, dh) return img, ratio, (dw, dh)
def random_perspective(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0, border=(0, 0)): def random_perspective(img, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
border=(0, 0)):
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10)) # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
# targets = [cls, xyxy] # targets = [cls, xyxy]
@ -877,37 +901,38 @@ def random_perspective(img, targets=(), degrees=10, translate=.1, scale=.1, shea
# Transform label coordinates # Transform label coordinates
n = len(targets) n = len(targets)
if n: if n:
# warp points use_segments = any(x.any() for x in segments)
new = np.zeros((n, 4))
if use_segments: # warp segments
segments = resample_segments(segments) # upsample
for i, segment in enumerate(segments):
xy = np.ones((len(segment), 3))
xy[:, :2] = segment
xy = xy @ M.T # transform
xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2] # perspective rescale or affine
# clip
new[i] = segment2box(xy, width, height)
else: # warp boxes
xy = np.ones((n * 4, 3)) xy = np.ones((n * 4, 3))
xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1 xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
xy = xy @ M.T # transform xy = xy @ M.T # transform
if perspective: xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale
else: # affine
xy = xy[:, :2].reshape(n, 8)
# create new boxes # create new boxes
x = xy[:, [0, 2, 4, 6]] x = xy[:, [0, 2, 4, 6]]
y = xy[:, [1, 3, 5, 7]] y = xy[:, [1, 3, 5, 7]]
xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
# # apply angle-based reduction of bounding boxes # clip
# radians = a * math.pi / 180 new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
# reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5 new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
# x = (xy[:, 2] + xy[:, 0]) / 2
# y = (xy[:, 3] + xy[:, 1]) / 2
# w = (xy[:, 2] - xy[:, 0]) * reduction
# h = (xy[:, 3] - xy[:, 1]) * reduction
# xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T
# clip boxes
xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
# filter candidates # filter candidates
i = box_candidates(box1=targets[:, 1:5].T * s, box2=xy.T) i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)
targets = targets[i] targets = targets[i]
targets[:, 1:5] = xy[i] targets[:, 1:5] = new[i]
return img, targets return img, targets
@ -1016,19 +1041,24 @@ def extract_boxes(path='../coco128/'): # from utils.datasets import *; extract_
assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}' assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
def autosplit(path='../coco128', weights=(0.9, 0.1, 0.0)): # from utils.datasets import *; autosplit('../coco128') def autosplit(path='../coco128', weights=(0.9, 0.1, 0.0), annotated_only=False):
""" Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
# Arguments Usage: from utils.datasets import *; autosplit('../coco128')
Arguments
path: Path to images directory path: Path to images directory
weights: Train, val, test weights (list) weights: Train, val, test weights (list)
annotated_only: Only use images with an annotated txt file
""" """
path = Path(path) # images dir path = Path(path) # images dir
files = list(path.rglob('*.*')) files = sum([list(path.rglob(f"*.{img_ext}")) for img_ext in img_formats], []) # image files only
n = len(files) # number of files n = len(files) # number of files
indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
[(path / x).unlink() for x in txt if (path / x).exists()] # remove existing [(path / x).unlink() for x in txt if (path / x).exists()] # remove existing
print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)
for i, img in tqdm(zip(indices, files), total=n): for i, img in tqdm(zip(indices, files), total=n):
if img.suffix[1:] in img_formats: if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
with open(path / txt[i], 'a') as f: with open(path / txt[i], 'a') as f:
f.write(str(img) + '\n') # add image to txt file f.write(str(img) + '\n') # add image to txt file

159
utils/general.py
View File

@ -1,9 +1,10 @@
# General utils # YOLOv3 general utils
import glob import glob
import logging import logging
import math import math
import os import os
import platform
import random import random
import re import re
import subprocess import subprocess
@ -12,6 +13,7 @@ from pathlib import Path
import cv2 import cv2
import numpy as np import numpy as np
import pandas as pd
import torch import torch
import torchvision import torchvision
import yaml import yaml
@ -23,6 +25,7 @@ from utils.torch_utils import init_torch_seeds
# Settings # Settings
torch.set_printoptions(linewidth=320, precision=5, profile='long') torch.set_printoptions(linewidth=320, precision=5, profile='long')
np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
pd.options.display.max_columns = 10
cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
@ -46,40 +49,75 @@ def get_latest_run(search_dir='.'):
return max(last_list, key=os.path.getctime) if last_list else '' return max(last_list, key=os.path.getctime) if last_list else ''
def isdocker():
# Is environment a Docker container
return Path('/workspace').exists() # or Path('/.dockerenv').exists()
def emojis(str=''):
# Return platform-dependent emoji-safe version of string
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
def check_online(): def check_online():
# Check internet connectivity # Check internet connectivity
import socket import socket
try: try:
socket.create_connection(("1.1.1.1", 53)) # check host accesability socket.create_connection(("1.1.1.1", 443), 5) # check host accesability
return True return True
except OSError: except OSError:
return False return False
def check_git_status(): def check_git_status():
# Suggest 'git pull' if YOLOv5 is out of date # Recommend 'git pull' if code is out of date
print(colorstr('github: '), end='') print(colorstr('github: '), end='')
try: try:
if Path('.git').exists() and check_online(): assert Path('.git').exists(), 'skipping check (not a git repository)'
url = subprocess.check_output( assert not isdocker(), 'skipping check (Docker image)'
'git fetch && git config --get remote.origin.url', shell=True).decode('utf-8')[:-1] assert check_online(), 'skipping check (offline)'
n = int(subprocess.check_output(
'git rev-list $(git rev-parse --abbrev-ref HEAD)..origin/master --count', shell=True)) # commits behind cmd = 'git fetch && git config --get remote.origin.url'
url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git') # github repo url
branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
if n > 0: if n > 0:
print(f"⚠️ WARNING: code is out of date by {n} {'commits' if n > 1 else 'commmit'}. " s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
f"Use 'git pull' to update or 'git clone {url}' to download latest.") f"Use 'git pull' to update or 'git clone {url}' to download latest."
else: else:
print(f'up to date with {url} ✅') s = f'up to date with {url} ✅'
print(emojis(s)) # emoji-safe
except Exception as e: except Exception as e:
print(e) print(e)
def check_requirements(file='requirements.txt'): def check_requirements(requirements='requirements.txt', exclude=()):
# Check installed dependencies meet requirements # Check installed dependencies meet requirements (pass *.txt file or list of packages)
import pkg_resources import pkg_resources as pkg
requirements = pkg_resources.parse_requirements(Path(file).open()) prefix = colorstr('red', 'bold', 'requirements:')
requirements = [x.name + ''.join(*x.specs) if len(x.specs) else x.name for x in requirements] if isinstance(requirements, (str, Path)): # requirements.txt file
pkg_resources.require(requirements) # DistributionNotFound or VersionConflict exception if requirements not met file = Path(requirements)
if not file.exists():
print(f"{prefix} {file.resolve()} not found, check failed.")
return
requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
else: # list or tuple of packages
requirements = [x for x in requirements if x not in exclude]
n = 0 # number of packages updates
for r in requirements:
try:
pkg.require(r)
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
n += 1
print(f"{prefix} {e.req} not found and is required by YOLOv3, attempting auto-update...")
print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode())
if n: # if packages updated
source = file.resolve() if 'file' in locals() else requirements
s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
print(emojis(s)) # emoji-safe
def check_img_size(img_size, s=32): def check_img_size(img_size, s=32):
@ -90,14 +128,28 @@ def check_img_size(img_size, s=32):
return new_size return new_size
def check_imshow():
# Check if environment supports image displays
try:
assert not isdocker(), 'cv2.imshow() is disabled in Docker environments'
cv2.imshow('test', np.zeros((1, 1, 3)))
cv2.waitKey(1)
cv2.destroyAllWindows()
cv2.waitKey(1)
return True
except Exception as e:
print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
return False
def check_file(file): def check_file(file):
# Search for file if not found # Search for file if not found
if os.path.isfile(file) or file == '': if Path(file).is_file() or file == '':
return file return file
else: else:
files = glob.glob('./**/' + file, recursive=True) # find file files = glob.glob('./**/' + file, recursive=True) # find file
assert len(files), 'File Not Found: %s' % file # assert file was found assert len(files), f'File Not Found: {file}' # assert file was found
assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files) # assert unique assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
return files[0] # return file return files[0] # return file
@ -220,6 +272,50 @@ def xywh2xyxy(x):
return y return y
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x
y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y
y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x
y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y
return y
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
# Convert normalized segments into pixel segments, shape (n,2)
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
y[:, 0] = w * x[:, 0] + padw # top left x
y[:, 1] = h * x[:, 1] + padh # top left y
return y
def segment2box(segment, width=640, height=640):
# Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
x, y = segment.T # segment xy
inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
x, y, = x[inside], y[inside]
return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # xyxy
def segments2boxes(segments):
# Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
boxes = []
for s in segments:
x, y = s.T # segment xy
boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy
return xyxy2xywh(np.array(boxes)) # cls, xywh
def resample_segments(segments, n=1000):
# Up-sample an (n,2) segment
for i, s in enumerate(segments):
x = np.linspace(0, len(s) - 1, n)
xp = np.arange(len(s))
segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy
return segments
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
# Rescale coords (xyxy) from img1_shape to img0_shape # Rescale coords (xyxy) from img1_shape to img0_shape
if ratio_pad is None: # calculate from img0_shape if ratio_pad is None: # calculate from img0_shape
@ -244,7 +340,7 @@ def clip_coords(boxes, img_shape):
boxes[:, 3].clamp_(0, img_shape[0]) # y2 boxes[:, 3].clamp_(0, img_shape[0]) # y2
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9): def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
box2 = box2.T box2 = box2.T
@ -280,7 +376,7 @@ def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
with torch.no_grad(): with torch.no_grad():
alpha = v / ((1 + eps) - iou + v) alpha = v / (v - iou + (1 + eps))
return iou - (rho2 / c2 + v * alpha) # CIoU return iou - (rho2 / c2 + v * alpha) # CIoU
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
c_area = cw * ch + eps # convex area c_area = cw * ch + eps # convex area
@ -322,11 +418,12 @@ def wh_iou(wh1, wh2):
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter) return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()): def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
"""Performs Non-Maximum Suppression (NMS) on inference results labels=()):
"""Runs Non-Maximum Suppression (NMS) on inference results
Returns: Returns:
detections with shape: nx6 (x1, y1, x2, y2, conf, cls) list of detections, on (n,6) tensor per image [xyxy, conf, cls]
""" """
nc = prediction.shape[2] - 5 # number of classes nc = prediction.shape[2] - 5 # number of classes
@ -338,7 +435,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms() max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
time_limit = 10.0 # seconds to quit after time_limit = 10.0 # seconds to quit after
redundant = True # require redundant detections redundant = True # require redundant detections
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
merge = False # use merge-NMS merge = False # use merge-NMS
t = time.time() t = time.time()
@ -412,18 +509,20 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
return output return output
def strip_optimizer(f='weights/best.pt', s=''): # from utils.general import *; strip_optimizer() def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
# Strip optimizer from 'f' to finalize training, optionally save as 's' # Strip optimizer from 'f' to finalize training, optionally save as 's'
x = torch.load(f, map_location=torch.device('cpu')) x = torch.load(f, map_location=torch.device('cpu'))
for key in 'optimizer', 'training_results', 'wandb_id': if x.get('ema'):
x[key] = None x['model'] = x['ema'] # replace model with ema
for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
x[k] = None
x['epoch'] = -1 x['epoch'] = -1
x['model'].half() # to FP16 x['model'].half() # to FP16
for p in x['model'].parameters(): for p in x['model'].parameters():
p.requires_grad = False p.requires_grad = False
torch.save(x, s or f) torch.save(x, s or f)
mb = os.path.getsize(s or f) / 1E6 # filesize mb = os.path.getsize(s or f) / 1E6 # filesize
print('Optimizer stripped from %s,%s %.1fMB' % (f, (' saved as %s,' % s) if s else '', mb)) print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''): def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):

10
utils/google_utils.py
View File

@ -21,11 +21,11 @@ def attempt_download(file, repo='ultralytics/yolov3'):
file = Path(str(file).strip().replace("'", '').lower()) file = Path(str(file).strip().replace("'", '').lower())
if not file.exists(): if not file.exists():
# try: try:
# response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
# assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
# tag = response['tag_name'] # i.e. 'v1.0' tag = response['tag_name'] # i.e. 'v1.0'
# except: # fallback plan except: # fallback plan
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt'] assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
tag = subprocess.check_output('git tag', shell=True).decode().split()[-1] tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]

61
utils/loss.py
View File

@ -85,26 +85,38 @@ class QFocalLoss(nn.Module):
return loss return loss
def compute_loss(p, targets, model): # predictions, targets, model class ComputeLoss:
device = targets.device # Compute losses
lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) def __init__(self, model, autobalance=False):
tcls, tbox, indices, anchors = build_targets(p, targets, model) # targets super(ComputeLoss, self).__init__()
device = next(model.parameters()).device # get model device
h = model.hyp # hyperparameters h = model.hyp # hyperparameters
# Define criteria # Define criteria
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) # weight=model.class_weights) BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device)) BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
cp, cn = smooth_BCE(eps=0.0) self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets
# Focal loss # Focal loss
g = h['fl_gamma'] # focal loss gamma g = h['fl_gamma'] # focal loss gamma
if g > 0: if g > 0:
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7
self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
for k in 'na', 'nc', 'nl', 'anchors':
setattr(self, k, getattr(det, k))
def __call__(self, p, targets): # predictions, targets, model
device = targets.device
lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets
# Losses # Losses
balance = [4.0, 1.0, 0.4, 0.1] # P3-P6
for i, pi in enumerate(p): # layer index, layer predictions for i, pi in enumerate(p): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi[..., 0], device=device) # target obj tobj = torch.zeros_like(pi[..., 0], device=device) # target obj
@ -121,33 +133,36 @@ def compute_loss(p, targets, model): # predictions, targets, model
lbox += (1.0 - iou).mean() # iou loss lbox += (1.0 - iou).mean() # iou loss
# Objectness # Objectness
tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
# Classification # Classification
if model.nc > 1: # cls loss (only if multiple classes) if self.nc > 1: # cls loss (only if multiple classes)
t = torch.full_like(ps[:, 5:], cn, device=device) # targets t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets
t[range(n), tcls[i]] = cp t[range(n), tcls[i]] = self.cp
lcls += BCEcls(ps[:, 5:], t) # BCE lcls += self.BCEcls(ps[:, 5:], t) # BCE
# Append targets to text file # Append targets to text file
# with open('targets.txt', 'a') as file: # with open('targets.txt', 'a') as file:
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss obji = self.BCEobj(pi[..., 4], tobj)
lobj += obji * self.balance[i] # obj loss
if self.autobalance:
self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
lbox *= h['box'] if self.autobalance:
lobj *= h['obj'] self.balance = [x / self.balance[self.ssi] for x in self.balance]
lcls *= h['cls'] lbox *= self.hyp['box']
lobj *= self.hyp['obj']
lcls *= self.hyp['cls']
bs = tobj.shape[0] # batch size bs = tobj.shape[0] # batch size
loss = lbox + lobj + lcls loss = lbox + lobj + lcls
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach() return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
def build_targets(self, p, targets):
def build_targets(p, targets, model):
# Build targets for compute_loss(), input targets(image,class,x,y,w,h) # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module na, nt = self.na, targets.shape[0] # number of anchors, targets
na, nt = det.na, targets.shape[0] # number of anchors, targets
tcls, tbox, indices, anch = [], [], [], [] tcls, tbox, indices, anch = [], [], [], []
gain = torch.ones(7, device=targets.device) # normalized to gridspace gain gain = torch.ones(7, device=targets.device) # normalized to gridspace gain
ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
@ -159,8 +174,8 @@ def build_targets(p, targets, model):
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
], device=targets.device).float() * g # offsets ], device=targets.device).float() * g # offsets
for i in range(det.nl): for i in range(self.nl):
anchors = det.anchors[i] anchors = self.anchors[i]
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
# Match targets to anchors # Match targets to anchors
@ -168,7 +183,7 @@ def build_targets(p, targets, model):
if nt: if nt:
# Matches # Matches
r = t[:, :, 4:6] / anchors[:, None] # wh ratio r = t[:, :, 4:6] / anchors[:, None] # wh ratio
j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t'] # compare j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
t = t[j] # filter t = t[j] # filter

61
utils/metrics.py
View File

@ -15,7 +15,7 @@ def fitness(x):
return (x[:, :4] * w).sum(1) return (x[:, :4] * w).sum(1)
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision-recall_curve.png', names=[]): def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()):
""" Compute the average precision, given the recall and precision curves. """ Compute the average precision, given the recall and precision curves.
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
# Arguments # Arguments
@ -35,12 +35,11 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision
# Find unique classes # Find unique classes
unique_classes = np.unique(target_cls) unique_classes = np.unique(target_cls)
nc = unique_classes.shape[0] # number of classes, number of detections
# Create Precision-Recall curve and compute AP for each class # Create Precision-Recall curve and compute AP for each class
px, py = np.linspace(0, 1, 1000), [] # for plotting px, py = np.linspace(0, 1, 1000), [] # for plotting
pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898 ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)
ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)
for ci, c in enumerate(unique_classes): for ci, c in enumerate(unique_classes):
i = pred_cls == c i = pred_cls == c
n_l = (target_cls == c).sum() # number of labels n_l = (target_cls == c).sum() # number of labels
@ -55,25 +54,28 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision
# Recall # Recall
recall = tpc / (n_l + 1e-16) # recall curve recall = tpc / (n_l + 1e-16) # recall curve
r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
# Precision # Precision
precision = tpc / (tpc + fpc) # precision curve precision = tpc / (tpc + fpc) # precision curve
p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
# AP from recall-precision curve # AP from recall-precision curve
for j in range(tp.shape[1]): for j in range(tp.shape[1]):
ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j]) ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
if plot and (j == 0): if plot and j == 0:
py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5 py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
# Compute F1 score (harmonic mean of precision and recall) # Compute F1 (harmonic mean of precision and recall)
f1 = 2 * p * r / (p + r + 1e-16) f1 = 2 * p * r / (p + r + 1e-16)
if plot: if plot:
plot_pr_curve(px, py, ap, save_dir, names) plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision')
plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall')
return p, r, ap, f1, unique_classes.astype('int32') i = f1.mean(0).argmax() # max F1 index
return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32')
def compute_ap(recall, precision): def compute_ap(recall, precision):
@ -145,12 +147,12 @@ class ConfusionMatrix:
if n and sum(j) == 1: if n and sum(j) == 1:
self.matrix[gc, detection_classes[m1[j]]] += 1 # correct self.matrix[gc, detection_classes[m1[j]]] += 1 # correct
else: else:
self.matrix[gc, self.nc] += 1 # background FP self.matrix[self.nc, gc] += 1 # background FP
if n: if n:
for i, dc in enumerate(detection_classes): for i, dc in enumerate(detection_classes):
if not any(m1 == i): if not any(m1 == i):
self.matrix[self.nc, dc] += 1 # background FN self.matrix[dc, self.nc] += 1 # background FN
def matrix(self): def matrix(self):
return self.matrix return self.matrix
@ -166,8 +168,8 @@ class ConfusionMatrix:
sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
xticklabels=names + ['background FN'] if labels else "auto", xticklabels=names + ['background FP'] if labels else "auto",
yticklabels=names + ['background FP'] if labels else "auto").set_facecolor((1, 1, 1)) yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
fig.axes[0].set_xlabel('True') fig.axes[0].set_xlabel('True')
fig.axes[0].set_ylabel('Predicted') fig.axes[0].set_ylabel('Predicted')
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
@ -181,13 +183,14 @@ class ConfusionMatrix:
# Plots ---------------------------------------------------------------------------------------------------------------- # Plots ----------------------------------------------------------------------------------------------------------------
def plot_pr_curve(px, py, ap, save_dir='.', names=()): def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
# Precision-recall curve
fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
py = np.stack(py, axis=1) py = np.stack(py, axis=1)
if 0 < len(names) < 21: # show mAP in legend if < 10 classes if 0 < len(names) < 21: # display per-class legend if < 21 classes
for i, y in enumerate(py.T): for i, y in enumerate(py.T):
ax.plot(px, y, linewidth=1, label=f'{names[i]} %.3f' % ap[i, 0]) # plot(recall, precision) ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}') # plot(recall, precision)
else: else:
ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision)
@ -197,4 +200,24 @@ def plot_pr_curve(px, py, ap, save_dir='.', names=()):
ax.set_xlim(0, 1) ax.set_xlim(0, 1)
ax.set_ylim(0, 1) ax.set_ylim(0, 1)
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
fig.savefig(Path(save_dir) / 'precision_recall_curve.png', dpi=250) fig.savefig(Path(save_dir), dpi=250)
def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
# Metric-confidence curve
fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
if 0 < len(names) < 21: # display per-class legend if < 21 classes
for i, y in enumerate(py):
ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric)
else:
ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric)
y = py.mean(0)
ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
fig.savefig(Path(save_dir), dpi=250)

56
utils/plots.py
View File

@ -15,7 +15,7 @@ import pandas as pd
import seaborn as sns import seaborn as sns
import torch import torch
import yaml import yaml
from PIL import Image, ImageDraw from PIL import Image, ImageDraw, ImageFont
from scipy.signal import butter, filtfilt from scipy.signal import butter, filtfilt
from utils.general import xywh2xyxy, xyxy2xywh from utils.general import xywh2xyxy, xyxy2xywh
@ -31,7 +31,7 @@ def color_list():
def hex2rgb(h): def hex2rgb(h):
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
return [hex2rgb(h) for h in plt.rcParams['axes.prop_cycle'].by_key()['color']] return [hex2rgb(h) for h in matplotlib.colors.TABLEAU_COLORS.values()] # or BASE_ (8), CSS4_ (148), XKCD_ (949)
def hist2d(x, y, n=100): def hist2d(x, y, n=100):
@ -54,7 +54,7 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
return filtfilt(b, a, data) # forward-backward filter return filtfilt(b, a, data) # forward-backward filter
def plot_one_box(x, img, color=None, label=None, line_thickness=None): def plot_one_box(x, img, color=None, label=None, line_thickness=3):
# Plots one bounding box on image img # Plots one bounding box on image img
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
color = color or [random.randint(0, 255) for _ in range(3)] color = color or [random.randint(0, 255) for _ in range(3)]
@ -68,6 +68,20 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None):
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None):
img = Image.fromarray(img)
draw = ImageDraw.Draw(img)
line_thickness = line_thickness or max(int(min(img.size) / 200), 2)
draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
if label:
fontsize = max(round(max(img.size) / 40), 12)
font = ImageFont.truetype("Arial.ttf", fontsize)
txt_width, txt_height = font.getsize(label)
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
return np.asarray(img)
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods() def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
# Compares the two methods for width-height anchor multiplication # Compares the two methods for width-height anchor multiplication
# https://github.com/ultralytics/yolov3/issues/168 # https://github.com/ultralytics/yolov3/issues/168
@ -223,38 +237,39 @@ def plot_targets_txt(): # from utils.plots import *; plot_targets_txt()
plt.savefig('targets.jpg', dpi=200) plt.savefig('targets.jpg', dpi=200)
def plot_study_txt(path='study/', x=None): # from utils.plots import *; plot_study_txt() def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_txt()
# Plot study.txt generated by test.py # Plot study.txt generated by test.py
fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True) fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
ax = ax.ravel() # ax = ax.ravel()
fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]: # for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov3-tiny', 'yolov3', 'yolov3-spp', 'yolov5l']]:
for f in sorted(Path(path).glob('study*.txt')):
y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
x = np.arange(y.shape[1]) if x is None else np.array(x) x = np.arange(y.shape[1]) if x is None else np.array(x)
s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)'] s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
for i in range(7): # for i in range(7):
ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8) # ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
ax[i].set_title(s[i]) # ax[i].set_title(s[i])
j = y[3].argmax() + 1 j = y[3].argmax() + 1
ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8, ax2.plot(y[6, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8,
label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO')) label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet') 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
ax2.grid() ax2.grid(alpha=0.2)
ax2.set_yticks(np.arange(30, 60, 5)) ax2.set_yticks(np.arange(20, 60, 5))
ax2.set_xlim(0, 30) ax2.set_xlim(0, 57)
ax2.set_ylim(29, 51) ax2.set_ylim(15, 55)
ax2.set_xlabel('GPU Speed (ms/img)') ax2.set_xlabel('GPU Speed (ms/img)')
ax2.set_ylabel('COCO AP val') ax2.set_ylabel('COCO AP val')
ax2.legend(loc='lower right') ax2.legend(loc='lower right')
plt.savefig('test_study.png', dpi=300) plt.savefig(str(Path(path).name) + '.png', dpi=300)
def plot_labels(labels, save_dir=Path(''), loggers=None): def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
# plot dataset labels # plot dataset labels
print('Plotting labels... ') print('Plotting labels... ')
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
@ -271,6 +286,11 @@ def plot_labels(labels, save_dir=Path(''), loggers=None):
matplotlib.use('svg') # faster matplotlib.use('svg') # faster
ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
ax[0].set_ylabel('instances')
if 0 < len(names) < 30:
ax[0].set_xticks(range(len(names)))
ax[0].set_xticklabels(names, rotation=90, fontsize=10)
else:
ax[0].set_xlabel('classes') ax[0].set_xlabel('classes')
sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9) sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9) sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
@ -295,13 +315,13 @@ def plot_labels(labels, save_dir=Path(''), loggers=None):
# loggers # loggers
for k, v in loggers.items() or {}: for k, v in loggers.items() or {}:
if k == 'wandb' and v: if k == 'wandb' and v:
v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}) v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}, commit=False)
def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution() def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution()
# Plot hyperparameter evolution results in evolve.txt # Plot hyperparameter evolution results in evolve.txt
with open(yaml_file) as f: with open(yaml_file) as f:
hyp = yaml.load(f, Loader=yaml.FullLoader) hyp = yaml.load(f, Loader=yaml.SafeLoader)
x = np.loadtxt('evolve.txt', ndmin=2) x = np.loadtxt('evolve.txt', ndmin=2)
f = fitness(x) f = fitness(x)
# weights = (f - f.min()) ** 2 # for weighted results # weights = (f - f.min()) ** 2 # for weighted results

31
utils/torch_utils.py
View File

@ -1,8 +1,10 @@
# PyTorch utils # YOLOv3 PyTorch utils
import datetime
import logging import logging
import math import math
import os import os
import platform
import subprocess import subprocess
import time import time
from contextlib import contextmanager from contextlib import contextmanager
@ -43,17 +45,24 @@ def init_torch_seeds(seed=0):
cudnn.benchmark, cudnn.deterministic = True, False cudnn.benchmark, cudnn.deterministic = True, False
def git_describe(): def date_modified(path=__file__):
# return human-readable file modification date, i.e. '2021-3-26'
t = datetime.datetime.fromtimestamp(Path(path).stat().st_mtime)
return f'{t.year}-{t.month}-{t.day}'
def git_describe(path=Path(__file__).parent): # path must be a directory
# return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
if Path('.git').exists(): s = f'git -C {path} describe --tags --long --always'
return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1] try:
else: return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1]
return '' except subprocess.CalledProcessError as e:
return '' # not a git repository
def select_device(device='', batch_size=None): def select_device(device='', batch_size=None):
# device = 'cpu' or '0' or '0,1,2,3' # device = 'cpu' or '0' or '0,1,2,3'
s = f'YOLOv3 {git_describe()} torch {torch.__version__} ' # string s = f'YOLOv3 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
cpu = device.lower() == 'cpu' cpu = device.lower() == 'cpu'
if cpu: if cpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
@ -73,7 +82,7 @@ def select_device(device='', batch_size=None):
else: else:
s += 'CPU\n' s += 'CPU\n'
logger.info(s) # skip a line logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
return torch.device('cuda:0' if cuda else 'cpu') return torch.device('cuda:0' if cuda else 'cpu')
@ -120,7 +129,7 @@ def profile(x, ops, n=100, device=None):
s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list' s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
print(f'{p:12.4g}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}') print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
def is_parallel(model): def is_parallel(model):
@ -182,7 +191,7 @@ def fuse_conv_and_bn(conv, bn):
# prepare filters # prepare filters
w_conv = conv.weight.clone().view(conv.out_channels, -1) w_conv = conv.weight.clone().view(conv.out_channels, -1)
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
# prepare spatial bias # prepare spatial bias
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
@ -205,7 +214,7 @@ def model_info(model, verbose=False, img_size=640):
try: # FLOPS try: # FLOPS
from thop import profile from thop import profile
stride = int(model.stride.max()) if hasattr(model, 'stride') else 32 stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS
img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float

0
utils/wandb_logging/__init__.py Normal file
View File

24
utils/wandb_logging/log_dataset.py Normal file
View File

@ -0,0 +1,24 @@
import argparse
import yaml
from wandb_utils import WandbLogger
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
def create_dataset_artifact(opt):
with open(opt.data) as f:
data = yaml.load(f, Loader=yaml.SafeLoader) # data dict
logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
opt = parser.parse_args()
opt.resume = False # Explicitly disallow resume check for dataset upload job
create_dataset_artifact(opt)

306
utils/wandb_logging/wandb_utils.py Normal file
View File

@ -0,0 +1,306 @@
import json
import sys
from pathlib import Path
import torch
import yaml
from tqdm import tqdm
sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
from utils.datasets import LoadImagesAndLabels
from utils.datasets import img2label_paths
from utils.general import colorstr, xywh2xyxy, check_dataset
try:
import wandb
from wandb import init, finish
except ImportError:
wandb = None
WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
return from_string[len(prefix):]
def check_wandb_config_file(data_config_file):
wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
if Path(wandb_config).is_file():
return wandb_config
return data_config_file
def get_run_info(run_path):
run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
run_id = run_path.stem
project = run_path.parent.stem
model_artifact_name = 'run_' + run_id + '_model'
return run_id, project, model_artifact_name
def check_wandb_resume(opt):
process_wandb_config_ddp_mode(opt) if opt.global_rank not in [-1, 0] else None
if isinstance(opt.resume, str):
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
if opt.global_rank not in [-1, 0]: # For resuming DDP runs
run_id, project, model_artifact_name = get_run_info(opt.resume)
api = wandb.Api()
artifact = api.artifact(project + '/' + model_artifact_name + ':latest')
modeldir = artifact.download()
opt.weights = str(Path(modeldir) / "last.pt")
return True
return None
def process_wandb_config_ddp_mode(opt):
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
train_dir, val_dir = None, None
if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
api = wandb.Api()
train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
train_dir = train_artifact.download()
train_path = Path(train_dir) / 'data/images/'
data_dict['train'] = str(train_path)
if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
api = wandb.Api()
val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
val_dir = val_artifact.download()
val_path = Path(val_dir) / 'data/images/'
data_dict['val'] = str(val_path)
if train_dir or val_dir:
ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
with open(ddp_data_path, 'w') as f:
yaml.dump(data_dict, f)
opt.data = ddp_data_path
class WandbLogger():
def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
# Pre-training routine --
self.job_type = job_type
self.wandb, self.wandb_run, self.data_dict = wandb, None if not wandb else wandb.run, data_dict
# It's more elegant to stick to 1 wandb.init call, but useful config data is overwritten in the WandbLogger's wandb.init call
if isinstance(opt.resume, str): # checks resume from artifact
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
run_id, project, model_artifact_name = get_run_info(opt.resume)
model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
assert wandb, 'install wandb to resume wandb runs'
# Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
self.wandb_run = wandb.init(id=run_id, project=project, resume='allow')
opt.resume = model_artifact_name
elif self.wandb:
self.wandb_run = wandb.init(config=opt,
resume="allow",
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
name=name,
job_type=job_type,
id=run_id) if not wandb.run else wandb.run
if self.wandb_run:
if self.job_type == 'Training':
if not opt.resume:
wandb_data_dict = self.check_and_upload_dataset(opt) if opt.upload_dataset else data_dict
# Info useful for resuming from artifacts
self.wandb_run.config.opt = vars(opt)
self.wandb_run.config.data_dict = wandb_data_dict
self.data_dict = self.setup_training(opt, data_dict)
if self.job_type == 'Dataset Creation':
self.data_dict = self.check_and_upload_dataset(opt)
else:
prefix = colorstr('wandb: ')
print(f"{prefix}Install Weights & Biases for YOLOv5 logging with 'pip install wandb' (recommended)")
def check_and_upload_dataset(self, opt):
assert wandb, 'Install wandb to upload dataset'
check_dataset(self.data_dict)
config_path = self.log_dataset_artifact(opt.data,
opt.single_cls,
'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
print("Created dataset config file ", config_path)
with open(config_path) as f:
wandb_data_dict = yaml.load(f, Loader=yaml.SafeLoader)
return wandb_data_dict
def setup_training(self, opt, data_dict):
self.log_dict, self.current_epoch, self.log_imgs = {}, 0, 16 # Logging Constants
self.bbox_interval = opt.bbox_interval
if isinstance(opt.resume, str):
modeldir, _ = self.download_model_artifact(opt)
if modeldir:
self.weights = Path(modeldir) / "last.pt"
config = self.wandb_run.config
opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str(
self.weights), config.save_period, config.total_batch_size, config.bbox_interval, config.epochs, \
config.opt['hyp']
data_dict = dict(self.wandb_run.config.data_dict) # eliminates the need for config file to resume
if 'val_artifact' not in self.__dict__: # If --upload_dataset is set, use the existing artifact, don't download
self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'),
opt.artifact_alias)
self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'),
opt.artifact_alias)
self.result_artifact, self.result_table, self.val_table, self.weights = None, None, None, None
if self.train_artifact_path is not None:
train_path = Path(self.train_artifact_path) / 'data/images/'
data_dict['train'] = str(train_path)
if self.val_artifact_path is not None:
val_path = Path(self.val_artifact_path) / 'data/images/'
data_dict['val'] = str(val_path)
self.val_table = self.val_artifact.get("val")
self.map_val_table_path()
if self.val_artifact is not None:
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
if opt.bbox_interval == -1:
self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
return data_dict
def download_dataset_artifact(self, path, alias):
if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
datadir = dataset_artifact.download()
return datadir, dataset_artifact
return None, None
def download_model_artifact(self, opt):
if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
modeldir = model_artifact.download()
epochs_trained = model_artifact.metadata.get('epochs_trained')
total_epochs = model_artifact.metadata.get('total_epochs')
assert epochs_trained < total_epochs, 'training to %g epochs is finished, nothing to resume.' % (
total_epochs)
return modeldir, model_artifact
return None, None
def log_model(self, path, opt, epoch, fitness_score, best_model=False):
model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={
'original_url': str(path),
'epochs_trained': epoch + 1,
'save period': opt.save_period,
'project': opt.project,
'total_epochs': opt.epochs,
'fitness_score': fitness_score
})
model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
wandb.log_artifact(model_artifact,
aliases=['latest', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
print("Saving model artifact on epoch ", epoch + 1)
def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
with open(data_file) as f:
data = yaml.load(f, Loader=yaml.SafeLoader) # data dict
nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
names = {k: v for k, v in enumerate(names)} # to index dictionary
self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(
data['train']), names, name='train') if data.get('train') else None
self.val_artifact = self.create_dataset_table(LoadImagesAndLabels(
data['val']), names, name='val') if data.get('val') else None
if data.get('train'):
data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
if data.get('val'):
data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
path = data_file if overwrite_config else '_wandb.'.join(data_file.rsplit('.', 1)) # updated data.yaml path
data.pop('download', None)
with open(path, 'w') as f:
yaml.dump(data, f)
if self.job_type == 'Training': # builds correct artifact pipeline graph
self.wandb_run.use_artifact(self.val_artifact)
self.wandb_run.use_artifact(self.train_artifact)
self.val_artifact.wait()
self.val_table = self.val_artifact.get('val')
self.map_val_table_path()
else:
self.wandb_run.log_artifact(self.train_artifact)
self.wandb_run.log_artifact(self.val_artifact)
return path
def map_val_table_path(self):
self.val_table_map = {}
print("Mapping dataset")
for i, data in enumerate(tqdm(self.val_table.data)):
self.val_table_map[data[3]] = data[0]
def create_dataset_table(self, dataset, class_to_id, name='dataset'):
# TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
artifact = wandb.Artifact(name=name, type="dataset")
img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
img_files = tqdm(dataset.img_files) if not img_files else img_files
for img_file in img_files:
if Path(img_file).is_dir():
artifact.add_dir(img_file, name='data/images')
labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
artifact.add_dir(labels_path, name='data/labels')
else:
artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
label_file = Path(img2label_paths([img_file])[0])
artifact.add_file(str(label_file),
name='data/labels/' + label_file.name) if label_file.exists() else None
table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
height, width = shapes[0]
labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4))) * torch.Tensor([width, height, width, height])
box_data, img_classes = [], {}
for cls, *xyxy in labels[:, 1:].tolist():
cls = int(cls)
box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": cls,
"box_caption": "%s" % (class_to_id[cls]),
"scores": {"acc": 1},
"domain": "pixel"})
img_classes[cls] = class_to_id[cls]
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes),
Path(paths).name)
artifact.add(table, name)
return artifact
def log_training_progress(self, predn, path, names):
if self.val_table and self.result_table:
class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
box_data = []
total_conf = 0
for *xyxy, conf, cls in predn.tolist():
if conf >= 0.25:
box_data.append(
{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf),
"scores": {"class_score": conf},
"domain": "pixel"})
total_conf = total_conf + conf
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
id = self.val_table_map[Path(path).name]
self.result_table.add_data(self.current_epoch,
id,
wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
total_conf / max(1, len(box_data))
)
def log(self, log_dict):
if self.wandb_run:
for key, value in log_dict.items():
self.log_dict[key] = value
def end_epoch(self, best_result=False):
if self.wandb_run:
wandb.log(self.log_dict)
self.log_dict = {}
if self.result_artifact:
train_results = wandb.JoinedTable(self.val_table, self.result_table, "id")
self.result_artifact.add(train_results, 'result')
wandb.log_artifact(self.result_artifact, aliases=['latest', 'epoch ' + str(self.current_epoch),
('best' if best_result else '')])
self.result_table = wandb.Table(["epoch", "id", "prediction", "avg_confidence"])
self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
def finish_run(self):
if self.wandb_run:
if self.log_dict:
wandb.log(self.log_dict)
wandb.run.finish()