YOLOv5 v5.0 release compatibility update for YOLOv3 (#1737)
* 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
GitHub
parent
5d8f03020c
commit
8eb4cde090
2
.gitattributes
vendored
2
.gitattributes
vendored
@ -1,2 +1,2 @@
|
||||
# remove notebooks from GitHub language stats
|
||||
# this drop notebooks from GitHub language stats
|
||||
*.ipynb linguist-vendored
|
||||
|
||||
9
.github/workflows/ci-testing.yml
vendored
9
.github/workflows/ci-testing.yml
vendored
@ -66,14 +66,15 @@ jobs:
|
||||
di=cpu # inference devices # define device
|
||||
|
||||
# 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
|
||||
python detect.py --weights weights/${{ matrix.model }}.pt --device $di
|
||||
python detect.py --weights runs/train/exp/weights/last.pt --device $di
|
||||
# test
|
||||
python test.py --img 256 --batch 8 --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 weights/${{ matrix.model }}.pt --device $di
|
||||
python test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
|
||||
|
||||
python hubconf.py # hub
|
||||
python 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
|
||||
|
||||
9
.github/workflows/greetings.yml
vendored
9
.github/workflows/greetings.yml
vendored
@ -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
|
||||
|
||||
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.
|
||||
|
||||
@ -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):
|
||||
|
||||
- **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>
|
||||
- **Kaggle Notebook** with free GPU: [https://www.kaggle.com/ultralytics/yolov3](https://www.kaggle.com/ultralytics/yolov3)
|
||||
- **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
|
||||
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov3. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) 
|
||||
- **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
|
||||
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
|
||||
|
||||
|
||||
## Status
|
||||
|
||||
|
||||
19
Dockerfile
19
Dockerfile
@ -1,14 +1,14 @@
|
||||
# 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
|
||||
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
|
||||
RUN pip install --upgrade pip
|
||||
COPY requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
RUN pip install gsutil
|
||||
RUN python -m pip install --upgrade pip
|
||||
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
|
||||
RUN mkdir -p /usr/src/app
|
||||
@ -17,6 +17,9 @@ WORKDIR /usr/src/app
|
||||
# Copy contents
|
||||
COPY . /usr/src/app
|
||||
|
||||
# Set environment variables
|
||||
ENV HOME=/usr/src/app
|
||||
|
||||
|
||||
# --------------------------------------------------- Extras Below ---------------------------------------------------
|
||||
|
||||
@ -34,13 +37,13 @@ COPY . /usr/src/app
|
||||
# sudo docker kill $(sudo docker ps -q)
|
||||
|
||||
# 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
|
||||
# sudo docker container exec -it ba65811811ab bash
|
||||
# sudo docker exec -it 5a9b5863d93d bash
|
||||
|
||||
# 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
|
||||
# python -c "from utils.general import *; strip_optimizer('runs/train/exp0_*/weights/best.pt', 'tmp.pt')" && gsutil cp tmp.pt gs://*.pt
|
||||
|
||||
116
README.md
116
README.md
@ -1,35 +1,61 @@
|
||||
<a 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>
|
||||
<a align="left" href="https://apps.apple.com/app/id1452689527" target="_blank">
|
||||
<img width="800" src="https://user-images.githubusercontent.com/26833433/99805965-8f2ca800-2b3d-11eb-8fad-13a96b222a23.jpg"></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:
|
||||
* [Master branch](https://github.com/ultralytics/yolov3/tree/master): Forward-compatible with all [YOLOv5](https://github.com/ultralytics/yolov5) models and methods (**recommended**).
|
||||
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.
|
||||
|
||||
<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
|
||||
$ 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
|
||||
$ 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
|
||||
|
||||
| Model | AP<sup>val</sup> | AP<sup>test</sup> | AP<sub>50</sub> | Speed<sub>GPU</sub> | FPS<sub>GPU</sub> || params | FLOPS |
|
||||
|---------- |------ |------ |------ | -------- | ------| ------ |------ | :------: |
|
||||
| [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
|
||||
| [YOLOv3-tiny](https://github.com/ultralytics/yolov3/releases) | 17.6 | 34.9 | 34.9 | **1.7ms** | **588** || 8.9M | 13.3B
|
||||
[assets3]: https://github.com/ultralytics/yolov3/releases
|
||||
[assets5]: https://github.com/ultralytics/yolov5/releases
|
||||
|
||||
Model |size<br><sup>(pixels) |mAP<sup>val<br>0.5:0.95 |mAP<sup>test<br>0.5:0.95 |mAP<sup>val<br>0.5 |Speed<br><sup>V100 (ms) | |params<br><sup>(M) |FLOPS<br><sup>640 (B)
|
||||
--- |--- |--- |--- |--- |--- |---|--- |---
|
||||
[YOLOv3-tiny][assets3] |640 |17.6 |17.6 |34.8 |**1.2** | |8.8 |13.2
|
||||
[YOLOv3][assets3] |640 |43.3 |43.3 |63.0 |4.1 | |61.9 |156.3
|
||||
[YOLOv3-SPP][assets3] |640 |44.3 |44.3 |64.6 |4.1 | |63.0 |157.1
|
||||
| | | | | | || |
|
||||
[YOLOv5l][assets5] |640 |**48.2** |**48.2** |**66.9** |3.7 | |47.0 |115.4
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Table Notes (click to expand)</summary>
|
||||
|
||||
* AP<sup>test</sup> denotes COCO [test-dev2017](http://cocodataset.org/#upload) server results, all other AP results denote val2017 accuracy.
|
||||
* AP values are for single-model single-scale unless otherwise noted. **Reproduce mAP** by `python test.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
|
||||
* Speed<sub>GPU</sub> averaged over 5000 COCO val2017 images using a GCP [n1-standard-16](https://cloud.google.com/compute/docs/machine-types#n1_standard_machine_types) V100 instance, and includes FP16 inference, postprocessing and NMS. **Reproduce speed** by `python test.py --data coco.yaml --img 640 --conf 0.25 --iou 0.45`
|
||||
* All checkpoints are trained to 300 epochs with default settings and hyperparameters (no autoaugmentation).
|
||||
</details>
|
||||
|
||||
** 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
|
||||
|
||||
@ -42,7 +68,9 @@ $ pip install -r requirements.txt
|
||||
## Tutorials
|
||||
|
||||
* [Train Custom Data](https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data) 🚀 RECOMMENDED
|
||||
* [Tips for Best Training Results](https://github.com/ultralytics/yolov5/wiki/Tips-for-Best-Training-Results) ☘️ RECOMMENDED
|
||||
* [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289) 🌟 NEW
|
||||
* [Supervisely Ecosystem](https://github.com/ultralytics/yolov5/issues/2518) 🌟 NEW
|
||||
* [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
|
||||
* [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) ⭐ NEW
|
||||
* [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):
|
||||
|
||||
- **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>
|
||||
- **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)
|
||||
- **Docker Image** https://hub.docker.com/r/ultralytics/yolov3. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) 
|
||||
- **Google Colab and Kaggle** notebooks with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov3/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> <a href="https://www.kaggle.com/ultralytics/yolov3"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
|
||||
- **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/GCP-Quickstart)
|
||||
- **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/AWS-Quickstart)
|
||||
- **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov3/wiki/Docker-Quickstart) <a href="https://hub.docker.com/r/ultralytics/yolov3"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov3?logo=docker" alt="Docker Pulls"></a>
|
||||
|
||||
|
||||
## 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
|
||||
$ python detect.py --source 0 # webcam
|
||||
file.jpg # image
|
||||
file.mp4 # video
|
||||
path/ # directory
|
||||
path/*.jpg # glob
|
||||
rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp stream
|
||||
rtmp://192.168.1.105/live/test # rtmp stream
|
||||
http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8 # http stream
|
||||
'https://youtu.be/NUsoVlDFqZg' # YouTube video
|
||||
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
||||
```
|
||||
|
||||
To run inference on example images in `data/images`:
|
||||
```bash
|
||||
$ 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
|
||||
|
||||
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
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
# 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
|
||||
img1 = Image.open('zidane.jpg')
|
||||
img2 = Image.open('bus.jpg')
|
||||
imgs = [img1, img2] # batched list of images
|
||||
dir = 'https://github.com/ultralytics/yolov3/raw/master/data/images/'
|
||||
imgs = [dir + f for f in ('zidane.jpg', 'bus.jpg')] # batch of images
|
||||
|
||||
# Inference
|
||||
prediction = model(imgs, size=640) # includes NMS
|
||||
results = model(imgs)
|
||||
results.print() # or .show(), .save()
|
||||
```
|
||||
|
||||
|
||||
## 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
|
||||
$ python train.py --data coco.yaml --cfg yolov3.yaml --weights '' --batch-size 24
|
||||
yolov3-spp.yaml 24
|
||||
yolov3-tiny.yaml 64
|
||||
$ python train.py --data coco.yaml --cfg yolov3.yaml --weights '' --batch-size 24
|
||||
yolov3-spp.yaml 24
|
||||
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
|
||||
|
||||
21
data/argoverse_hd.yaml
Normal file
21
data/argoverse_hd.yaml
Normal 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' ]
|
||||
@ -18,15 +18,15 @@ test: ../coco/test-dev2017.txt # 20288 of 40670 images, submit to https://compe
|
||||
nc: 80
|
||||
|
||||
# class names
|
||||
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||
'hair drier', 'toothbrush']
|
||||
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||
'hair drier', 'toothbrush' ]
|
||||
|
||||
# Print classes
|
||||
# with open('data/coco.yaml') as f:
|
||||
|
||||
@ -17,12 +17,12 @@ val: ../coco128/images/train2017/ # 128 images
|
||||
nc: 80
|
||||
|
||||
# class names
|
||||
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||
'hair drier', 'toothbrush']
|
||||
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
||||
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||||
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
||||
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||||
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
||||
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
||||
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
||||
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||||
'hair drier', 'toothbrush' ]
|
||||
|
||||
62
data/scripts/get_argoverse_hd.sh
Normal file
62
data/scripts/get_argoverse_hd.sh
Normal 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/
|
||||
@ -10,8 +10,9 @@
|
||||
# Download/unzip labels
|
||||
d='../' # unzip directory
|
||||
url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
||||
f='coco2017labels.zip' # 68 MB
|
||||
echo 'Downloading' $url$f ' ...' && curl -L $url$f -o $f && unzip -q $f -d $d && rm $f # download, unzip, remove
|
||||
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 in background
|
||||
|
||||
# Download/unzip images
|
||||
d='../coco/images' # unzip directory
|
||||
@ -20,5 +21,7 @@ f1='train2017.zip' # 19G, 118k images
|
||||
f2='val2017.zip' # 1G, 5k images
|
||||
f3='test2017.zip' # 7G, 41k images (optional)
|
||||
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
|
||||
wait # finish background tasks
|
||||
|
||||
@ -17,9 +17,11 @@ url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
|
||||
f1=VOCtrainval_06-Nov-2007.zip # 446MB, 5012 images
|
||||
f2=VOCtest_06-Nov-2007.zip # 438MB, 4953 images
|
||||
f3=VOCtrainval_11-May-2012.zip # 1.95GB, 17126 images
|
||||
for f in $f1 $f2 $f3; do
|
||||
echo 'Downloading' $url$f ' ...' && curl -L $url$f -o $f && unzip -q $f -d $d && rm $f # download, unzip, remove
|
||||
for f in $f3 $f2 $f1; do
|
||||
echo 'Downloading' $url$f '...'
|
||||
curl -L $url$f -o $f && unzip -q $f -d $d && rm $f & # download, unzip, remove in background
|
||||
done
|
||||
wait # finish background tasks
|
||||
|
||||
end=$(date +%s)
|
||||
runtime=$((end - start))
|
||||
|
||||
@ -17,5 +17,5 @@ val: ../VOC/images/val/ # 4952 images
|
||||
nc: 20
|
||||
|
||||
# class names
|
||||
names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
|
||||
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
|
||||
names: [ 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
|
||||
'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]
|
||||
|
||||
43
detect.py
43
detect.py
@ -9,16 +9,17 @@ from numpy import random
|
||||
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import LoadStreams, LoadImages
|
||||
from utils.general import check_img_size, check_requirements, non_max_suppression, apply_classifier, scale_coords, \
|
||||
xyxy2xywh, strip_optimizer, set_logging, increment_path
|
||||
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
|
||||
scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
|
||||
from utils.plots import plot_one_box
|
||||
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
||||
|
||||
|
||||
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
|
||||
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
|
||||
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
|
||||
('rtsp://', 'rtmp://', 'http://'))
|
||||
('rtsp://', 'rtmp://', 'http://', 'https://'))
|
||||
|
||||
# Directories
|
||||
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
|
||||
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:
|
||||
model.half() # to FP16
|
||||
|
||||
@ -44,21 +46,20 @@ def detect(save_img=False):
|
||||
# Set Dataloader
|
||||
vid_path, vid_writer = None, None
|
||||
if webcam:
|
||||
view_img = True
|
||||
view_img = check_imshow()
|
||||
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:
|
||||
save_img = True
|
||||
dataset = LoadImages(source, img_size=imgsz)
|
||||
dataset = LoadImages(source, img_size=imgsz, stride=stride)
|
||||
|
||||
# Get names and colors
|
||||
names = model.module.names if hasattr(model, 'module') else model.names
|
||||
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
|
||||
|
||||
# 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()
|
||||
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:
|
||||
img = torch.from_numpy(img).to(device)
|
||||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||||
@ -97,7 +98,7 @@ def detect(save_img=False):
|
||||
# Print results
|
||||
for c in det[:, -1].unique():
|
||||
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
|
||||
for *xyxy, conf, cls in reversed(det):
|
||||
@ -117,22 +118,25 @@ def detect(save_img=False):
|
||||
# Stream results
|
||||
if view_img:
|
||||
cv2.imshow(str(p), im0)
|
||||
cv2.waitKey(1) # 1 millisecond
|
||||
|
||||
# Save results (image with detections)
|
||||
if save_img:
|
||||
if dataset.mode == 'image':
|
||||
cv2.imwrite(save_path, im0)
|
||||
else: # 'video'
|
||||
else: # 'video' or 'stream'
|
||||
if vid_path != save_path: # new video
|
||||
vid_path = save_path
|
||||
if isinstance(vid_writer, cv2.VideoWriter):
|
||||
vid_writer.release() # release previous video writer
|
||||
|
||||
fourcc = 'mp4v' # output video codec
|
||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
|
||||
if vid_cap: # video
|
||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
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)
|
||||
|
||||
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('--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('--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('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
||||
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')
|
||||
opt = parser.parse_args()
|
||||
print(opt)
|
||||
check_requirements()
|
||||
check_requirements(exclude=('pycocotools', 'thop'))
|
||||
|
||||
with torch.no_grad():
|
||||
if opt.update: # update all models (to fix SourceChangeWarning)
|
||||
|
||||
101
hubconf.py
101
hubconf.py
@ -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:
|
||||
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
|
||||
@ -10,10 +10,12 @@ from pathlib import Path
|
||||
import torch
|
||||
|
||||
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.torch_utils import select_device
|
||||
|
||||
dependencies = ['torch', 'yaml']
|
||||
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
|
||||
set_logging()
|
||||
|
||||
|
||||
@ -21,7 +23,7 @@ def create(name, pretrained, channels, classes, autoshape):
|
||||
"""Creates a specified YOLOv3 model
|
||||
|
||||
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
|
||||
channels (int): number of input channels
|
||||
classes (int): number of model classes
|
||||
@ -29,21 +31,23 @@ def create(name, pretrained, channels, classes, autoshape):
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
config = Path(__file__).parent / 'models' / f'{name}.yaml' # model.yaml path
|
||||
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:
|
||||
fname = f'{name}.pt' # checkpoint filename
|
||||
attempt_download(fname) # download if not found locally
|
||||
ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
|
||||
state_dict = ckpt['model'].float().state_dict() # to FP32
|
||||
state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
|
||||
model.load_state_dict(state_dict, strict=False) # load
|
||||
msd = model.state_dict() # model state_dict
|
||||
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
|
||||
csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter
|
||||
model.load_state_dict(csd, strict=False) # load
|
||||
if len(ckpt['model'].names) == classes:
|
||||
model.names = ckpt['model'].names # set class names attribute
|
||||
if autoshape:
|
||||
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:
|
||||
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
|
||||
|
||||
|
||||
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):
|
||||
"""YOLOv3-custom model from https://github.com/ultralytics/yolov3
|
||||
"""YOLOv3-custom model https://github.com/ultralytics/yolov3
|
||||
|
||||
Arguments (3 options):
|
||||
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
|
||||
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.load_state_dict(model.float().state_dict()) # load state_dict
|
||||
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__':
|
||||
@ -119,9 +99,14 @@ if __name__ == '__main__':
|
||||
# model = custom(path_or_model='path/to/model.pt') # custom example
|
||||
|
||||
# Verify inference
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
imgs = [Image.open(x) for x in Path('data/images').glob('*.jpg')]
|
||||
results = model(imgs)
|
||||
results.show()
|
||||
imgs = [Image.open('data/images/bus.jpg'), # PIL
|
||||
'data/images/zidane.jpg', # filename
|
||||
'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.save()
|
||||
|
||||
163
models/common.py
163
models/common.py
@ -1,16 +1,21 @@
|
||||
# This file contains modules common to various models
|
||||
# YOLOv3 common modules
|
||||
|
||||
import math
|
||||
from copy import copy
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import requests
|
||||
import torch
|
||||
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.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
|
||||
from utils.plots import color_list
|
||||
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh
|
||||
from utils.plots import color_list, plot_one_box
|
||||
from utils.torch_utils import time_synchronized
|
||||
|
||||
|
||||
def autopad(k, p=None): # kernel, padding
|
||||
@ -40,6 +45,52 @@ class Conv(nn.Module):
|
||||
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):
|
||||
# Standard bottleneck
|
||||
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))
|
||||
|
||||
|
||||
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):
|
||||
# Spatial pyramid pooling layer used in YOLOv3-SPP
|
||||
def __init__(self, c1, c2, k=(5, 9, 13)):
|
||||
@ -166,7 +225,6 @@ class NMS(nn.Module):
|
||||
|
||||
class autoShape(nn.Module):
|
||||
# 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
|
||||
iou = 0.45 # NMS IoU threshold
|
||||
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()
|
||||
return self
|
||||
|
||||
@torch.no_grad()
|
||||
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'
|
||||
# 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)
|
||||
# PIL: = Image.open('image.jpg') # HWC x(720,1280,3)
|
||||
# numpy: = np.zeros((720,1280,3)) # HWC
|
||||
# torch: = torch.zeros(16,3,720,1280) # BCHW
|
||||
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
||||
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
|
||||
# numpy: = np.zeros((640,1280,3)) # HWC
|
||||
# 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
|
||||
|
||||
t = [time_synchronized()]
|
||||
p = next(self.model.parameters()) # for device and type
|
||||
if isinstance(imgs, torch.Tensor): # torch
|
||||
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
|
||||
with amp.autocast(enabled=p.device.type != 'cpu'):
|
||||
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
|
||||
|
||||
# Pre-process
|
||||
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):
|
||||
f = f'image{i}' # filename
|
||||
if isinstance(im, str): # filename or uri
|
||||
im = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im) # open
|
||||
im = np.array(im) # to numpy
|
||||
im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
|
||||
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
|
||||
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
|
||||
@ -213,82 +277,101 @@ class autoShape(nn.Module):
|
||||
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 = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
|
||||
t.append(time_synchronized())
|
||||
|
||||
# Inference
|
||||
with torch.no_grad():
|
||||
with amp.autocast(enabled=p.device.type != 'cpu'):
|
||||
# Inference
|
||||
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
|
||||
for i in range(n):
|
||||
scale_coords(shape1, y[i][:, :4], shape0[i])
|
||||
# Post-process
|
||||
y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
|
||||
for i in range(n):
|
||||
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:
|
||||
# detections class for YOLOv5 inference results
|
||||
def __init__(self, imgs, pred, names=None):
|
||||
# detections class for YOLOv3 inference results
|
||||
def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
|
||||
super(Detections, self).__init__()
|
||||
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
|
||||
self.imgs = imgs # list of images as numpy arrays
|
||||
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
|
||||
self.names = names # class names
|
||||
self.files = files # image filenames
|
||||
self.xyxy = pred # xyxy 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.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()
|
||||
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:
|
||||
for c in pred[:, -1].unique():
|
||||
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:
|
||||
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
|
||||
# str += '%s %.2f, ' % (names[int(cls)], conf) # label
|
||||
ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
|
||||
label = f'{self.names[int(cls)]} {conf:.2f}'
|
||||
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:
|
||||
print(str)
|
||||
print(str.rstrip(', '))
|
||||
if show:
|
||||
img.show(f'Image {i}') # show
|
||||
img.show(self.files[i]) # show
|
||||
if save:
|
||||
f = f'results{i}.jpg'
|
||||
str += f"saved to '{f}'"
|
||||
img.save(f) # save
|
||||
f = self.files[i]
|
||||
img.save(Path(save_dir) / f) # save
|
||||
print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
|
||||
if render:
|
||||
self.imgs[i] = np.asarray(img)
|
||||
|
||||
def print(self):
|
||||
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):
|
||||
self.display(show=True) # show results
|
||||
|
||||
def save(self):
|
||||
self.display(save=True) # save results
|
||||
def save(self, save_dir='runs/hub/exp'):
|
||||
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):
|
||||
self.display(render=True) # render results
|
||||
return self.imgs
|
||||
|
||||
def __len__(self):
|
||||
return self.n
|
||||
def pandas(self):
|
||||
# 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):
|
||||
# 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 k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
|
||||
setattr(d, k, getattr(d, k)[0]) # pop out of list
|
||||
return x
|
||||
|
||||
def __len__(self):
|
||||
return self.n
|
||||
|
||||
|
||||
class Classify(nn.Module):
|
||||
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
# This file contains experimental modules
|
||||
# YOLOv3 experimental modules
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@ -58,7 +58,7 @@ class GhostConv(nn.Module):
|
||||
|
||||
class GhostBottleneck(nn.Module):
|
||||
# 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__()
|
||||
c_ = c2 // 2
|
||||
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
|
||||
@ -115,11 +115,12 @@ def attempt_load(weights, map_location=None):
|
||||
model = Ensemble()
|
||||
for w in weights if isinstance(weights, list) else [weights]:
|
||||
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
|
||||
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
|
||||
elif type(m) is Conv:
|
||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
|
||||
|
||||
@ -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:
|
||||
$ 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 utils.activations import Hardswish, SiLU
|
||||
from utils.general import set_logging, check_img_size
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', type=str, default='./yolov3.pt', help='weights path') # from yolov3/models/
|
||||
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('--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.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
|
||||
print(opt)
|
||||
@ -30,7 +34,8 @@ if __name__ == '__main__':
|
||||
t = time.time()
|
||||
|
||||
# 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
|
||||
|
||||
# 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
|
||||
|
||||
# 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
|
||||
for k, m in model.named_modules():
|
||||
@ -50,14 +55,14 @@ if __name__ == '__main__':
|
||||
m.act = SiLU()
|
||||
# elif isinstance(m, models.yolo.Detect):
|
||||
# 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
|
||||
|
||||
# TorchScript export
|
||||
try:
|
||||
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
||||
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)
|
||||
print('TorchScript export success, saved as %s' % f)
|
||||
except Exception as e:
|
||||
@ -70,7 +75,9 @@ if __name__ == '__main__':
|
||||
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
|
||||
f = opt.weights.replace('.pt', '.onnx') # filename
|
||||
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
|
||||
onnx_model = onnx.load(f) # load onnx model
|
||||
|
||||
@ -1,14 +1,15 @@
|
||||
# YOLOv3 YOLO-specific modules
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from models.common import *
|
||||
from models.experimental import MixConv2d, CrossConv
|
||||
from models.experimental import *
|
||||
from utils.autoanchor import check_anchor_order
|
||||
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, \
|
||||
@ -50,7 +51,7 @@ class Detect(nn.Module):
|
||||
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
|
||||
|
||||
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
|
||||
z.append(y.view(bs, -1, self.no))
|
||||
|
||||
@ -63,7 +64,7 @@ class Detect(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__()
|
||||
if isinstance(cfg, dict):
|
||||
self.yaml = cfg # model dict
|
||||
@ -71,13 +72,16 @@ class Model(nn.Module):
|
||||
import yaml # for torch hub
|
||||
self.yaml_file = Path(cfg).name
|
||||
with open(cfg) as f:
|
||||
self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict
|
||||
|
||||
# Define model
|
||||
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
||||
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
|
||||
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.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))])
|
||||
@ -107,7 +111,7 @@ class Model(nn.Module):
|
||||
for si, fi in zip(s, f):
|
||||
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
|
||||
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
|
||||
if fi == 2:
|
||||
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
|
||||
|
||||
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]
|
||||
|
||||
# Normal
|
||||
# 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
|
||||
if c2 != no: # if not output
|
||||
c2 = make_divisible(c2 * gw, 8)
|
||||
|
||||
args = [c1, c2, *args[1:]]
|
||||
if m in [BottleneckCSP, C3]:
|
||||
args.insert(2, n)
|
||||
if m in [BottleneckCSP, C3, C3TR]:
|
||||
args.insert(2, n) # number of repeats
|
||||
n = 1
|
||||
elif m is nn.BatchNorm2d:
|
||||
args = [ch[f]]
|
||||
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:
|
||||
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
|
||||
args[1] = [list(range(args[1] * 2))] * len(f)
|
||||
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:
|
||||
c2 = ch[f if f < 0 else f + 1] // args[0] ** 2
|
||||
c2 = ch[f] // args[0] ** 2
|
||||
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
|
||||
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
|
||||
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
||||
layers.append(m_)
|
||||
if i == 0:
|
||||
ch = []
|
||||
ch.append(c2)
|
||||
return nn.Sequential(*layers), sorted(save)
|
||||
|
||||
|
||||
@ -1,19 +1,18 @@
|
||||
# pip install -r requirements.txt
|
||||
|
||||
# base ----------------------------------------
|
||||
Cython
|
||||
matplotlib>=3.2.2
|
||||
numpy>=1.18.5
|
||||
opencv-python>=4.1.2
|
||||
Pillow
|
||||
PyYAML>=5.3
|
||||
PyYAML>=5.3.1
|
||||
scipy>=1.4.1
|
||||
tensorboard>=2.2
|
||||
torch>=1.7.0
|
||||
torchvision>=0.8.1
|
||||
tqdm>=4.41.0
|
||||
|
||||
# logging -------------------------------------
|
||||
tensorboard>=2.4.1
|
||||
# wandb
|
||||
|
||||
# plotting ------------------------------------
|
||||
@ -21,8 +20,8 @@ seaborn>=0.11.0
|
||||
pandas
|
||||
|
||||
# export --------------------------------------
|
||||
# coremltools==4.0
|
||||
# onnx>=1.8.0
|
||||
# coremltools>=4.1
|
||||
# onnx>=1.8.1
|
||||
# scikit-learn==0.19.2 # for coreml quantization
|
||||
|
||||
# extras --------------------------------------
|
||||
|
||||
111
test.py
111
test.py
@ -13,7 +13,6 @@ from models.experimental import attempt_load
|
||||
from utils.datasets import create_dataloader
|
||||
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
|
||||
box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
|
||||
from utils.loss import compute_loss
|
||||
from utils.metrics import ap_per_class, ConfusionMatrix
|
||||
from utils.plots import plot_images, output_to_target, plot_study_txt
|
||||
from utils.torch_utils import select_device, time_synchronized
|
||||
@ -36,8 +35,10 @@ def test(data,
|
||||
save_hybrid=False, # for hybrid auto-labelling
|
||||
save_conf=False, # save auto-label confidences
|
||||
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
|
||||
training = model is not None
|
||||
if training: # called by train.py
|
||||
@ -53,47 +54,46 @@ def test(data,
|
||||
|
||||
# Load 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
|
||||
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
||||
# model = nn.DataParallel(model)
|
||||
|
||||
# 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:
|
||||
model.half()
|
||||
|
||||
# Configure
|
||||
model.eval()
|
||||
is_coco = data.endswith('coco.yaml') # is COCO dataset
|
||||
with open(data) as f:
|
||||
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
if isinstance(data, str):
|
||||
is_coco = data.endswith('coco.yaml')
|
||||
with open(data) as f:
|
||||
data = yaml.load(f, Loader=yaml.SafeLoader)
|
||||
check_dataset(data) # check
|
||||
nc = 1 if single_cls else int(data['nc']) # number of classes
|
||||
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
||||
niou = iouv.numel()
|
||||
|
||||
# Logging
|
||||
log_imgs, 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
|
||||
if not training:
|
||||
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
|
||||
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
|
||||
dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True,
|
||||
prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
|
||||
if device.type != 'cpu':
|
||||
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
|
||||
task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images
|
||||
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True,
|
||||
prefix=colorstr(f'{task}: '))[0]
|
||||
|
||||
seen = 0
|
||||
confusion_matrix = ConfusionMatrix(nc=nc)
|
||||
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
|
||||
coco91class = coco80_to_coco91_class()
|
||||
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', '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.
|
||||
loss = torch.zeros(3, device=device)
|
||||
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
|
||||
@ -107,22 +107,22 @@ def test(data,
|
||||
with torch.no_grad():
|
||||
# Run model
|
||||
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
|
||||
|
||||
# Compute loss
|
||||
if training:
|
||||
loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # box, obj, cls
|
||||
if compute_loss:
|
||||
loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
|
||||
|
||||
# Run NMS
|
||||
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
||||
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
||||
t = time_synchronized()
|
||||
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
|
||||
|
||||
# Statistics per image
|
||||
for si, pred in enumerate(output):
|
||||
for si, pred in enumerate(out):
|
||||
labels = targets[targets[:, 0] == si, 1:]
|
||||
nl = len(labels)
|
||||
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:
|
||||
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
||||
|
||||
# W&B logging
|
||||
if plots and len(wandb_images) < log_imgs:
|
||||
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
|
||||
# W&B logging - Media Panel Plots
|
||||
if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0: # Check for test operation
|
||||
if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
|
||||
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
|
||||
"class_id": int(cls),
|
||||
"box_caption": "%s %.3f" % (names[cls], conf),
|
||||
"scores": {"class_score": conf},
|
||||
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
||||
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
||||
wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))
|
||||
wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None
|
||||
|
||||
# Append to pycocotools JSON dictionary
|
||||
if save_json:
|
||||
@ -179,7 +181,7 @@ def test(data,
|
||||
tbox = xywh2xyxy(labels[:, 1:5])
|
||||
scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
|
||||
if plots:
|
||||
confusion_matrix.process_batch(pred, torch.cat((labels[:, 0:1], tbox), 1))
|
||||
confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
|
||||
|
||||
# Per target class
|
||||
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
|
||||
Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()
|
||||
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
|
||||
Thread(target=plot_images, args=(img, output_to_target(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
|
||||
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
|
||||
if len(stats) and stats[0].any():
|
||||
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
|
||||
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()
|
||||
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
|
||||
else:
|
||||
nt = torch.zeros(1)
|
||||
|
||||
# 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 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):
|
||||
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
|
||||
|
||||
@ -239,9 +241,11 @@ def test(data,
|
||||
# Plots
|
||||
if plots:
|
||||
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
|
||||
if wandb and wandb.run:
|
||||
wandb.log({"Images": wandb_images})
|
||||
wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]})
|
||||
if wandb_logger and wandb_logger.wandb:
|
||||
val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
|
||||
wandb_logger.log({"Validation": val_batches})
|
||||
if wandb_images:
|
||||
wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
|
||||
|
||||
# Save JSON
|
||||
if save_json and len(jdict):
|
||||
@ -269,10 +273,10 @@ def test(data,
|
||||
print(f'pycocotools unable to run: {e}')
|
||||
|
||||
# Return results
|
||||
model.float() # for training
|
||||
if not training:
|
||||
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
||||
print(f"Results saved to {save_dir}{s}")
|
||||
model.float() # for training
|
||||
maps = np.zeros(nc) + map
|
||||
for i, c in enumerate(ap_class):
|
||||
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('--conf-thres', type=float, default=0.001, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
|
||||
parser.add_argument('--task', default='val', help="'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('--single-cls', action='store_true', help='treat as single-class dataset')
|
||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||
@ -305,7 +309,7 @@ if __name__ == '__main__':
|
||||
print(opt)
|
||||
check_requirements()
|
||||
|
||||
if opt.task in ['val', 'test']: # run normally
|
||||
if opt.task in ('train', 'val', 'test'): # run normally
|
||||
test(opt.data,
|
||||
opt.weights,
|
||||
opt.batch_size,
|
||||
@ -321,16 +325,21 @@ if __name__ == '__main__':
|
||||
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
|
||||
for weights in ['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(320, 800, 64)) # x axis
|
||||
# python test.py --task study --data coco.yaml --iou 0.7 --weights yolov3.pt yolov3-spp.pt yolov3-tiny.pt
|
||||
x = list(range(256, 1536 + 128, 128)) # x axis (image sizes)
|
||||
for w in opt.weights:
|
||||
f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt' # filename to save to
|
||||
y = [] # y axis
|
||||
for i in x: # img-size
|
||||
print('\nRunning %s point %s...' % (f, i))
|
||||
r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
|
||||
print(f'\nRunning {f} point {i}...')
|
||||
r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
|
||||
plots=False)
|
||||
y.append(r + t) # results and times
|
||||
np.savetxt(f, y, fmt='%10.4g') # save
|
||||
os.system('zip -r study.zip study_*.txt')
|
||||
plot_study_txt(f, x) # plot
|
||||
plot_study_txt(x=x) # plot
|
||||
|
||||
225
train.py
225
train.py
@ -4,6 +4,7 @@ import math
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
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, \
|
||||
check_requirements, print_mutation, set_logging, one_cycle, colorstr
|
||||
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.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__)
|
||||
|
||||
|
||||
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()))
|
||||
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
|
||||
@ -60,10 +62,19 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
|
||||
init_seeds(2 + rank)
|
||||
with open(opt.data) as f:
|
||||
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
|
||||
with torch_distributed_zero_first(rank):
|
||||
check_dataset(data_dict) # check
|
||||
train_path = data_dict['train']
|
||||
test_path = data_dict['val']
|
||||
is_coco = opt.data.endswith('coco.yaml')
|
||||
|
||||
# Logging- Doing this before checking the dataset. Might update data_dict
|
||||
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
|
||||
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
|
||||
@ -74,16 +85,18 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
|
||||
with torch_distributed_zero_first(rank):
|
||||
attempt_download(weights) # download if not found locally
|
||||
ckpt = torch.load(weights, map_location=device) # load checkpoint
|
||||
if hyp.get('anchors'):
|
||||
ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor
|
||||
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
|
||||
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
|
||||
exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys
|
||||
state_dict = ckpt['model'].float().state_dict() # to FP32
|
||||
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect
|
||||
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
|
||||
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 = [] # 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
|
||||
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
|
||||
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
|
||||
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']
|
||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
|
||||
# plot_lr_scheduler(optimizer, scheduler, epochs)
|
||||
|
||||
# Logging
|
||||
if rank in [-1, 0] and wandb and wandb.run is 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
|
||||
# EMA
|
||||
ema = ModelEMA(model) if rank in [-1, 0] else None
|
||||
|
||||
# Resume
|
||||
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'])
|
||||
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
|
||||
if ckpt.get('training_results') is not None:
|
||||
with open(results_file, 'w') as file:
|
||||
file.write(ckpt['training_results']) # write results.txt
|
||||
results_file.write_text(ckpt['training_results']) # write results.txt
|
||||
|
||||
# Epochs
|
||||
start_epoch = ckpt['epoch'] + 1
|
||||
@ -158,7 +172,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
|
||||
del ckpt, state_dict
|
||||
|
||||
# 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'])
|
||||
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)
|
||||
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
|
||||
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
|
||||
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
|
||||
if rank in [-1, 0]:
|
||||
ema.updates = start_epoch * nb // accumulate # set EMA updates
|
||||
testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, # testloader
|
||||
testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader
|
||||
hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
|
||||
world_size=opt.world_size, workers=opt.workers,
|
||||
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
|
||||
# model._initialize_biases(cf.to(device))
|
||||
if plots:
|
||||
plot_labels(labels, save_dir, loggers)
|
||||
plot_labels(labels, names, save_dir, loggers)
|
||||
if tb_writer:
|
||||
tb_writer.add_histogram('classes', c, 0)
|
||||
|
||||
# Anchors
|
||||
if not opt.noautoanchor:
|
||||
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
|
||||
hyp['box'] *= 3. / nl # scale to layers
|
||||
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
|
||||
hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers
|
||||
hyp['label_smoothing'] = opt.label_smoothing
|
||||
model.nc = nc # attach number of classes to model
|
||||
model.hyp = hyp # attach hyperparameters to model
|
||||
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)
|
||||
scheduler.last_epoch = start_epoch - 1 # do not move
|
||||
scaler = amp.GradScaler(enabled=cuda)
|
||||
compute_loss = ComputeLoss(model) # init loss class
|
||||
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
|
||||
f'Using {dataloader.num_workers} dataloader workers\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:
|
||||
dataloader.sampler.set_epoch(epoch)
|
||||
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]:
|
||||
pbar = tqdm(pbar, total=nb) # progress bar
|
||||
optimizer.zero_grad()
|
||||
@ -286,7 +301,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
|
||||
# Forward
|
||||
with amp.autocast(enabled=cuda):
|
||||
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:
|
||||
loss *= opt.world_size # gradient averaged between devices in DDP mode
|
||||
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()
|
||||
# if tb_writer:
|
||||
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
|
||||
# tb_writer.add_graph(model, imgs) # add model to tensorboard
|
||||
elif plots and ni == 3 and wandb:
|
||||
wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')]})
|
||||
# tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
|
||||
elif plots and ni == 10 and wandb_logger.wandb:
|
||||
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
|
||||
save_dir.glob('train*.jpg') if x.exists()]})
|
||||
|
||||
# end batch ------------------------------------------------------------------------------------------------
|
||||
# end epoch ----------------------------------------------------------------------------------------------------
|
||||
@ -331,23 +347,26 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
|
||||
# DDP process 0 or single-GPU
|
||||
if rank in [-1, 0]:
|
||||
# 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
|
||||
if not opt.notest or final_epoch: # Calculate mAP
|
||||
results, maps, times = test.test(opt.data,
|
||||
batch_size=total_batch_size,
|
||||
wandb_logger.current_epoch = epoch + 1
|
||||
results, maps, times = test.test(data_dict,
|
||||
batch_size=batch_size * 2,
|
||||
imgsz=imgsz_test,
|
||||
model=ema.ema,
|
||||
single_cls=opt.single_cls,
|
||||
dataloader=testloader,
|
||||
save_dir=save_dir,
|
||||
verbose=nc < 50 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
|
||||
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:
|
||||
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):
|
||||
if tb_writer:
|
||||
tb_writer.add_scalar(tag, x, epoch) # tensorboard
|
||||
if wandb:
|
||||
wandb.log({tag: x}) # W&B
|
||||
if wandb_logger.wandb:
|
||||
wandb_logger.log({tag: x}) # W&B
|
||||
|
||||
# Update best mAP
|
||||
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
|
||||
if fi > best_fitness:
|
||||
best_fitness = fi
|
||||
wandb_logger.end_epoch(best_result=best_fitness == fi)
|
||||
|
||||
# Save model
|
||||
save = (not opt.nosave) or (final_epoch and not opt.evolve)
|
||||
if save:
|
||||
with open(results_file, 'r') as f: # create checkpoint
|
||||
ckpt = {'epoch': epoch,
|
||||
'best_fitness': best_fitness,
|
||||
'training_results': f.read(),
|
||||
'model': ema.ema,
|
||||
'optimizer': None if final_epoch else optimizer.state_dict(),
|
||||
'wandb_id': wandb_run.id if wandb else None}
|
||||
if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
|
||||
ckpt = {'epoch': epoch,
|
||||
'best_fitness': best_fitness,
|
||||
'training_results': results_file.read_text(),
|
||||
'model': deepcopy(model.module if is_parallel(model) else model).half(),
|
||||
'ema': deepcopy(ema.ema).half(),
|
||||
'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
|
||||
torch.save(ckpt, last)
|
||||
if best_fitness == fi:
|
||||
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
|
||||
|
||||
# end epoch ----------------------------------------------------------------------------------------------------
|
||||
# end training
|
||||
|
||||
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
|
||||
final = best if best.exists() else last # final model
|
||||
for f in [last, best]:
|
||||
for f in last, best:
|
||||
if f.exists():
|
||||
strip_optimizer(f) # strip optimizers
|
||||
if opt.bucket:
|
||||
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
|
||||
|
||||
# Plots
|
||||
if plots:
|
||||
plot_results(save_dir=save_dir) # save as results.png
|
||||
if wandb:
|
||||
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)
|
||||
|
||||
if wandb_logger.wandb and not opt.evolve: # Log the stripped model
|
||||
wandb_logger.wandb.log_artifact(str(final), type='model',
|
||||
name='run_' + wandb_logger.wandb_run.id + '_model',
|
||||
aliases=['last', 'best', 'stripped'])
|
||||
wandb_logger.finish_run()
|
||||
else:
|
||||
dist.destroy_process_group()
|
||||
|
||||
wandb.run.finish() if wandb and wandb.run else None
|
||||
torch.cuda.empty_cache()
|
||||
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('--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('--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('--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('--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('--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()
|
||||
|
||||
# Set DDP variables
|
||||
@ -471,13 +500,14 @@ if __name__ == '__main__':
|
||||
check_requirements()
|
||||
|
||||
# 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
|
||||
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
|
||||
apriori = opt.global_rank, opt.local_rank
|
||||
with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
|
||||
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)
|
||||
else:
|
||||
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
|
||||
@ -504,18 +534,13 @@ if __name__ == '__main__':
|
||||
|
||||
# Train
|
||||
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:
|
||||
tb_writer = None # init loggers
|
||||
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
|
||||
train(hyp, opt, device, tb_writer, wandb)
|
||||
train(hyp, opt, device, tb_writer)
|
||||
|
||||
# Evolve hyperparameters (optional)
|
||||
else:
|
||||
@ -589,7 +614,7 @@ if __name__ == '__main__':
|
||||
hyp[k] = round(hyp[k], 5) # significant digits
|
||||
|
||||
# Train mutation
|
||||
results = train(hyp.copy(), opt, device, wandb=wandb)
|
||||
results = train(hyp.copy(), opt, device)
|
||||
|
||||
# Write mutation results
|
||||
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
|
||||
|
||||
450
tutorial.ipynb
vendored
450
tutorial.ipynb
vendored
File diff suppressed because one or more lines are too long
@ -37,17 +37,21 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
||||
bpr = (best > 1. / thr).float().mean() # best possible recall
|
||||
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='')
|
||||
if bpr < 0.98: # threshold to recompute
|
||||
print('. Attempting to improve anchors, please wait...')
|
||||
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)
|
||||
new_bpr = metric(new_anchors.reshape(-1, 2))[0]
|
||||
try:
|
||||
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
|
||||
new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors)
|
||||
m.anchor_grid[:] = new_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
|
||||
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
||||
m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
|
||||
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
||||
check_anchor_order(m)
|
||||
print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
||||
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
|
||||
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
|
||||
dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
|
||||
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...')
|
||||
s = wh.std(0) # sigmas for whitening
|
||||
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
|
||||
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
||||
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
||||
|
||||
0
utils/aws/__init__.py
Normal file
0
utils/aws/__init__.py
Normal file
26
utils/aws/mime.sh
Normal file
26
utils/aws/mime.sh
Normal 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
37
utils/aws/resume.py
Normal 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
27
utils/aws/userdata.sh
Normal 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
|
||||
@ -20,12 +20,13 @@ from PIL import Image, ExifTags
|
||||
from torch.utils.data import Dataset
|
||||
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
|
||||
|
||||
# Parameters
|
||||
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
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -119,9 +120,8 @@ class _RepeatSampler(object):
|
||||
|
||||
|
||||
class LoadImages: # for inference
|
||||
def __init__(self, path, img_size=640):
|
||||
p = str(Path(path)) # os-agnostic
|
||||
p = os.path.abspath(p) # absolute path
|
||||
def __init__(self, path, img_size=640, stride=32):
|
||||
p = str(Path(path).absolute()) # os-agnostic absolute path
|
||||
if '*' in p:
|
||||
files = sorted(glob.glob(p, recursive=True)) # glob
|
||||
elif os.path.isdir(p):
|
||||
@ -136,6 +136,7 @@ class LoadImages: # for inference
|
||||
ni, nv = len(images), len(videos)
|
||||
|
||||
self.img_size = img_size
|
||||
self.stride = stride
|
||||
self.files = images + videos
|
||||
self.nf = ni + nv # number of files
|
||||
self.video_flag = [False] * ni + [True] * nv
|
||||
@ -181,7 +182,7 @@ class LoadImages: # for inference
|
||||
print(f'image {self.count}/{self.nf} {path}: ', end='')
|
||||
|
||||
# Padded resize
|
||||
img = letterbox(img0, new_shape=self.img_size)[0]
|
||||
img = letterbox(img0, self.img_size, stride=self.stride)[0]
|
||||
|
||||
# Convert
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||
@ -199,8 +200,9 @@ class LoadImages: # 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.stride = stride
|
||||
|
||||
if pipe.isnumeric():
|
||||
pipe = eval(pipe) # local camera
|
||||
@ -243,7 +245,7 @@ class LoadWebcam: # for inference
|
||||
print(f'webcam {self.count}: ', end='')
|
||||
|
||||
# Padded resize
|
||||
img = letterbox(img0, new_shape=self.img_size)[0]
|
||||
img = letterbox(img0, self.img_size, stride=self.stride)[0]
|
||||
|
||||
# Convert
|
||||
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
|
||||
def __init__(self, sources='streams.txt', img_size=640):
|
||||
def __init__(self, sources='streams.txt', img_size=640, stride=32):
|
||||
self.mode = 'stream'
|
||||
self.img_size = img_size
|
||||
self.stride = stride
|
||||
|
||||
if os.path.isfile(sources):
|
||||
with open(sources, 'r') as f:
|
||||
@ -272,19 +275,25 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
for i, s in enumerate(sources):
|
||||
# Start the thread to read frames from the video stream
|
||||
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}'
|
||||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
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
|
||||
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()
|
||||
print('') # newline
|
||||
|
||||
# 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
|
||||
if not self.rect:
|
||||
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()
|
||||
cap.grab()
|
||||
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
|
||||
time.sleep(0.01) # wait time
|
||||
time.sleep(1 / self.fps) # wait time
|
||||
|
||||
def __iter__(self):
|
||||
self.count = -1
|
||||
@ -313,7 +323,7 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
raise StopIteration
|
||||
|
||||
# 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
|
||||
img = np.stack(img, 0)
|
||||
@ -331,7 +341,7 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||
def img2label_paths(img_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
|
||||
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
|
||||
@ -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_border = [-img_size // 2, -img_size // 2]
|
||||
self.stride = stride
|
||||
self.path = path
|
||||
|
||||
try:
|
||||
f = [] # image files
|
||||
@ -352,37 +363,42 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
p = Path(p) # os-agnostic
|
||||
if p.is_dir(): # dir
|
||||
f += glob.glob(str(p / '**' / '*.*'), recursive=True)
|
||||
# f = list(p.rglob('**/*.*')) # pathlib
|
||||
elif p.is_file(): # file
|
||||
with open(p, 'r') as t:
|
||||
t = t.read().strip().splitlines()
|
||||
parent = str(p.parent) + os.sep
|
||||
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:
|
||||
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 for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
|
||||
assert self.img_files, f'{prefix}No images found'
|
||||
except Exception as e:
|
||||
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')
|
||||
|
||||
# Check cache
|
||||
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():
|
||||
cache = torch.load(cache_path) # load
|
||||
if cache['hash'] != get_hash(self.label_files + self.img_files) or 'results' not in cache: # changed
|
||||
cache = self.cache_labels(cache_path, prefix) # re-cache
|
||||
cache, exists = torch.load(cache_path), True # load
|
||||
if cache['hash'] != get_hash(self.label_files + self.img_files) or 'version' not in cache: # changed
|
||||
cache, exists = self.cache_labels(cache_path, prefix), False # re-cache
|
||||
else:
|
||||
cache = self.cache_labels(cache_path, prefix) # cache
|
||||
cache, exists = self.cache_labels(cache_path, prefix), False # cache
|
||||
|
||||
# Display cache
|
||||
[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"
|
||||
tqdm(None, desc=prefix + desc, total=n, initial=n)
|
||||
nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
|
||||
if exists:
|
||||
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}'
|
||||
|
||||
# Read cache
|
||||
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.shapes = np.array(shapes, dtype=np.float64)
|
||||
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)
|
||||
gb += self.imgs[i].nbytes
|
||||
pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)'
|
||||
pbar.close()
|
||||
|
||||
def cache_labels(self, path=Path('./labels.cache'), prefix=''):
|
||||
# Cache dataset labels, check images and read shapes
|
||||
@ -445,13 +462,20 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
im = Image.open(im_file)
|
||||
im.verify() # PIL verify
|
||||
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
|
||||
if os.path.isfile(lb_file):
|
||||
nf += 1 # label found
|
||||
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):
|
||||
assert l.shape[1] == 5, 'labels require 5 columns each'
|
||||
assert (l >= 0).all(), 'negative labels'
|
||||
@ -463,19 +487,21 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
else:
|
||||
nm += 1 # label missing
|
||||
l = np.zeros((0, 5), dtype=np.float32)
|
||||
x[im_file] = [l, shape]
|
||||
x[im_file] = [l, shape, segments]
|
||||
except Exception as e:
|
||||
nc += 1
|
||||
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"
|
||||
pbar.close()
|
||||
|
||||
if nf == 0:
|
||||
print(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
|
||||
|
||||
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
|
||||
logging.info(f'{prefix}New cache created: {path}')
|
||||
return x
|
||||
@ -515,16 +541,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
|
||||
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
|
||||
|
||||
# Load labels
|
||||
labels = []
|
||||
x = self.labels[index]
|
||||
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]
|
||||
labels = self.labels[index].copy()
|
||||
if labels.size: # normalized xywh to pixel xyxy format
|
||||
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
|
||||
|
||||
if self.augment:
|
||||
# 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)
|
||||
cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed
|
||||
|
||||
# Histogram equalization
|
||||
# if random.random() < 0.2:
|
||||
# for i in range(3):
|
||||
# img[:, :, i] = cv2.equalizeHist(img[:, :, i])
|
||||
|
||||
def hist_equalize(img, clahe=True, bgr=False):
|
||||
# Equalize histogram on BGR image 'img' with img.shape(n,m,3) and range 0-255
|
||||
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):
|
||||
# loads images in a 4-mosaic
|
||||
|
||||
labels4 = []
|
||||
labels4, segments4 = [], []
|
||||
s = self.img_size
|
||||
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):
|
||||
# Load image
|
||||
img, _, (h, w) = load_image(self, index)
|
||||
@ -674,23 +699,21 @@ def load_mosaic(self, index):
|
||||
padh = y1a - y1b
|
||||
|
||||
# Labels
|
||||
x = self.labels[index]
|
||||
labels = x.copy()
|
||||
if x.size > 0: # Normalized xywh to pixel xyxy format
|
||||
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw
|
||||
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
|
||||
labels, segments = self.labels[index].copy(), self.segments[index].copy()
|
||||
if labels.size:
|
||||
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
|
||||
segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
|
||||
labels4.append(labels)
|
||||
segments4.extend(segments)
|
||||
|
||||
# Concat/clip labels
|
||||
if len(labels4):
|
||||
labels4 = np.concatenate(labels4, 0)
|
||||
np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_perspective
|
||||
# img4, labels4 = replicate(img4, labels4) # replicate
|
||||
labels4 = np.concatenate(labels4, 0)
|
||||
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
|
||||
|
||||
# Augment
|
||||
img4, labels4 = random_perspective(img4, labels4,
|
||||
img4, labels4 = random_perspective(img4, labels4, segments4,
|
||||
degrees=self.hyp['degrees'],
|
||||
translate=self.hyp['translate'],
|
||||
scale=self.hyp['scale'],
|
||||
@ -704,9 +727,9 @@ def load_mosaic(self, index):
|
||||
def load_mosaic9(self, index):
|
||||
# loads images in a 9-mosaic
|
||||
|
||||
labels9 = []
|
||||
labels9, segments9 = [], []
|
||||
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):
|
||||
# Load image
|
||||
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
|
||||
|
||||
# Labels
|
||||
x = self.labels[index]
|
||||
labels = x.copy()
|
||||
if x.size > 0: # Normalized xywh to pixel xyxy format
|
||||
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padx
|
||||
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
|
||||
labels, segments = self.labels[index].copy(), self.segments[index].copy()
|
||||
if labels.size:
|
||||
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
|
||||
segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
|
||||
labels9.append(labels)
|
||||
segments9.extend(segments)
|
||||
|
||||
# Image
|
||||
img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
|
||||
hp, wp = h, w # height, width previous
|
||||
|
||||
# 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]
|
||||
|
||||
# Concat/clip labels
|
||||
if len(labels9):
|
||||
labels9 = np.concatenate(labels9, 0)
|
||||
labels9[:, [1, 3]] -= xc
|
||||
labels9[:, [2, 4]] -= yc
|
||||
labels9 = np.concatenate(labels9, 0)
|
||||
labels9[:, [1, 3]] -= xc
|
||||
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
|
||||
# img9, labels9 = replicate(img9, labels9) # replicate
|
||||
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
|
||||
|
||||
# Augment
|
||||
img9, labels9 = random_perspective(img9, labels9,
|
||||
img9, labels9 = random_perspective(img9, labels9, segments9,
|
||||
degrees=self.hyp['degrees'],
|
||||
translate=self.hyp['translate'],
|
||||
scale=self.hyp['scale'],
|
||||
@ -792,8 +815,8 @@ def replicate(img, labels):
|
||||
return img, labels
|
||||
|
||||
|
||||
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
|
||||
# Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
|
||||
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
||||
# Resize and pad image while meeting stride-multiple constraints
|
||||
shape = img.shape[:2] # current shape [height, width]
|
||||
if isinstance(new_shape, int):
|
||||
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))
|
||||
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
||||
if auto: # minimum rectangle
|
||||
dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding
|
||||
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
||||
elif scaleFill: # stretch
|
||||
dw, dh = 0.0, 0.0
|
||||
new_unpad = (new_shape[1], new_shape[0])
|
||||
@ -825,7 +848,8 @@ def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale
|
||||
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))
|
||||
# targets = [cls, xyxy]
|
||||
|
||||
@ -877,37 +901,38 @@ def random_perspective(img, targets=(), degrees=10, translate=.1, scale=.1, shea
|
||||
# Transform label coordinates
|
||||
n = len(targets)
|
||||
if n:
|
||||
# warp points
|
||||
xy = np.ones((n * 4, 3))
|
||||
xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
|
||||
xy = xy @ M.T # transform
|
||||
if perspective:
|
||||
xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale
|
||||
else: # affine
|
||||
xy = xy[:, :2].reshape(n, 8)
|
||||
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
|
||||
|
||||
# create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
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
|
||||
# clip
|
||||
new[i] = segment2box(xy, width, height)
|
||||
|
||||
# # apply angle-based reduction of bounding boxes
|
||||
# radians = a * math.pi / 180
|
||||
# reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5
|
||||
# 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
|
||||
else: # warp boxes
|
||||
xy = np.ones((n * 4, 3))
|
||||
xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
|
||||
xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
|
||||
|
||||
# clip boxes
|
||||
xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
|
||||
xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
|
||||
# create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||
|
||||
# clip
|
||||
new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
|
||||
new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
|
||||
|
||||
# filter candidates
|
||||
i = box_candidates(box1=targets[:, 1:5].T * s, box2=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[:, 1:5] = xy[i]
|
||||
targets[:, 1:5] = new[i]
|
||||
|
||||
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}'
|
||||
|
||||
|
||||
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
|
||||
# Arguments
|
||||
path: Path to images directory
|
||||
weights: Train, val, test weights (list)
|
||||
Usage: from utils.datasets import *; autosplit('../coco128')
|
||||
Arguments
|
||||
path: Path to images directory
|
||||
weights: Train, val, test weights (list)
|
||||
annotated_only: Only use images with an annotated txt file
|
||||
"""
|
||||
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
|
||||
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
|
||||
[(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):
|
||||
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:
|
||||
f.write(str(img) + '\n') # add image to txt file
|
||||
|
||||
163
utils/general.py
163
utils/general.py
@ -1,9 +1,10 @@
|
||||
# General utils
|
||||
# YOLOv3 general utils
|
||||
|
||||
import glob
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import platform
|
||||
import random
|
||||
import re
|
||||
import subprocess
|
||||
@ -12,6 +13,7 @@ from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torchvision
|
||||
import yaml
|
||||
@ -23,6 +25,7 @@ from utils.torch_utils import init_torch_seeds
|
||||
# Settings
|
||||
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
|
||||
pd.options.display.max_columns = 10
|
||||
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
|
||||
|
||||
@ -46,40 +49,75 @@ def get_latest_run(search_dir='.'):
|
||||
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():
|
||||
# Check internet connectivity
|
||||
import socket
|
||||
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
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
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='')
|
||||
try:
|
||||
if Path('.git').exists() and check_online():
|
||||
url = subprocess.check_output(
|
||||
'git fetch && git config --get remote.origin.url', shell=True).decode('utf-8')[:-1]
|
||||
n = int(subprocess.check_output(
|
||||
'git rev-list $(git rev-parse --abbrev-ref HEAD)..origin/master --count', shell=True)) # commits behind
|
||||
if n > 0:
|
||||
print(f"⚠️ WARNING: code is out of date by {n} {'commits' if n > 1 else 'commmit'}. "
|
||||
f"Use 'git pull' to update or 'git clone {url}' to download latest.")
|
||||
else:
|
||||
print(f'up to date with {url} ✅')
|
||||
assert Path('.git').exists(), 'skipping check (not a git repository)'
|
||||
assert not isdocker(), 'skipping check (Docker image)'
|
||||
assert check_online(), 'skipping check (offline)'
|
||||
|
||||
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:
|
||||
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."
|
||||
else:
|
||||
s = f'up to date with {url} ✅'
|
||||
print(emojis(s)) # emoji-safe
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
def check_requirements(file='requirements.txt'):
|
||||
# Check installed dependencies meet requirements
|
||||
import pkg_resources
|
||||
requirements = pkg_resources.parse_requirements(Path(file).open())
|
||||
requirements = [x.name + ''.join(*x.specs) if len(x.specs) else x.name for x in requirements]
|
||||
pkg_resources.require(requirements) # DistributionNotFound or VersionConflict exception if requirements not met
|
||||
def check_requirements(requirements='requirements.txt', exclude=()):
|
||||
# Check installed dependencies meet requirements (pass *.txt file or list of packages)
|
||||
import pkg_resources as pkg
|
||||
prefix = colorstr('red', 'bold', 'requirements:')
|
||||
if isinstance(requirements, (str, Path)): # requirements.txt file
|
||||
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):
|
||||
@ -90,14 +128,28 @@ def check_img_size(img_size, s=32):
|
||||
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):
|
||||
# Search for file if not found
|
||||
if os.path.isfile(file) or file == '':
|
||||
if Path(file).is_file() or file == '':
|
||||
return file
|
||||
else:
|
||||
files = glob.glob('./**/' + file, recursive=True) # find file
|
||||
assert len(files), 'File Not Found: %s' % file # assert file was found
|
||||
assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files) # assert unique
|
||||
assert len(files), f'File Not Found: {file}' # assert file was found
|
||||
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
|
||||
return files[0] # return file
|
||||
|
||||
|
||||
@ -220,6 +272,50 @@ def xywh2xyxy(x):
|
||||
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):
|
||||
# Rescale coords (xyxy) from img1_shape to 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
|
||||
|
||||
|
||||
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
|
||||
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
|
||||
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
||||
with torch.no_grad():
|
||||
alpha = v / ((1 + eps) - iou + v)
|
||||
alpha = v / (v - iou + (1 + eps))
|
||||
return iou - (rho2 / c2 + v * alpha) # CIoU
|
||||
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
||||
c_area = cw * ch + eps # convex area
|
||||
@ -322,11 +418,12 @@ def wh_iou(wh1, wh2):
|
||||
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=()):
|
||||
"""Performs Non-Maximum Suppression (NMS) on inference results
|
||||
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
||||
labels=()):
|
||||
"""Runs Non-Maximum Suppression (NMS) on inference results
|
||||
|
||||
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
|
||||
@ -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()
|
||||
time_limit = 10.0 # seconds to quit after
|
||||
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
|
||||
|
||||
t = time.time()
|
||||
@ -412,18 +509,20 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
|
||||
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'
|
||||
x = torch.load(f, map_location=torch.device('cpu'))
|
||||
for key in 'optimizer', 'training_results', 'wandb_id':
|
||||
x[key] = None
|
||||
if x.get('ema'):
|
||||
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['model'].half() # to FP16
|
||||
for p in x['model'].parameters():
|
||||
p.requires_grad = False
|
||||
torch.save(x, s or f)
|
||||
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=''):
|
||||
|
||||
@ -21,13 +21,13 @@ def attempt_download(file, repo='ultralytics/yolov3'):
|
||||
file = Path(str(file).strip().replace("'", '').lower())
|
||||
|
||||
if not file.exists():
|
||||
# try:
|
||||
# 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', ...]
|
||||
# tag = response['tag_name'] # i.e. 'v1.0'
|
||||
# except: # fallback plan
|
||||
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
||||
tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
|
||||
try:
|
||||
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', ...]
|
||||
tag = response['tag_name'] # i.e. 'v1.0'
|
||||
except: # fallback plan
|
||||
assets = ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']
|
||||
tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
|
||||
|
||||
name = file.name
|
||||
if name in assets:
|
||||
|
||||
201
utils/loss.py
201
utils/loss.py
@ -85,117 +85,132 @@ class QFocalLoss(nn.Module):
|
||||
return loss
|
||||
|
||||
|
||||
def compute_loss(p, targets, model): # 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 = build_targets(p, targets, model) # targets
|
||||
h = model.hyp # hyperparameters
|
||||
class ComputeLoss:
|
||||
# Compute losses
|
||||
def __init__(self, model, autobalance=False):
|
||||
super(ComputeLoss, self).__init__()
|
||||
device = next(model.parameters()).device # get model device
|
||||
h = model.hyp # hyperparameters
|
||||
|
||||
# Define criteria
|
||||
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) # weight=model.class_weights)
|
||||
BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
|
||||
# Define criteria
|
||||
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_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
|
||||
cp, cn = smooth_BCE(eps=0.0)
|
||||
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
|
||||
self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets
|
||||
|
||||
# Focal loss
|
||||
g = h['fl_gamma'] # focal loss gamma
|
||||
if g > 0:
|
||||
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
|
||||
# Focal loss
|
||||
g = h['fl_gamma'] # focal loss gamma
|
||||
if g > 0:
|
||||
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
|
||||
|
||||
# Losses
|
||||
balance = [4.0, 1.0, 0.4, 0.1] # P3-P6
|
||||
for i, pi in enumerate(p): # layer index, layer predictions
|
||||
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
||||
tobj = torch.zeros_like(pi[..., 0], device=device) # target obj
|
||||
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))
|
||||
|
||||
n = b.shape[0] # number of targets
|
||||
if n:
|
||||
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
|
||||
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
|
||||
|
||||
# Regression
|
||||
pxy = ps[:, :2].sigmoid() * 2. - 0.5
|
||||
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
|
||||
pbox = torch.cat((pxy, pwh), 1) # predicted box
|
||||
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
|
||||
lbox += (1.0 - iou).mean() # iou loss
|
||||
# Losses
|
||||
for i, pi in enumerate(p): # layer index, layer predictions
|
||||
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
||||
tobj = torch.zeros_like(pi[..., 0], device=device) # target obj
|
||||
|
||||
# Objectness
|
||||
tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
|
||||
n = b.shape[0] # number of targets
|
||||
if n:
|
||||
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
|
||||
|
||||
# Classification
|
||||
if model.nc > 1: # cls loss (only if multiple classes)
|
||||
t = torch.full_like(ps[:, 5:], cn, device=device) # targets
|
||||
t[range(n), tcls[i]] = cp
|
||||
lcls += BCEcls(ps[:, 5:], t) # BCE
|
||||
# Regression
|
||||
pxy = ps[:, :2].sigmoid() * 2. - 0.5
|
||||
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
|
||||
pbox = torch.cat((pxy, pwh), 1) # predicted box
|
||||
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
|
||||
lbox += (1.0 - iou).mean() # iou loss
|
||||
|
||||
# Append targets to text 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)]
|
||||
# Objectness
|
||||
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
|
||||
|
||||
lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss
|
||||
# Classification
|
||||
if self.nc > 1: # cls loss (only if multiple classes)
|
||||
t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets
|
||||
t[range(n), tcls[i]] = self.cp
|
||||
lcls += self.BCEcls(ps[:, 5:], t) # BCE
|
||||
|
||||
lbox *= h['box']
|
||||
lobj *= h['obj']
|
||||
lcls *= h['cls']
|
||||
bs = tobj.shape[0] # batch size
|
||||
# Append targets to text 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)]
|
||||
|
||||
loss = lbox + lobj + lcls
|
||||
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
|
||||
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()
|
||||
|
||||
if self.autobalance:
|
||||
self.balance = [x / self.balance[self.ssi] for x in self.balance]
|
||||
lbox *= self.hyp['box']
|
||||
lobj *= self.hyp['obj']
|
||||
lcls *= self.hyp['cls']
|
||||
bs = tobj.shape[0] # batch size
|
||||
|
||||
def build_targets(p, targets, model):
|
||||
# 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 = det.na, targets.shape[0] # number of anchors, targets
|
||||
tcls, tbox, indices, anch = [], [], [], []
|
||||
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)
|
||||
targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices
|
||||
loss = lbox + lobj + lcls
|
||||
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
|
||||
|
||||
g = 0.5 # bias
|
||||
off = torch.tensor([[0, 0],
|
||||
# [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
|
||||
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
|
||||
], device=targets.device).float() * g # offsets
|
||||
def build_targets(self, p, targets):
|
||||
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
|
||||
na, nt = self.na, targets.shape[0] # number of anchors, targets
|
||||
tcls, tbox, indices, anch = [], [], [], []
|
||||
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)
|
||||
targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices
|
||||
|
||||
for i in range(det.nl):
|
||||
anchors = det.anchors[i]
|
||||
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
|
||||
g = 0.5 # bias
|
||||
off = torch.tensor([[0, 0],
|
||||
# [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
|
||||
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
|
||||
], device=targets.device).float() * g # offsets
|
||||
|
||||
# Match targets to anchors
|
||||
t = targets * gain
|
||||
if nt:
|
||||
# Matches
|
||||
r = t[:, :, 4:6] / anchors[:, None] # wh ratio
|
||||
j = torch.max(r, 1. / r).max(2)[0] < model.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))
|
||||
t = t[j] # filter
|
||||
for i in range(self.nl):
|
||||
anchors = self.anchors[i]
|
||||
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
|
||||
|
||||
# Offsets
|
||||
# Match targets to anchors
|
||||
t = targets * gain
|
||||
if nt:
|
||||
# Matches
|
||||
r = t[:, :, 4:6] / anchors[:, None] # wh ratio
|
||||
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))
|
||||
t = t[j] # filter
|
||||
|
||||
# Offsets
|
||||
gxy = t[:, 2:4] # grid xy
|
||||
gxi = gain[[2, 3]] - gxy # inverse
|
||||
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
|
||||
l, m = ((gxi % 1. < g) & (gxi > 1.)).T
|
||||
j = torch.stack((torch.ones_like(j),))
|
||||
t = t.repeat((off.shape[0], 1, 1))[j]
|
||||
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
|
||||
else:
|
||||
t = targets[0]
|
||||
offsets = 0
|
||||
|
||||
# Define
|
||||
b, c = t[:, :2].long().T # image, class
|
||||
gxy = t[:, 2:4] # grid xy
|
||||
gxi = gain[[2, 3]] - gxy # inverse
|
||||
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
|
||||
l, m = ((gxi % 1. < g) & (gxi > 1.)).T
|
||||
j = torch.stack((torch.ones_like(j),))
|
||||
t = t.repeat((off.shape[0], 1, 1))[j]
|
||||
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
|
||||
else:
|
||||
t = targets[0]
|
||||
offsets = 0
|
||||
gwh = t[:, 4:6] # grid wh
|
||||
gij = (gxy - offsets).long()
|
||||
gi, gj = gij.T # grid xy indices
|
||||
|
||||
# Define
|
||||
b, c = t[:, :2].long().T # image, class
|
||||
gxy = t[:, 2:4] # grid xy
|
||||
gwh = t[:, 4:6] # grid wh
|
||||
gij = (gxy - offsets).long()
|
||||
gi, gj = gij.T # grid xy indices
|
||||
# Append
|
||||
a = t[:, 6].long() # anchor indices
|
||||
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices
|
||||
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
|
||||
anch.append(anchors[a]) # anchors
|
||||
tcls.append(c) # class
|
||||
|
||||
# Append
|
||||
a = t[:, 6].long() # anchor indices
|
||||
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices
|
||||
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
|
||||
anch.append(anchors[a]) # anchors
|
||||
tcls.append(c) # class
|
||||
|
||||
return tcls, tbox, indices, anch
|
||||
return tcls, tbox, indices, anch
|
||||
|
||||
@ -15,7 +15,7 @@ def fitness(x):
|
||||
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.
|
||||
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
|
||||
# Arguments
|
||||
@ -35,12 +35,11 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision
|
||||
|
||||
# Find unique classes
|
||||
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
|
||||
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
|
||||
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)
|
||||
ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
|
||||
for ci, c in enumerate(unique_classes):
|
||||
i = pred_cls == c
|
||||
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 = 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 = 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
|
||||
for j in range(tp.shape[1]):
|
||||
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
|
||||
|
||||
# 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)
|
||||
|
||||
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):
|
||||
@ -145,12 +147,12 @@ class ConfusionMatrix:
|
||||
if n and sum(j) == 1:
|
||||
self.matrix[gc, detection_classes[m1[j]]] += 1 # correct
|
||||
else:
|
||||
self.matrix[gc, self.nc] += 1 # background FP
|
||||
self.matrix[self.nc, gc] += 1 # background FP
|
||||
|
||||
if n:
|
||||
for i, dc in enumerate(detection_classes):
|
||||
if not any(m1 == i):
|
||||
self.matrix[self.nc, dc] += 1 # background FN
|
||||
self.matrix[dc, self.nc] += 1 # background FN
|
||||
|
||||
def matrix(self):
|
||||
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
|
||||
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,
|
||||
xticklabels=names + ['background FN'] if labels else "auto",
|
||||
yticklabels=names + ['background FP'] if labels else "auto").set_facecolor((1, 1, 1))
|
||||
xticklabels=names + ['background FP'] if labels else "auto",
|
||||
yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
|
||||
fig.axes[0].set_xlabel('True')
|
||||
fig.axes[0].set_ylabel('Predicted')
|
||||
fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
|
||||
@ -181,13 +183,14 @@ class ConfusionMatrix:
|
||||
|
||||
# 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)
|
||||
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):
|
||||
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:
|
||||
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_ylim(0, 1)
|
||||
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)
|
||||
|
||||
@ -15,7 +15,7 @@ import pandas as pd
|
||||
import seaborn as sns
|
||||
import torch
|
||||
import yaml
|
||||
from PIL import Image, ImageDraw
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from scipy.signal import butter, filtfilt
|
||||
|
||||
from utils.general import xywh2xyxy, xyxy2xywh
|
||||
@ -31,7 +31,7 @@ def color_list():
|
||||
def hex2rgb(h):
|
||||
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):
|
||||
@ -54,7 +54,7 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
||||
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
|
||||
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)]
|
||||
@ -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)
|
||||
|
||||
|
||||
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()
|
||||
# Compares the two methods for width-height anchor multiplication
|
||||
# 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)
|
||||
|
||||
|
||||
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
|
||||
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)
|
||||
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
|
||||
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)']
|
||||
for i in range(7):
|
||||
ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
|
||||
ax[i].set_title(s[i])
|
||||
# for i in range(7):
|
||||
# ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
|
||||
# ax[i].set_title(s[i])
|
||||
|
||||
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'))
|
||||
|
||||
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')
|
||||
|
||||
ax2.grid()
|
||||
ax2.set_yticks(np.arange(30, 60, 5))
|
||||
ax2.set_xlim(0, 30)
|
||||
ax2.set_ylim(29, 51)
|
||||
ax2.grid(alpha=0.2)
|
||||
ax2.set_yticks(np.arange(20, 60, 5))
|
||||
ax2.set_xlim(0, 57)
|
||||
ax2.set_ylim(15, 55)
|
||||
ax2.set_xlabel('GPU Speed (ms/img)')
|
||||
ax2.set_ylabel('COCO AP val')
|
||||
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
|
||||
print('Plotting labels... ')
|
||||
c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
|
||||
@ -271,7 +286,12 @@ def plot_labels(labels, save_dir=Path(''), loggers=None):
|
||||
matplotlib.use('svg') # faster
|
||||
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].set_xlabel('classes')
|
||||
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')
|
||||
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)
|
||||
|
||||
@ -295,13 +315,13 @@ def plot_labels(labels, save_dir=Path(''), loggers=None):
|
||||
# loggers
|
||||
for k, v in loggers.items() or {}:
|
||||
if k == 'wandb' and v:
|
||||
v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]})
|
||||
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()
|
||||
# Plot hyperparameter evolution results in evolve.txt
|
||||
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)
|
||||
f = fitness(x)
|
||||
# weights = (f - f.min()) ** 2 # for weighted results
|
||||
|
||||
@ -1,8 +1,10 @@
|
||||
# PyTorch utils
|
||||
# YOLOv3 PyTorch utils
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
@ -43,17 +45,24 @@ def init_torch_seeds(seed=0):
|
||||
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
|
||||
if Path('.git').exists():
|
||||
return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1]
|
||||
else:
|
||||
return ''
|
||||
s = f'git -C {path} describe --tags --long --always'
|
||||
try:
|
||||
return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1]
|
||||
except subprocess.CalledProcessError as e:
|
||||
return '' # not a git repository
|
||||
|
||||
|
||||
def select_device(device='', batch_size=None):
|
||||
# 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'
|
||||
if cpu:
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
|
||||
@ -73,7 +82,7 @@ def select_device(device='', batch_size=None):
|
||||
else:
|
||||
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')
|
||||
|
||||
|
||||
@ -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_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
|
||||
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
|
||||
print(f'{p:12.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):
|
||||
@ -182,7 +191,7 @@ def fuse_conv_and_bn(conv, bn):
|
||||
# prepare filters
|
||||
w_conv = conv.weight.clone().view(conv.out_channels, -1)
|
||||
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
|
||||
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
|
||||
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
|
||||
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
|
||||
|
||||
0
utils/wandb_logging/__init__.py
Normal file
0
utils/wandb_logging/__init__.py
Normal file
24
utils/wandb_logging/log_dataset.py
Normal file
24
utils/wandb_logging/log_dataset.py
Normal 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
306
utils/wandb_logging/wandb_utils.py
Normal 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()
|
||||
Loading…
x
Reference in New Issue
Block a user