Skip to content

Commit

Permalink
Update code for yolov4 model
Browse files Browse the repository at this point in the history
  • Loading branch information
@hunglc007
hunglc007 committed Apr 26, 2020
1 parent d8d6b54 commit b707010
Show file tree
Hide file tree
Showing 12 changed files with 302 additions and 88 deletions.
60 changes: 60 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
@@ -1,2 +1,62 @@
# tensorflow-yolov4-tflite
YoloV4 Implemented in Tensorflow 2.0. Convert .weights to .tflite format for tensorflow lite.

Download yolov4.weights file: https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT


### Prerequisites
Tensorflow 2.1.0
tensorflow_addons 0.9.1 (require for mish activation)
requirements.txt

###Performance
<p align="center"><img src="data/performance.png" width="640"\></p>

### Demo

```bash
# yolov4
python detect.py --weights ./data/yolov4.weights --framework tf --size 608 --image ./data/kite.jpg
```

####Output
<p align="center"><img src="result.png" width="640"\></p>



### Evaluate on COCO 2017 Dataset
```bash
# preprocess coco dataset
cd data
mkdir dataset
cd ..
cd scripts
python coco_convert.py --input COCO_ANOTATION_DATA_PATH --output val2017.pkl
python coco_annotation.py --coco_path COCO_DATA_PATH
cd ..

# evaluate yolov4 model
python evaluate.py --weights ./data/yolov4.weights
cd mAP/extra
python remove_space.py
cd ..
python main.py --output results_yolov4_tf

```

### TODO
* [ ] Training code
* [ ] greedy NMS
* [ ] Update scale xy
* [ ] ciou
* [ ] Mosaic data augmentation
* [ ] yolov4 tflite version
* [ ] yolov4 in8 tflite version for mobile

### References

* YOLOv4: Optimal Speed and Accuracy of Object Detection [YOLOv4](https://arxiv.org/abs/2004.10934).
* [darknet](https://github.com/AlexeyAB/darknet)
* [Yolov3 tensorflow](https://github.com/YunYang1994/tensorflow-yolov3)
* [Yolov3 tf2](https://github.com/zzh8829/yolov3-tf2)

93 changes: 62 additions & 31 deletions core/backbone.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,63 +4,94 @@
import tensorflow as tf
import core.common as common


def cspdarknet53(input_data):
def darknet53(input_data):

input_data = common.convolutional(input_data, (3, 3, 3, 32))
input_data = common.convolutional(input_data, (3, 3, 32, 64), downsample=True)

for i in range(1):
input_data = common.residual_block(input_data, 64, 32, 64)

input_data = common.convolutional(input_data, (3, 3, 64, 128), downsample=True)

for i in range(2):
input_data = common.residual_block(input_data, 128, 64, 128)

input_data = common.convolutional(input_data, (3, 3, 128, 256), downsample=True)

for i in range(8):
input_data = common.residual_block(input_data, 256, 128, 256)

route_1 = input_data
input_data = common.convolutional(input_data, (3, 3, 256, 512), downsample=True)

for i in range(8):
input_data = common.residual_block(input_data, 512, 256, 512)

route_2 = input_data
input_data = common.convolutional(input_data, (3, 3, 512, 1024), downsample=True)

for i in range(4):
input_data = common.residual_block(input_data, 1024, 512, 1024)

return route_1, route_2, input_data

def cspdarknet53(input_data):

input_data = common.convolutional(input_data, (3, 3, 3, 32), activate_type="mish")
input_data = common.convolutional(input_data, (3, 3, 32, 64), downsample=True, activate_type="mish")

route = input_data
input_data = common.convolutional(input_data, (1, 1, 64, 64))
route = common.convolutional(route, (1, 1, 64, 64), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 64, 64), activate_type="mish")
for i in range(1):
input_data = common.residual_block(input_data, 64, 32, 64)
input_data = common.convolutional(input_data, (1, 1, 64, 64))
route = common.convolutional(route, (1, 1, 64, 64))
input_data = tf.concat([route, input_data], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 64, 64), activate_type="mish")

input_data = common.convolutional(input_data, (1, 1, 128, 64))
input_data = common.convolutional(input_data, (3, 3, 64, 128), downsample=True)
input_data = tf.concat([input_data, route], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 128, 64), activate_type="mish")
input_data = common.convolutional(input_data, (3, 3, 64, 128), downsample=True, activate_type="mish")
route = input_data
input_data = common.convolutional(input_data, (1, 1, 128, 64))
route = common.convolutional(route, (1, 1, 128, 64), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 128, 64), activate_type="mish")
for i in range(2):
input_data = common.residual_block(input_data, 64, 64, 64)
input_data = common.convolutional(input_data, (1, 1, 64, 64))
route = common.convolutional(route, (1, 1, 128, 64))
input_data = tf.concat([route, input_data], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 64, 64), activate_type="mish")
input_data = tf.concat([input_data, route], axis=-1)

input_data = common.convolutional(input_data, (1, 1, 128, 128))
input_data = common.convolutional(input_data, (3, 3, 128, 256), downsample=True)
input_data = common.convolutional(input_data, (1, 1, 128, 128), activate_type="mish")
input_data = common.convolutional(input_data, (3, 3, 128, 256), downsample=True, activate_type="mish")
route = input_data
input_data = common.convolutional(input_data, (1, 1, 256, 128))
route = common.convolutional(route, (1, 1, 256, 128), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 256, 128), activate_type="mish")
for i in range(8):
input_data = common.residual_block(input_data, 128, 128, 128)
input_data = common.convolutional(input_data, (1, 1, 128, 128))
route = common.convolutional(route, (1, 1, 256, 128))
input_data = tf.concat([route, input_data], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 128, 128), activate_type="mish")
input_data = tf.concat([input_data, route], axis=-1)

input_data = common.convolutional(input_data, (1, 1, 256, 256))
input_data = common.convolutional(input_data, (1, 1, 256, 256), activate_type="mish")
route_1 = input_data
input_data = common.convolutional(input_data, (3, 3, 256, 512), downsample=True)
input_data = common.convolutional(input_data, (3, 3, 256, 512), downsample=True, activate_type="mish")
route = input_data
input_data = common.convolutional(route, (1, 1, 512, 256))
route = common.convolutional(route, (1, 1, 512, 256), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 512, 256), activate_type="mish")
for i in range(8):
input_data = common.residual_block(input_data, 256, 256, 256)
input_data = common.convolutional(input_data, (1, 1, 256, 256))
route = common.convolutional(route, (1, 1, 512, 256))
input_data = tf.concat([route, input_data], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 256, 256), activate_type="mish")
input_data = tf.concat([input_data, route], axis=-1)

input_data = common.convolutional(input_data, (1, 1, 512, 512))
input_data = common.convolutional(input_data, (1, 1, 512, 512), activate_type="mish")
route_2 = input_data
input_data = common.convolutional(input_data, (3, 3, 512, 1024), downsample=True)
input_data = common.convolutional(input_data, (3, 3, 512, 1024), downsample=True, activate_type="mish")
route = input_data
input_data = common.convolutional(input_data, (1, 1, 1024, 512))
route = common.convolutional(route, (1, 1, 1024, 512), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 1024, 512), activate_type="mish")
for i in range(4):
input_data = common.residual_block(input_data, 512, 512, 512)
input_data = common.convolutional(input_data, (1, 1, 512, 512))
route = common.convolutional(route, (1, 1, 1024, 512))
input_data = tf.concat([route, input_data], axis=-1)
input_data = common.convolutional(input_data, (1, 1, 512, 512), activate_type="mish")
input_data = tf.concat([input_data, route], axis=-1)

input_data = common.convolutional(input_data, (1, 1, 1024, 1024))
input_data = common.convolutional(input_data, (1, 1, 1024, 1024), activate_type="mish")
input_data = common.convolutional(input_data, (1, 1, 1024, 512))
input_data = common.convolutional(input_data, (3, 3, 512, 1024))
input_data = common.convolutional(input_data, (1, 1, 1024, 512))
Expand Down
20 changes: 15 additions & 5 deletions core/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
# coding=utf-8

import tensorflow as tf

