Skip to content

Commit

Permalink
Merge pull request #81 from ar90n/add-yolov3-spp-support
Browse files Browse the repository at this point in the history
Add yolov3-spp support
  • Loading branch information
mystic123 authored Aug 30, 2019
2 parents 74b2f46 + 142150c commit 136fb66
Show file tree
Hide file tree
Showing 5 changed files with 53 additions and 6 deletions.
11 changes: 8 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ To run demo type this in the command line:
1. Download binary file with desired weights:
1. Full weights: `wget https://pjreddie.com/media/files/yolov3.weights`
1. Tiny weights: `wget https://pjreddie.com/media/files/yolov3-tiny.weights`
1. SPP weights: `wget https://pjreddie.com/media/files/yolov3-spp.weights`
2. Run `python ./convert_weights.py` and `python ./convert_weights_pb.py`
3. Run `python ./demo.py --input_img <path-to-image> --output_img <name-of-output-image> --frozen_model <path-to-frozen-model>`

Expand All @@ -33,7 +34,9 @@ To run demo type this in the command line:
1. `NCHW` (gpu only) or `NHWC`
4. `--tiny`
1. Use yolov3-tiny
5. `--ckpt_file`
5. `--spp`
1. Use yolov3-spp
6. `--ckpt_file`
1. Output checkpoint file
2. convert_weights_pb.py:
1. `--class_names`
Expand All @@ -44,7 +47,9 @@ To run demo type this in the command line:
1. `NCHW` (gpu only) or `NHWC`
4. `--tiny`
1. Use yolov3-tiny
5. `--output_graph`
5. `--spp`
1. Use yolov3-spp
6. `--output_graph`
1. Location to write the output .pb graph to
3. demo.py
1. `--class_names`
Expand All @@ -62,4 +67,4 @@ To run demo type this in the command line:
7. `--iou_threshold`
1. Desired iou threshold
8. `--gpu_memory_fraction`
1. Fraction of gpu memory to work with
1. Fraction of gpu memory to work with
4 changes: 4 additions & 0 deletions convert_weights.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,17 @@
'data_format', 'NCHW', 'Data format: NCHW (gpu only) / NHWC')
tf.app.flags.DEFINE_bool(
'tiny', False, 'Use tiny version of YOLOv3')
tf.app.flags.DEFINE_bool(
'spp', False, 'Use SPP version of YOLOv3')
tf.app.flags.DEFINE_string(
'ckpt_file', './saved_model/model.ckpt', 'Chceckpoint file')


def main(argv=None):
if FLAGS.tiny:
model = yolo_v3_tiny.yolo_v3_tiny
elif FLAGS.spp:
model = yolo_v3.yolo_v3_spp
else:
model = yolo_v3.yolo_v3

Expand Down
4 changes: 4 additions & 0 deletions convert_weights_pb.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@

tf.app.flags.DEFINE_bool(
'tiny', False, 'Use tiny version of YOLOv3')
tf.app.flags.DEFINE_bool(
'spp', False, 'Use SPP version of YOLOv3')
tf.app.flags.DEFINE_integer(
'size', 416, 'Image size')

Expand All @@ -29,6 +31,8 @@
def main(argv=None):
if FLAGS.tiny:
model = yolo_v3_tiny.yolo_v3_tiny
elif FLAGS.spp:
model = yolo_v3.yolo_v3_spp
else:
model = yolo_v3.yolo_v3

Expand Down
4 changes: 4 additions & 0 deletions demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
'frozen_model', '', 'Frozen tensorflow protobuf model')
tf.app.flags.DEFINE_bool(
'tiny', False, 'Use tiny version of YOLOv3')
tf.app.flags.DEFINE_bool(
'spp', False, 'Use SPP version of YOLOv3')

tf.app.flags.DEFINE_integer(
'size', 416, 'Image size')
Expand Down Expand Up @@ -71,6 +73,8 @@ def main(argv=None):
else:
if FLAGS.tiny:
model = yolo_v3_tiny.yolo_v3_tiny
elif FLAGS.spp:
model = yolo_v3.yolo_v3_spp
else:
model = yolo_v3.yolo_v3

Expand Down
36 changes: 33 additions & 3 deletions yolo_v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,14 @@ def _darknet53_block(inputs, filters):
return inputs


def _spp_block(inputs, data_format='NCHW'):
return tf.concat([slim.max_pool2d(inputs, 13, 1, 'SAME'),
slim.max_pool2d(inputs, 9, 1, 'SAME'),
slim.max_pool2d(inputs, 5, 1, 'SAME'),
inputs],
axis=1 if data_format == 'NCHW' else 3)


@tf.contrib.framework.add_arg_scope
def _fixed_padding(inputs, kernel_size, *args, mode='CONSTANT', **kwargs):
"""
Expand Down Expand Up @@ -95,10 +103,15 @@ def _fixed_padding(inputs, kernel_size, *args, mode='CONSTANT', **kwargs):
return padded_inputs


def _yolo_block(inputs, filters):
def _yolo_block(inputs, filters, data_format='NCHW', with_spp=False):
inputs = _conv2d_fixed_padding(inputs, filters, 1)
inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
inputs = _conv2d_fixed_padding(inputs, filters, 1)

if with_spp:
inputs = _spp_block(inputs, data_format)
inputs = _conv2d_fixed_padding(inputs, filters, 1)

inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
inputs = _conv2d_fixed_padding(inputs, filters, 1)
route = inputs
Expand Down Expand Up @@ -187,7 +200,7 @@ def _upsample(inputs, out_shape, data_format='NCHW'):
return inputs


def yolo_v3(inputs, num_classes, is_training=False, data_format='NCHW', reuse=False):
def yolo_v3(inputs, num_classes, is_training=False, data_format='NCHW', reuse=False, with_spp=False):
"""
Creates YOLO v3 model.
Expand All @@ -197,6 +210,7 @@ def yolo_v3(inputs, num_classes, is_training=False, data_format='NCHW', reuse=Fa
:param is_training: whether is training or not.
:param data_format: data format NCHW or NHWC.
:param reuse: whether or not the network and its variables should be reused.
:param with_spp: whether or not is using spp layer.
:return:
"""
# it will be needed later on
Expand Down Expand Up @@ -228,7 +242,8 @@ def yolo_v3(inputs, num_classes, is_training=False, data_format='NCHW', reuse=Fa
route_1, route_2, inputs = darknet53(inputs)

with tf.variable_scope('yolo-v3'):
route, inputs = _yolo_block(inputs, 512)
route, inputs = _yolo_block(inputs, 512, data_format, with_spp)

detect_1 = _detection_layer(
inputs, num_classes, _ANCHORS[6:9], img_size, data_format)
detect_1 = tf.identity(detect_1, name='detect_1')
Expand Down Expand Up @@ -260,3 +275,18 @@ def yolo_v3(inputs, num_classes, is_training=False, data_format='NCHW', reuse=Fa
detections = tf.concat([detect_1, detect_2, detect_3], axis=1)
detections = tf.identity(detections, name='detections')
return detections


def yolo_v3_spp(inputs, num_classes, is_training=False, data_format='NCHW', reuse=False):
"""
Creates YOLO v3 with SPP model.
:param inputs: a 4-D tensor of size [batch_size, height, width, channels].
Dimension batch_size may be undefined. The channel order is RGB.
:param num_classes: number of predicted classes.
:param is_training: whether is training or not.
:param data_format: data format NCHW or NHWC.
:param reuse: whether or not the network and its variables should be reused.
:return:
"""
return yolo_v3(inputs, num_classes, is_training=is_training, data_format=data_format, reuse=reuse, with_spp=True)

0 comments on commit 136fb66

Please sign in to comment.