Showing preview only (234K chars total). Download the full file or copy to clipboard to get everything.
Repository: yjh0410/yolov2-yolov3_PyTorch
Branch: master
Commit: 49dbf4bbcaba
Files: 35
Total size: 223.1 KB
Directory structure:
gitextract_dmpt53ff/
├── .gitignore
├── LICENSE
├── README.md
├── backbone/
│ ├── __init__.py
│ ├── darknet19.py
│ ├── darknet53.py
│ ├── darknet_tiny.py
│ ├── resnet.py
│ └── weights/
│ └── README.md
├── data/
│ ├── __init__.py
│ ├── coco2017.py
│ ├── config.py
│ ├── scripts/
│ │ ├── COCO2017.sh
│ │ ├── VOC2007.sh
│ │ └── VOC2012.sh
│ └── voc0712.py
├── demo.py
├── eval.py
├── models/
│ ├── yolov2_d19.py
│ ├── yolov2_r50.py
│ ├── yolov3.py
│ ├── yolov3_spp.py
│ └── yolov3_tiny.py
├── test.py
├── tools.py
├── train.py
├── utils/
│ ├── __init__.py
│ ├── augmentations.py
│ ├── cocoapi_evaluator.py
│ ├── com_paras_flops.py
│ ├── distributed_utils.py
│ ├── kmeans_anchor.py
│ ├── modules.py
│ └── vocapi_evaluator.py
└── weights/
└── README.md
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
*.pt
*.pth
*.txt
*.pkl
__pycache__
.vscode
det_results
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
================================================
FILE: README.md
================================================
# Update
Recently, I have released a new YOLO project:
https://github.com/yjh0410/PyTorch_YOLO_Tutorial
In my new YOLO project, you can enjoy:
- a new and stronger YOLOv1
- a new and stronger YOLOv2
- YOLOv3
- YOLOv4
- YOLOv5
- YOLOv7
- YOLOX
- RTCDet
# This project
In this project, you can enjoy:
- YOLOv2 with DarkNet-19
- YOLOv2 with ResNet-50
- YOLOv2Slim
- YOLOv3
- YOLOv3-Spp
- YOLOv3-Tiny
I just want to provide a good YOLO project for everyone who is interested in Object Detection.
# Weights
Google Drive: https://drive.google.com/drive/folders/1T5hHyGICbFSdu6u2_vqvxn_puotvPsbd?usp=sharing
BaiDuYunDisk: https://pan.baidu.com/s/1tSylvzOVFReUAvaAxKRSwg
Password d266
You can download all my models from the above links.
# YOLOv2
## YOLOv2 with DarkNet-19
### Tricks
Tricks in official paper:
- [x] batch norm
- [x] hi-res classifier
- [x] convolutional
- [x] anchor boxes
- [x] new network
- [x] dimension priors
- [x] location prediction
- [x] passthrough
- [x] multi-scale
- [x] hi-red detector
## VOC2007
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> size </td><td bgcolor=white> Original (darknet) </td><td bgcolor=white> Ours (pytorch) 160peochs </td><td bgcolor=white> Ours (pytorch) 250epochs </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> VOC07 test</th><td bgcolor=white> 416 </td><td bgcolor=white> 76.8 </td><td bgcolor=white> 76.0 </td><td bgcolor=white> 77.1 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> VOC07 test</th><td bgcolor=white> 544 </td><td bgcolor=white> 78.6 </td><td bgcolor=white> 77.0 </td><td bgcolor=white> 78.1 </td></tr>
</table></tbody>
## COCO
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Original (darknet)</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 21.6 </td><td bgcolor=white> 44.0 </td><td bgcolor=white> 19.2 </td><td bgcolor=white> 5.0 </td><td bgcolor=white> 22.4 </td><td bgcolor=white> 35.5 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Ours (pytorch)</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 26.8 </td><td bgcolor=white> 46.6 </td><td bgcolor=white> 26.8 </td><td bgcolor=white> 5.8 </td><td bgcolor=white> 27.4 </td><td bgcolor=white> 45.2 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Ours (pytorch)</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 26.6 </td><td bgcolor=white> 46.0 </td><td bgcolor=white> 26.7 </td><td bgcolor=white> 5.9 </td><td bgcolor=white> 27.8 </td><td bgcolor=white> 47.1 </td></tr>
</table></tbody>
## YOLOv2 with ResNet-50
I replace darknet-19 with resnet-50 and get a better result on COCO-val
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Our YOLOv2-320</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 25.8 </td><td bgcolor=white> 44.6 </td><td bgcolor=white> 25.9 </td><td bgcolor=white> 4.6 </td><td bgcolor=white> 26.8 </td><td bgcolor=white> 47.9 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Our YOLOv2-416</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 29.0 </td><td bgcolor=white> 48.8 </td><td bgcolor=white> 29.7 </td><td bgcolor=white> 7.4 </td><td bgcolor=white> 31.9 </td><td bgcolor=white> 48.3 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Our YOLOv2-512</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 30.4 </td><td bgcolor=white> 51.6 </td><td bgcolor=white> 30.9 </td><td bgcolor=white> 10.1 </td><td bgcolor=white> 34.9 </td><td bgcolor=white> 46.6 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Our YOLOv2-544</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 30.4 </td><td bgcolor=white> 51.9 </td><td bgcolor=white> 30.9 </td><td bgcolor=white> 11.1 </td><td bgcolor=white> 35.8 </td><td bgcolor=white> 45.5 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> Our YOLOv2-608</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 29.2 </td><td bgcolor=white> 51.6 </td><td bgcolor=white> 29.1 </td><td bgcolor=white> 13.6 </td><td bgcolor=white> 36.8 </td><td bgcolor=white> 40.5 </td></tr>
</table></tbody>
# YOLOv3
## VOC2007
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> size </td><td bgcolor=white> Original (darknet) </td><td bgcolor=white> Ours (pytorch) 250epochs </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> VOC07 test</th><td bgcolor=white> 416 </td><td bgcolor=white> 80.25 </td><td bgcolor=white> 81.4 </td></tr>
</table></tbody>
# COCO
Official YOLOv3:
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-320</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 28.2 </td><td bgcolor=white> 51.5 </td><td bgcolor=white> - </td><td bgcolor=white> - </td><td bgcolor=white> - </td><td bgcolor=white> - </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-416</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 31.0 </td><td bgcolor=white> 55.3 </td><td bgcolor=white> - </td><td bgcolor=white> - </td><td bgcolor=white> - </td><td bgcolor=white> - </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-608</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 33.0 </td><td bgcolor=white> 57.0 </td><td bgcolor=white> 34.4 </td><td bgcolor=white> 18.3 </td><td bgcolor=white> 35.4 </td><td bgcolor=white> 41.9 </td></tr>
</table></tbody>
Our YOLOv3:
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-320</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 33.1 </td><td bgcolor=white> 54.1 </td><td bgcolor=white> 34.5 </td><td bgcolor=white> 12.1 </td><td bgcolor=white> 34.5 </td><td bgcolor=white> 49.6 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-416</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 36.0 </td><td bgcolor=white> 57.4 </td><td bgcolor=white> 37.0 </td><td bgcolor=white> 16.3 </td><td bgcolor=white> 37.5 </td><td bgcolor=white> 51.1 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3-608</th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> 37.6 </td><td bgcolor=white> 59.4 </td><td bgcolor=white> 39.9 </td><td bgcolor=white> 20.4 </td><td bgcolor=white> 39.9 </td><td bgcolor=white> 48.2 </td></tr>
</table></tbody>
# YOLOv3SPP
## COCO:
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3Spp-320</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 32.78 </td><td bgcolor=white> 53.79 </td><td bgcolor=white> 33.9 </td><td bgcolor=white> 12.4 </td><td bgcolor=white> 35.5 </td><td bgcolor=white> 50.6 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3Spp-416</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 35.66 </td><td bgcolor=white> 57.09 </td><td bgcolor=white> 37.4 </td><td bgcolor=white> 16.8 </td><td bgcolor=white> 38.1 </td><td bgcolor=white> 50.7 </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> YOLOv3Spp-608</th><td bgcolor=white> COCO eval </td><td bgcolor=white> 37.52 </td><td bgcolor=white> 59.44 </td><td bgcolor=white> 39.3 </td><td bgcolor=white> 21.5 </td><td bgcolor=white> 40.6 </td><td bgcolor=white> 49.6 </td></tr>
</table></tbody>
# YOLOv3Tiny
<table><tbody>
<tr><th align="left" bgcolor=#f8f8f8> </th> <td bgcolor=white> data </td><td bgcolor=white> AP </td><td bgcolor=white> AP50 </td><td bgcolor=white> AP75 </td><td bgcolor=white> AP_S </td><td bgcolor=white> AP_M </td><td bgcolor=white> AP_L </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> (official) YOLOv3Tiny </th><td bgcolor=white> COCO test-dev </td><td bgcolor=white> - </td><td bgcolor=white> 33.1 </td><td bgcolor=white> - </td><td bgcolor=white>- </td><td bgcolor=white> - </td><td bgcolor=white> - </td></tr>
<tr><th align="left" bgcolor=#f8f8f8> (Our) YOLOv3Tiny </th><td bgcolor=white> COCO val </td><td bgcolor=white> 15.9 </td><td bgcolor=white> 33.8 </td><td bgcolor=white> 12.8 </td><td bgcolor=white> 7.6 </td><td bgcolor=white> 17.7 </td><td bgcolor=white> 22.4 </td></tr>
</table></tbody>
# Installation
- Pytorch-gpu 1.1.0/1.2.0/1.3.0
- Tensorboard 1.14.
- opencv-python, python3.6/3.7
# Dataset
## VOC Dataset
I copy the download files from the following excellent project:
https://github.com/amdegroot/ssd.pytorch
I have uploaded the VOC2007 and VOC2012 to BaiDuYunDisk, so for researchers in China, you can download them from BaiDuYunDisk:
Link：https://pan.baidu.com/s/1tYPGCYGyC0wjpC97H-zzMQ
Password：4la9
You will get a ```VOCdevkit.zip```, then what you need to do is just to unzip it and put it into ```data/```. After that, the whole path to VOC dataset is ```data/VOCdevkit/VOC2007``` and ```data/VOCdevkit/VOC2012```.
### Download VOC2007 trainval & test
```Shell
# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/VOC2007.sh # <directory>
```
### Download VOC2012 trainval
```Shell
# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/VOC2012.sh # <directory>
```
## MSCOCO Dataset
I copy the download files from the following excellent project:
https://github.com/DeNA/PyTorch_YOLOv3
### Download MSCOCO 2017 dataset
Just run ```sh data/scripts/COCO2017.sh```. You will get COCO train2017, val2017, test2017.
# Train
## VOC
```Shell
python train.py -d voc --cuda -v [select a model] -hr -ms --ema
```
You can run ```python train.py -h``` to check all optional argument.
## COCO
If you have only one gpu:
```Shell
python train.py -d coco --cuda -v [select a model] -hr -ms --ema
```
If you have multi gpus like 8, and you put 4 images on each gpu:
```Shell
python -m torch.distributed.launch --nproc_per_node=8 train.py -d coco --cuda -v [select a model] -hr -ms --ema \
-dist \
--sybn \
--num_gpu 8\
--batch_size 4
```
# Test
## VOC
```Shell
python test.py -d voc --cuda -v [select a model] --trained_model [ Please input the path to model dir. ]
```
## COCO
```Shell
python test.py -d coco-val --cuda -v [select a model] --trained_model [ Please input the path to model dir. ]
```
# Evaluation
## VOC
```Shell
python eval.py -d voc --cuda -v [select a model] --train_model [ Please input the path to model dir. ]
```
## COCO
To run on COCO_val:
```Shell
python eval.py -d coco-val --cuda -v [select a model] --train_model [ Please input the path to model dir. ]
```
To run on COCO_test-dev(You must be sure that you have downloaded test2017):
```Shell
python eval.py -d coco-test --cuda -v [select a model] --train_model [ Please input the path to model dir. ]
```
You will get a .json file which can be evaluated on COCO test server.
================================================
FILE: backbone/__init__.py
================================================
from .resnet import build_resnet
from .darknet19 import build_darknet19
from .darknet53 import build_darknet53
from .darknet_tiny import build_darknet_tiny
def build_backbone(model_name='resnet18', pretrained=False):
if 'resnet' in model_name:
backbone = build_resnet(model_name, pretrained)
elif model_name == 'darknet19':
backbone = build_darknet19(pretrained)
elif model_name == 'darknet53':
backbone = build_darknet53(pretrained)
elif model_name == 'darknet19':
backbone = build_darknet_tiny(pretrained)
return backbone
================================================
FILE: backbone/darknet19.py
================================================
import torch
import torch.nn as nn
import os
model_urls = {
"darknet19": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet19.pth",
}
__all__ = ['darknet19']
class Conv_BN_LeakyReLU(nn.Module):
def __init__(self, in_channels, out_channels, ksize, padding=0, stride=1, dilation=1):
super(Conv_BN_LeakyReLU, self).__init__()
self.convs = nn.Sequential(
nn.Conv2d(in_channels, out_channels, ksize, padding=padding, stride=stride, dilation=dilation),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.1, inplace=True)
)
def forward(self, x):
return self.convs(x)
class DarkNet_19(nn.Module):
def __init__(self):
super(DarkNet_19, self).__init__()
# backbone network : DarkNet-19
# output : stride = 2, c = 32
self.conv_1 = nn.Sequential(
Conv_BN_LeakyReLU(3, 32, 3, 1),
nn.MaxPool2d((2,2), 2),
)
# output : stride = 4, c = 64
self.conv_2 = nn.Sequential(
Conv_BN_LeakyReLU(32, 64, 3, 1),
nn.MaxPool2d((2,2), 2)
)
# output : stride = 8, c = 128
self.conv_3 = nn.Sequential(
Conv_BN_LeakyReLU(64, 128, 3, 1),
Conv_BN_LeakyReLU(128, 64, 1),
Conv_BN_LeakyReLU(64, 128, 3, 1),
nn.MaxPool2d((2,2), 2)
)
# output : stride = 8, c = 256
self.conv_4 = nn.Sequential(
Conv_BN_LeakyReLU(128, 256, 3, 1),
Conv_BN_LeakyReLU(256, 128, 1),
Conv_BN_LeakyReLU(128, 256, 3, 1),
)
# output : stride = 16, c = 512
self.maxpool_4 = nn.MaxPool2d((2, 2), 2)
self.conv_5 = nn.Sequential(
Conv_BN_LeakyReLU(256, 512, 3, 1),
Conv_BN_LeakyReLU(512, 256, 1),
Conv_BN_LeakyReLU(256, 512, 3, 1),
Conv_BN_LeakyReLU(512, 256, 1),
Conv_BN_LeakyReLU(256, 512, 3, 1),
)
# output : stride = 32, c = 1024
self.maxpool_5 = nn.MaxPool2d((2, 2), 2)
self.conv_6 = nn.Sequential(
Conv_BN_LeakyReLU(512, 1024, 3, 1),
Conv_BN_LeakyReLU(1024, 512, 1),
Conv_BN_LeakyReLU(512, 1024, 3, 1),
Conv_BN_LeakyReLU(1024, 512, 1),
Conv_BN_LeakyReLU(512, 1024, 3, 1)
)
def forward(self, x):
c1 = self.conv_1(x)
c2 = self.conv_2(c1)
c3 = self.conv_3(c2)
c3 = self.conv_4(c3)
c4 = self.conv_5(self.maxpool_4(c3))
c5 = self.conv_6(self.maxpool_5(c4))
output = {
'layer1': c3,
'layer2': c4,
'layer3': c5
}
return output
def build_darknet19(pretrained=False):
# model
model = DarkNet_19()
# load weight
if pretrained:
print('Loading pretrained weight ...')
url = model_urls['darknet19']
# checkpoint state dict
checkpoint_state_dict = torch.hub.load_state_dict_from_url(
url=url, map_location="cpu", check_hash=True)
# model state dict
model_state_dict = model.state_dict()
# check
for k in list(checkpoint_state_dict.keys()):
if k in model_state_dict:
shape_model = tuple(model_state_dict[k].shape)
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
if shape_model != shape_checkpoint:
checkpoint_state_dict.pop(k)
else:
checkpoint_state_dict.pop(k)
print(k)
model.load_state_dict(checkpoint_state_dict)
return model
if __name__ == '__main__':
import time
net = build_darknet19(pretrained=True)
x = torch.randn(1, 3, 224, 224)
t0 = time.time()
output = net(x)
t1 = time.time()
print('Time: ', t1 - t0)
for k in output.keys():
print('{} : {}'.format(k, output[k].shape))
================================================
FILE: backbone/darknet53.py
================================================
import torch
import torch.nn as nn
model_urls = {
"darknet53": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet53.pth",
}
__all__ = ['darknet53']
class Conv_BN_LeakyReLU(nn.Module):
def __init__(self, in_channels, out_channels, ksize, padding=0, stride=1, dilation=1):
super(Conv_BN_LeakyReLU, self).__init__()
self.convs = nn.Sequential(
nn.Conv2d(in_channels, out_channels, ksize, padding=padding, stride=stride, dilation=dilation),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.1, inplace=True)
)
def forward(self, x):
return self.convs(x)
class ResBlock(nn.Module):
def __init__(self, ch, nblocks=1):
super().__init__()
self.module_list = nn.ModuleList()
for _ in range(nblocks):
resblock_one = nn.Sequential(
Conv_BN_LeakyReLU(ch, ch//2, 1),
Conv_BN_LeakyReLU(ch//2, ch, 3, padding=1)
)
self.module_list.append(resblock_one)
def forward(self, x):
for module in self.module_list:
x = module(x) + x
return x
class DarkNet_53(nn.Module):
"""
DarkNet-53.
"""
def __init__(self):
super(DarkNet_53, self).__init__()
# stride = 2
self.layer_1 = nn.Sequential(
Conv_BN_LeakyReLU(3, 32, 3, padding=1),
Conv_BN_LeakyReLU(32, 64, 3, padding=1, stride=2),
ResBlock(64, nblocks=1)
)
# stride = 4
self.layer_2 = nn.Sequential(
Conv_BN_LeakyReLU(64, 128, 3, padding=1, stride=2),
ResBlock(128, nblocks=2)
)
# stride = 8
self.layer_3 = nn.Sequential(
Conv_BN_LeakyReLU(128, 256, 3, padding=1, stride=2),
ResBlock(256, nblocks=8)
)
# stride = 16
self.layer_4 = nn.Sequential(
Conv_BN_LeakyReLU(256, 512, 3, padding=1, stride=2),
ResBlock(512, nblocks=8)
)
# stride = 32
self.layer_5 = nn.Sequential(
Conv_BN_LeakyReLU(512, 1024, 3, padding=1, stride=2),
ResBlock(1024, nblocks=4)
)
def forward(self, x, targets=None):
c1 = self.layer_1(x)
c2 = self.layer_2(c1)
c3 = self.layer_3(c2)
c4 = self.layer_4(c3)
c5 = self.layer_5(c4)
output = {
'layer1': c3,
'layer2': c4,
'layer3': c5
}
return output
def build_darknet53(pretrained=False):
# model
model = DarkNet_53()
# load weight
if pretrained:
print('Loading pretrained weight ...')
url = model_urls['darknet53']
# checkpoint state dict
checkpoint_state_dict = torch.hub.load_state_dict_from_url(
url=url, map_location="cpu", check_hash=True)
# model state dict
model_state_dict = model.state_dict()
# check
for k in list(checkpoint_state_dict.keys()):
if k in model_state_dict:
shape_model = tuple(model_state_dict[k].shape)
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
if shape_model != shape_checkpoint:
checkpoint_state_dict.pop(k)
else:
checkpoint_state_dict.pop(k)
print(k)
model.load_state_dict(checkpoint_state_dict)
return model
if __name__ == '__main__':
import time
net = build_darknet53(pretrained=True)
x = torch.randn(1, 3, 224, 224)
t0 = time.time()
output = net(x)
t1 = time.time()
print('Time: ', t1 - t0)
for k in output.keys():
print('{} : {}'.format(k, output[k].shape))
================================================
FILE: backbone/darknet_tiny.py
================================================
import torch
import torch.nn as nn
model_urls = {
"darknet_tiny": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet_tiny.pth",
}
__all__ = ['darknet_tiny']
class Conv_BN_LeakyReLU(nn.Module):
def __init__(self, in_channels, out_channels, ksize, padding=0, stride=1, dilation=1):
super(Conv_BN_LeakyReLU, self).__init__()
self.convs = nn.Sequential(
nn.Conv2d(in_channels, out_channels, ksize, padding=padding, stride=stride, dilation=dilation),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.1, inplace=True)
)
def forward(self, x):
return self.convs(x)
class DarkNet_Tiny(nn.Module):
def __init__(self):
super(DarkNet_Tiny, self).__init__()
# backbone network : DarkNet_Tiny
self.conv_1 = Conv_BN_LeakyReLU(3, 16, 3, 1)
self.maxpool_1 = nn.MaxPool2d((2, 2), 2) # stride = 2
self.conv_2 = Conv_BN_LeakyReLU(16, 32, 3, 1)
self.maxpool_2 = nn.MaxPool2d((2, 2), 2) # stride = 4
self.conv_3 = Conv_BN_LeakyReLU(32, 64, 3, 1)
self.maxpool_3 = nn.MaxPool2d((2, 2), 2) # stride = 8
self.conv_4 = Conv_BN_LeakyReLU(64, 128, 3, 1)
self.maxpool_4 = nn.MaxPool2d((2, 2), 2) # stride = 16
self.conv_5 = Conv_BN_LeakyReLU(128, 256, 3, 1)
self.maxpool_5 = nn.MaxPool2d((2, 2), 2) # stride = 32
self.conv_6 = Conv_BN_LeakyReLU(256, 512, 3, 1)
self.maxpool_6 = nn.Sequential(
nn.ZeroPad2d((0, 1, 0, 1)),
nn.MaxPool2d((2, 2), 1) # stride = 32
)
self.conv_7 = Conv_BN_LeakyReLU(512, 1024, 3, 1)
def forward(self, x):
x = self.conv_1(x)
c1 = self.maxpool_1(x)
c1 = self.conv_2(c1)
c2 = self.maxpool_2(c1)
c2 = self.conv_3(c2)
c3 = self.maxpool_3(c2)
c3 = self.conv_4(c3)
c4 = self.maxpool_4(c3)
c4 = self.conv_5(c4) # stride = 16
c5 = self.maxpool_5(c4)
c5 = self.conv_6(c5)
c5 = self.maxpool_6(c5)
c5 = self.conv_7(c5) # stride = 32
output = {
'layer1': c3,
'layer2': c4,
'layer3': c5
}
return output
def build_darknet_tiny(pretrained=False):
# model
model = DarkNet_Tiny()
# load weight
if pretrained:
print('Loading pretrained weight ...')
url = model_urls['darknet_tiny']
# checkpoint state dict
checkpoint_state_dict = torch.hub.load_state_dict_from_url(
url=url, map_location="cpu", check_hash=True)
# model state dict
model_state_dict = model.state_dict()
# check
for k in list(checkpoint_state_dict.keys()):
if k in model_state_dict:
shape_model = tuple(model_state_dict[k].shape)
shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
if shape_model != shape_checkpoint:
checkpoint_state_dict.pop(k)
else:
checkpoint_state_dict.pop(k)
print(k)
model.load_state_dict(checkpoint_state_dict)
return model
if __name__ == '__main__':
import time
net = build_darknet_tiny(pretrained=True)
x = torch.randn(1, 3, 224, 224)
t0 = time.time()
output = net(x)
t1 = time.time()
print('Time: ', t1 - t0)
for k in output.keys():
print('{} : {}'.format(k, output[k].shape))
================================================
FILE: backbone/resnet.py
================================================
import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
'resnet152']
model_urls = {
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
def conv1x1(in_planes, out_planes, stride=1):
"""1x1 convolution"""
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = conv1x1(inplanes, planes)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = conv3x3(planes, planes, stride)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = conv1x1(planes, planes * self.expansion)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, layers, zero_init_residual=False):
super(ResNet, self).__init__()
self.inplanes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
# Zero-initialize the last BN in each residual branch,
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
nn.init.constant_(m.bn3.weight, 0)
elif isinstance(m, BasicBlock):
nn.init.constant_(m.bn2.weight, 0)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv1x1(self.inplanes, planes * block.expansion, stride),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
c1 = self.conv1(x)
c1 = self.bn1(c1)
c1 = self.relu(c1)
c1 = self.maxpool(c1)
c2 = self.layer1(c1)
c3 = self.layer2(c2)
c4 = self.layer3(c3)
c5 = self.layer4(c4)
output = {
'layer1': c3,
'layer2': c4,
'layer3': c5
}
return output
def resnet18(pretrained=False, **kwargs):
"""Constructs a ResNet-18 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
if pretrained:
# strict = False as we don't need fc layer params.
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
return model
def resnet34(pretrained=False, **kwargs):
"""Constructs a ResNet-34 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']), strict=False)
return model
def resnet50(pretrained=False, **kwargs):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']), strict=False)
return model
def resnet101(pretrained=False, **kwargs):
"""Constructs a ResNet-101 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']), strict=False)
return model
def resnet152(pretrained=False, **kwargs):
"""Constructs a ResNet-152 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
return model
def build_resnet(model_name='resnet18', pretrained=False):
if model_name == 'resnet18':
model = resnet18(pretrained=pretrained)
elif model_name == 'resnet34':
model = resnet34(pretrained=pretrained)
elif model_name == 'resnet50':
model = resnet50(pretrained=pretrained)
elif model_name == 'resnet101':
model = resnet101(pretrained=pretrained)
elif model_name == 'resnet152':
model = resnet152(pretrained=pretrained)
return model
if __name__ == "__main__":
import time
model = build_resnet(model_name='resnet18', pretrained=True)
x = torch.randn(1, 3, 224, 224)
t0 = time.time()
output = model(x)
t1 = time.time()
print('Time: ', t1 - t0)
for k in output.keys():
print('{} : {}'.format(k, output[k].shape))
================================================
FILE: backbone/weights/README.md
================================================
# darknet19, darknet53, darknet-tiny, darknet-light
darknet-tiny is designed by myself. It is a very simple and lightweight backbone.
darknet-light is same to the backbone used in official TinyYOLOv3.
For researchers in China, you can download them from BaiduYunDisk:
link：https://pan.baidu.com/s/1Rm87Fcj1RXZFmeTUrDWANA
password：qgzn
Also, you can download them from Google Drive:
link: https://drive.google.com/drive/folders/15saMtvYiz3yfFNu5EnC7GSltEAvTImMB?usp=sharing
================================================
FILE: data/__init__.py
================================================
from .voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES
from .coco2017 import COCODataset, coco_class_labels, coco_class_index
from .config import *
import torch
import cv2
import numpy as np
def detection_collate(batch):
"""Custom collate fn for dealing with batches of images that have a different
number of associated object annotations (bounding boxes).
Arguments:
batch: (tuple) A tuple of tensor images and lists of annotations
Return:
A tuple containing:
1) (tensor) batch of images stacked on their 0 dim
2) (list of tensors) annotations for a given image are stacked on
0 dim
"""
targets = []
imgs = []
for sample in batch:
imgs.append(sample[0])
targets.append(torch.FloatTensor(sample[1]))
return torch.stack(imgs, 0), targets
def base_transform(image, size, mean, std):
x = cv2.resize(image, (size, size)).astype(np.float32)
x /= 255.
x -= mean
x /= std
return x
class BaseTransform:
def __init__(self, size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)):
self.size = size
self.mean = np.array(mean, dtype=np.float32)
self.std = np.array(std, dtype=np.float32)
def __call__(self, image, boxes=None, labels=None):
return base_transform(image, self.size, self.mean, self.std), boxes, labels
================================================
FILE: data/coco2017.py
================================================
import os
import numpy as np
import random
import torch
from torch.utils.data import Dataset
import cv2
from pycocotools.coco import COCO
coco_class_labels = ('background',
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
'boat', 'traffic light', 'fire hydrant', 'street sign', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack', 'umbrella',
'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'mirror', 'dining table', 'window', 'desk',
'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')
coco_class_index = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20,
21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67,
70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
class COCODataset(Dataset):
"""
COCO dataset class.
"""
def __init__(self,
data_dir=None,
transform=None,
json_file='instances_train2017.json',
name='train2017'):
"""
COCO dataset initialization. Annotation data are read into memory by COCO API.
Args:
data_dir (str): dataset root directory
json_file (str): COCO json file name
name (str): COCO data name (e.g. 'train2017' or 'val2017')
img_size (int): target image size after pre-processing
min_size (int): bounding boxes smaller than this are ignored
debug (bool): if True, only one data id is selected from the dataset
"""
self.data_dir = data_dir
self.json_file = json_file
self.coco = COCO(os.path.join(self.data_dir, 'annotations', self.json_file))
self.ids = self.coco.getImgIds()
self.class_ids = sorted(self.coco.getCatIds())
self.name = name
self.transform = transform
def __len__(self):
return len(self.ids)
def pull_image(self, index):
id_ = self.ids[index]
img_file = os.path.join(self.data_dir, self.name,
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
if self.json_file == 'instances_val5k.json' and img is None:
img_file = os.path.join(self.data_dir, 'train2017',
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
return img, id_
def pull_anno(self, index):
id_ = self.ids[index]
anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
annotations = self.coco.loadAnns(anno_ids)
target = []
for anno in annotations:
if 'bbox' in anno:
xmin = np.max((0, anno['bbox'][0]))
ymin = np.max((0, anno['bbox'][1]))
xmax = xmin + anno['bbox'][2]
ymax = ymin + anno['bbox'][3]
if anno['area'] > 0 and xmax >= xmin and ymax >= ymin:
label_ind = anno['category_id']
cls_id = self.class_ids.index(label_ind)
target.append([xmin, ymin, xmax, ymax, cls_id]) # [xmin, ymin, xmax, ymax, label_ind]
else:
print('No bbox !!')
return target
def __getitem__(self, index):
img, gt, h, w = self.pull_item(index)
return img, gt
def pull_item(self, index):
id_ = self.ids[index]
anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
annotations = self.coco.loadAnns(anno_ids)
# load an image
img_file = os.path.join(self.data_dir, self.name,
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
if self.json_file == 'instances_val5k.json' and img is None:
img_file = os.path.join(self.data_dir, 'train2017',
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
assert img is not None
height, width, channels = img.shape
# load a target
target = []
for anno in annotations:
if 'bbox' in anno and anno['area'] > 0:
xmin = np.max((0, anno['bbox'][0]))
ymin = np.max((0, anno['bbox'][1]))
xmax = np.min((width - 1, xmin + np.max((0, anno['bbox'][2] - 1))))
ymax = np.min((height - 1, ymin + np.max((0, anno['bbox'][3] - 1))))
if xmax > xmin and ymax > ymin:
label_ind = anno['category_id']
cls_id = self.class_ids.index(label_ind)
xmin /= width
ymin /= height
xmax /= width
ymax /= height
target.append([xmin, ymin, xmax, ymax, cls_id]) # [xmin, ymin, xmax, ymax, label_ind]
else:
print('No bbox !!!')
# check target
if len(target) == 0:
target = np.zeros([1, 5])
else:
target = np.array(target)
# transform
if self.transform is not None:
img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])
# to rgb
img = img[:, :, (2, 1, 0)]
# to tensor
img = torch.from_numpy(img).permute(2, 0, 1).float()
target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
return img, target, height, width
if __name__ == "__main__":
def base_transform(image, size, mean):
x = cv2.resize(image, (size, size)).astype(np.float32)
x -= mean
x = x.astype(np.float32)
return x
class BaseTransform:
def __init__(self, size, mean):
self.size = size
self.mean = np.array(mean, dtype=np.float32)
def __call__(self, image, boxes=None, labels=None):
return base_transform(image, self.size, self.mean), boxes, labels
img_size = 640
dataset = COCODataset(
data_dir='/mnt/share/ssd2/dataset/COCO/',
transform=BaseTransform(img_size, (0, 0, 0)))
for i in range(1000):
im, gt, h, w = dataset.pull_item(i)
img = im.permute(1,2,0).numpy()[:, :, (2, 1, 0)].astype(np.uint8)
img = img.copy()
for box in gt:
xmin, ymin, xmax, ymax, _ = box
xmin *= img_size
ymin *= img_size
xmax *= img_size
ymax *= img_size
img = cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0,0,255), 2)
cv2.imshow('gt', img)
# cv2.imwrite(str(i)+'.jpg', img)
cv2.waitKey(0)
================================================
FILE: data/config.py
================================================
# config.py
# YOLOv2 with darknet-19
yolov2_d19_cfg = {
# network
'backbone': 'd19',
# for multi-scale trick
'train_size': 640,
'val_size': 416,
'random_size_range': [10, 19],
# anchor size
'anchor_size_voc': [[1.19, 1.98], [2.79, 4.59], [4.53, 8.92], [8.06, 5.29], [10.32, 10.65]],
'anchor_size_coco': [[0.53, 0.79], [1.71, 2.36], [2.89, 6.44], [6.33, 3.79], [9.03, 9.74]],
# train
'lr_epoch': (150, 200),
'max_epoch': 250,
'ignore_thresh': 0.5
}
# YOLOv2 with resnet-50
yolov2_r50_cfg = {
# network
'backbone': 'r50',
# for multi-scale trick
'train_size': 640,
'val_size': 416,
'random_size_range': [10, 19],
# anchor size
'anchor_size_voc': [[1.19, 1.98], [2.79, 4.59], [4.53, 8.92], [8.06, 5.29], [10.32, 10.65]],
'anchor_size_coco': [[0.53, 0.79], [1.71, 2.36], [2.89, 6.44], [6.33, 3.79], [9.03, 9.74]],
# train
'lr_epoch': (150, 200),
'max_epoch': 250,
'ignore_thresh': 0.5
}
# YOLOv3 / YOLOv3Spp
yolov3_d53_cfg = {
# network
'backbone': 'd53',
# for multi-scale trick
'train_size': 640,
'val_size': 416,
'random_size_range': [10, 19],
# anchor size
'anchor_size_voc': [[32.64, 47.68], [50.24, 108.16], [126.72, 96.32],
[78.4, 201.92], [178.24, 178.56], [129.6, 294.72],
[331.84, 194.56], [227.84, 325.76], [365.44, 358.72]],
'anchor_size_coco': [[12.48, 19.2], [31.36, 46.4],[46.4, 113.92],
[97.28, 55.04], [133.12, 127.36], [79.04, 224.],
[301.12, 150.4 ], [172.16, 285.76], [348.16, 341.12]],
# train
'lr_epoch': (150, 200),
'max_epoch': 250,
'ignore_thresh': 0.5
}
# YOLOv3Tiny
yolov3_tiny_cfg = {
# network
'backbone': 'd-light',
# for multi-scale trick
'train_size': 640,
'val_size': 416,
'random_size_range':[10, 19],
# anchor size
'anchor_size_voc': [[34.01, 61.79], [86.94, 109.68], [93.49, 227.46],
[246.38, 163.33], [178.68, 306.55], [344.89, 337.14]],
'anchor_size_coco': [[15.09, 23.25], [46.36, 61.47], [68.41, 161.84],
[168.88, 93.59], [154.96, 257.45], [334.74, 302.47]],
# train
'lr_epoch': (150, 200),
'max_epoch': 250,
'ignore_thresh': 0.5
}
================================================
FILE: data/scripts/COCO2017.sh
================================================
mkdir COCO
cd COCO
wget http://images.cocodataset.org/zips/train2017.zip
wget http://images.cocodataset.org/zips/val2017.zip
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip
wget http://images.cocodataset.org/zips/test2017.zip
wget http://images.cocodataset.org/annotations/image_info_test2017.zip
unzip train2017.zip
unzip val2017.zip
unzip annotations_trainval2017.zip
unzip test2017.zip
unzip image_info_test2017.zip
# rm -f train2017.zip
# rm -f val2017.zip
# rm -f annotations_trainval2017.zip
# rm -f test2017.zip
# rm -f image_info_test2017.zip
================================================
FILE: data/scripts/VOC2007.sh
================================================
#!/bin/bash
# Ellis Brown
start=`date +%s`
# handle optional download dir
if [ -z "$1" ]
then
# navigate to ~/data
echo "navigating to ~/data/ ..."
mkdir -p ~/data
cd ~/data/
else
# check if is valid directory
if [ ! -d $1 ]; then
echo $1 "is not a valid directory"
exit 0
fi
echo "navigating to" $1 "..."
cd $1
fi
echo "Downloading VOC2007 trainval ..."
# Download the data.
curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
echo "Downloading VOC2007 test data ..."
curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
echo "Done downloading."
# Extract data
echo "Extracting trainval ..."
tar -xvf VOCtrainval_06-Nov-2007.tar
echo "Extracting test ..."
tar -xvf VOCtest_06-Nov-2007.tar
echo "removing tars ..."
rm VOCtrainval_06-Nov-2007.tar
rm VOCtest_06-Nov-2007.tar
end=`date +%s`
runtime=$((end-start))
echo "Completed in" $runtime "seconds"
================================================
FILE: data/scripts/VOC2012.sh
================================================
#!/bin/bash
# Ellis Brown
start=`date +%s`
# handle optional download dir
if [ -z "$1" ]
then
# navigate to ~/data
echo "navigating to ~/data/ ..."
mkdir -p ~/data
cd ~/data/
else
# check if is valid directory
if [ ! -d $1 ]; then
echo $1 "is not a valid directory"
exit 0
fi
echo "navigating to" $1 "..."
cd $1
fi
echo "Downloading VOC2012 trainval ..."
# Download the data.
curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
echo "Done downloading."
# Extract data
echo "Extracting trainval ..."
tar -xvf VOCtrainval_11-May-2012.tar
echo "removing tar ..."
rm VOCtrainval_11-May-2012.tar
end=`date +%s`
runtime=$((end-start))
echo "Completed in" $runtime "seconds"
================================================
FILE: data/voc0712.py
================================================
"""VOC Dataset Classes
Original author: Francisco Massa
https://github.com/fmassa/vision/blob/voc_dataset/torchvision/datasets/voc.py
Updated by: Ellis Brown, Max deGroot
"""
import os.path as osp
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
import random
import xml.etree.ElementTree as ET
VOC_CLASSES = ( # always index 0
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor')
class VOCAnnotationTransform(object):
"""Transforms a VOC annotation into a Tensor of bbox coords and label index
Initilized with a dictionary lookup of classnames to indexes
Arguments:
class_to_ind (dict, optional): dictionary lookup of classnames -> indexes
(default: alphabetic indexing of VOC's 20 classes)
keep_difficult (bool, optional): keep difficult instances or not
(default: False)
height (int): height
width (int): width
"""
def __init__(self, class_to_ind=None, keep_difficult=False):
self.class_to_ind = class_to_ind or dict(
zip(VOC_CLASSES, range(len(VOC_CLASSES))))
self.keep_difficult = keep_difficult
def __call__(self, target, width, height):
"""
Arguments:
target (annotation) : the target annotation to be made usable
will be an ET.Element
Returns:
a list containing lists of bounding boxes [bbox coords, class name]
"""
res = []
for obj in target.iter('object'):
difficult = int(obj.find('difficult').text) == 1
if not self.keep_difficult and difficult:
continue
name = obj.find('name').text.lower().strip()
bbox = obj.find('bndbox')
pts = ['xmin', 'ymin', 'xmax', 'ymax']
bndbox = []
for i, pt in enumerate(pts):
cur_pt = int(bbox.find(pt).text) - 1
# scale height or width
cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height
bndbox.append(cur_pt)
label_idx = self.class_to_ind[name]
bndbox.append(label_idx)
res += [bndbox] # [xmin, ymin, xmax, ymax, label_ind]
# img_id = target.find('filename').text[:-4]
return res # [[xmin, ymin, xmax, ymax, label_ind], ... ]
class VOCDetection(data.Dataset):
"""VOC Detection Dataset Object
input is image, target is annotation
Arguments:
root (string): filepath to VOCdevkit folder.
image_set (string): imageset to use (eg. 'train', 'val', 'test')
transform (callable, optional): transformation to perform on the
input image
target_transform (callable, optional): transformation to perform on the
target `annotation`
(eg: take in caption string, return tensor of word indices)
dataset_name (string, optional): which dataset to load
(default: 'VOC2007')
"""
def __init__(self,
data_dir=None,
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
transform=None,
target_transform=VOCAnnotationTransform(),
dataset_name='VOC0712'):
self.root = data_dir
self.image_set = image_sets
self.transform = transform
self.target_transform = target_transform
self.name = dataset_name
self._annopath = osp.join('%s', 'Annotations', '%s.xml')
self._imgpath = osp.join('%s', 'JPEGImages', '%s.jpg')
self.ids = list()
for (year, name) in image_sets:
rootpath = osp.join(self.root, 'VOC' + year)
for line in open(osp.join(rootpath, 'ImageSets', 'Main', name + '.txt')):
self.ids.append((rootpath, line.strip()))
def __getitem__(self, index):
im, gt, h, w = self.pull_item(index)
return im, gt
def __len__(self):
return len(self.ids)
def pull_item(self, index):
# load an image
img_id = self.ids[index]
img = cv2.imread(self._imgpath % img_id)
height, width, channels = img.shape
# load a target
target = ET.parse(self._annopath % img_id).getroot()
if self.target_transform is not None:
target = self.target_transform(target, width, height)
# check target
if len(target) == 0:
target = np.zeros([1, 5])
else:
target = np.array(target)
# transform
if self.transform is not None:
img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])
# to rgb
img = img[:, :, (2, 1, 0)]
# to tensor
img = torch.from_numpy(img).permute(2, 0, 1).float()
# target
target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
return img, target, height, width
def pull_image(self, index):
'''Returns the original image object at index in PIL form
Note: not using self.__getitem__(), as any transformations passed in
could mess up this functionality.
Argument:
index (int): index of img to show
Return:
PIL img
'''
img_id = self.ids[index]
return cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR), img_id
def pull_anno(self, index):
'''Returns the original annotation of image at index
Note: not using self.__getitem__(), as any transformations passed in
could mess up this functionality.
Argument:
index (int): index of img to get annotation of
Return:
list: [img_id, [(label, bbox coords),...]]
eg: ('001718', [('dog', (96, 13, 438, 332))])
'''
img_id = self.ids[index]
anno = ET.parse(self._annopath % img_id).getroot()
gt = self.target_transform(anno, 1, 1)
return img_id[1], gt
def pull_tensor(self, index):
'''Returns the original image at an index in tensor form
Note: not using self.__getitem__(), as any transformations passed in
could mess up this functionality.
Argument:
index (int): index of img to show
Return:
tensorized version of img, squeezed
'''
return torch.Tensor(self.pull_image(index)).unsqueeze_(0)
if __name__ == "__main__":
def base_transform(image, size, mean):
x = cv2.resize(image, (size, size)).astype(np.float32)
x -= mean
x = x.astype(np.float32)
return x
class BaseTransform:
def __init__(self, size, mean):
self.size = size
self.mean = np.array(mean, dtype=np.float32)
def __call__(self, image, boxes=None, labels=None):
return base_transform(image, self.size, self.mean), boxes, labels
img_size = 640
# dataset
dataset = VOCDetection(data_dir='/mnt/share/ssd2/dataset/VOCdevkit/',
image_sets=[('2007', 'trainval')],
transform=BaseTransform(img_size, (0, 0, 0)))
for i in range(1000):
im, gt, h, w = dataset.pull_item(i)
img = im.permute(1,2,0).numpy()[:, :, (2, 1, 0)].astype(np.uint8)
img = img.copy()
for box in gt:
xmin, ymin, xmax, ymax, _ = box
xmin *= img_size
ymin *= img_size
xmax *= img_size
ymax *= img_size
img = cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0,0,255), 2)
cv2.imshow('gt', img)
cv2.waitKey(0)
================================================
FILE: demo.py
================================================
import argparse
import os
import numpy as np
import cv2
import time
import torch
from data.coco2017 import coco_class_index, coco_class_labels
from data import config, BaseTransform
def parse_args():
parser = argparse.ArgumentParser(description='YOLO Demo Detection')
# basic
parser.add_argument('--mode', default='image',
type=str, help='Use the data from image, video or camera')
parser.add_argument('-size', '--input_size', default=416, type=int,
help='input_size')
parser.add_argument('--cuda', action='store_true', default=False,
help='Use cuda')
parser.add_argument('--path_to_img', default='data/demo/images/',
type=str, help='The path to image files')
parser.add_argument('--path_to_vid', default='data/demo/videos/',
type=str, help='The path to video files')
parser.add_argument('--path_to_save', default='det_results/',
type=str, help='The path to save the detection results')
parser.add_argument('-vs', '--visual_threshold', default=0.3,
type=float, help='visual threshold')
# model
parser.add_argument('-v', '--version', default='yolo_v2',
help='yolov2_d19, yolov2_r50, yolov2_slim, yolov3, yolov3_spp, yolov3_tiny')
parser.add_argument('--conf_thresh', default=0.1, type=float,
help='NMS threshold')
parser.add_argument('--nms_thresh', default=0.45, type=float,
help='NMS threshold')
parser.add_argument('--trained_model', default='weights/',
type=str, help='Trained state_dict file path to open')
return parser.parse_args()
def plot_bbox_labels(img, bbox, label, cls_color, test_scale=0.4):
x1, y1, x2, y2 = bbox
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
t_size = cv2.getTextSize(label, 0, fontScale=1, thickness=2)[0]
# plot bbox
cv2.rectangle(img, (x1, y1), (x2, y2), cls_color, 2)
# plot title bbox
cv2.rectangle(img, (x1, y1-t_size[1]), (int(x1 + t_size[0] * test_scale), y1), cls_color, -1)
# put the test on the title bbox
cv2.putText(img, label, (int(x1), int(y1 - 5)), 0, test_scale, (0, 0, 0), 1, lineType=cv2.LINE_AA)
return img
def visualize(img, bboxes, scores, cls_inds, class_colors, vis_thresh=0.3):
ts = 0.4
for i, bbox in enumerate(bboxes):
if scores[i] > vis_thresh:
cls_color = class_colors[int(cls_inds[i])]
cls_id = coco_class_index[int(cls_inds[i])]
mess = '%s: %.2f' % (coco_class_labels[cls_id], scores[i])
img = plot_bbox_labels(img, bbox, mess, cls_color, test_scale=ts)
return img
def detect(net,
device,
transform,
vis_thresh,
mode='image',
path_to_img=None,
path_to_vid=None,
path_to_save=None):
# class color
class_colors = [(np.random.randint(255),
np.random.randint(255),
np.random.randint(255)) for _ in range(80)]
save_path = os.path.join(path_to_save, mode)
os.makedirs(save_path, exist_ok=True)
# ------------------------- Camera ----------------------------
if mode == 'camera':
print('use camera !!!')
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
while True:
ret, frame = cap.read()
if ret:
if cv2.waitKey(1) == ord('q'):
break
img_h, img_w = frame.shape[:2]
scale = np.array([[img_w, img_h, img_w, img_h]])
# prepare
x = torch.from_numpy(transform(frame)[0][:, :, ::-1]).permute(2, 0, 1)
x = x.unsqueeze(0).to(device)
# inference
t0 = time.time()
bboxes, scores, cls_inds = net(x)
t1 = time.time()
print("detection time used ", t1-t0, "s")
# rescale
bboxes *= scale
frame_processed = visualize(img=frame,
bboxes=bboxes,
scores=scores,
cls_inds=cls_inds,
class_colors=class_colors,
vis_thresh=vis_thresh)
cv2.imshow('detection result', frame_processed)
cv2.waitKey(1)
else:
break
cap.release()
cv2.destroyAllWindows()
# ------------------------- Image ----------------------------
elif mode == 'image':
for i, img_id in enumerate(os.listdir(path_to_img)):
img = cv2.imread(path_to_img + '/' + img_id, cv2.IMREAD_COLOR)
img_h, img_w = img.shape[:2]
scale = np.array([[img_w, img_h, img_w, img_h]])
# prepare
x = torch.from_numpy(transform(img)[0][:, :, ::-1]).permute(2, 0, 1)
x = x.unsqueeze(0).to(device)
# inference
t0 = time.time()
bboxes, scores, cls_inds = net(x)
t1 = time.time()
print("detection time used ", t1-t0, "s")
# rescale
bboxes *= scale
img_processed = visualize(img=img,
bboxes=bboxes,
scores=scores,
cls_inds=cls_inds,
class_colors=class_colors,
vis_thresh=vis_thresh)
cv2.imshow('detection', img_processed)
cv2.imwrite(os.path.join(save_path, str(i).zfill(6)+'.jpg'), img_processed)
cv2.waitKey(0)
# ------------------------- Video ---------------------------
elif mode == 'video':
video = cv2.VideoCapture(path_to_vid)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
save_size = (640, 480)
save_path = os.path.join(save_path, 'det.avi')
fps = 15.0
out = cv2.VideoWriter(save_path, fourcc, fps, save_size)
while(True):
ret, frame = video.read()
if ret:
# ------------------------- Detection ---------------------------
img_h, img_w = frame.shape[:2]
scale = np.array([[img_w, img_h, img_w, img_h]])
# prepare
x = torch.from_numpy(transform(frame)[0][:, :, ::-1]).permute(2, 0, 1)
x = x.unsqueeze(0).to(device)
# inference
t0 = time.time()
bboxes, scores, cls_inds = net(x)
t1 = time.time()
print("detection time used ", t1-t0, "s")
# rescale
bboxes *= scale
frame_processed = visualize(img=frame,
bboxes=bboxes,
scores=scores,
cls_inds=cls_inds,
class_colors=class_colors,
vis_thresh=vis_thresh)
frame_processed_resize = cv2.resize(frame_processed, save_size)
out.write(frame_processed_resize)
cv2.imshow('detection', frame_processed)
cv2.waitKey(1)
else:
break
video.release()
out.release()
cv2.destroyAllWindows()
def run():
args = parse_args()
# use cuda
if args.cuda:
device = torch.device("cuda")
else:
device = torch.device("cpu")
# model
model_name = args.version
print('Model: ', model_name)
# load model and config file
if model_name == 'yolov2_d19':
from models.yolov2_d19 import YOLOv2D19 as yolo_net
cfg = config.yolov2_d19_cfg
elif model_name == 'yolov2_r50':
from models.yolov2_r50 import YOLOv2R50 as yolo_net
cfg = config.yolov2_r50_cfg
elif model_name == 'yolov2_slim':
from models.yolov2_slim import YOLOv2Slim as yolo_net
cfg = config.yolov2_slim_cfg
elif model_name == 'yolov3':
from models.yolov3 import YOLOv3 as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_spp':
from models.yolov3_spp import YOLOv3Spp as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_tiny':
from models.yolov3_tiny import YOLOv3tiny as yolo_net
cfg = config.yolov3_tiny_cfg
else:
print('Unknown model name...')
exit(0)
input_size = [args.input_size, args.input_size]
# build model
anchor_size = cfg['anchor_size_coco']
net = yolo_net(device=device,
input_size=input_size,
num_classes=80,
trainable=False,
conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh,
anchor_size=anchor_size)
# load weight
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.to(device).eval()
print('Finished loading model!')
# run
detect(net=net,
device=device,
transform=BaseTransform(input_size),
mode=args.mode,
path_to_img=args.path_to_img,
path_to_vid=args.path_to_vid,
path_to_save=args.path_to_save,
thresh=args.visual_threshold
)
if __name__ == '__main__':
run()
================================================
FILE: eval.py
================================================
import argparse
import os
import torch
from utils.vocapi_evaluator import VOCAPIEvaluator
from utils.cocoapi_evaluator import COCOAPIEvaluator
from data import BaseTransform, config
parser = argparse.ArgumentParser(description='YOLO Detector Evaluation')
parser.add_argument('-v', '--version', default='yolo_v2',
help='yolov2_d19, yolov2_r50, yolov2_slim, yolov3, yolov3_spp, yolov3_tiny')
parser.add_argument('--trained_model', type=str, default='weights/',
help='Trained state_dict file path to open')
parser.add_argument('-size', '--input_size', default=416, type=int,
help='input_size')
parser.add_argument('--cuda', action='store_true', default=False,
help='Use cuda')
# dataset
parser.add_argument('--root', default='/mnt/share/ssd2/dataset',
help='data root')
parser.add_argument('-d', '--dataset', default='coco-val',
help='voc, coco-val, coco-test.')
args = parser.parse_args()
def voc_test(model, data_dir, device, input_size):
evaluator = VOCAPIEvaluator(data_root=data_dir,
img_size=input_size,
device=device,
transform=BaseTransform(input_size),
display=True)
# VOC evaluation
evaluator.evaluate(model)
def coco_test(model, data_dir, device, input_size, test=False):
if test:
# test-dev
print('test on test-dev 2017')
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=input_size,
device=device,
testset=True,
transform=BaseTransform(input_size)
)
else:
# eval
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=input_size,
device=device,
testset=False,
transform=BaseTransform(input_size)
)
# COCO evaluation
evaluator.evaluate(model)
if __name__ == '__main__':
# dataset
if args.dataset == 'voc':
print('eval on voc ...')
num_classes = 20
data_dir = os.path.join(args.root, 'VOCdevkit')
elif args.dataset == 'coco-val':
print('eval on coco-val ...')
num_classes = 80
data_dir = os.path.join(args.root, 'COCO')
elif args.dataset == 'coco-test':
print('eval on coco-test-dev ...')
num_classes = 80
data_dir = os.path.join(args.root, 'COCO')
else:
print('unknow dataset !! we only support voc, coco-val, coco-test !!!')
exit(0)
# cuda
if args.cuda:
print('use cuda')
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# model
model_name = args.version
print('Model: ', model_name)
# load model and config file
if model_name == 'yolov2_d19':
from models.yolov2_d19 import YOLOv2D19 as yolo_net
cfg = config.yolov2_d19_cfg
elif model_name == 'yolov2_r50':
from models.yolov2_r50 import YOLOv2R50 as yolo_net
cfg = config.yolov2_r50_cfg
elif model_name == 'yolov2_slim':
from models.yolov2_slim import YOLOv2Slim as yolo_net
cfg = config.yolov2_slim_cfg
elif model_name == 'yolov3':
from models.yolov3 import YOLOv3 as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_spp':
from models.yolov3_spp import YOLOv3Spp as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_tiny':
from models.yolov3_tiny import YOLOv3tiny as yolo_net
cfg = config.yolov3_tiny_cfg
else:
print('Unknown model name...')
exit(0)
# input size
input_size = args.input_size
# build model
anchor_size = cfg['anchor_size_voc'] if args.dataset == 'voc' else cfg['anchor_size_coco']
net = yolo_net(device=device,
input_size=input_size,
num_classes=num_classes,
trainable=False,
anchor_size=anchor_size)
# load net
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.eval()
print('Finished loading model!')
net = net.to(device)
# evaluation
with torch.no_grad():
if args.dataset == 'voc':
voc_test(net, data_dir, device, input_size)
elif args.dataset == 'coco-val':
coco_test(net, data_dir, device, input_size, test=False)
elif args.dataset == 'coco-test':
coco_test(net, data_dir, device, input_size, test=True)
================================================
FILE: models/yolov2_d19.py
================================================
import numpy as np
import torch
import torch.nn as nn
from utils.modules import Conv, reorg_layer
from backbone import build_backbone
import tools
class YOLOv2D19(nn.Module):
def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.001, nms_thresh=0.5, anchor_size=None):
super(YOLOv2D19, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.anchor_size = torch.tensor(anchor_size)
self.num_anchors = len(anchor_size)
self.stride = 32
self.grid_cell, self.all_anchor_wh = self.create_grid(input_size)
# backbone darknet-19
self.backbone = build_backbone(model_name='darknet19', pretrained=trainable)
# detection head
self.convsets_1 = nn.Sequential(
Conv(1024, 1024, k=3, p=1),
Conv(1024, 1024, k=3, p=1)
)
self.route_layer = Conv(512, 64, k=1)
self.reorg = reorg_layer(stride=2)
self.convsets_2 = Conv(1280, 1024, k=3, p=1)
# prediction layer
self.pred = nn.Conv2d(1024, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
def create_grid(self, input_size):
w, h = input_size, input_size
# generate grid cells
ws, hs = w // self.stride, h // self.stride
grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
grid_xy = grid_xy.view(1, hs*ws, 1, 2).to(self.device)
# generate anchor_wh tensor
anchor_wh = self.anchor_size.repeat(hs*ws, 1, 1).unsqueeze(0).to(self.device)
return grid_xy, anchor_wh
def set_grid(self, input_size):
self.input_size = input_size
self.grid_cell, self.all_anchor_wh = self.create_grid(input_size)
def decode_xywh(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
xywh_pred : [B, H*W*anchor_n, 4] \n
"""
B, HW, ab_n, _ = txtytwth_pred.size()
# b_x = sigmoid(tx) + gride_x
# b_y = sigmoid(ty) + gride_y
xy_pred = torch.sigmoid(txtytwth_pred[..., :2]) + self.grid_cell
# b_w = anchor_w * exp(tw)
# b_h = anchor_h * exp(th)
wh_pred = torch.exp(txtytwth_pred[..., 2:]) * self.all_anchor_wh
# [B, H*W, anchor_n, 4] -> [B, H*W*anchor_n, 4]
xywh_pred = torch.cat([xy_pred, wh_pred], -1).view(B, -1, 4) * self.stride
return xywh_pred
def decode_boxes(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
x1y1x2y2_pred : [B, H*W*anchor_n, 4] \n
"""
# txtytwth -> cxcywh
xywh_pred = self.decode_xywh(txtytwth_pred)
# cxcywh -> x1y1x2y2
x1y1x2y2_pred = torch.zeros_like(xywh_pred)
x1y1_pred = xywh_pred[..., :2] - xywh_pred[..., 2:] * 0.5
x2y2_pred = xywh_pred[..., :2] + xywh_pred[..., 2:] * 0.5
x1y1x2y2_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
return x1y1x2y2_pred
def nms(self, dets, scores):
""""Pure Python NMS baseline."""
x1 = dets[:, 0] #xmin
y1 = dets[:, 1] #ymin
x2 = dets[:, 2] #xmax
y2 = dets[:, 3] #ymax
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(1e-10, xx2 - xx1)
h = np.maximum(1e-10, yy2 - yy1)
inter = w * h
# Cross Area / (bbox + particular area - Cross Area)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#reserve all the boundingbox whose ovr less than thresh
inds = np.where(ovr <= self.nms_thresh)[0]
order = order[inds + 1]
return keep
def postprocess(self, bboxes, scores):
"""
bboxes: (HxW, 4), bsize = 1
scores: (HxW, num_classes), bsize = 1
"""
cls_inds = np.argmax(scores, axis=1)
scores = scores[(np.arange(scores.shape[0]), cls_inds)]
# threshold
keep = np.where(scores >= self.conf_thresh)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# NMS
keep = np.zeros(len(bboxes), dtype=np.int)
for i in range(self.num_classes):
inds = np.where(cls_inds == i)[0]
if len(inds) == 0:
continue
c_bboxes = bboxes[inds]
c_scores = scores[inds]
c_keep = self.nms(c_bboxes, c_scores)
keep[inds[c_keep]] = 1
keep = np.where(keep > 0)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
return bboxes, scores, cls_inds
@ torch.no_grad()
def inference(self, x):
# backbone
feats = self.backbone(x)
# reorg layer
p5 = self.convsets_1(feats['layer3'])
p4 = self.reorg(self.route_layer(feats['layer2']))
p5 = torch.cat([p4, p5], dim=1)
# head
p5 = self.convsets_2(p5)
# pred
pred = self.pred(p5)
B, abC, H, W = pred.size()
# [B, num_anchor * C, H, W] -> [B, H, W, num_anchor * C] -> [B, H*W, num_anchor*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, H*W, abC)
# [B, H*W*num_anchor, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, 1)
# [B, H*W, num_anchor, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, self.num_classes)
# [B, H*W, num_anchor, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
# decode box
reg_pred = reg_pred.view(B, H*W, self.num_anchors, 4)
box_pred = self.decode_boxes(reg_pred)
# batch size = 1
conf_pred = conf_pred[0]
cls_pred = cls_pred[0]
box_pred = box_pred[0]
# score
scores = torch.sigmoid(conf_pred) * torch.softmax(cls_pred, dim=-1)
# normalize bbox
bboxes = torch.clamp(box_pred / self.input_size, 0., 1.)
# to cpu
scores = scores.to('cpu').numpy()
bboxes = bboxes.to('cpu').numpy()
# post-process
bboxes, scores, cls_inds = self.postprocess(bboxes, scores)
return bboxes, scores, cls_inds
def forward(self, x, target=None):
if not self.trainable:
return self.inference(x)
else:
# backbone
feats = self.backbone(x)
# reorg layer
p5 = self.convsets_1(feats['layer3'])
p4 = self.reorg(self.route_layer(feats['layer2']))
p5 = torch.cat([p4, p5], dim=1)
# head
p5 = self.convsets_2(p5)
# pred
pred = self.pred(p5)
B, abC, H, W = pred.size()
# [B, num_anchor * C, H, W] -> [B, H, W, num_anchor * C] -> [B, H*W, num_anchor*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, H*W, abC)
# [B, H*W*num_anchor, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, 1)
# [B, H*W, num_anchor, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, self.num_classes)
# [B, H*W, num_anchor, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
reg_pred = reg_pred.view(B, H*W, self.num_anchors, 4)
# decode bbox
x1y1x2y2_pred = (self.decode_boxes(reg_pred) / self.input_size).view(-1, 4)
x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)
reg_pred = reg_pred.view(B, H*W*self.num_anchors, 4)
# set conf target
iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt).view(B, -1, 1)
gt_conf = iou_pred.clone().detach()
# [obj, cls, txtytwth, x1y1x2y2] -> [conf, obj, cls, txtytwth]
target = torch.cat([gt_conf, target[:, :, :7]], dim=2)
# loss
(
conf_loss,
cls_loss,
bbox_loss,
iou_loss
) = tools.loss(pred_conf=conf_pred,
pred_cls=cls_pred,
pred_txtytwth=reg_pred,
pred_iou=iou_pred,
label=target
)
return conf_loss, cls_loss, bbox_loss, iou_loss
================================================
FILE: models/yolov2_r50.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.modules import Conv, reorg_layer
from backbone import build_backbone
import numpy as np
import tools
class YOLOv2R50(nn.Module):
def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.001, nms_thresh=0.6, anchor_size=None, hr=False):
super(YOLOv2R50, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.anchor_size = torch.tensor(anchor_size)
self.num_anchors = len(anchor_size)
self.stride = 32
self.grid_cell, self.all_anchor_wh = self.create_grid(input_size)
# backbone
self.backbone = build_backbone(model_name='resnet50', pretrained=trainable)
# head
self.convsets_1 = nn.Sequential(
Conv(2048, 1024, k=1),
Conv(1024, 1024, k=3, p=1),
Conv(1024, 1024, k=3, p=1)
)
# reorg
self.route_layer = Conv(1024, 128, k=1)
self.reorg = reorg_layer(stride=2)
# head
self.convsets_2 = Conv(1024+128*4, 1024, k=3, p=1)
# pred
self.pred = nn.Conv2d(1024, self.num_anchors*(1 + 4 + self.num_classes), 1)
if self.trainable:
# init bias
self.init_bias()
def init_bias(self):
# init bias
init_prob = 0.01
bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob))
nn.init.constant_(self.pred.bias[..., :self.num_anchors], bias_value)
def create_grid(self, input_size):
w, h = input_size, input_size
# generate grid cells
ws, hs = w // self.stride, h // self.stride
grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
grid_xy = grid_xy.view(1, hs*ws, 1, 2).to(self.device)
# generate anchor_wh tensor
anchor_wh = self.anchor_size.repeat(hs*ws, 1, 1).unsqueeze(0).to(self.device)
return grid_xy, anchor_wh
def set_grid(self, input_size):
self.input_size = input_size
self.grid_cell, self.all_anchor_wh = self.create_grid(input_size)
def decode_xywh(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
xywh_pred : [B, H*W*anchor_n, 4] \n
"""
B, HW, ab_n, _ = txtytwth_pred.size()
# b_x = sigmoid(tx) + gride_x
# b_y = sigmoid(ty) + gride_y
xy_pred = torch.sigmoid(txtytwth_pred[:, :, :, :2]) + self.grid_cell
# b_w = anchor_w * exp(tw)
# b_h = anchor_h * exp(th)
wh_pred = torch.exp(txtytwth_pred[:, :, :, 2:]) * self.all_anchor_wh
# [H*W, anchor_n, 4] -> [H*W*anchor_n, 4]
xywh_pred = torch.cat([xy_pred, wh_pred], -1).view(B, -1, 4) * self.stride
return xywh_pred
def decode_boxes(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
x1y1x2y2_pred : [B, H*W*anchor_n, 4] \n
"""
# txtytwth -> cxcywh
xywh_pred = self.decode_xywh(txtytwth_pred)
# cxcywh -> x1y1x2y2
x1y1x2y2_pred = torch.zeros_like(xywh_pred)
x1y1_pred = xywh_pred[..., :2] - xywh_pred[..., 2:] * 0.5
x2y2_pred = xywh_pred[..., :2] + xywh_pred[..., 2:] * 0.5
x1y1x2y2_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
return x1y1x2y2_pred
def nms(self, dets, scores):
""""Pure Python NMS baseline."""
x1 = dets[:, 0] #xmin
y1 = dets[:, 1] #ymin
x2 = dets[:, 2] #xmax
y2 = dets[:, 3] #ymax
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(1e-10, xx2 - xx1)
h = np.maximum(1e-10, yy2 - yy1)
inter = w * h
# Cross Area / (bbox + particular area - Cross Area)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#reserve all the boundingbox whose ovr less than thresh
inds = np.where(ovr <= self.nms_thresh)[0]
order = order[inds + 1]
return keep
def postprocess(self, bboxes, scores):
"""
bboxes: (HxW, 4), bsize = 1
scores: (HxW, num_classes), bsize = 1
"""
cls_inds = np.argmax(scores, axis=1)
scores = scores[(np.arange(scores.shape[0]), cls_inds)]
# threshold
keep = np.where(scores >= self.conf_thresh)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# NMS
keep = np.zeros(len(bboxes), dtype=np.int)
for i in range(self.num_classes):
inds = np.where(cls_inds == i)[0]
if len(inds) == 0:
continue
c_bboxes = bboxes[inds]
c_scores = scores[inds]
c_keep = self.nms(c_bboxes, c_scores)
keep[inds[c_keep]] = 1
keep = np.where(keep > 0)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
return bboxes, scores, cls_inds
@ torch.no_grad()
def inference(self, x):
# backbone
feats = self.backbone(x)
# reorg layer
p5 = self.convsets_1(feats['layer3'])
p4 = self.reorg(self.route_layer(feats['layer2']))
p5 = torch.cat([p4, p5], dim=1)
# head
p5 = self.convsets_2(p5)
# pred
pred = self.pred(p5)
B, abC, H, W = pred.size()
# [B, num_anchor * C, H, W] -> [B, H, W, num_anchor * C] -> [B, H*W, num_anchor*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, H*W, abC)
# [B, H*W*num_anchor, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, 1)
# [B, H*W, num_anchor, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, self.num_classes)
# [B, H*W, num_anchor, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
# decode box
reg_pred = reg_pred.view(B, H*W, self.num_anchors, 4)
box_pred = self.decode_boxes(reg_pred)
# batch size = 1
conf_pred = conf_pred[0]
cls_pred = cls_pred[0]
box_pred = box_pred[0]
# score
scores = torch.sigmoid(conf_pred) * torch.softmax(cls_pred, dim=-1)
# normalize bbox
bboxes = torch.clamp(box_pred / self.input_size, 0., 1.)
# to cpu
scores = scores.to('cpu').numpy()
bboxes = bboxes.to('cpu').numpy()
# post-process
bboxes, scores, cls_inds = self.postprocess(bboxes, scores)
return bboxes, scores, cls_inds
def forward(self, x, target=None):
if not self.trainable:
return self.inference(x)
else:
# backbone
feats = self.backbone(x)
# reorg layer
p5 = self.convsets_1(feats['layer3'])
p4 = self.reorg(self.route_layer(feats['layer2']))
p5 = torch.cat([p4, p5], dim=1)
# head
p5 = self.convsets_2(p5)
# pred
pred = self.pred(p5)
B, abC, H, W = pred.size()
# [B, num_anchor * C, H, W] -> [B, H, W, num_anchor * C] -> [B, H*W, num_anchor*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, H*W, abC)
# [B, H*W*num_anchor, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, 1)
# [B, H*W, num_anchor, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, H*W*self.num_anchors, self.num_classes)
# [B, H*W, num_anchor, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
reg_pred = reg_pred.view(B, H*W, self.num_anchors, 4)
# decode bbox
x1y1x2y2_pred = (self.decode_boxes(reg_pred) / self.input_size).view(-1, 4)
x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)
reg_pred = reg_pred.view(B, H*W*self.num_anchors, 4)
# set conf target
iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt).view(B, -1, 1)
gt_conf = iou_pred.clone().detach()
# [obj, cls, txtytwth, x1y1x2y2] -> [conf, obj, cls, txtytwth]
target = torch.cat([gt_conf, target[:, :, :7]], dim=2)
# loss
(
conf_loss,
cls_loss,
bbox_loss,
iou_loss
) = tools.loss(pred_conf=conf_pred,
pred_cls=cls_pred,
pred_txtytwth=reg_pred,
pred_iou=iou_pred,
label=target
)
return conf_loss, cls_loss, bbox_loss, iou_loss
================================================
FILE: models/yolov3.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.modules import Conv
from backbone import build_backbone
import numpy as np
import tools
class YOLOv3(nn.Module):
def __init__(self,
device,
input_size=None,
num_classes=20,
trainable=False,
conf_thresh=0.001,
nms_thresh=0.50,
anchor_size=None):
super(YOLOv3, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.topk = 3000
self.stride = [8, 16, 32]
self.anchor_size = torch.tensor(anchor_size).view(3, len(anchor_size) // 3, 2)
self.num_anchors = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
# backbone
self.backbone = build_backbone(model_name='darknet53', pretrained=trainable)
# s = 32
self.conv_set_3 = nn.Sequential(
Conv(1024, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1)
)
self.conv_1x1_3 = Conv(512, 256, k=1)
self.extra_conv_3 = Conv(512, 1024, k=3, p=1)
self.pred_3 = nn.Conv2d(1024, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
# s = 16
self.conv_set_2 = nn.Sequential(
Conv(768, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1)
)
self.conv_1x1_2 = Conv(256, 128, k=1)
self.extra_conv_2 = Conv(256, 512, k=3, p=1)
self.pred_2 = nn.Conv2d(512, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
# s = 8
self.conv_set_1 = nn.Sequential(
Conv(384, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1)
)
self.extra_conv_1 = Conv(128, 256, k=3, p=1)
self.pred_1 = nn.Conv2d(256, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
self.init_yolo()
def init_yolo(self):
# Init head
init_prob = 0.01
bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob))
# init obj&cls pred
for pred in [self.pred_1, self.pred_2, self.pred_3]:
nn.init.constant_(pred.bias[..., :self.num_anchors], bias_value)
nn.init.constant_(pred.bias[..., self.num_anchors : (1 + self.num_classes) * self.num_anchors], bias_value)
def create_grid(self, input_size):
total_grid_xy = []
total_stride = []
total_anchor_wh = []
w, h = input_size, input_size
for ind, s in enumerate(self.stride):
# generate grid cells
ws, hs = w // s, h // s
grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
grid_xy = grid_xy.view(1, hs*ws, 1, 2)
# generate stride tensor
stride_tensor = torch.ones([1, hs*ws, self.num_anchors, 2]) * s
# generate anchor_wh tensor
anchor_wh = self.anchor_size[ind].repeat(hs*ws, 1, 1)
total_grid_xy.append(grid_xy)
total_stride.append(stride_tensor)
total_anchor_wh.append(anchor_wh)
total_grid_xy = torch.cat(total_grid_xy, dim=1).to(self.device)
total_stride = torch.cat(total_stride, dim=1).to(self.device)
total_anchor_wh = torch.cat(total_anchor_wh, dim=0).to(self.device).unsqueeze(0)
return total_grid_xy, total_stride, total_anchor_wh
def set_grid(self, input_size):
self.input_size = input_size
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
def decode_xywh(self, txtytwth_pred):
"""
Input:
txtytwth_pred : [B, H*W, anchor_n, 4] containing [tx, ty, tw, th]
Output:
xywh_pred : [B, H*W*anchor_n, 4] containing [x, y, w, h]
"""
# b_x = sigmoid(tx) + gride_x, b_y = sigmoid(ty) + gride_y
B, HW, ab_n, _ = txtytwth_pred.size()
c_xy_pred = (torch.sigmoid(txtytwth_pred[..., :2]) + self.grid_cell) * self.stride_tensor
# b_w = anchor_w * exp(tw), b_h = anchor_h * exp(th)
b_wh_pred = torch.exp(txtytwth_pred[..., 2:]) * self.all_anchors_wh
# [B, H*W, anchor_n, 4] -> [B, H*W*anchor_n, 4]
xywh_pred = torch.cat([c_xy_pred, b_wh_pred], -1).view(B, HW*ab_n, 4)
return xywh_pred
def decode_boxes(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
x1y1x2y2_pred : [B, H*W*anchor_n, 4] \n
"""
# txtytwth -> cxcywh
xywh_pred = self.decode_xywh(txtytwth_pred)
# cxcywh -> x1y1x2y2
x1y1x2y2_pred = torch.zeros_like(xywh_pred)
x1y1_pred = xywh_pred[..., :2] - xywh_pred[..., 2:] * 0.5
x2y2_pred = xywh_pred[..., :2] + xywh_pred[..., 2:] * 0.5
x1y1x2y2_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
return x1y1x2y2_pred
def nms(self, dets, scores):
""""Pure Python NMS baseline."""
x1 = dets[:, 0] #xmin
y1 = dets[:, 1] #ymin
x2 = dets[:, 2] #xmax
y2 = dets[:, 3] #ymax
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(1e-10, xx2 - xx1)
h = np.maximum(1e-10, yy2 - yy1)
inter = w * h
# Cross Area / (bbox + particular area - Cross Area)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#reserve all the boundingbox whose ovr less than thresh
inds = np.where(ovr <= self.nms_thresh)[0]
order = order[inds + 1]
return keep
def postprocess(self, bboxes, scores):
"""
bboxes: (HxW, 4), bsize = 1
scores: (HxW, num_classes), bsize = 1
"""
cls_inds = np.argmax(scores, axis=1)
scores = scores[(np.arange(scores.shape[0]), cls_inds)]
# threshold
keep = np.where(scores >= self.conf_thresh)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# NMS
keep = np.zeros(len(bboxes), dtype=np.int)
for i in range(self.num_classes):
inds = np.where(cls_inds == i)[0]
if len(inds) == 0:
continue
c_bboxes = bboxes[inds]
c_scores = scores[inds]
c_keep = self.nms(c_bboxes, c_scores)
keep[inds[c_keep]] = 1
keep = np.where(keep > 0)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# topk
scores_sorted, scores_sorted_inds = np.sort(scores), np.argsort(scores)
topk_scores, topk_scores_inds = scores_sorted[:self.topk], scores_sorted_inds[:self.topk]
topk_bboxes = bboxes[topk_scores_inds]
topk_cls_inds = cls_inds[topk_scores_inds]
return topk_bboxes, topk_scores, topk_cls_inds
@torch.no_grad()
def inference(self, x):
B = x.size(0)
# backbone
feats = self.backbone(x)
c3, c4, c5 = feats['layer1'], feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_3(c5)
p5_up = F.interpolate(self.conv_1x1_3(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], 1)
p4 = self.conv_set_2(p4)
p4_up = F.interpolate(self.conv_1x1_2(p4), scale_factor=2.0, mode='bilinear', align_corners=True)
p3 = torch.cat([c3, p4_up], 1)
p3 = self.conv_set_1(p3)
# head
# s = 32
p5 = self.extra_conv_3(p5)
pred_3 = self.pred_3(p5)
# s = 16
p4 = self.extra_conv_2(p4)
pred_2 = self.pred_2(p4)
# s = 8
p3 = self.extra_conv_1(p3)
pred_1 = self.pred_1(p3)
preds = [pred_1, pred_2, pred_3]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
box_pred = self.decode_boxes(reg_pred)
# batch size = 1
conf_pred = conf_pred[0]
cls_pred = cls_pred[0]
box_pred = box_pred[0]
# score
scores = torch.sigmoid(conf_pred) * torch.softmax(cls_pred, dim=-1)
# normalize bbox
bboxes = torch.clamp(box_pred / self.input_size, 0., 1.)
# to cpu
scores = scores.to('cpu').numpy()
bboxes = bboxes.to('cpu').numpy()
# post-process
bboxes, scores, cls_inds = self.postprocess(bboxes, scores)
return bboxes, scores, cls_inds
def forward(self, x, target=None):
if not self.trainable:
return self.inference(x)
else:
# backbone
B = x.size(0)
# backbone
feats = self.backbone(x)
c3, c4, c5 = feats['layer1'], feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_3(c5)
p5_up = F.interpolate(self.conv_1x1_3(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], 1)
p4 = self.conv_set_2(p4)
p4_up = F.interpolate(self.conv_1x1_2(p4), scale_factor=2.0, mode='bilinear', align_corners=True)
p3 = torch.cat([c3, p4_up], 1)
p3 = self.conv_set_1(p3)
# head
# s = 32
p5 = self.extra_conv_3(p5)
pred_3 = self.pred_3(p5)
# s = 16
p4 = self.extra_conv_2(p4)
pred_2 = self.pred_2(p4)
# s = 8
p3 = self.extra_conv_1(p3)
pred_1 = self.pred_1(p3)
preds = [pred_1, pred_2, pred_3]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
x1y1x2y2_pred = (self.decode_boxes(reg_pred) / self.input_size).view(-1, 4)
reg_pred = reg_pred.view(B, -1, 4)
x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)
# set conf target
iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt).view(B, -1, 1)
gt_conf = iou_pred.clone().detach()
# [obj, cls, txtytwth, scale_weight, x1y1x2y2] -> [conf, obj, cls, txtytwth, scale_weight]
target = torch.cat([gt_conf, target[:, :, :7]], dim=2)
# loss
(
conf_loss,
cls_loss,
bbox_loss,
iou_loss
) = tools.loss(pred_conf=conf_pred,
pred_cls=cls_pred,
pred_txtytwth=reg_pred,
pred_iou=iou_pred,
label=target
)
return conf_loss, cls_loss, bbox_loss, iou_loss
================================================
FILE: models/yolov3_spp.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from utils.modules import Conv, SPP
from backbone import build_backbone
import tools
# YOLOv3 SPP
class YOLOv3Spp(nn.Module):
def __init__(self,
device,
input_size=None,
num_classes=20,
trainable=False,
conf_thresh=0.001,
nms_thresh=0.50,
anchor_size=None):
super(YOLOv3Spp, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = [8, 16, 32]
self.anchor_size = torch.tensor(anchor_size).view(3, len(anchor_size) // 3, 2)
self.num_anchors = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
# backbone
self.backbone = build_backbone(model_name='darknet53', pretrained=trainable)
# s = 32
self.conv_set_3 = nn.Sequential(
SPP(),
Conv(1024*4, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1)
)
self.conv_1x1_3 = Conv(512, 256, k=1)
self.extra_conv_3 = Conv(512, 1024, k=3, p=1)
self.pred_3 = nn.Conv2d(1024, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
# s = 16
self.conv_set_2 = nn.Sequential(
Conv(768, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1)
)
self.conv_1x1_2 = Conv(256, 128, k=1)
self.extra_conv_2 = Conv(256, 512, k=3, p=1)
self.pred_2 = nn.Conv2d(512, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
# s = 8
self.conv_set_1 = nn.Sequential(
Conv(384, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1)
)
self.extra_conv_1 = Conv(128, 256, k=3, p=1)
self.pred_1 = nn.Conv2d(256, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
self.init_yolo()
def init_yolo(self):
# Init head
init_prob = 0.01
bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob))
# init obj&cls pred
for pred in [self.pred_1, self.pred_2, self.pred_3]:
nn.init.constant_(pred.bias[..., :self.num_anchors], bias_value)
nn.init.constant_(pred.bias[..., self.num_anchors : (1 + self.num_classes) * self.num_anchors], bias_value)
def create_grid(self, input_size):
total_grid_xy = []
total_stride = []
total_anchor_wh = []
w, h = input_size, input_size
for ind, s in enumerate(self.stride):
# generate grid cells
ws, hs = w // s, h // s
grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
grid_xy = grid_xy.view(1, hs*ws, 1, 2)
# generate stride tensor
stride_tensor = torch.ones([1, hs*ws, self.num_anchors, 2]) * s
# generate anchor_wh tensor
anchor_wh = self.anchor_size[ind].repeat(hs*ws, 1, 1)
total_grid_xy.append(grid_xy)
total_stride.append(stride_tensor)
total_anchor_wh.append(anchor_wh)
total_grid_xy = torch.cat(total_grid_xy, dim=1).to(self.device)
total_stride = torch.cat(total_stride, dim=1).to(self.device)
total_anchor_wh = torch.cat(total_anchor_wh, dim=0).to(self.device).unsqueeze(0)
return total_grid_xy, total_stride, total_anchor_wh
def set_grid(self, input_size):
self.input_size = input_size
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
def decode_xywh(self, txtytwth_pred):
"""
Input:
txtytwth_pred : [B, H*W, anchor_n, 4] containing [tx, ty, tw, th]
Output:
xywh_pred : [B, H*W*anchor_n, 4] containing [x, y, w, h]
"""
# b_x = sigmoid(tx) + gride_x, b_y = sigmoid(ty) + gride_y
B, HW, ab_n, _ = txtytwth_pred.size()
c_xy_pred = (torch.sigmoid(txtytwth_pred[:, :, :, :2]) + self.grid_cell) * self.stride_tensor
# b_w = anchor_w * exp(tw), b_h = anchor_h * exp(th)
b_wh_pred = torch.exp(txtytwth_pred[:, :, :, 2:]) * self.all_anchors_wh
# [B, H*W, anchor_n, 4] -> [B, H*W*anchor_n, 4]
xywh_pred = torch.cat([c_xy_pred, b_wh_pred], -1).view(B, HW*ab_n, 4)
return xywh_pred
def decode_boxes(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
x1y1x2y2_pred : [B, H*W*anchor_n, 4] \n
"""
# txtytwth -> cxcywh
xywh_pred = self.decode_xywh(txtytwth_pred)
# cxcywh -> x1y1x2y2
x1y1x2y2_pred = torch.zeros_like(xywh_pred)
x1y1_pred = xywh_pred[..., :2] - xywh_pred[..., 2:] * 0.5
x2y2_pred = xywh_pred[..., :2] + xywh_pred[..., 2:] * 0.5
x1y1x2y2_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
return x1y1x2y2_pred
def nms(self, dets, scores):
""""Pure Python NMS baseline."""
x1 = dets[:, 0] #xmin
y1 = dets[:, 1] #ymin
x2 = dets[:, 2] #xmax
y2 = dets[:, 3] #ymax
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(1e-10, xx2 - xx1)
h = np.maximum(1e-10, yy2 - yy1)
inter = w * h
# Cross Area / (bbox + particular area - Cross Area)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#reserve all the boundingbox whose ovr less than thresh
inds = np.where(ovr <= self.nms_thresh)[0]
order = order[inds + 1]
return keep
def postprocess(self, bboxes, scores):
"""
bboxes: (HxW, 4), bsize = 1
scores: (HxW, num_classes), bsize = 1
"""
cls_inds = np.argmax(scores, axis=1)
scores = scores[(np.arange(scores.shape[0]), cls_inds)]
# threshold
keep = np.where(scores >= self.conf_thresh)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# NMS
keep = np.zeros(len(bboxes), dtype=np.int)
for i in range(self.num_classes):
inds = np.where(cls_inds == i)[0]
if len(inds) == 0:
continue
c_bboxes = bboxes[inds]
c_scores = scores[inds]
c_keep = self.nms(c_bboxes, c_scores)
keep[inds[c_keep]] = 1
keep = np.where(keep > 0)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
return bboxes, scores, cls_inds
@torch.no_grad()
def inference(self, x):
B = x.size(0)
# backbone
feats = self.backbone(x)
c3, c4, c5 = feats['layer1'], feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_3(c5)
p5_up = F.interpolate(self.conv_1x1_3(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], 1)
p4 = self.conv_set_2(p4)
p4_up = F.interpolate(self.conv_1x1_2(p4), scale_factor=2.0, mode='bilinear', align_corners=True)
p3 = torch.cat([c3, p4_up], 1)
p3 = self.conv_set_1(p3)
# head
# s = 32
p5 = self.extra_conv_3(p5)
pred_3 = self.pred_3(p5)
# s = 16
p4 = self.extra_conv_2(p4)
pred_2 = self.pred_2(p4)
# s = 8
p3 = self.extra_conv_1(p3)
pred_1 = self.pred_1(p3)
preds = [pred_1, pred_2, pred_3]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
box_pred = self.decode_boxes(reg_pred)
# batch size = 1
conf_pred = conf_pred[0]
cls_pred = cls_pred[0]
box_pred = box_pred[0]
# score
scores = torch.sigmoid(conf_pred) * torch.softmax(cls_pred, dim=-1)
# normalize bbox
bboxes = torch.clamp(box_pred / self.input_size, 0., 1.)
# to cpu
scores = scores.to('cpu').numpy()
bboxes = bboxes.to('cpu').numpy()
# post-process
bboxes, scores, cls_inds = self.postprocess(bboxes, scores)
return bboxes, scores, cls_inds
def forward(self, x, target=None):
if not self.trainable:
return self.inference(x)
else:
# backbone
B = x.size(0)
# backbone
feats = self.backbone(x)
c3, c4, c5 = feats['layer1'], feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_3(c5)
p5_up = F.interpolate(self.conv_1x1_3(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], 1)
p4 = self.conv_set_2(p4)
p4_up = F.interpolate(self.conv_1x1_2(p4), scale_factor=2.0, mode='bilinear', align_corners=True)
p3 = torch.cat([c3, p4_up], 1)
p3 = self.conv_set_1(p3)
# head
# s = 32
p5 = self.extra_conv_3(p5)
pred_3 = self.pred_3(p5)
# s = 16
p4 = self.extra_conv_2(p4)
pred_2 = self.pred_2(p4)
# s = 8
p3 = self.extra_conv_1(p3)
pred_1 = self.pred_1(p3)
preds = [pred_1, pred_2, pred_3]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
x1y1x2y2_pred = (self.decode_boxes(reg_pred) / self.input_size).view(-1, 4)
reg_pred = reg_pred.view(B, -1, 4)
x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)
# set conf target
iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt).view(B, -1, 1)
gt_conf = iou_pred.clone().detach()
# [obj, cls, txtytwth, scale_weight, x1y1x2y2] -> [conf, obj, cls, txtytwth, scale_weight]
target = torch.cat([gt_conf, target[:, :, :7]], dim=2)
# loss
(
conf_loss,
cls_loss,
bbox_loss,
iou_loss
) = tools.loss(pred_conf=conf_pred,
pred_cls=cls_pred,
pred_txtytwth=reg_pred,
pred_iou=iou_pred,
label=target
)
return conf_loss, cls_loss, bbox_loss, iou_loss
================================================
FILE: models/yolov3_tiny.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from utils.modules import Conv
from backbone import build_backbone
import tools
# YOLOv3 Tiny
class YOLOv3tiny(nn.Module):
def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.01, nms_thresh=0.50, anchor_size=None, hr=False):
super(YOLOv3tiny, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = [16, 32]
self.anchor_size = torch.tensor(anchor_size).view(2, len(anchor_size) // 2, 2)
self.num_anchors = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
# backbone
self.backbone = build_backbone(model_name='darknet_tiny', pretrained=trainable)
# s = 32
self.conv_set_2 = Conv(1024, 256, k=3, p=1)
self.conv_1x1_2 = Conv(256, 128, k=1)
self.extra_conv_2 = Conv(256, 512, k=3, p=1)
self.pred_2 = nn.Conv2d(512, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
# s = 16
self.conv_set_1 = Conv(384, 256, k=3, p=1)
self.pred_1 = nn.Conv2d(256, self.num_anchors*(1 + 4 + self.num_classes), kernel_size=1)
self.init_yolo()
def init_yolo(self):
# Init head
init_prob = 0.01
bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob))
# init obj&cls pred
for pred in [self.pred_1, self.pred_2, self.pred_3]:
nn.init.constant_(pred.bias[..., :self.num_anchors], bias_value)
nn.init.constant_(pred.bias[..., self.num_anchors : (1 + self.num_classes) * self.num_anchors], bias_value)
def create_grid(self, input_size):
total_grid_xy = []
total_stride = []
total_anchor_wh = []
w, h = input_size, input_size
for ind, s in enumerate(self.stride):
# generate grid cells
ws, hs = w // s, h // s
grid_y, grid_x = torch.meshgrid([torch.arange(hs), torch.arange(ws)])
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
grid_xy = grid_xy.view(1, hs*ws, 1, 2)
# generate stride tensor
stride_tensor = torch.ones([1, hs*ws, self.num_anchors, 2]) * s
# generate anchor_wh tensor
anchor_wh = self.anchor_size[ind].repeat(hs*ws, 1, 1)
total_grid_xy.append(grid_xy)
total_stride.append(stride_tensor)
total_anchor_wh.append(anchor_wh)
total_grid_xy = torch.cat(total_grid_xy, dim=1).to(self.device)
total_stride = torch.cat(total_stride, dim=1).to(self.device)
total_anchor_wh = torch.cat(total_anchor_wh, dim=0).to(self.device).unsqueeze(0)
return total_grid_xy, total_stride, total_anchor_wh
def set_grid(self, input_size):
self.input_size = input_size
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
def decode_xywh(self, txtytwth_pred):
"""
Input:
txtytwth_pred : [B, H*W, anchor_n, 4] containing [tx, ty, tw, th]
Output:
xywh_pred : [B, H*W*anchor_n, 4] containing [x, y, w, h]
"""
# b_x = sigmoid(tx) + gride_x, b_y = sigmoid(ty) + gride_y
B, HW, ab_n, _ = txtytwth_pred.size()
c_xy_pred = (torch.sigmoid(txtytwth_pred[:, :, :, :2]) + self.grid_cell) * self.stride_tensor
# b_w = anchor_w * exp(tw), b_h = anchor_h * exp(th)
b_wh_pred = torch.exp(txtytwth_pred[:, :, :, 2:]) * self.all_anchors_wh
# [B, H*W, anchor_n, 4] -> [B, H*W*anchor_n, 4]
xywh_pred = torch.cat([c_xy_pred, b_wh_pred], -1).view(B, HW*ab_n, 4)
return xywh_pred
def decode_boxes(self, txtytwth_pred):
"""
Input: \n
txtytwth_pred : [B, H*W, anchor_n, 4] \n
Output: \n
x1y1x2y2_pred : [B, H*W*anchor_n, 4] \n
"""
# txtytwth -> cxcywh
xywh_pred = self.decode_xywh(txtytwth_pred)
# cxcywh -> x1y1x2y2
x1y1x2y2_pred = torch.zeros_like(xywh_pred)
x1y1_pred = xywh_pred[..., :2] - xywh_pred[..., 2:] * 0.5
x2y2_pred = xywh_pred[..., :2] + xywh_pred[..., 2:] * 0.5
x1y1x2y2_pred = torch.cat([x1y1_pred, x2y2_pred], dim=-1)
return x1y1x2y2_pred
def nms(self, dets, scores):
""""Pure Python NMS baseline."""
x1 = dets[:, 0] #xmin
y1 = dets[:, 1] #ymin
x2 = dets[:, 2] #xmax
y2 = dets[:, 3] #ymax
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(1e-10, xx2 - xx1)
h = np.maximum(1e-10, yy2 - yy1)
inter = w * h
# Cross Area / (bbox + particular area - Cross Area)
ovr = inter / (areas[i] + areas[order[1:]] - inter)
#reserve all the boundingbox whose ovr less than thresh
inds = np.where(ovr <= self.nms_thresh)[0]
order = order[inds + 1]
return keep
def postprocess(self, bboxes, scores):
"""
bboxes: (HxW, 4), bsize = 1
scores: (HxW, num_classes), bsize = 1
"""
cls_inds = np.argmax(scores, axis=1)
scores = scores[(np.arange(scores.shape[0]), cls_inds)]
# threshold
keep = np.where(scores >= self.conf_thresh)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
# NMS
keep = np.zeros(len(bboxes), dtype=np.int)
for i in range(self.num_classes):
inds = np.where(cls_inds == i)[0]
if len(inds) == 0:
continue
c_bboxes = bboxes[inds]
c_scores = scores[inds]
c_keep = self.nms(c_bboxes, c_scores)
keep[inds[c_keep]] = 1
keep = np.where(keep > 0)
bboxes = bboxes[keep]
scores = scores[keep]
cls_inds = cls_inds[keep]
return bboxes, scores, cls_inds
@torch.no_grad()
def inference(self, x):
B = x.size(0)
# backbone
feats = self.backbone(x)
c4, c5 = feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_2(c5)
p5_up = F.interpolate(self.conv_1x1_2(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], dim=1)
p4 = self.conv_set_1(p4)
# head
# s = 32
p5 = self.extra_conv_2(p5)
pred_2 = self.pred_2(p5)
# s = 16
pred_1 = self.pred_1(p4)
preds = [pred_1, pred_2]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# Divide prediction to obj_pred, xywh_pred and cls_pred
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
box_pred = self.decode_boxes(reg_pred)
# batch size = 1
conf_pred = conf_pred[0]
cls_pred = cls_pred[0]
box_pred = box_pred[0]
# score
scores = torch.sigmoid(conf_pred) * torch.softmax(cls_pred, dim=-1)
# normalize bbox
bboxes = torch.clamp(box_pred / self.input_size, 0., 1.)
# to cpu
scores = scores.to('cpu').numpy()
bboxes = bboxes.to('cpu').numpy()
# post-process
bboxes, scores, cls_inds = self.postprocess(bboxes, scores)
return bboxes, scores, cls_inds
def forward(self, x, target=None):
if not self.trainable:
return self.inference(x)
else:
# backbone
B = x.size(0)
# backbone
feats = self.backbone(x)
c4, c5 = feats['layer2'], feats['layer3']
# FPN
p5 = self.conv_set_2(c5)
p5_up = F.interpolate(self.conv_1x1_2(p5), scale_factor=2.0, mode='bilinear', align_corners=True)
p4 = torch.cat([c4, p5_up], dim=1)
p4 = self.conv_set_1(p4)
# head
# s = 32
p5 = self.extra_conv_2(p5)
pred_2 = self.pred_2(p5)
# s = 16
pred_1 = self.pred_1(p4)
preds = [pred_1, pred_2]
total_conf_pred = []
total_cls_pred = []
total_reg_pred = []
for pred in preds:
C = pred.size(1)
# [B, anchor_n * C, H, W] -> [B, H, W, anchor_n * C] -> [B, H*W, anchor_n*C]
pred = pred.permute(0, 2, 3, 1).contiguous().view(B, -1, C)
# Divide prediction to obj_pred, xywh_pred and cls_pred
# [B, H*W*anchor_n, 1]
conf_pred = pred[:, :, :1 * self.num_anchors].contiguous().view(B, -1, 1)
# [B, H*W*anchor_n, num_cls]
cls_pred = pred[:, :, 1 * self.num_anchors : (1 + self.num_classes) * self.num_anchors].contiguous().view(B, -1, self.num_classes)
# [B, H*W*anchor_n, 4]
reg_pred = pred[:, :, (1 + self.num_classes) * self.num_anchors:].contiguous()
total_conf_pred.append(conf_pred)
total_cls_pred.append(cls_pred)
total_reg_pred.append(reg_pred)
conf_pred = torch.cat(total_conf_pred, dim=1)
cls_pred = torch.cat(total_cls_pred, dim=1)
reg_pred = torch.cat(total_reg_pred, dim=1)
# decode bbox
reg_pred = reg_pred.view(B, -1, self.num_anchors, 4)
x1y1x2y2_pred = (self.decode_boxes(reg_pred) / self.input_size).view(-1, 4)
reg_pred = reg_pred.view(B, -1, 4)
x1y1x2y2_gt = target[:, :, 7:].view(-1, 4)
# set conf target
iou_pred = tools.iou_score(x1y1x2y2_pred, x1y1x2y2_gt).view(B, -1, 1)
gt_conf = iou_pred.clone().detach()
# [obj, cls, txtytwth, scale_weight, x1y1x2y2] -> [conf, obj, cls, txtytwth, scale_weight]
target = torch.cat([gt_conf, target[:, :, :7]], dim=2)
# loss
(
conf_loss,
cls_loss,
bbox_loss,
iou_loss
) = tools.loss(pred_conf=conf_pred,
pred_cls=cls_pred,
pred_txtytwth=reg_pred,
pred_iou=iou_pred,
label=target
)
return conf_loss, cls_loss, bbox_loss, iou_loss
================================================
FILE: test.py
================================================
import os
import argparse
import torch
import torch.backends.cudnn as cudnn
from data.voc0712 import VOC_CLASSES, VOCDetection
from data.coco2017 import COCODataset, coco_class_index, coco_class_labels
from data import config, BaseTransform
import numpy as np
import cv2
import time
parser = argparse.ArgumentParser(description='YOLO Detection')
# basic
parser.add_argument('-size', '--input_size', default=416, type=int,
help='input_size')
parser.add_argument('--cuda', action='store_true', default=False,
help='use cuda.')
# model
parser.add_argument('-v', '--version', default='yolo_v2',
help='yolov2_d19, yolov2_r50, yolov2_slim, yolov3, yolov3_spp, yolov3_tiny')
parser.add_argument('--trained_model', default='weight/',
type=str, help='Trained state_dict file path to open')
parser.add_argument('--conf_thresh', default=0.1, type=float,
help='Confidence threshold')
parser.add_argument('--nms_thresh', default=0.50, type=float,
help='NMS threshold')
# dataset
parser.add_argument('-root', '--data_root', default='/mnt/share/ssd2/dataset',
help='dataset root')
parser.add_argument('-d', '--dataset', default='voc',
help='voc or coco')
# visualize
parser.add_argument('-vs', '--visual_threshold', default=0.25, type=float,
help='Final confidence threshold')
parser.add_argument('--show', action='store_true', default=False,
help='show the visulization results.')
args = parser.parse_args()
def plot_bbox_labels(img, bbox, label=None, cls_color=None, text_scale=0.4):
x1, y1, x2, y2 = bbox
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
t_size = cv2.getTextSize(label, 0, fontScale=1, thickness=2)[0]
# plot bbox
cv2.rectangle(img, (x1, y1), (x2, y2), cls_color, 2)
if label is not None:
# plot title bbox
cv2.rectangle(img, (x1, y1-t_size[1]), (int(x1 + t_size[0] * text_scale), y1), cls_color, -1)
# put the test on the title bbox
cv2.putText(img, label, (int(x1), int(y1 - 5)), 0, text_scale, (0, 0, 0), 1, lineType=cv2.LINE_AA)
return img
def visualize(img,
bboxes,
scores,
cls_inds,
vis_thresh,
class_colors,
class_names,
class_indexs=None,
dataset_name='voc'):
ts = 0.4
for i, bbox in enumerate(bboxes):
if scores[i] > vis_thresh:
cls_id = int(cls_inds[i])
if dataset_name == 'coco':
cls_color = class_colors[cls_id]
cls_id = class_indexs[cls_id]
else:
cls_color = class_colors[cls_id]
if len(class_names) > 1:
mess = '%s: %.2f' % (class_names[cls_id], scores[i])
else:
cls_color = [255, 0, 0]
mess = None
img = plot_bbox_labels(img, bbox, mess, cls_color, text_scale=ts)
return img
def test(net,
device,
dataset,
transform,
vis_thresh,
class_colors=None,
class_names=None,
class_indexs=None,
dataset_name='voc'):
num_images = len(dataset)
save_path = os.path.join('det_results/', args.dataset, args.version)
os.makedirs(save_path, exist_ok=True)
for index in range(num_images):
print('Testing image {:d}/{:d}....'.format(index+1, num_images))
image, _ = dataset.pull_image(index)
h, w, _ = image.shape
scale = np.array([[w, h, w, h]])
# to tensor
x = torch.from_numpy(transform(image)[0][:, :, (2, 1, 0)]).permute(2, 0, 1)
x = x.unsqueeze(0).to(device)
t0 = time.time()
# forward
bboxes, scores, cls_inds = net(x)
print("detection time used ", time.time() - t0, "s")
# rescale
bboxes *= scale
# vis detection
img_processed = visualize(
img=image,
bboxes=bboxes,
scores=scores,
cls_inds=cls_inds,
vis_thresh=vis_thresh,
class_colors=class_colors,
class_names=class_names,
class_indexs=class_indexs,
dataset_name=dataset_name
)
if args.show:
cv2.imshow('detection', img_processed)
cv2.waitKey(0)
# save result
cv2.imwrite(os.path.join(save_path, str(index).zfill(6) +'.jpg'), img_processed)
if __name__ == '__main__':
# cuda
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# input size
input_size = args.input_size
# dataset
if args.dataset == 'voc':
print('test on voc ...')
data_dir = os.path.join(args.data_root, 'VOCdevkit')
class_names = VOC_CLASSES
class_indexs = None
num_classes = 20
dataset = VOCDetection(root=data_dir,
image_sets=[('2007', 'test')])
elif args.dataset == 'coco':
print('test on coco-val ...')
data_dir = os.path.join(args.data_root, 'COCO')
class_names = coco_class_labels
class_indexs = coco_class_index
num_classes = 80
dataset = COCODataset(
data_dir=data_dir,
json_file='instances_val2017.json',
name='val2017')
class_colors = [(np.random.randint(255),
np.random.randint(255),
np.random.randint(255)) for _ in range(num_classes)]
# model
model_name = args.version
print('Model: ', model_name)
# load model and config file
if model_name == 'yolov2_d19':
from models.yolov2_d19 import YOLOv2D19 as yolo_net
cfg = config.yolov2_d19_cfg
elif model_name == 'yolov2_r50':
from models.yolov2_r50 import YOLOv2R50 as yolo_net
cfg = config.yolov2_r50_cfg
elif model_name == 'yolov3':
from models.yolov3 import YOLOv3 as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_spp':
from models.yolov3_spp import YOLOv3Spp as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_tiny':
from models.yolov3_tiny import YOLOv3tiny as yolo_net
cfg = config.yolov3_tiny_cfg
else:
print('Unknown model name...')
exit(0)
# build model
anchor_size = cfg['anchor_size_voc'] if args.dataset == 'voc' else cfg['anchor_size_coco']
net = yolo_net(device=device,
input_size=input_size,
num_classes=num_classes,
trainable=False,
conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh,
anchor_size=anchor_size)
# load weight
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test(net=net,
device=device,
dataset=dataset,
transform=BaseTransform(input_size),
vis_thresh=args.visual_threshold,
class_colors=class_colors,
class_names=class_names,
class_indexs=class_indexs,
dataset_name=args.dataset
)
================================================
FILE: tools.py
================================================
import numpy as np
from data import *
import torch.nn as nn
import torch.nn.functional as F
# We use ignore thresh to decide which anchor box can be kept.
ignore_thresh = 0.5
class MSEWithLogitsLoss(nn.Module):
def __init__(self, reduction='mean'):
super(MSEWithLogitsLoss, self).__init__()
self.reduction = reduction
def forward(self, logits, targets, mask):
inputs = torch.sigmoid(logits)
# We ignore those whose tarhets == -1.0.
pos_id = (mask==1.0).float()
neg_id = (mask==0.0).float()
pos_loss = pos_id * (inputs - targets)**2
neg_loss = neg_id * (inputs)**2
loss = 5.0*pos_loss + 1.0*neg_loss
if self.reduction == 'mean':
batch_size = logits.size(0)
loss = torch.sum(loss) / batch_size
return loss
else:
return loss
def compute_iou(anchor_boxes, gt_box):
"""
Input:
anchor_boxes : ndarray -> [[c_x_s, c_y_s, anchor_w, anchor_h], ..., [c_x_s, c_y_s, anchor_w, anchor_h]].
gt_box : ndarray -> [c_x_s, c_y_s, anchor_w, anchor_h].
Output:
iou : ndarray -> [iou_1, iou_2, ..., iou_m], and m is equal to the number of anchor boxes.
"""
# compute the iou between anchor box and gt box
# First, change [c_x_s, c_y_s, anchor_w, anchor_h] -> [xmin, ymin, xmax, ymax]
# anchor box :
ab_x1y1_x2y2 = np.zeros([len(anchor_boxes), 4])
ab_x1y1_x2y2[:, 0] = anchor_boxes[:, 0] - anchor_boxes[:, 2] / 2 # xmin
ab_x1y1_x2y2[:, 1] = anchor_boxes[:, 1] - anchor_boxes[:, 3] / 2 # ymin
ab_x1y1_x2y2[:, 2] = anchor_boxes[:, 0] + anchor_boxes[:, 2] / 2 # xmax
ab_x1y1_x2y2[:, 3] = anchor_boxes[:, 1] + anchor_boxes[:, 3] / 2 # ymax
w_ab, h_ab = anchor_boxes[:, 2], anchor_boxes[:, 3]
# gt_box :
# We need to expand gt_box(ndarray) to the shape of anchor_boxes(ndarray), in order to compute IoU easily.
gt_box_expand = np.repeat(gt_box, len(anchor_boxes), axis=0)
gb_x1y1_x2y2 = np.zeros([len(anchor_boxes), 4])
gb_x1y1_x2y2[:, 0] = gt_box_expand[:, 0] - gt_box_expand[:, 2] / 2 # xmin
gb_x1y1_x2y2[:, 1] = gt_box_expand[:, 1] - gt_box_expand[:, 3] / 2 # ymin
gb_x1y1_x2y2[:, 2] = gt_box_expand[:, 0] + gt_box_expand[:, 2] / 2 # xmax
gb_x1y1_x2y2[:, 3] = gt_box_expand[:, 1] + gt_box_expand[:, 3] / 2 # ymin
w_gt, h_gt = gt_box_expand[:, 2], gt_box_expand[:, 3]
# Then we compute IoU between anchor_box and gt_box
S_gt = w_gt * h_gt
S_ab = w_ab * h_ab
I_w = np.minimum(gb_x1y1_x2y2[:, 2], ab_x1y1_x2y2[:, 2]) - np.maximum(gb_x1y1_x2y2[:, 0], ab_x1y1_x2y2[:, 0])
I_h = np.minimum(gb_x1y1_x2y2[:, 3], ab_x1y1_x2y2[:, 3]) - np.maximum(gb_x1y1_x2y2[:, 1], ab_x1y1_x2y2[:, 1])
S_I = I_h * I_w
U = S_gt + S_ab - S_I + 1e-20
IoU = S_I / U
return IoU
def set_anchors(anchor_size):
"""
Input:
anchor_size : list -> [[h_1, w_1], [h_2, w_2], ..., [h_n, w_n]].
Output:
anchor_boxes : ndarray -> [[0, 0, anchor_w, anchor_h],
[0, 0, anchor_w, anchor_h],
...
[0, 0, anchor_w, anchor_h]].
"""
anchor_number = len(anchor_size)
anchor_boxes = np.zeros([anchor_number, 4])
for index, size in enumerate(anchor_size):
anchor_w, anchor_h = size
anchor_boxes[index] = np.array([0, 0, anchor_w, anchor_h])
return anchor_boxes
def generate_txtytwth(gt_label, w, h, s, all_anchor_size):
xmin, ymin, xmax, ymax = gt_label[:-1]
# compute the center, width and height
c_x = (xmax + xmin) / 2 * w
c_y = (ymax + ymin) / 2 * h
box_w = (xmax - xmin) * w
box_h = (ymax - ymin) * h
if box_w < 1. or box_h < 1.:
# print('A dirty data !!!')
return False
# map the center, width and height to the feature map size
c_x_s = c_x / s
c_y_s = c_y / s
box_ws = box_w / s
box_hs = box_h / s
# the grid cell location
grid_x = int(c_x_s)
grid_y = int(c_y_s)
# generate anchor boxes
anchor_boxes = set_anchors(all_anchor_size)
gt_box = np.array([[0, 0, box_ws, box_hs]])
# compute the IoU
iou = compute_iou(anchor_boxes, gt_box)
# We consider those anchor boxes whose IoU is more than ignore thresh,
iou_mask = (iou > ignore_thresh)
result = []
if iou_mask.sum() == 0:
# We assign the anchor box with highest IoU score.
index = np.argmax(iou)
p_w, p_h = all_anchor_size[index]
tx = c_x_s - grid_x
ty = c_y_s - grid_y
tw = np.log(box_ws / p_w)
th = np.log(box_hs / p_h)
weight = 2.0 - (box_w / w) * (box_h / h)
result.append([index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax])
return result
else:
# There are more than one anchor boxes whose IoU are higher than ignore thresh.
# But we only assign only one anchor box whose IoU is the best(objectness target is 1) and ignore other
# anchor boxes whose(we set their objectness as -1 which means we will ignore them during computing obj loss )
# iou_ = iou * iou_mask
# We get the index of the best IoU
best_index = np.argmax(iou)
for index, iou_m in enumerate(iou_mask):
if iou_m:
if index == best_index:
p_w, p_h = all_anchor_size[index]
tx = c_x_s - grid_x
ty = c_y_s - grid_y
tw = np.log(box_ws / p_w)
th = np.log(box_hs / p_h)
weight = 2.0 - (box_w / w) * (box_h / h)
result.append([index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax])
else:
# we ignore other anchor boxes even if their iou scores all higher than ignore thresh
result.append([index, grid_x, grid_y, 0., 0., 0., 0., -1.0, 0., 0., 0., 0.])
return result
def gt_creator(input_size, stride, label_lists, anchor_size):
"""
Input:
input_size : list -> the size of image in the training stage.
stride : int or list -> the downSample of the CNN, such as 32, 64 and so on.
label_list : list -> [[[xmin, ymin, xmax, ymax, cls_ind], ... ], [[xmin, ymin, xmax, ymax, cls_ind], ... ]],
and len(label_list) = batch_size;
len(label_list[i]) = the number of class instance in a image;
(xmin, ymin, xmax, ymax) : the coords of a bbox whose valus is between 0 and 1;
cls_ind : the corresponding class label.
Output:
gt_tensor : ndarray -> shape = [batch_size, anchor_number, 1+1+4, grid_cell number ]
"""
# prepare the all empty gt datas
batch_size = len(label_lists)
h = w = input_size
# We make gt labels by anchor-free method and anchor-based method.
ws = w // stride
hs = h // stride
s = stride
# We use anchor boxes to build training target.
all_anchor_size = anchor_size
anchor_number = len(all_anchor_size)
gt_tensor = np.zeros([batch_size, hs, ws, anchor_number, 1+1+4+1+4])
for batch_index in range(batch_size):
for gt_label in label_lists[batch_index]:
# get a bbox coords
gt_class = int(gt_label[-1])
results = generate_txtytwth(gt_label, w, h, s, all_anchor_size)
if results:
for result in results:
index, grid_x, grid_y, tx, ty, tw, th, weight, xmin, ymin, xmax, ymax = result
if weight > 0.:
if grid_y < gt_tensor.shape[1] and grid_x < gt_tensor.shape[2]:
gt_tensor[batch_index, grid_y, grid_x, index, 0] = 1.0
gt_tensor[batch_index, grid_y, grid_x, index, 1] = gt_class
gt_tensor[batch_index, grid_y, grid_x, index, 2:6] = np.array([tx, ty, tw, th])
gt_tensor[batch_index, grid_y, grid_x, index, 6] = weight
gt_tensor[batch_index, grid_y, grid_x, index, 7:] = np.array([xmin, ymin, xmax, ymax])
else:
gt_tensor[batch_index, grid_y, grid_x, index, 0] = -1.0
gt_tensor[batch_index, grid_y, grid_x, index, 6] = -1.0
gt_tensor = gt_tensor.reshape(batch_size, hs * ws * anchor_number, 1+1+4+1+4)
return gt_tensor
def multi_gt_creator(input_size, strides, label_lists, anchor_size):
"""creator multi scales gt"""
# prepare the all empty gt datas
batch_size = len(label_lists)
h = w = input_size
num_scale = len(strides)
gt_tensor = []
all_anchor_size = anchor_size
anchor_number = len(all_anchor_size) // num_scale
for s in strides:
gt_tensor.append(np.zeros([batch_size, h//s, w//s, anchor_number, 1+1+4+1+4]))
# generate gt datas
for batch_index in range(batch_size):
for gt_label in label_lists[batch_index]:
# get a bbox coords
gt_class = int(gt_label[-1])
xmin, ymin, xmax, ymax = gt_label[:-1]
# compute the center, width and height
c_x = (xmax + xmin) / 2 * w
c_y = (ymax + ymin) / 2 * h
box_w = (xmax - xmin) * w
box_h = (ymax - ymin) * h
if box_w < 1. or box_h < 1.:
# print('A dirty data !!!')
continue
# compute the IoU
anchor_boxes = set_anchors(all_anchor_size)
gt_box = np.array([[0, 0, box_w, box_h]])
iou = compute_iou(anchor_boxes, gt_box)
# We only consider those anchor boxes whose IoU is more than ignore thresh,
iou_mask = (iou > ignore_thresh)
if iou_mask.sum() == 0:
# We assign the anchor box with highest IoU score.
index = np.argmax(iou)
# s_indx, ab_ind = index // num_scale, index % num_scale
s_indx = index // anchor_number
ab_ind = index - s_indx * anchor_number
# get the corresponding stride
s = strides[s_indx]
# get the corresponding anchor box
p_w, p_h = anchor_boxes[index, 2], anchor_boxes[index, 3]
# compute the gride cell location
c_x_s = c_x / s
c_y_s = c_y / s
grid_x = int(c_x_s)
grid_y = int(c_y_s)
# compute gt labels
tx = c_x_s - grid_x
ty = c_y_s - grid_y
tw = np.log(box_w / p_w)
th = np.log(box_h / p_h)
weight = 2.0 - (box_w / w) * (box_h / h)
if grid_y < gt_tensor[s_indx].shape[1] and grid_x < gt_tensor[s_indx].shape[2]:
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = 1.0
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 1] = gt_class
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 2:6] = np.array([tx, ty, tw, th])
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = weight
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 7:] = np.array([xmin, ymin, xmax, ymax])
else:
# There are more than one anchor boxes whose IoU are higher than ignore thresh.
# But we only assign only one anchor box whose IoU is the best(objectness target is 1) and ignore other
# anchor boxes whose(we set their objectness as -1 which means we will ignore them during computing obj loss )
# iou_ = iou * iou_mask
# We get the index of the best IoU
best_index = np.argmax(iou)
for index, iou_m in enumerate(iou_mask):
if iou_m:
if index == best_index:
# s_indx, ab_ind = index // num_scale, index % num_scale
s_indx = index // anchor_number
ab_ind = index - s_indx * anchor_number
# get the corresponding stride
s = strides[s_indx]
# get the corresponding anchor box
p_w, p_h = anchor_boxes[index, 2], anchor_boxes[index, 3]
# compute the gride cell location
c_x_s = c_x / s
c_y_s = c_y / s
grid_x = int(c_x_s)
grid_y = int(c_y_s)
# compute gt labels
tx = c_x_s - grid_x
ty = c_y_s - grid_y
tw = np.log(box_w / p_w)
th = np.log(box_h / p_h)
weight = 2.0 - (box_w / w) * (box_h / h)
if grid_y < gt_tensor[s_indx].shape[1] and grid_x < gt_tensor[s_indx].shape[2]:
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = 1.0
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 1] = gt_class
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 2:6] = np.array([tx, ty, tw, th])
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = weight
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 7:] = np.array([xmin, ymin, xmax, ymax])
else:
# we ignore other anchor boxes even if their iou scores are higher than ignore thresh
# s_indx, ab_ind = index // num_scale, index % num_scale
s_indx = index // anchor_number
ab_ind = index - s_indx * anchor_number
s = strides[s_indx]
c_x_s = c_x / s
c_y_s = c_y / s
grid_x = int(c_x_s)
grid_y = int(c_y_s)
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 0] = -1.0
gt_tensor[s_indx][batch_index, grid_y, grid_x, ab_ind, 6] = -1.0
gt_tensor = [gt.reshape(batch_size, -1, 1+1+4+1+4) for gt in gt_tensor]
gt_tensor = np.concatenate(gt_tensor, 1)
return gt_tensor
def iou_score(bboxes_a, bboxes_b):
"""
bbox_1 : [B*N, 4] = [x1, y1, x2, y2]
bbox_2 : [B*N, 4] = [x1, y1, x2, y2]
"""
tl = torch.max(bboxes_a[:, :2], bboxes_b[:, :2])
br = torch.min(bboxes_a[:, 2:], bboxes_b[:, 2:])
area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)
en = (tl < br).type(tl.type()).prod(dim=1)
area_i = torch.prod(br - tl, 1) * en # * ((tl < br).all())
return area_i / (area_a + area_b - area_i + 1e-14)
def loss(pred_conf, pred_cls, pred_txtytwth, pred_iou, label):
# loss func
conf_loss_function = MSEWithLogitsLoss(reduction='mean')
cls_loss_function = nn.CrossEntropyLoss(reduction='none')
txty_loss_function = nn.BCEWithLogitsLoss(reduction='none')
twth_loss_function = nn.MSELoss(reduction='none')
iou_loss_function = nn.SmoothL1Loss(reduction='none')
# pred
pred_conf = pred_conf[:, :, 0]
pred_cls = pred_cls.permute(0, 2, 1)
pred_txty = pred_txtytwth[:, :, :2]
pred_twth = pred_txtytwth[:, :, 2:]
pred_iou = pred_iou[:, :, 0]
# gt
gt_conf = label[:, :, 0].float()
gt_obj = label[:, :, 1].float()
gt_cls = label[:, :, 2].long()
gt_txty = label[:, :, 3:5].float()
gt_twth = label[:, :, 5:7].float()
gt_box_scale_weight = label[:, :, 7].float()
gt_iou = (gt_box_scale_weight > 0.).float()
gt_mask = (gt_box_scale_weight > 0.).float()
batch_size = pred_conf.size(0)
# objectness loss
conf_loss = conf_loss_function(pred_conf, gt_conf, gt_obj)
# class loss
cls_loss = torch.sum(cls_loss_function(pred_cls, gt_cls) * gt_mask) / batch_size
# box loss
txty_loss = torch.sum(torch.sum(txty_loss_function(pred_txty, gt_txty), dim=-1) * gt_box_scale_weight * gt_mask) / batch_size
twth_loss = torch.sum(torch.sum(twth_loss_function(pred_twth, gt_twth), dim=-1) * gt_box_scale_weight * gt_mask) / batch_size
bbox_loss = txty_loss + twth_loss
# iou loss
iou_loss = torch.sum(iou_loss_function(pred_iou, gt_iou) * gt_mask) / batch_size
return conf_loss, cls_loss, bbox_loss, iou_loss
if __name__ == "__main__":
gt_box = np.array([[0.0, 0.0, 10, 10]])
anchor_boxes = np.array([[0.0, 0.0, 10, 10],
[0.0, 0.0, 4, 4],
[0.0, 0.0, 8, 8],
[0.0, 0.0, 16, 16]
])
iou = compute_iou(anchor_boxes, gt_box)
print(iou)
================================================
FILE: train.py
================================================
from __future__ import division
import os
import random
import argparse
import time
import cv2
import numpy as np
from copy import deepcopy
import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from data.voc0712 import VOCDetection
from data.coco2017 import COCODataset
from data import config
from data import BaseTransform, detection_collate
import tools
from utils import distributed_utils
from utils.com_paras_flops import FLOPs_and_Params
from utils.augmentations import SSDAugmentation, ColorAugmentation
from utils.cocoapi_evaluator import COCOAPIEvaluator
from utils.vocapi_evaluator import VOCAPIEvaluator
from utils.modules import ModelEMA
def parse_args():
parser = argparse.ArgumentParser(description='YOLO Detection')
# basic
parser.add_argument('--cuda', action='store_true', default=False,
help='use cuda.')
parser.add_argument('-bs', '--batch_size', default=16, type=int,
help='Batch size for training')
parser.add_argument('--lr', default=1e-3, type=float,
help='initial learning rate')
parser.add_argument('--wp_epoch', type=int, default=2,
help='The upper bound of warm-up')
parser.add_argument('--start_epoch', type=int, default=0,
help='start epoch to train')
parser.add_argument('-r', '--resume', default=None, type=str,
help='keep training')
parser.add_argument('--momentum', default=0.9, type=float,
help='Momentum value for optim')
parser.add_argument('--weight_decay', default=5e-4, type=float,
help='Weight decay for SGD')
parser.add_argument('--num_workers', default=8, type=int,
help='Number of workers used in dataloading')
parser.add_argument('--num_gpu', default=1, type=int,
help='Number of GPUs to train')
parser.add_argument('--eval_epoch', type=int,
default=10, help='interval between evaluations')
parser.add_argument('--tfboard', action='store_true', default=False,
help='use tensorboard')
parser.add_argument('--save_folder', default='weights/', type=str,
help='Gamma update for SGD')
parser.add_argument('--vis', action='store_true', default=False,
help='visualize target.')
# model
parser.add_argument('-v', '--version', default='yolo_v2',
help='yolov2_d19, yolov2_r50, yolov2_slim, yolov3, yolov3_spp, yolov3_tiny')
# dataset
parser.add_argument('-root', '--data_root', default='/mnt/share/ssd2/dataset',
help='dataset root')
parser.add_argument('-d', '--dataset', default='voc',
help='voc or coco')
# train trick
parser.add_argument('--no_warmup', action='store_true', default=False,
help='do not use warmup')
parser.add_argument('-ms', '--multi_scale', action='store_true', default=False,
help='use multi-scale trick')
parser.add_argument('--mosaic', action='store_true', default=False,
help='use mosaic augmentation')
parser.add_argument('--ema', action='store_true', default=False,
help='use ema training trick')
# DDP train
parser.add_argument('-dist', '--distributed', action='store_true', default=False,
help='distributed training')
parser.add_argument('--dist_url', default='env://',
help='url used to set up distributed training')
parser.add_argument('--world_size', default=1, type=int,
help='number of distributed processes')
parser.add_argument('--sybn', action='store_true', default=False,
help='use sybn.')
return parser.parse_args()
def train():
args = parse_args()
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
# set distributed
print('World size: {}'.format(distributed_utils.get_world_size()))
if args.distributed:
distributed_utils.init_distributed_mode(args)
print("git:\n {}\n".format(distributed_utils.get_sha()))
# cuda
if args.cuda:
print('use cuda')
# cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
model_name = args.version
print('Model: ', model_name)
# load model and config file
if model_name == 'yolov2_d19':
from models.yolov2_d19 import YOLOv2D19 as yolo_net
cfg = config.yolov2_d19_cfg
elif model_name == 'yolov2_r50':
from models.yolov2_r50 import YOLOv2R50 as yolo_net
cfg = config.yolov2_r50_cfg
elif model_name == 'yolov3':
from models.yolov3 import YOLOv3 as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_spp':
from models.yolov3_spp import YOLOv3Spp as yolo_net
cfg = config.yolov3_d53_cfg
elif model_name == 'yolov3_tiny':
from models.yolov3_tiny import YOLOv3tiny as yolo_net
cfg = config.yolov3_tiny_cfg
else:
print('Unknown model name...')
exit(0)
# path to save model
path_to_save = os.path.join(args.save_folder, args.dataset, args.version)
os.makedirs(path_to_save, exist_ok=True)
# multi-scale
if args.multi_scale:
print('use the multi-scale trick ...')
train_size = cfg['train_size']
val_size = cfg['val_size']
else:
train_size = val_size = cfg['train_size']
# Model ENA
if args.ema:
print('use EMA trick ...')
# dataset and evaluator
if args.dataset == 'voc':
data_dir = os.path.join(args.data_root, 'VOCdevkit')
num_classes = 20
dataset = VOCDetection(data_dir=data_dir,
transform=SSDAugmentation(train_size))
evaluator = VOCAPIEvaluator(data_root=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size))
elif args.dataset == 'coco':
data_dir = os.path.join(args.data_root, 'COCO')
num_classes = 80
dataset = COCODataset(
data_dir=data_dir,
transform=SSDAugmentation(train_size))
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size))
else:
print('unknow dataset !! Only support voc and coco !!')
exit(0)
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# build model
anchor_size = cfg['anchor_size_voc'] if args.dataset == 'voc' else cfg['anchor_size_coco']
net = yolo_net(device=device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size)
model = net
model = model.to(device).train()
# SyncBatchNorm
if args.sybn and args.distributed:
print('use SyncBatchNorm ...')
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
# DDP
model_without_ddp = model
if args.distributed:
model = DDP(model, device_ids=[args.gpu])
model_without_ddp = model.module
# compute FLOPs and Params
if distributed_utils.is_main_process:
model_copy = deepcopy(model_without_ddp)
model_copy.trainable = False
model_copy.eval()
FLOPs_and_Params(model=model_copy,
size=train_size,
device=device)
model_copy.trainable = True
model_copy.train()
if args.distributed:
# wait for all processes to synchronize
dist.barrier()
# dataloader
batch_size = args.batch_size * distributed_utils.get_world_size()
if args.distributed and args.num_gpu > 1:
dataloader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True,
sampler=torch.utils.data.distributed.DistributedSampler(dataset)
)
else:
# dataloader
dataloader = torch.utils.data.DataLoader(
dataset=dataset,
shuffle=True,
batch_size=batch_size,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True
)
# keep training
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
# EMA
ema = ModelEMA(model) if args.ema else None
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = os.path.join('log/', args.dataset, c_time)
os.makedirs(log_path, exist_ok=True)
tblogger = SummaryWriter(log_path)
# optimizer setup
base_lr = (args.lr / 16) * batch_size
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay
)
max_epoch = cfg['max_epoch']
epoch_size = len(dataloader)
best_map = -1.
warmup = not args.no_warmup
t0 = time.time()
# start training loop
for epoch in range(args.start_epoch, max_epoch):
if args.distributed:
dataloader.sampler.set_epoch(epoch)
# use step lr
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
ni = iter_i + epoch * epoch_size
# warmup
if epoch < args.wp_epoch and warmup:
nw = args.wp_epoch * epoch_size
tmp_lr = base_lr * pow(ni / nw, 4)
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0 and warmup:
# warmup is over
warmup = False
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
r = cfg['random_size_range']
train_size = random.randint(r[0], r[1]) * 32
model.set_grid(train_size)
if args.multi_scale:
# interpolate
images = torch.nn.functional.interpolate(images, size=train_size, mode='bilinear', align_corners=False)
targets = [label.tolist() for label in targets]
# visualize labels
if args.vis:
vis_data(images, targets, train_size)
continue
# label assignment
if model_name in ['yolov2_d19', 'yolov2_r50']:
targets = tools.gt_creator(input_size=train_size,
stride=net.stride,
label_lists=targets,
anchor_size=anchor_size
)
else:
targets = tools.multi_gt_creator(input_size=train_size,
strides=net.stride,
label_lists=targets,
anchor_size=anchor_size
)
# to device
images = images.float().to(device)
targets = torch.tensor(targets).float().to(device)
# forward
conf_loss, cls_loss, box_loss, iou_loss = model(images, target=targets)
# compute loss
total_loss = conf_loss + cls_loss + box_loss + iou_loss
loss_dict = dict(conf_loss=conf_loss,
cls_loss=cls_loss,
box_loss=box_loss,
iou_loss=iou_loss,
total_loss=total_loss
)
loss_dict_reduced = distributed_utils.reduce_dict(loss_dict)
# check NAN for loss
if torch.isnan(total_loss):
print('loss is nan !!')
continue
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
# ema
if args.ema:
ema.update(model)
# display
if distributed_utils.is_main_process() and iter_i % 10 == 0:
if args.tfboard:
# viz loss
tblogger.add_scalar('conf loss', loss_dict_reduced['conf_loss'].item(), iter_i + epoch * epoch_size)
tblogger.add_scalar('cls loss', loss_dict_reduced['cls_loss'].item(), iter_i + epoch * epoch_size)
tblogger.add_scalar('box loss', loss_dict_reduced['box_loss'].item(), iter_i + epoch * epoch_size)
tblogger.add_scalar('iou loss', loss_dict_reduced['iou_loss'].item(), iter_i + epoch * epoch_size)
t1 = time.time()
cur_lr = [param_group['lr'] for param_group in optimizer.param_groups]
# basic infor
log = '[Epoch: {}/{}]'.format(epoch+1, max_epoch)
log += '[Iter: {}/{}]'.format(iter_i, epoch_size)
log += '[lr: {:.6f}]'.format(cur_lr[0])
# loss infor
for k in loss_dict_reduced.keys():
log += '[{}: {:.2f}]'.format(k, loss_dict[k])
# other infor
log += '[time: {:.2f}]'.format(t1 - t0)
log += '[size: {}]'.format(train_size)
# print log infor
print(log, flush=True)
t0 = time.time()
if distributed_utils.is_main_process():
# evaluation
if (epoch % args.eval_epoch) == 0 or (epoch == max_epoch - 1):
if args.ema:
model_eval = ema.ema
else:
model_eval = model_without_ddp
# check evaluator
if evaluator is None:
print('No evaluator ... save model and go on training.')
print('Saving state, epoch: {}'.format(epoch + 1))
weight_name = '{}_epoch_{}.pth'.format(args.version, epoch + 1)
checkpoint_path = os.path.join(path_to_save, weight_name)
torch.save(model_eval.state_dict(), checkpoint_path)
else:
print('eval ...')
# set eval mode
model_eval.trainable = False
model_eval.set_grid(val_size)
model_eval.eval()
# evaluate
evaluator.evaluate(model_eval)
cur_map = evaluator.map
if cur_map > best_map:
# update best-map
best_map = cur_map
# save model
print('Saving state, epoch:', epoch + 1)
weight_name = '{}_epoch_{}_{:.2f}.pth'.format(args.version, epoch + 1, best_map*100)
checkpoint_path = os.path.join(path_to_save, weight_name)
torch.save(model_eval.state_dict(), checkpoint_path)
if args.tfboard:
if args.dataset == 'voc':
tblogger.add_scalar('07test/mAP', evaluator.map, epoch)
elif args.dataset == 'coco':
tblogger.add_scalar('val/AP50_95', evaluator.ap50_95, epoch)
tblogger.add_scalar('val/AP50', evaluator.ap50, epoch)
# set train mode.
model_eval.trainable = True
model_eval.set_grid(train_size)
model_eval.train()
# wait for all processes to synchronize
if args.distributed:
dist.barrier()
if args.tfboard:
tblogger.close()
def set_lr(optimizer, lr):
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def vis_data(images, targets, input_size):
# vis data
mean=(0.406, 0.456, 0.485)
std=(0.225, 0.224, 0.229)
mean = np.array(mean, dtype=np.float32)
std = np.array(std, dtype=np.float32)
img = images[0].permute(1, 2, 0).cpu().numpy()[:, :, ::-1]
img = ((img * std + mean)*255).astype(np.uint8)
img = img.copy()
for box in targets[0]:
xmin, ymin, xmax, ymax = box[:-1]
# print(xmin, ymin, xmax, ymax)
xmin *= input_size
ymin *= input_size
xmax *= input_size
ymax *= input_size
cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255), 2)
cv2.imshow('img', img)
cv2.waitKey(0)
if __name__ == '__main__':
train()
================================================
FILE: utils/__init__.py
================================================
================================================
FILE: utils/augmentations.py
================================================
import cv2
import numpy as np
from numpy import random
def intersect(box_a, box_b):
max_xy = np.minimum(box_a[:, 2:], box_b[2:])
min_xy = np.maximum(box_a[:, :2], box_b[:2])
inter = np.clip((max_xy - min_xy), a_min=0, a_max=np.inf)
return inter[:, 0] * inter[:, 1]
def jaccard_numpy(box_a, box_b):
"""Compute the jaccard overlap of two sets of boxes. The jaccard overlap
is simply the intersection over union of two boxes.
E.g.:
A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
Args:
box_a: Multiple bounding boxes, Shape: [num_boxes,4]
box_b: Single bounding box, Shape: [4]
Return:
jaccard overlap: Shape: [box_a.shape[0], box_a.shape[1]]
"""
inter = intersect(box_a, box_b)
area_a = ((box_a[:, 2]-box_a[:, 0]) *
(box_a[:, 3]-box_a[:, 1])) # [A,B]
area_b = ((box_b[2]-box_b[0]) *
(box_b[3]-box_b[1])) # [A,B]
union = area_a + area_b - inter
return inter / union # [A,B]
class Compose(object):
"""Composes several augmentations together.
Args:
transforms (List[Transform]): list of transforms to compose.
Example:
>>> augmentations.Compose([
>>> transforms.CenterCrop(10),
>>> transforms.ToTensor(),
>>> ])
"""
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, img, boxes=None, labels=None):
for t in self.transforms:
img, boxes, labels = t(img, boxes, labels)
return img, boxes, labels
class ConvertFromInts(object):
def __call__(self, image, boxes=None, labels=None):
return image.astype(np.float32), boxes, labels
class Normalize(object):
def __init__(self, mean=None, std=None):
self.mean = np.array(mean, dtype=np.float32)
self.std = np.array(std, dtype=np.float32)
def __call__(self, image, boxes=None, labels=None):
image = image.astype(np.float32)
image /= 255.
image -= self.mean
image /= self.std
return image, boxes, labels
class ToAbsoluteCoords(object):
def __call__(self, image, boxes=None, labels=None):
height, width, channels = image.shape
boxes[:, 0] *= width
boxes[:, 2] *= width
boxes[:, 1] *= height
boxes[:, 3] *= height
return image, boxes, labels
class ToPercentCoords(object):
def __call__(self, image, boxes=None, labels=None):
height, width, channels = image.shape
boxes[:, 0] /= width
boxes[:, 2] /= width
boxes[:, 1] /= height
boxes[:, 3] /= height
return image, boxes, labels
class Resize(object):
def __init__(self, size=416):
self.size = size
def __call__(self, image, boxes=None, labels=None):
image = cv2.resize(image, (self.size, self.size))
return image, boxes, labels
class RandomSaturation(object):
def __init__(self, lower=0.5, upper=1.5):
self.lower = lower
self.upper = upper
assert self.upper >= self.lower, "contrast upper must be >= lower."
assert self.lower >= 0, "contrast lower must be non-negative."
def __call__(self, image, boxes=None, labels=None):
if random.randint(2):
image[:, :, 1] *= random.uniform(self.lower, self.upper)
return image, boxes, labels
class RandomHue(object):
def __init__(self, delta=18.0):
assert delta >= 0.0 and delta <= 360.0
self.delta = delta
def __call__(self, image, boxes=None, labels=None):
if random.randint(2):
image[:, :, 0] += random.uniform(-self.delta, self.delta)
image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0
image[:, :, 0][image[:, :, 0] < 0.0] += 360.0
return image, boxes, labels
class RandomLightingNoise(object):
def __init__(self):
self.perms = ((0, 1, 2), (0, 2, 1),
(1, 0, 2), (1, 2, 0),
(2, 0, 1), (2, 1, 0))
def __call__(self, image, boxes=None, labels=None):
if random.randint(2):
swap = self.perms[random.randint(len(self.perms))]
shuffle = SwapChannels(swap) # shuffle channels
image = shuffle(image)
return image, boxes, labels
class ConvertColor(object):
def __init__(self, current='BGR', transform='HSV'):
self.transform = transform
self.current = current
def __call__(self, image, boxes=None, labels=None):
if self.current == 'BGR' and self.transform == 'HSV':
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
elif self.current == 'HSV' and self.transform == 'BGR':
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
else:
raise NotImplementedError
return image, boxes, labels
class RandomContrast(object):
def __init__(self, lower=0.5, upper=1.5):
self.lower = lower
self.upper = upper
assert self.upper >= self.lower, "contrast upper must be >= lower."
assert self.lower >= 0, "contrast lower must be non-negative."
# expects float image
def __call__(self, image, boxes=None, labels=None):
if random.randint(2):
alpha = random.uniform(self.lower, self.upper)
image *= alpha
return image, boxes, labels
class RandomBrightness(object):
def __init__(self, delta=32):
assert delta >= 0.0
assert delta <= 255.0
self.delta = delta
def __call__(self, image, boxes=None, labels=None):
if random.randint(2):
delta = random.uniform(-self.delta, self.delta)
image += delta
return image, boxes, labels
class RandomSampleCrop(object):
"""Crop
Arguments:
img (Image): the image being input during training
boxes (Tensor): the original bounding boxes in pt form
labels (Tensor): the class labels for each bbox
mode (float tuple): the min and max jaccard overlaps
Return:
(img, boxes, classes)
img (Image): the cropped image
boxes (Tensor): the adjusted bounding boxes in pt form
labels (Tensor): the class labels for each bbox
"""
def __init__(self):
self.sample_options = (
# using entire original input image
None,
# sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9
(0.1, None),
(0.3, None),
(0.7, None),
(0.9, None),
# randomly sample a patch
(None, None),
)
def __call__(self, image, boxes=None, labels=None):
height, width, _ = image.shape
while True:
# randomly choose a mode
sample_id = np.random.randint(len(self.sample_options))
mode = self.sample_options[sample_id]
if mode is None:
return image, boxes, labels
min_iou, max_iou = mode
if min_iou is None:
min_iou = float('-inf')
if max_iou is None:
max_iou = float('inf')
# max trails (50)
for _ in range(50):
current_image = image
w = random.uniform(0.3 * width, width)
h = random.uniform(0.3 * height, height)
# aspect ratio constraint b/t .5 & 2
if h / w < 0.5 or h / w > 2:
continue
left = random.uniform(width - w)
top = random.uniform(height - h)
# convert to integer rect x1,y1,x2,y2
rect = np.array([int(left), int(top), int(left+w), int(top+h)])
# calculate IoU (jaccard overlap) b/t the cropped and gt boxes
overlap = jaccard_numpy(boxes, rect)
# is min and max overlap constraint satisfied? if not try again
if overlap.min() < min_iou and max_iou < overlap.max():
continue
# cut the crop from the image
current_image = current_image[rect[1]:rect[3], rect[0]:rect[2],
:]
# keep overlap with gt box IF center in sampled patch
centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0
# mask in all gt boxes that above and to the left of centers
m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1])
# mask in all gt boxes that under and to the right of centers
m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1])
# mask in that both m1 and m2 are true
mask = m1 * m2
# have any valid boxes? try again if not
if not mask.any():
continue
# take only matching gt boxes
current_boxes = boxes[mask, :].copy()
# take only matching gt labels
current_labels = labels[mask]
# should we use the box left and top corner or the crop's
current_boxes[:, :2] = np.maximum(current_boxes[:, :2],
rect[:2])
# adjust to crop (by substracting crop's left,top)
current_boxes[:, :2] -= rect[:2]
current_boxes[:, 2:] = np.minimum(current_boxes[:, 2:],
rect[2:])
# adjust to crop (by substracting crop's left,top)
current_boxes[:, 2:] -= rect[:2]
return current_image, current_boxes, current_labels
class RandomMirror(object):
def __call__(self, image, boxes, classes):
_, width, _ = image.shape
if random.randint(2):
image = image[:, ::-1]
boxes = boxes.copy()
boxes[:, 0::2] = width - boxes[:, 2::-2]
return image, boxes, classes
class SwapChannels(object):
"""Transforms a tensorized image by swapping the channels in the order
specified in the swap tuple.
Args:
swaps (int triple): final order of channels
eg: (2, 1, 0)
"""
def __init__(self, swaps):
self.swaps = swaps
def __call__(self, image):
"""
Args:
image (Tensor): image tensor to be transformed
Return:
a tensor with channels swapped according to swap
"""
# if torch.is_tensor(image):
# image = image.data.cpu().numpy()
# else:
# image = np.array(image)
image = image[:, :, self.swaps]
return image
class PhotometricDistort(object):
def __init__(self):
self.pd = [
RandomContrast(),
ConvertColor(transform='HSV'),
RandomSaturation(),
RandomHue(),
ConvertColor(current='HSV', transform='BGR'),
RandomContrast()
]
self.rand_brightness = RandomBrightness()
# self.rand_light_noise = RandomLightingNoise()
def __call__(self, image, boxes, labels):
im = image.copy()
im, boxes, labels = self.rand_brightness(im, boxes, labels)
if random.randint(2):
distort = Compose(self.pd[:-1])
else:
distort = Compose(self.pd[1:])
im, boxes, labels = distort(im, boxes, labels)
return im, boxes, labels
# return self.rand_light_noise(im, boxes, labels)
class SSDAugmentation(object):
def __init__(self, size=416, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)):
self.mean = mean
self.size = size
self.std = std
self.augment = Compose([
ConvertFromInts(),
ToAbsoluteCoords(),
PhotometricDistort(),
RandomSampleCrop(),
RandomMirror(),
ToPercentCoords(),
Resize(self.size),
Normalize(self.mean, self.std)
])
def __call__(self, img, boxes, labels):
return self.augment(img, boxes, labels)
class ColorAugmentation(object):
def __init__(self, size=416, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)):
self.mean = mean
self.size = size
self.std = std
self.augment = Compose([
ConvertFromInts(),
ToAbsoluteCoords(),
PhotometricDistort(),
RandomMirror(),
ToPercentCoords(),
Resize(self.size),
Normalize(self.mean, self.std)
])
def __call__(self, img, boxes, labels):
return self.augment(img, boxes, labels)
================================================
FILE: utils/cocoapi_evaluator.py
================================================
import json
import tempfile
from pycocotools.cocoeval import COCOeval
from torch.autograd import Variable
from data.coco2017 import *
from data import *
class COCOAPIEvaluator():
"""
COCO AP Evaluation class.
All the data in the val2017 dataset are processed \
and evaluated by COCO API.
"""
def __init__(self, data_dir, img_size, device, testset=False, transform=None):
"""
Args:
data_dir (str): dataset root directory
img_size (int): image size after preprocess. images are resized \
to squares whose shape is (img_size, img_size).
confthre (float):
confidence threshold ranging from 0 to 1, \
which is defined in the config file.
nmsthre (float):
IoU threshold of non-max supression ranging from 0 to 1.
"""
self.testset = testset
if self.testset:
json_file='image_info_test-dev2017.json'
name = 'test2017'
else:
json_file='instances_val2017.json'
name='val2017'
self.dataset = COCODataset(data_dir=data_dir,
json_file=json_file,
name=name)
self.img_size = img_size
self.transform = transform
self.device = device
self.map = 0.
self.ap50_95 = 0.
self.ap50 = 0.
def evaluate(self, model):
"""
COCO average precision (AP) Evaluation. Iterate inference on the test dataset
and the results are evaluated by COCO API.
Args:
model : model object
Returns:
ap50_95 (float) : calculated COCO AP for IoU=50:95
ap50 (float) : calculated COCO AP for IoU=50
"""
model.eval()
ids = []
data_dict = []
num_images = len(self.dataset)
print('total number of images: %d' % (num_images))
# start testing
for index in range(num_images): # all the data in val2017
if index % 500 == 0:
print('[Eval: %d / %d]'%(index, num_images))
img, id_ = self.dataset.pull_im
gitextract_dmpt53ff/
├── .gitignore
├── LICENSE
├── README.md
├── backbone/
│ ├── __init__.py
│ ├── darknet19.py
│ ├── darknet53.py
│ ├── darknet_tiny.py
│ ├── resnet.py
│ └── weights/
│ └── README.md
├── data/
│ ├── __init__.py
│ ├── coco2017.py
│ ├── config.py
│ ├── scripts/
│ │ ├── COCO2017.sh
│ │ ├── VOC2007.sh
│ │ └── VOC2012.sh
│ └── voc0712.py
├── demo.py
├── eval.py
├── models/
│ ├── yolov2_d19.py
│ ├── yolov2_r50.py
│ ├── yolov3.py
│ ├── yolov3_spp.py
│ └── yolov3_tiny.py
├── test.py
├── tools.py
├── train.py
├── utils/
│ ├── __init__.py
│ ├── augmentations.py
│ ├── cocoapi_evaluator.py
│ ├── com_paras_flops.py
│ ├── distributed_utils.py
│ ├── kmeans_anchor.py
│ ├── modules.py
│ └── vocapi_evaluator.py
└── weights/
└── README.md
SYMBOL INDEX (251 symbols across 25 files)
FILE: backbone/__init__.py
function build_backbone (line 7) | def build_backbone(model_name='resnet18', pretrained=False):
FILE: backbone/darknet19.py
class Conv_BN_LeakyReLU (line 14) | class Conv_BN_LeakyReLU(nn.Module):
method __init__ (line 15) | def __init__(self, in_channels, out_channels, ksize, padding=0, stride...
method forward (line 23) | def forward(self, x):
class DarkNet_19 (line 27) | class DarkNet_19(nn.Module):
method __init__ (line 28) | def __init__(self):
method forward (line 78) | def forward(self, x):
function build_darknet19 (line 95) | def build_darknet19(pretrained=False):
FILE: backbone/darknet53.py
class Conv_BN_LeakyReLU (line 13) | class Conv_BN_LeakyReLU(nn.Module):
method __init__ (line 14) | def __init__(self, in_channels, out_channels, ksize, padding=0, stride...
method forward (line 22) | def forward(self, x):
class ResBlock (line 26) | class ResBlock(nn.Module):
method __init__ (line 27) | def __init__(self, ch, nblocks=1):
method forward (line 37) | def forward(self, x):
class DarkNet_53 (line 43) | class DarkNet_53(nn.Module):
method __init__ (line 47) | def __init__(self):
method forward (line 77) | def forward(self, x, targets=None):
function build_darknet53 (line 93) | def build_darknet53(pretrained=False):
FILE: backbone/darknet_tiny.py
class Conv_BN_LeakyReLU (line 13) | class Conv_BN_LeakyReLU(nn.Module):
method __init__ (line 14) | def __init__(self, in_channels, out_channels, ksize, padding=0, stride...
method forward (line 22) | def forward(self, x):
class DarkNet_Tiny (line 26) | class DarkNet_Tiny(nn.Module):
method __init__ (line 27) | def __init__(self):
method forward (line 55) | def forward(self, x):
function build_darknet_tiny (line 79) | def build_darknet_tiny(pretrained=False):
FILE: backbone/resnet.py
function conv3x3 (line 19) | def conv3x3(in_planes, out_planes, stride=1):
function conv1x1 (line 24) | def conv1x1(in_planes, out_planes, stride=1):
class BasicBlock (line 28) | class BasicBlock(nn.Module):
method __init__ (line 31) | def __init__(self, inplanes, planes, stride=1, downsample=None):
method forward (line 41) | def forward(self, x):
class Bottleneck (line 59) | class Bottleneck(nn.Module):
method __init__ (line 62) | def __init__(self, inplanes, planes, stride=1, downsample=None):
method forward (line 74) | def forward(self, x):
class ResNet (line 96) | class ResNet(nn.Module):
method __init__ (line 98) | def __init__(self, block, layers, zero_init_residual=False):
method _make_layer (line 128) | def _make_layer(self, block, planes, blocks, stride=1):
method forward (line 144) | def forward(self, x):
function resnet18 (line 164) | def resnet18(pretrained=False, **kwargs):
function resnet34 (line 176) | def resnet34(pretrained=False, **kwargs):
function resnet50 (line 187) | def resnet50(pretrained=False, **kwargs):
function resnet101 (line 198) | def resnet101(pretrained=False, **kwargs):
function resnet152 (line 209) | def resnet152(pretrained=False, **kwargs):
function build_resnet (line 221) | def build_resnet(model_name='resnet18', pretrained=False):
FILE: data/__init__.py
function detection_collate (line 9) | def detection_collate(batch):
function base_transform (line 30) | def base_transform(image, size, mean, std):
class BaseTransform (line 38) | class BaseTransform:
method __init__ (line 39) | def __init__(self, size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224...
method __call__ (line 44) | def __call__(self, image, boxes=None, labels=None):
FILE: data/coco2017.py
class COCODataset (line 32) | class COCODataset(Dataset):
method __init__ (line 36) | def __init__(self,
method __len__ (line 60) | def __len__(self):
method pull_image (line 64) | def pull_image(self, index):
method pull_anno (line 78) | def pull_anno(self, index):
method __getitem__ (line 102) | def __getitem__(self, index):
method pull_item (line 108) | def pull_item(self, index):
function base_transform (line 166) | def base_transform(image, size, mean):
class BaseTransform (line 172) | class BaseTransform:
method __init__ (line 173) | def __init__(self, size, mean):
method __call__ (line 177) | def __call__(self, image, boxes=None, labels=None):
FILE: data/voc0712.py
class VOCAnnotationTransform (line 26) | class VOCAnnotationTransform(object):
method __init__ (line 39) | def __init__(self, class_to_ind=None, keep_difficult=False):
method __call__ (line 44) | def __call__(self, target, width, height):
class VOCDetection (line 75) | class VOCDetection(data.Dataset):
method __init__ (line 92) | def __init__(self,
method __getitem__ (line 112) | def __getitem__(self, index):
method __len__ (line 118) | def __len__(self):
method pull_item (line 122) | def pull_item(self, index):
method pull_image (line 151) | def pull_image(self, index):
method pull_anno (line 166) | def pull_anno(self, index):
method pull_tensor (line 184) | def pull_tensor(self, index):
function base_transform (line 199) | def base_transform(image, size, mean):
class BaseTransform (line 205) | class BaseTransform:
method __init__ (line 206) | def __init__(self, size, mean):
method __call__ (line 210) | def __call__(self, image, boxes=None, labels=None):
FILE: demo.py
function parse_args (line 12) | def parse_args():
function plot_bbox_labels (line 42) | def plot_bbox_labels(img, bbox, label, cls_color, test_scale=0.4):
function visualize (line 56) | def visualize(img, bboxes, scores, cls_inds, class_colors, vis_thresh=0.3):
function detect (line 68) | def detect(net,
function run (line 196) | def run():
FILE: eval.py
function voc_test (line 30) | def voc_test(model, data_dir, device, input_size):
function coco_test (line 41) | def coco_test(model, data_dir, device, input_size, test=False):
FILE: models/yolov2_d19.py
class YOLOv2D19 (line 10) | class YOLOv2D19(nn.Module):
method __init__ (line 11) | def __init__(self, device, input_size=None, num_classes=20, trainable=...
method create_grid (line 42) | def create_grid(self, input_size):
method set_grid (line 56) | def set_grid(self, input_size):
method decode_xywh (line 61) | def decode_xywh(self, txtytwth_pred):
method decode_boxes (line 81) | def decode_boxes(self, txtytwth_pred):
method nms (line 100) | def nms(self, dets, scores):
method postprocess (line 132) | def postprocess(self, bboxes, scores):
method inference (line 167) | def inference(self, x):
method forward (line 218) | def forward(self, x, target=None):
FILE: models/yolov2_r50.py
class YOLOv2R50 (line 10) | class YOLOv2R50(nn.Module):
method __init__ (line 11) | def __init__(self, device, input_size=None, num_classes=20, trainable=...
method init_bias (line 50) | def init_bias(self):
method create_grid (line 57) | def create_grid(self, input_size):
method set_grid (line 72) | def set_grid(self, input_size):
method decode_xywh (line 77) | def decode_xywh(self, txtytwth_pred):
method decode_boxes (line 97) | def decode_boxes(self, txtytwth_pred):
method nms (line 116) | def nms(self, dets, scores):
method postprocess (line 148) | def postprocess(self, bboxes, scores):
method inference (line 183) | def inference(self, x):
method forward (line 234) | def forward(self, x, target=None):
FILE: models/yolov3.py
class YOLOv3 (line 10) | class YOLOv3(nn.Module):
method __init__ (line 11) | def __init__(self,
method init_yolo (line 74) | def init_yolo(self):
method create_grid (line 84) | def create_grid(self, input_size):
method set_grid (line 113) | def set_grid(self, input_size):
method decode_xywh (line 118) | def decode_xywh(self, txtytwth_pred):
method decode_boxes (line 136) | def decode_boxes(self, txtytwth_pred):
method nms (line 155) | def nms(self, dets, scores):
method postprocess (line 187) | def postprocess(self, bboxes, scores):
method inference (line 228) | def inference(self, x):
method forward (line 307) | def forward(self, x, target=None):
FILE: models/yolov3_spp.py
class YOLOv3Spp (line 12) | class YOLOv3Spp(nn.Module):
method __init__ (line 13) | def __init__(self,
method init_yolo (line 77) | def init_yolo(self):
method create_grid (line 87) | def create_grid(self, input_size):
method set_grid (line 116) | def set_grid(self, input_size):
method decode_xywh (line 121) | def decode_xywh(self, txtytwth_pred):
method decode_boxes (line 139) | def decode_boxes(self, txtytwth_pred):
method nms (line 158) | def nms(self, dets, scores):
method postprocess (line 190) | def postprocess(self, bboxes, scores):
method inference (line 225) | def inference(self, x):
method forward (line 304) | def forward(self, x, target=None):
FILE: models/yolov3_tiny.py
class YOLOv3tiny (line 12) | class YOLOv3tiny(nn.Module):
method __init__ (line 13) | def __init__(self, device, input_size=None, num_classes=20, trainable=...
method init_yolo (line 46) | def init_yolo(self):
method create_grid (line 56) | def create_grid(self, input_size):
method set_grid (line 85) | def set_grid(self, input_size):
method decode_xywh (line 90) | def decode_xywh(self, txtytwth_pred):
method decode_boxes (line 108) | def decode_boxes(self, txtytwth_pred):
method nms (line 127) | def nms(self, dets, scores):
method postprocess (line 159) | def postprocess(self, bboxes, scores):
method inference (line 194) | def inference(self, x):
method forward (line 266) | def forward(self, x, target=None):
FILE: test.py
function plot_bbox_labels (line 43) | def plot_bbox_labels(img, bbox, label=None, cls_color=None, text_scale=0...
function visualize (line 59) | def visualize(img,
function test (line 88) | def test(net,
FILE: tools.py
class MSEWithLogitsLoss (line 10) | class MSEWithLogitsLoss(nn.Module):
method __init__ (line 11) | def __init__(self, reduction='mean'):
method forward (line 15) | def forward(self, logits, targets, mask):
function compute_iou (line 35) | def compute_iou(anchor_boxes, gt_box):
function set_anchors (line 76) | def set_anchors(anchor_size):
function generate_txtytwth (line 95) | def generate_txtytwth(gt_label, w, h, s, all_anchor_size):
function gt_creator (line 165) | def gt_creator(input_size, stride, label_lists, anchor_size):
function multi_gt_creator (line 218) | def multi_gt_creator(input_size, strides, label_lists, anchor_size):
function iou_score (line 340) | def iou_score(bboxes_a, bboxes_b):
function loss (line 355) | def loss(pred_conf, pred_cls, pred_txtytwth, pred_iou, label):
FILE: train.py
function parse_args (line 32) | def parse_args():
function train (line 97) | def train():
function set_lr (line 457) | def set_lr(optimizer, lr):
function vis_data (line 462) | def vis_data(images, targets, input_size):
FILE: utils/augmentations.py
function intersect (line 6) | def intersect(box_a, box_b):
function jaccard_numpy (line 13) | def jaccard_numpy(box_a, box_b):
class Compose (line 33) | class Compose(object):
method __init__ (line 44) | def __init__(self, transforms):
method __call__ (line 47) | def __call__(self, img, boxes=None, labels=None):
class ConvertFromInts (line 53) | class ConvertFromInts(object):
method __call__ (line 54) | def __call__(self, image, boxes=None, labels=None):
class Normalize (line 58) | class Normalize(object):
method __init__ (line 59) | def __init__(self, mean=None, std=None):
method __call__ (line 63) | def __call__(self, image, boxes=None, labels=None):
class ToAbsoluteCoords (line 72) | class ToAbsoluteCoords(object):
method __call__ (line 73) | def __call__(self, image, boxes=None, labels=None):
class ToPercentCoords (line 83) | class ToPercentCoords(object):
method __call__ (line 84) | def __call__(self, image, boxes=None, labels=None):
class Resize (line 94) | class Resize(object):
method __init__ (line 95) | def __init__(self, size=416):
method __call__ (line 98) | def __call__(self, image, boxes=None, labels=None):
class RandomSaturation (line 103) | class RandomSaturation(object):
method __init__ (line 104) | def __init__(self, lower=0.5, upper=1.5):
method __call__ (line 110) | def __call__(self, image, boxes=None, labels=None):
class RandomHue (line 117) | class RandomHue(object):
method __init__ (line 118) | def __init__(self, delta=18.0):
method __call__ (line 122) | def __call__(self, image, boxes=None, labels=None):
class RandomLightingNoise (line 130) | class RandomLightingNoise(object):
method __init__ (line 131) | def __init__(self):
method __call__ (line 136) | def __call__(self, image, boxes=None, labels=None):
class ConvertColor (line 144) | class ConvertColor(object):
method __init__ (line 145) | def __init__(self, current='BGR', transform='HSV'):
method __call__ (line 149) | def __call__(self, image, boxes=None, labels=None):
class RandomContrast (line 159) | class RandomContrast(object):
method __init__ (line 160) | def __init__(self, lower=0.5, upper=1.5):
method __call__ (line 167) | def __call__(self, image, boxes=None, labels=None):
class RandomBrightness (line 174) | class RandomBrightness(object):
method __init__ (line 175) | def __init__(self, delta=32):
method __call__ (line 180) | def __call__(self, image, boxes=None, labels=None):
class RandomSampleCrop (line 187) | class RandomSampleCrop(object):
method __init__ (line 200) | def __init__(self):
method __call__ (line 213) | def __call__(self, image, boxes=None, labels=None):
class RandomMirror (line 292) | class RandomMirror(object):
method __call__ (line 293) | def __call__(self, image, boxes, classes):
class SwapChannels (line 302) | class SwapChannels(object):
method __init__ (line 310) | def __init__(self, swaps):
method __call__ (line 313) | def __call__(self, image):
class PhotometricDistort (line 328) | class PhotometricDistort(object):
method __init__ (line 329) | def __init__(self):
method __call__ (line 341) | def __call__(self, image, boxes, labels):
class SSDAugmentation (line 353) | class SSDAugmentation(object):
method __init__ (line 354) | def __init__(self, size=416, mean=(0.406, 0.456, 0.485), std=(0.225, 0...
method __call__ (line 369) | def __call__(self, img, boxes, labels):
class ColorAugmentation (line 373) | class ColorAugmentation(object):
method __init__ (line 374) | def __init__(self, size=416, mean=(0.406, 0.456, 0.485), std=(0.225, 0...
method __call__ (line 388) | def __call__(self, img, boxes, labels):
FILE: utils/cocoapi_evaluator.py
class COCOAPIEvaluator (line 11) | class COCOAPIEvaluator():
method __init__ (line 17) | def __init__(self, data_dir, img_size, device, testset=False, transfor...
method evaluate (line 48) | def evaluate(self, model):
FILE: utils/com_paras_flops.py
function FLOPs_and_Params (line 5) | def FLOPs_and_Params(model, size, device):
FILE: utils/distributed_utils.py
function all_gather (line 10) | def all_gather(data):
function reduce_dict (line 53) | def reduce_dict(input_dict, average=True):
function get_sha (line 80) | def get_sha():
function setup_for_distributed (line 100) | def setup_for_distributed(is_master):
function is_dist_avail_and_initialized (line 115) | def is_dist_avail_and_initialized():
function get_world_size (line 123) | def get_world_size():
function get_rank (line 129) | def get_rank():
function is_main_process (line 135) | def is_main_process():
function save_on_master (line 139) | def save_on_master(*args, **kwargs):
function init_distributed_mode (line 144) | def init_distributed_mode(args):
FILE: utils/kmeans_anchor.py
function parse_args (line 12) | def parse_args():
class Box (line 30) | class Box():
method __init__ (line 31) | def __init__(self, x, y, w, h):
function iou (line 38) | def iou(box1, box2):
function init_centroids (line 61) | def init_centroids(boxes, n_anchors):
function do_kmeans (line 98) | def do_kmeans(n_anchors, boxes, centroids):
function anchor_box_kmeans (line 129) | def anchor_box_kmeans(total_gt_boxes, n_anchors, loss_convergence, iters...
FILE: utils/modules.py
class Conv (line 7) | class Conv(nn.Module):
method __init__ (line 8) | def __init__(self, in_ch, out_ch, k=1, p=0, s=1, d=1, g=1, act=True):
method forward (line 22) | def forward(self, x):
class UpSample (line 26) | class UpSample(nn.Module):
method __init__ (line 27) | def __init__(self, size=None, scale_factor=None, mode='nearest', align...
method forward (line 34) | def forward(self, x):
class reorg_layer (line 39) | class reorg_layer(nn.Module):
method __init__ (line 40) | def __init__(self, stride):
method forward (line 44) | def forward(self, x):
class SPP (line 56) | class SPP(nn.Module):
method __init__ (line 60) | def __init__(self):
method forward (line 63) | def forward(self, x):
class ModelEMA (line 72) | class ModelEMA(object):
method __init__ (line 73) | def __init__(self, model, decay=0.9999, updates=0):
method update (line 81) | def update(self, model):
FILE: utils/vocapi_evaluator.py
class VOCAPIEvaluator (line 17) | class VOCAPIEvaluator():
method __init__ (line 19) | def __init__(self, data_root, img_size, device, transform, set_type='t...
method evaluate (line 42) | def evaluate(self, net):
method parse_rec (line 89) | def parse_rec(self, filename):
method get_output_dir (line 109) | def get_output_dir(self, name, phase):
method get_voc_results_file_template (line 121) | def get_voc_results_file_template(self, cls):
method write_voc_results_file (line 131) | def write_voc_results_file(self, all_boxes):
method do_python_eval (line 149) | def do_python_eval(self, use_07=True):
method voc_ap (line 188) | def voc_ap(self, rec, prec, use_07_metric=True):
method voc_eval (line 222) | def voc_eval(self, detpath, classname, cachedir, ovthresh=0.5, use_07_...
method evaluate_detections (line 330) | def evaluate_detections(self, box_list):
Condensed preview — 35 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (237K chars).
[
{
"path": ".gitignore",
"chars": 54,
"preview": "*.pt\n*.pth\n*.txt\n*.pkl\n__pycache__\n.vscode\ndet_results"
},
{
"path": "LICENSE",
"chars": 11323,
"preview": "Apache License\n Version 2.0, January 2004\n http://www.apache.org/licens"
},
{
"path": "README.md",
"chars": 11943,
"preview": "# Update\nRecently, I have released a new YOLO project:\n\nhttps://github.com/yjh0410/PyTorch_YOLO_Tutorial\n\nIn my new YOLO"
},
{
"path": "backbone/__init__.py",
"chars": 606,
"preview": "from .resnet import build_resnet\nfrom .darknet19 import build_darknet19\nfrom .darknet53 import build_darknet53\nfrom .dar"
},
{
"path": "backbone/darknet19.py",
"chars": 3978,
"preview": "import torch\nimport torch.nn as nn\nimport os\n\n\nmodel_urls = {\n \"darknet19\": \"https://github.com/yjh0410/image_classif"
},
{
"path": "backbone/darknet53.py",
"chars": 3759,
"preview": "import torch\nimport torch.nn as nn\n\n\nmodel_urls = {\n \"darknet53\": \"https://github.com/yjh0410/image_classification_py"
},
{
"path": "backbone/darknet_tiny.py",
"chars": 3601,
"preview": "import torch\nimport torch.nn as nn\n\n\nmodel_urls = {\n \"darknet_tiny\": \"https://github.com/yjh0410/image_classification"
},
{
"path": "backbone/resnet.py",
"chars": 7966,
"preview": "import torch\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n\n__all__ = ['ResNet', 'resnet18', 'resnet3"
},
{
"path": "backbone/weights/README.md",
"chars": 481,
"preview": "# darknet19, darknet53, darknet-tiny, darknet-light\ndarknet-tiny is designed by myself. It is a very simple and lightwei"
},
{
"path": "data/__init__.py",
"chars": 1419,
"preview": "from .voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES\nfrom .coco2017 import COCODataset, coco_class_lab"
},
{
"path": "data/coco2017.py",
"chars": 7759,
"preview": "import os\nimport numpy as np\nimport random\n\nimport torch\nfrom torch.utils.data import Dataset\nimport cv2\nfrom pycocotool"
},
{
"path": "data/config.py",
"chars": 2348,
"preview": "# config.py\n\n# YOLOv2 with darknet-19\nyolov2_d19_cfg = {\n # network\n 'backbone': 'd19',\n # for multi-scale tric"
},
{
"path": "data/scripts/COCO2017.sh",
"chars": 583,
"preview": "mkdir COCO\ncd COCO\n\nwget http://images.cocodataset.org/zips/train2017.zip\nwget http://images.cocodataset.org/zips/val201"
},
{
"path": "data/scripts/VOC2007.sh",
"chars": 971,
"preview": "#!/bin/bash\n# Ellis Brown\n\nstart=`date +%s`\n\n# handle optional download dir\nif [ -z \"$1\" ]\n then\n # navigate to ~/da"
},
{
"path": "data/scripts/VOC2012.sh",
"chars": 763,
"preview": "#!/bin/bash\n# Ellis Brown\n\nstart=`date +%s`\n\n# handle optional download dir\nif [ -z \"$1\" ]\n then\n # navigate to ~/da"
},
{
"path": "data/voc0712.py",
"chars": 7821,
"preview": "\"\"\"VOC Dataset Classes\n\nOriginal author: Francisco Massa\nhttps://github.com/fmassa/vision/blob/voc_dataset/torchvision/d"
},
{
"path": "demo.py",
"chars": 9757,
"preview": "import argparse\nimport os\nimport numpy as np\nimport cv2\nimport time\nimport torch\nfrom data.coco2017 import coco_class_in"
},
{
"path": "eval.py",
"chars": 4840,
"preview": "import argparse\nimport os\nimport torch\n\nfrom utils.vocapi_evaluator import VOCAPIEvaluator\nfrom utils.cocoapi_evaluator "
},
{
"path": "models/yolov2_d19.py",
"chars": 9253,
"preview": "import numpy as np\nimport torch\nimport torch.nn as nn\nfrom utils.modules import Conv, reorg_layer\n\nfrom backbone import "
},
{
"path": "models/yolov2_r50.py",
"chars": 9648,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom utils.modules import Conv, reorg_layer\nfrom back"
},
{
"path": "models/yolov3.py",
"chars": 13798,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom utils.modules import Conv\nfrom backbone import b"
},
{
"path": "models/yolov3_spp.py",
"chars": 13514,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\n\nfrom utils.modules import Conv, S"
},
{
"path": "models/yolov3_tiny.py",
"chars": 12005,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\n\nfrom utils.modules import Conv\nfr"
},
{
"path": "test.py",
"chars": 7649,
"preview": "import os\nimport argparse\nimport torch\nimport torch.backends.cudnn as cudnn\nfrom data.voc0712 import VOC_CLASSES, VOCDet"
},
{
"path": "tools.py",
"chars": 17322,
"preview": "import numpy as np\nfrom data import *\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# We use ignore thresh to d"
},
{
"path": "train.py",
"chars": 18388,
"preview": "from __future__ import division\n\nimport os\nimport random\nimport argparse\nimport time\nimport cv2\nimport numpy as np\nfrom "
},
{
"path": "utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "utils/augmentations.py",
"chars": 12671,
"preview": "import cv2\nimport numpy as np\nfrom numpy import random\n\n\ndef intersect(box_a, box_b):\n max_xy = np.minimum(box_a[:, 2"
},
{
"path": "utils/cocoapi_evaluator.py",
"chars": 4552,
"preview": "import json\nimport tempfile\n\nfrom pycocotools.cocoeval import COCOeval\nfrom torch.autograd import Variable\n\nfrom data.co"
},
{
"path": "utils/com_paras_flops.py",
"chars": 395,
"preview": "import torch\nfrom thop import profile\n\n\ndef FLOPs_and_Params(model, size, device):\n x = torch.randn(1, 3, size, size)"
},
{
"path": "utils/distributed_utils.py",
"chars": 5065,
"preview": "# from github: https://github.com/ruinmessi/ASFF/blob/master/utils/distributed_util.py\n\nimport torch\nimport torch.distri"
},
{
"path": "utils/kmeans_anchor.py",
"chars": 7496,
"preview": "import numpy as np\nimport random\nimport argparse\nimport os\nimport sys\nsys.path.append('..')\n\nfrom data.voc0712 import VO"
},
{
"path": "utils/modules.py",
"chars": 3072,
"preview": "import math\nimport torch\nimport torch.nn as nn\nfrom copy import deepcopy\n\n\nclass Conv(nn.Module):\n def __init__(self,"
},
{
"path": "utils/vocapi_evaluator.py",
"chars": 13214,
"preview": "\"\"\"Adapted from:\n @longcw faster_rcnn_pytorch: https://github.com/longcw/faster_rcnn_pytorch\n @rbgirshick py-faste"
},
{
"path": "weights/README.md",
"chars": 424,
"preview": "# yolo-v2-v3 and tiny model\nHi, guys ! \n\nFor researchers in China, you can download them from BaiduYunDisk. \nThere are 5"
}
]
About this extraction
This page contains the full source code of the yjh0410/yolov2-yolov3_PyTorch GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 35 files (223.1 KB), approximately 62.2k tokens, and a symbol index with 251 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.