import tensorflow_addons as tfa
class BatchNormalization(tf.keras.layers.BatchNormalization):
"""
"Frozen state" and "inference mode" are two separate concepts.
Expand All @@ -16,7 +16,7 @@ def call(self, x, training=False):
training = tf.logical_and(training, self.trainable)
return super().call(x, training)

def convolutional(input_layer, filters_shape, downsample=False, activate=True, bn=True):
def convolutional(input_layer, filters_shape, downsample=False, activate=True, bn=True, activate_type='leaky'):
if downsample:
input_layer = tf.keras.layers.ZeroPadding2D(((1, 0), (1, 0)))(input_layer)
padding = 'valid'
Expand All @@ -29,17 +29,27 @@ def convolutional(input_layer, filters_shape, downsample=False, activate=True, b
use_bias=not bn, kernel_regularizer=tf.keras.regularizers.l2(0.0005),
kernel_initializer=tf.random_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(0.))(input_layer)
# conv = tf.keras.layers.Conv2D(filters=filters_shape[-1], kernel_size = filters_shape[0], strides=strides, padding=padding
# , use_bias=not bn)(input_layer)

if bn: conv = BatchNormalization()(conv)
if activate == True: conv = tf.nn.leaky_relu(conv, alpha=0.1)
if activate == True:
if activate_type == "leaky":
conv = tf.nn.leaky_relu(conv, alpha=0.1)
elif activate_type == "mish":
conv = tfa.activations.mish(conv)
# conv = conv * tf.nn.tanh(tf.keras.activations.relu(tf.nn.softplus(conv), max_value=20))
# conv = tf.nn.softplus(conv)
# conv = tf.keras.activations.relu(tf.nn.softplus(conv), max_value=20)
# if activate == True: conv = tf.keras.layers.ReLU()(conv)

return conv


def residual_block(input_layer, input_channel, filter_num1, filter_num2):
short_cut = input_layer
conv = convolutional(input_layer, filters_shape=(1, 1, input_channel, filter_num1))
conv = convolutional(conv , filters_shape=(3, 3, filter_num1, filter_num2))
conv = convolutional(input_layer, filters_shape=(1, 1, input_channel, filter_num1), activate_type="mish")
conv = convolutional(conv , filters_shape=(3, 3, filter_num1, filter_num2), activate_type="mish")

residual_output = short_cut + conv
return residual_output
Expand Down
50 changes: 45 additions & 5 deletions core/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ def load_weights_tiny(model, weights_file):
major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

j = 0
for i in range(13):
for i in range(1):
conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'

Expand Down Expand Up @@ -43,12 +43,12 @@ def load_weights_tiny(model, weights_file):
assert len(wf.read()) == 0, 'failed to read all data'
wf.close()

def load_weights(model, weights_file):
def load_weights_v3(model, weights_file):
wf = open(weights_file, 'rb')
major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

j = 0
for i in range(78):
for i in range(75):
conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'

Expand All @@ -57,7 +57,7 @@ def load_weights(model, weights_file):
k_size = conv_layer.kernel_size[0]
in_dim = conv_layer.input_shape[-1]

if i not in [99, 100, 111]:
if i not in [58, 66, 74]:
# darknet weights: [beta, gamma, mean, variance]
bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
# tf weights: [gamma, beta, mean, variance]
Expand All @@ -73,7 +73,7 @@ def load_weights(model, weights_file):
# tf shape (height, width, in_dim, out_dim)
conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])

if i not in [99, 100, 111]:
if i not in [58, 66, 74]:
conv_layer.set_weights([conv_weights])
bn_layer.set_weights(bn_weights)
else:
Expand All @@ -82,6 +82,45 @@ def load_weights(model, weights_file):
# assert len(wf.read()) == 0, 'failed to read all data'
wf.close()

def load_weights(model, weights_file):
wf = open(weights_file, 'rb')
major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

j = 0
for i in range(110):
conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'

conv_layer = model.get_layer(conv_layer_name)
filters = conv_layer.filters
k_size = conv_layer.kernel_size[0]
in_dim = conv_layer.input_shape[-1]

if i not in [93, 101, 109]:
# darknet weights: [beta, gamma, mean, variance]
bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
# tf weights: [gamma, beta, mean, variance]
bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]
bn_layer = model.get_layer(bn_layer_name)
j += 1
else:
conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)

# darknet shape (out_dim, in_dim, height, width)
conv_shape = (filters, in_dim, k_size, k_size)
conv_weights = np.fromfile(wf, dtype=np.float32, count=np.product(conv_shape))
# tf shape (height, width, in_dim, out_dim)
conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])

if i not in [93, 101, 109]:
conv_layer.set_weights([conv_weights])
bn_layer.set_weights(bn_weights)
else:
conv_layer.set_weights([conv_weights, conv_bias])

assert len(wf.read()) == 0, 'failed to read all data'
wf.close()


def read_class_names(class_file_name):
'''loads class name from a file'''
Expand Down Expand Up @@ -255,6 +294,7 @@ def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
# # (5) discard some boxes with low scores
classes = np.argmax(pred_prob, axis=-1)
scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
# scores = pred_prob[np.arange(len(pred_coor)), classes]
score_mask = scores > score_threshold
mask = np.logical_and(scale_mask, score_mask)
coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]
Expand Down
Loading

0 comments on commit b707010

Please sign in to comment.