def pan_down(backbone_output_layers):
output_layers = []
pan_down_output_layer = scaled_yolov4_csp_block(
backbone_output_layers[-1],
backbone_output_layers[-1].shape[-1] // 2,
3,
type='spp')
output_layers.append(pan_down_output_layer)
for layer_index in range(2, len(backbone_output_layers) + 1):
pan_down_output_layer = conv2d_bn_mish(
pan_down_output_layer, pan_down_output_layer.shape[-1] // 2,
(1, 1))
pan_down_output_layer = tf.keras.layers.UpSampling2D(
size=(2, 2))(pan_down_output_layer)
backbone_output_layer = conv2d_bn_mish(
backbone_output_layers[-layer_index],
backbone_output_layers[-layer_index].shape[3] // 2, (1, 1))
pan_down_output_layer = tf.keras.layers.Concatenate()(
[backbone_output_layer, pan_down_output_layer])
pan_down_output_layer = scaled_yolov4_csp_block(
pan_down_output_layer,
pan_down_output_layer.shape[-1] // 2,
3,
type='head')
output_layers.append(pan_down_output_layer)
return output_layers
def scaled_yolov4_csp_darknet53(x, mode='p5'):
darknet53_filters = [64 * 2**i for i in range(5)]
if mode == 'p5':
loop_nums = [1, 3, 15, 15, 7]
elif mode == 'p6':
loop_nums = [1, 3, 15, 15, 7, 7]
darknet53_filters += [1024]
elif mode == 'p7':
loop_nums = [1, 3, 15, 15, 7, 7, 7]
darknet53_filters += [1024] * 2
x = conv2d_bn_mish(x, 32, (3, 3), name="first_block")
output_layers = []
for block_index in range(len(loop_nums)):
x = tf.keras.layers.ZeroPadding2D(((1, 0), (1, 0)))(x)
x = conv2d_bn_mish(x,
darknet53_filters[block_index], (3, 3),
strides=(2, 2),
padding='valid',
name="backbone_block_{}_0".format(block_index))
x = scaled_yolov4_csp_block(
x,
darknet53_filters[block_index],
loop_nums[block_index],
type="backbone",
name="backbone_block_{}_1".format(block_index))
output_layers.append(x)
return output_layers[2:]
def head_up(head_down_output_layers):
output_layers = []
head_up_output_layers = head_down_output_layers[-1]
output_layers.append(head_up_output_layers)
head_up_block_index = 0
for layer_index in range(2, len(head_down_output_layers) + 1):
if layer_index >= 4:
head_up_out_channel = head_up_output_layers.shape[-1]
else:
head_up_out_channel = head_up_output_layers.shape[-1] * 2
head_up_output_layers = tf.keras.layers.ZeroPadding2D(
((1, 0), (1, 0)))(head_up_output_layers)
head_up_output_layers = conv2d_bn_mish(
head_up_output_layers,
head_up_out_channel, (3, 3),
strides=(2, 2),
padding='valid',
name="head_up_block_{}_0".format(head_up_block_index))
head_up_output_layers = tf.keras.layers.Concatenate()(
[head_up_output_layers, head_down_output_layers[-layer_index]])
head_up_output_layers = scaled_yolov4_csp_block(
head_up_output_layers,
head_up_output_layers.shape[-1] // 2,
3,
type='head',
name="head_up_block_{}_1".format(head_up_block_index))
head_up_block_index += 1
output_layers.append(head_up_output_layers)
return output_layers
def spp(x):
pool_sizes = [5, 9, 13]
pooling_results = [
tf.keras.layers.MaxPooling2D((pool_size, pool_size),
strides=(1, 1),
padding='same')(x)
for pool_size in pool_sizes
]
spp_result = tf.keras.layers.Concatenate()(pooling_results + [x])
spp_result = conv2d_bn_mish(spp_result, x.shape[3], (1, 1))
return spp_result
def head_down(backbone_output_layers):
output_layers = []
head_down_output_layer = scaled_yolov4_csp_block(
backbone_output_layers[-1],
backbone_output_layers[-1].shape[-1] // 2,
3,
type='spp',
name='spp')
output_layers.append(head_down_output_layer)
head_down_block_index = 0
for layer_index in range(2, len(backbone_output_layers) + 1):
if layer_index < len(backbone_output_layers) - 1:
spp_out_channel = head_down_output_layer.shape[-1]
else:
spp_out_channel = head_down_output_layer.shape[-1] // 2
head_down_output_layer = conv2d_bn_mish(
head_down_output_layer,
spp_out_channel, (1, 1),
name="head_down_block_{}_0".format(head_down_block_index))
head_down_output_layer = tf.keras.layers.UpSampling2D(
size=(2, 2))(head_down_output_layer)
backbone_output_layer = conv2d_bn_mish(
backbone_output_layers[-layer_index],
backbone_output_layers[-layer_index].shape[3] // 2, (1, 1),
name="head_down_block_{}_1".format(head_down_block_index))
head_down_output_layer = tf.keras.layers.Concatenate()(
[backbone_output_layer, head_down_output_layer])
head_down_output_layer = scaled_yolov4_csp_block(
head_down_output_layer,
head_down_output_layer.shape[-1] // 2,
3,
type='head',
name="head_down_block_{}_2".format(head_down_block_index))
head_down_block_index += 1
output_layers.append(head_down_output_layer)
return output_layers
def pan_up(pan_down_output_layers):
output_layers = []
pan_up_output_layers = pan_down_output_layers[-1]
output_layers.append(pan_up_output_layers)
for layer_index in range(2, len(pan_down_output_layers) + 1):
pan_up_output_layers = tf.keras.layers.ZeroPadding2D(
((1, 0), (1, 0)))(pan_up_output_layers)
pan_up_output_layers = conv2d_bn_mish(pan_up_output_layers,
pan_up_output_layers.shape[-1] *
2, (3, 3),
strides=(2, 2),
padding='valid')
pan_up_output_layers = tf.keras.layers.Concatenate()(
[pan_up_output_layers, pan_down_output_layers[-layer_index]])
pan_down_output_layer = scaled_yolov4_csp_block(
pan_up_output_layers,
pan_up_output_layers.shape[-1] // 2,
3,
type='head')
output_layers.append(pan_down_output_layer)
return output_layers
def yolov3_head(inputs, args):
class_num = int(args.num_classes)
if class_num == 1:
class_num = 0
anchors = yolo_anchors[args.model_type]
output_layers = []
for index, x in enumerate(inputs):
x = conv2d_bn_mish(x,
x.shape[-1] * 2, (3, 3),
name='yolov3_head_%d_0' % index)
x = tf.keras.layers.Conv2D(len(anchors[index]) * (class_num + 5),
(1, 1),
use_bias=True,
name='yolov3_head_%d_1_conv2d' % index)(x)
x = tf.reshape(x, [
tf.shape(x)[0],
tf.shape(x)[1],
tf.shape(x)[2], -1, class_num + 5
])
output_layers.append(x)
return output_layers
def csp_darknet_block(x, loop_num, filters, is_half_filters=True):
x = tf.keras.layers.ZeroPadding2D(((1, 0), (1, 0)))(x)
x = conv2d_bn_mish(x, filters, (3, 3), strides=(2, 2), padding='valid')
csp_branch = conv2d_bn_mish(x,
filters // 2 if is_half_filters else filters,
(1, 1))
darknet_branch = conv2d_bn_mish(
x, filters // 2 if is_half_filters else filters, (1, 1))
for i in range(loop_num):
x = conv2d_bn_mish(darknet_branch, filters // 2, (1, 1))
x = conv2d_bn_mish(x, filters // 2 if is_half_filters else filters,
(3, 3))
darknet_branch = tf.keras.layers.Add()([darknet_branch, x])
darknet_branch = conv2d_bn_mish(
darknet_branch, filters // 2 if is_half_filters else filters, (1, 1))
x = tf.keras.layers.Concatenate()([darknet_branch, csp_branch])
return conv2d_bn_mish(x, filters, (1, 1))