import tensorflow as tf
from tensorflow.keras import layers, initializers, models
def conv(x, filters, kernel_size, downsampling=False, activation='leaky', batch_norm=True):
def mish(x):
return x * tf.math.tanh(tf.math.softplus(x))
if downsampling:
x = layers.ZeroPadding2D(padding=((1, 0), (1, 0)))(x) # top & left padding
padding = 'valid'
strides = 2
else:
padding = 'same'
strides = 1
x = layers.Conv2D(filters,
kernel_size,
strides=strides,
padding=padding,
use_bias=not batch_norm,
# kernel_regularizer=regularizers.l2(0.0005),
kernel_initializer=initializers.RandomNormal(mean=0.0, stddev=0.01),
# bias_initializer=initializers.Zeros()
)(x)
if batch_norm:
x = layers.BatchNormalization()(x)
if activation == 'mish':
x = mish(x)
elif activation == 'leaky':
x = layers.LeakyReLU(alpha=0.1)(x)
return x
def residual_block(x, filters1, filters2, activation='leaky'):
"""
:param x: input tensor
:param filters1: num of filter for 1x1 conv
:param filters2: num of filter for 3x3 conv
:param activation: default activation function: leaky relu
:return:
"""
y = conv(x, filters1, kernel_size=1, activation=activation)
y = conv(y, filters2, kernel_size=3, activation=activation)
return layers.Add()([x, y])
def csp_block(x, residual_out, repeat, residual_bottleneck=False):
"""
Cross Stage Partial Network (CSPNet)
transition_bottleneck_dims: 1x1 bottleneck
output_dims: 3x3
:param x:
:param residual_out:
:param repeat:
:param residual_bottleneck:
:return:
"""
route = x
route = conv(route, residual_out, 1, activation="mish")
x = conv(x, residual_out, 1, activation="mish")
for i in range(repeat):
x = residual_block(x,
residual_out // 2 if residual_bottleneck else residual_out,
residual_out,
activation="mish")
x = conv(x, residual_out, 1, activation="mish")
x = layers.Concatenate()([x, route])
return x
def darknet53(x):
x = conv(x, 32, 3)
x = conv(x, 64, 3, downsampling=True)
for i in range(1):
x = residual_block(x, 32, 64)
x = conv(x, 128, 3, downsampling=True)
for i in range(2):
x = residual_block(x, 64, 128)
x = conv(x, 256, 3, downsampling=True)
for i in range(8):
x = residual_block(x, 128, 256)
route_1 = x
x = conv(x, 512, 3, downsampling=True)
for i in range(8):
x = residual_block(x, 256, 512)
route_2 = x
x = conv(x, 1024, 3, downsampling=True)
for i in range(4):
x = residual_block(x, 512, 1024)
return route_1, route_2, x
def cspdarknet53(input):
x = conv(input, 32, 3)
x = conv(x, 64, 3, downsampling=True)
x = csp_block(x, residual_out=64, repeat=1, residual_bottleneck=True)
x = conv(x, 64, 1, activation='mish')
x = conv(x, 128, 3, activation='mish', downsampling=True)
x = csp_block(x, residual_out=64, repeat=2)
x = conv(x, 128, 1, activation='mish')
x = conv(x, 256, 3, activation='mish', downsampling=True)
x = csp_block(x, residual_out=128, repeat=8)
x = conv(x, 256, 1, activation='mish')
route0 = x
x = conv(x, 512, 3, activation='mish', downsampling=True)
x = csp_block(x, residual_out=256, repeat=8)
x = conv(x, 512, 1, activation='mish')
route1 = x
x = conv(x, 1024, 3, activation='mish', downsampling=True)
x = csp_block(x, residual_out=512, repeat=4)
x = conv(x, 1024, 1, activation="mish")
x = conv(x, 512, 1)
x = conv(x, 1024, 3)
x = conv(x, 512, 1)
x = layers.Concatenate()([layers.MaxPooling2D(pool_size=13, strides=1, padding='same')(x),
layers.MaxPooling2D(pool_size=9, strides=1, padding='same')(x),
layers.MaxPooling2D(pool_size=5, strides=1, padding='same')(x),
x
])
x = conv(x, 512, 1)
x = conv(x, 1024, 3)
route2 = conv(x, 512, 1)
return models.Model(input, [route0, route1, route2])
def yolov4_neck(x, num_classes):
backbone_model = cspdarknet53(x)
route0, route1, route2 = backbone_model.output
route_input = route2
x = conv(route2, 256, 1)
x = layers.UpSampling2D()(x)
route1 = conv(route1, 256, 1)
x = layers.Concatenate()([route1, x])
x = conv(x, 256, 1)
x = conv(x, 512, 3)
x = conv(x, 256, 1)
x = conv(x, 512, 3)
x = conv(x, 256, 1)
route1 = x
x = conv(x, 128, 1)
x = layers.UpSampling2D()(x)
route0 = conv(route0, 128, 1)
x = layers.Concatenate()([route0, x])
x = conv(x, 128, 1)
x = conv(x, 256, 3)
x = conv(x, 128, 1)
x = conv(x, 256, 3)
x = conv(x, 128, 1)
route0 = x
x = conv(x, 256, 3)
conv_sbbox = conv(x, 3 * (num_classes + 5), 1, activation=None, batch_norm=False)
x = conv(route0, 256, 3, downsampling=True)
x = layers.Concatenate()([x, route1])
x = conv(x, 256, 1)
x = conv(x, 512, 3)
x = conv(x, 256, 1)
x = conv(x, 512, 3)
x = conv(x, 256, 1)
route1 = x
x = conv(x, 512, 3)
conv_mbbox = conv(x, 3 * (num_classes + 5), 1, activation=None, batch_norm=False)
x = conv(route1, 512, 3, downsampling=True)
x = layers.Concatenate()([x, route_input])
x = conv(x, 512, 1)
x = conv(x, 1024, 3)
x = conv(x, 512, 1)
x = conv(x, 1024, 3)
x = conv(x, 512, 1)
x = conv(x, 1024, 3)
conv_lbbox = conv(x, 3 * (num_classes + 5), 1, activation=None, batch_norm=False)
return [conv_sbbox, conv_mbbox, conv_lbbox]
def yolov4_head(yolo_neck_outputs, classes, anchors, xyscale):
bbox0, object_probability0, class_probabilities0, pred_box0 = get_boxes(yolo_neck_outputs[0],
anchors=anchors[0, :, :], classes=classes,
grid_size=52, strides=8,
xyscale=xyscale[0])
bbox1, object_probability1, class_probabilities1, pred_box1 = get_boxes(yolo_neck_outputs[1],
anchors=anchors[1, :, :], classes=classes,
grid_size=26, strides=16,
xyscale=xyscale[1])
bbox2, object_probability2, class_probabilities2, pred_box2 = get_boxes(yolo_neck_outputs[2],
anchors=anchors[2, :, :], classes=classes,
grid_size=13, strides=32,
xyscale=xyscale[2])
x = [bbox0, object_probability0, class_probabilities0, pred_box0,
bbox1, object_probability1, class_probabilities1, pred_box1,
bbox2, object_probability2, class_probabilities2, pred_box2]
return x
def get_boxes(pred, anchors, classes, grid_size, strides, xyscale):
"""
:param pred:
:param anchors:
:param classes:
:param grid_size:
:param strides:
:param xyscale:
:return:
"""
pred = tf.reshape(pred,
(tf.shape(pred)[0],
grid_size,
grid_size,
3,
5 + classes)) # (batch_size, grid_size, grid_size, 3, 5+classes)
box_xy, box_wh, obj_prob, class_prob = tf.split(
pred, (2, 2, 1, classes), axis=-1
) # (?, 52, 52, 3, 2) (?, 52, 52, 3, 2) (?, 52, 52, 3, 1) (?, 52, 52, 3, 80)
box_xy = tf.sigmoid(box_xy) # (?, 52, 52, 3, 2)
obj_prob = tf.sigmoid(obj_prob) # (?, 52, 52, 3, 1)
class_prob = tf.sigmoid(class_prob) # (?, 52, 52, 3, 80)
pred_box_xywh = tf.concat((box_xy, box_wh), axis=-1) # (?, 52, 52, 3, 4)
grid = tf.meshgrid(tf.range(grid_size), tf.range(grid_size)) # (52, 52) (52, 52)
grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2) # (52, 52, 1, 2)
grid = tf.cast(grid, dtype=tf.float32)
box_xy = ((box_xy * xyscale) - 0.5 * (xyscale - 1) + grid) * strides # (?, 52, 52, 1, 4)
box_wh = tf.exp(box_wh) * anchors # (?, 52, 52, 3, 2)
box_x1y1 = box_xy - box_wh / 2 # (?, 52, 52, 3, 2)
box_x2y2 = box_xy + box_wh / 2 # (?, 52, 52, 3, 2)
pred_box_x1y1x2y2 = tf.concat([box_x1y1, box_x2y2], axis=-1) # (?, 52, 52, 3, 4)
return pred_box_x1y1x2y2, obj_prob, class_prob, pred_box_xywh
# pred_box_x1y1x2y2: absolute xy value
def nms(model_ouputs, input_shape, num_class, iou_threshold=0.413, score_threshold=0.3):
"""
Apply Non-Maximum suppression
ref: https://www.tensorflow.org/api_docs/python/tf/image/combined_non_max_suppression
:param model_ouputs: yolo model model_ouputs
:param input_shape: size of input image
:return: nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections
"""
bs = tf.shape(model_ouputs[0])[0]
boxes = tf.zeros((bs, 0, 4))
confidence = tf.zeros((bs, 0, 1))
class_probabilities = tf.zeros((bs, 0, num_class))
for output_idx in range(0, len(model_ouputs), 4):
output_xy = model_ouputs[output_idx]
output_conf = model_ouputs[output_idx + 1]
output_classes = model_ouputs[output_idx + 2]
boxes = tf.concat([boxes, tf.reshape(output_xy, (bs, -1, 4))], axis=1)
confidence = tf.concat([confidence, tf.reshape(output_conf, (bs, -1, 1))], axis=1)
class_probabilities = tf.concat([class_probabilities, tf.reshape(output_classes, (bs, -1, num_class))], axis=1)
scores = confidence * class_probabilities
boxes = tf.expand_dims(boxes, axis=-2)
boxes = boxes / input_shape[0] # box normalization: relative img size
print(f'nms iou: {iou_threshold} score: {score_threshold}')
(nmsed_boxes, # [bs, max_detections, 4]
nmsed_scores, # [bs, max_detections]
nmsed_classes, # [bs, max_detections]
valid_detections # [batch_size]
) = tf.image.combined_non_max_suppression(
boxes=boxes, # y1x1, y2x2 [0~1]
scores=scores,
max_output_size_per_class=100,
max_total_size=100, # max_boxes: Maximum nmsed_boxes in a single img.
iou_threshold=iou_threshold, # iou_threshold: Minimum overlap that counts as a valid detection.
score_threshold=score_threshold, # # Minimum confidence that counts as a valid detection.
)
return nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections