def create_yolo_model(inputs, anchors_count, classes_count):
route_1, route_2, x = darknet53(inputs)
x = _conv_block_body(x, 1024)
conv_large_branch, conv_large_box = _conv_box(x, 1024, anchors_count,
classes_count)
x = conv_block(x, filters=256, kernel_size=1)
x = UpSampling2D()(x)
x = tf.concat([x, route_2], axis=-1)
x = _conv_block_body(x, 512)
conv_medium_branch, conv_medium_box = _conv_box(x, 512, anchors_count,
classes_count)
x = conv_block(x, filters=128, kernel_size=1)
x = UpSampling2D()(x)
x = tf.concat([x, route_1], axis=-1)
x = _conv_block_body(x, 256)
conv_small_branch, conv_small_box = _conv_box(x, 256, anchors_count,
classes_count)
# return Model(inputs, [conv_large_box, conv_medium_box, conv_small_box])
return [conv_large_box, conv_medium_box, conv_small_box]
def _conv_box(inputs, filters, num_anchors, num_classes):
branch = conv_block(inputs, filters=filters, kernel_size=3)
box = conv_block(branch,
filters=num_anchors * (num_classes + 5),
kernel_size=1,
activate=False,
batch_norm=False)
return branch, box
def _conv_block_body(inputs, filters):
half = filters // 2
x = conv_block(inputs, filters=half, kernel_size=1)
x = conv_block(x, filters=filters, kernel_size=3)
x = conv_block(x, filters=half, kernel_size=1)
x = conv_block(x, filters=filters, kernel_size=3)
x = conv_block(x, filters=half, kernel_size=1)
return x
def darknet53(inputs):
x = conv_block(inputs, filters=32, kernel_size=3)
x = _res_block_body(x, 64, 1)
x = _res_block_body(x, 128, 2)
x = _res_block_body(x, 256, 8)
route_1 = x
x = _res_block_body(x, 512, 8)
route_2 = x
x = _res_block_body(x, 1024, 4)
return route_1, route_2, x
def _res_block_body(inputs, filters, res_blocks_count):
x = conv_block(inputs, filters=filters, kernel_size=3, down_sample=True)
for i in range(res_blocks_count):
x = res_block(x, filters // 2, filters)
return x