Skip to content

Commit ff1ea1c

Browse files
wizyoungwizyoung
committed
first commit
1 parent 746e952 commit ff1ea1c

27 files changed

+2187
-1
lines changed

‎.gitignore

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# latex
2+
## Core latex/pdflatex auxiliary files:
3+
*.aux
4+
*.lof
5+
*.log
6+
*.lot
7+
*.fls
8+
*.out
9+
*.toc
10+
*.fmt
11+
*.fot
12+
*.cb
13+
*.cb2
14+
.*.lb
15+
*.bbl
16+
*.blg
17+
*.synctex.gz
18+
## Intermediate documents:
19+
*.dvi
20+
*.xdv
21+
*-converted-to.*
22+
# these rules might exclude image files for figures etc.
23+
# *.ps
24+
# *.eps
25+
# *.pdf
26+
27+
# python
28+
# Byte-compiled / optimized / DLL files
29+
__pycache__/
30+
*.py[cod]
31+
*$py.class
32+
# C extensions
33+
*.so
34+
35+
# c++
36+
# Prerequisites
37+
*.d
38+
# Compiled Object files
39+
*.slo
40+
*.lo
41+
*.o
42+
*.obj
43+
# Precompiled Headers
44+
*.gch
45+
*.pch
46+
# Compiled Dynamic libraries
47+
*.so
48+
*.dylib
49+
*.dll
50+
51+
52+
# tfrecord
53+
*.tfrecord
54+
55+
# macOS
56+
.DS_Store
57+
58+
# folders
59+
*.vscode/*
60+
*.texpadtmp/*
61+
62+
# darknet
63+
*.weights

‎README.md

Lines changed: 213 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,213 @@
1-
# YOLOv3_TensorFlow
1+
# YOLOv3_TensorFlow
2+
3+
### 1. Introduction
4+
5+
This is my implementation of [YOLOv3](https://pjreddie.com/media/files/papers/YOLOv3.pdf) in pure TensorFlow. It contains the full pipeline of training and evaluation on your own dataset. The keys features of this repo are:
6+
7+
- Efficient tf.data pipeline
8+
9+
- Weights converter (converting pretrained darknet weights on COCO dataset to TensorFlow checkpoint.)
10+
- Extremely fast GPU non maximum supression.
11+
- Full training pipeline.
12+
- Kmeans algorithm to select prior anchor boxes.
13+
14+
### 2. Requirements
15+
16+
- tensorflow >= 1.8.0 (lower versions may work too)
17+
- opencv-python
18+
19+
### 3. Weights convertion
20+
21+
The pretrained darknet weights file can be downloaded [here](https://pjreddie.com/media/files/yolov3.weights). Place this weights file under directory `./data/darknet_weights/` and then run:
22+
23+
```shell
24+
python convert_weight.py
25+
```
26+
27+
Then the converted TensorFlow checkpoint file will be saved to `./data/darknet_weights/` directory.
28+
29+
You can also download the converted TensorFlow checkpoint file by me via [Google Drive link](https://drive.google.com/drive/folders/1mXbNgNxyXPi7JNsnBaxEv1-nWr7SVoQt?usp=sharing) and then place it to the same directory.
30+
31+
### 4. Running demos
32+
33+
There are some demo images and videos under the `./data/demo_data/`. You can run the demo by:
34+
35+
Single image test demo:
36+
37+
```shell
38+
python test_single_image.py ./data/demo_data/messi.jpg
39+
```
40+
41+
Video test demo:
42+
43+
```shell
44+
python video_test.py ./data/demo_data/video.mp4
45+
```
46+
47+
Some results:
48+
49+
![](https://github.com/wizyoung/YOLOv3_TensorFlow/blob/master/data/demo_data/results/dog.jpg?raw=true)
50+
51+
![](https://github.com/wizyoung/YOLOv3_TensorFlow/blob/master/data/demo_data/results/messi.jpg?raw=true)
52+
53+
![](https://github.com/wizyoung/YOLOv3_TensorFlow/blob/master/data/demo_data/results/kite.jpg?raw=true)
54+
55+
(The kite result is under image resolution 1344x896)
56+
57+
#### 5. Inference speed
58+
59+
How fast is the inference speed? With images scaled to 416*416:
60+
61+
62+
| Backbone | GPU | Time(ms) |
63+
| :-------------------- | :------: | :------: |
64+
| Darknet-53 (paper) | Titan X | 29 |
65+
| Darknet-53 (my impl.) | Titan XP | ~23 |
66+
67+
### 6. Training
68+
69+
#### 6.1 Data preparation
70+
71+
(1) annotation file
72+
73+
Generate `train.txt/val.txt/test.txt` files under `./data/my_data/` directory. One line for one image, in the format like `image_absolute_path box_1 box_2 ... box_n`. Box_format: `label_index x_min y_min x_max y_max`.
74+
75+
For example:
76+
77+
```
78+
xxx/xxx/1.jpg 0 453 369 473 391 1 588 245 608 268
79+
xxx/xxx/2.jpg 1 466 403 485 422 2 793 300 809 320
80+
...
81+
```
82+
83+
**NOTE**: **You should leave a blank line at the end of each txt file.**
84+
85+
(2) class_names file:
86+
87+
Generate the `data.names` file under `./data/my_data/` directory. Each line represents a class name.
88+
89+
For example:
90+
91+
```
92+
bird
93+
person
94+
bike
95+
...
96+
```
97+
98+
The COCO dataset class names file is placed at `./data/coco.names`.
99+
100+
(3) prior anchor file:
101+
102+
Using the kmeans algorithm to get the prior anchors:
103+
104+
```
105+
python get_kmeans.py
106+
```
107+
108+
Then you will get 9 anchors and the average IOU. Save the anchors to a txt file.
109+
110+
The COCO dataset anchors offered by YOLO v3 author is placed at `./data/yolo_anchors.txt`, you can use that one too.
111+
112+
**NOTE: The yolo anchors should be scaled to the rescaled new image size. Suppose your image size is [W, H], and the image will be rescale to 416*416 as input, for each generated anchor [anchor_w, anchor_h], you should apply the transformation anchor_w = anchor_w / W * 416, anchor_h = anchor_g / H * 416.**
113+
114+
#### 6.2 Training
115+
116+
Using `train.py`. The parameters are as following:
117+
118+
```shell
119+
$ python train.py -h
120+
usage: train.py [-h] [--train_file TRAIN_FILE]
121+
[--val_file VAL_FILE]
122+
[--restore_path RESTORE_PATH]
123+
[--save_dir SAVE_DIR]
124+
[--log_dir LOG_DIR]
125+
[--progress_log_path PROGRESS_LOG_PATH]
126+
[--anchor_path ANCHOR_PATH]
127+
[--class_name_path CLASS_NAME_PATH] [--batch_size BATCH_SIZE]
128+
[--img_size [IMG_SIZE [IMG_SIZE ...]]]
129+
[--total_epoches TOTAL_EPOCHES]
130+
[--train_evaluation_freq TRAIN_EVALUATION_FREQ]
131+
[--val_evaluation_freq VAL_EVALUATION_FREQ]
132+
[--save_freq SAVE_FREQ] [--num_threads NUM_THREADS]
133+
[--prefetech_buffer PREFETECH_BUFFER]
134+
[--optimizer_name OPTIMIZER_NAME]
135+
[--save_optimizer SAVE_OPTIMIZER]
136+
[--learning_rate_init LEARNING_RATE_INIT] [--lr_type LR_TYPE]
137+
[--lr_decay_freq LR_DECAY_FREQ]
138+
[--lr_decay_factor LR_DECAY_FACTOR]
139+
[--lr_lower_bound LR_LOWER_BOUND]
140+
[--restore_part [RESTORE_PART [RESTORE_PART ...]]]
141+
[--update_part [UPDATE_PART [UPDATE_PART ...]]]
142+
```
143+
144+
Check the `train.py` for more details. You should set the parameters yourself.
145+
146+
Some training tricks in my experiment:
147+
148+
(1) Apply the two-stage training strategy:
149+
150+
First stage: Restore `darknet53_body` part weights from COCO checkpoints, train the `yolov3_head` with big learning rate like 1e-3 until the loss reaches to a low level, like less than 1.
151+
152+
Second stage: Restore the weights from the first stage, then train the whole model with small learning rate like 1e-4 or smaller. At this stage remember to restore the optimizer parameters if you use optimizers like adam.
153+
154+
(2) Quick train:
155+
156+
If you want to obtain good results in a short time like in 10 minutes. You can use the coco names but substitute several with real class names in your dataset. In this way you restore the whole pretrained COCO model and get a 80 class classification model, but you only care the class names from your dataset.
157+
158+
### 7. Evaluation
159+
160+
Using `eval.py` to evaluate the validation or test dataset. The parameters are as following:
161+
162+
```shell
163+
$ python eval.py -h
164+
usage: eval.py [-h] [--eval_file EVAL_FILE]
165+
[--restore_path RESTORE_PATH]
166+
[--anchor_path ANCHOR_PATH]
167+
[--class_name_path CLASS_NAME_PATH]
168+
[--batch_size BATCH_SIZE]
169+
[--img_size [IMG_SIZE [IMG_SIZE ...]]]
170+
[--num_threads NUM_THREADS]
171+
[--prefetech_buffer PREFETECH_BUFFER]
172+
```
173+
174+
Check the `eval.py` for more details. You should set the parameters yourself.
175+
176+
You will get the loss, recall and precision metrics results, like:
177+
178+
```shell
179+
recall: 0.927, precision: 0.945
180+
total_loss: 0.210, loss_xy: 0.010, loss_wh: 0.025, loss_conf: 0.125, loss_class: 0.050
181+
```
182+
183+
### 8. Other skills
184+
185+
There are many skills you can try during training:
186+
187+
(1) Data augmentation: You can implement your data augmentation like color jittering under `data_augmentation` method in `./utils/data_utils.py`.
188+
189+
(2) Mutil-scale training: You can change the input image scales periodically like the author does in the original paper.
190+
191+
### 9. TODO
192+
193+
- [ ] Multi-GPU training with sync batch norm.
194+
195+
-------
196+
197+
### Credits:
198+
199+
I refer to many fantastic repos during the implementation:
200+
201+
https://github.com/YunYang1994/tensorflow-yolov3
202+
203+
https://github.com/qqwweee/keras-yolo3
204+
205+
https://github.com/eriklindernoren/PyTorch-YOLOv3
206+
207+
https://github.com/pjreddie/darknet
208+
209+
210+
211+
212+
213+

‎convert_weight.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# coding: utf-8
2+
# for more details about the yolo darknet weights file, refer to
3+
# https://itnext.io/implementing-yolo-v3-in-tensorflow-tf-slim-c3c55ff59dbe
4+
5+
import os
6+
import sys
7+
import tensorflow as tf
8+
import numpy as np
9+
10+
from model import yolov3
11+
from utils.misc_utils import parse_anchors, load_weights
12+
13+
num_class = 80
14+
img_size = 416
15+
weight_path = './data/darknet_weights/yolov3.weights'
16+
save_path = './data/darknet_weights/yolov3.ckpt'
17+
anchors = parse_anchors('./data/yolo_anchors.txt')
18+
19+
model = yolov3(80, anchors)
20+
with tf.Session() as sess:
21+
inputs = tf.placeholder(tf.float32, [1, img_size, img_size, 3])
22+
23+
with tf.variable_scope('yolov3'):
24+
feature_map = model.forward(inputs)
25+
26+
saver = tf.train.Saver(var_list=tf.global_variables(scope='yolov3'))
27+
28+
load_ops = load_weights(tf.global_variables(scope='yolov3'), weight_path)
29+
sess.run(load_ops)
30+
saver.save(sess, save_path=save_path)
31+
print('TensorFlow model checkpoint has been saved to {}'.format(save_path))
32+
33+
34+

‎data/coco.names

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
person
2+
bicycle
3+
car
4+
motorbike
5+
aeroplane
6+
bus
7+
train
8+
truck
9+
boat
10+
traffic light
11+
fire hydrant
12+
stop sign
13+
parking meter
14+
bench
15+
bird
16+
cat
17+
dog
18+
horse
19+
sheep
20+
cow
21+
elephant
22+
bear
23+
zebra
24+
giraffe
25+
backpack
26+
umbrella
27+
handbag
28+
tie
29+
suitcase
30+
frisbee
31+
skis
32+
snowboard
33+
sports ball
34+
kite
35+
baseball bat
36+
baseball glove
37+
skateboard
38+
surfboard
39+
tennis racket
40+
bottle
41+
wine glass
42+
cup
43+
fork
44+
knife
45+
spoon
46+
bowl
47+
banana
48+
apple
49+
sandwich
50+
orange
51+
broccoli
52+
carrot
53+
hot dog
54+
pizza
55+
donut
56+
cake
57+
chair
58+
sofa
59+
pottedplant
60+
bed
61+
diningtable
62+
toilet
63+
tvmonitor
64+
laptop
65+
mouse
66+
remote
67+
keyboard
68+
cell phone
69+
microwave
70+
oven
71+
toaster
72+
sink
73+
refrigerator
74+
book
75+
clock
76+
vase
77+
scissors
78+
teddy bear
79+
hair drier
80+
toothbrush

0 commit comments

Comments
 (0)