def Process(inputs, value, type='up'):
if type == 'up':
x = L.UpSampling2D()(inputs[0])
add = wBiFPNAdd(name='bifpn-add-{0}-{1}'.format(index, value))(
[x, inputs[1]])
if type == 'max':
x = L.MaxPooling2D(pool_size=3, strides=2,
padding='same')(inputs[0])
add = wBiFPNAdd(name='bifpn-add-{0}-{1}'.format(index, value))(
[x, inputs[1], inputs[2]])
if type == 'out':
x = L.MaxPooling2D(pool_size=3, strides=2,
padding='same')(inputs[0])
add = wBiFPNAdd(name='bifpn-add-{0}-{1}'.format(index, value))(
[x, inputs[1]])
out = L.Activation(lambda y: tf.nn.swish(y))(add)
out = ConvBlock(value=value, kernel_size=3, strides=1)(out)
return out
def single_BiFPN_merge_step(feature_map_other_level,
feature_maps_current_level, upsampling,
num_channels, idx_BiFPN_layer, node_idx, op_idx):
"""
Merges two feature maps of different levels in the BiFPN
Args:
feature_map_other_level: Input feature map of a different level. Needs to be resized before merging.
feature_maps_current_level: Input feature map of the current level
upsampling: Boolean indicating wheter to upsample or downsample the feature map of the different level to match the shape of the current level
num_channels: Number of channels used in the BiFPN
idx_BiFPN_layer: The index of the BiFPN layer to build
node_idx, op_idx: Integers needed to set the correct layer names
Returns:
The merged feature map
"""
if upsampling:
feature_map_resampled = layers.UpSampling2D()(feature_map_other_level)
else:
feature_map_resampled = layers.MaxPooling2D(
pool_size=3, strides=2, padding='same')(feature_map_other_level)
merged_feature_map = wBiFPNAdd(
name=f'fpn_cells/cell_{idx_BiFPN_layer}/fnode{node_idx}/add')(
feature_maps_current_level + [feature_map_resampled])
merged_feature_map = layers.Activation(lambda x: tf.nn.swish(x))(
merged_feature_map)
merged_feature_map = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=
f'fpn_cells/cell_{idx_BiFPN_layer}/fnode{node_idx}/op_after_combine{op_idx}'
)(merged_feature_map)
return merged_feature_map
def build_wBiFPN(features, num_channels, id, freeze_bn=False):
if id == 0:
_, _, C3, C4, C5 = features
P3_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P3'.format(id))(C3)
P4_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P4'.format(id))(C4)
P5_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P5'.format(id))(C5)
P6_in = ConvBlock(num_channels,
kernel_size=3,
strides=2,
freeze_bn=freeze_bn,
name='BiFPN_{}_P6'.format(id))(C5)
P7_in = ConvBlock(num_channels,
kernel_size=3,
strides=2,
freeze_bn=freeze_bn,
name='BiFPN_{}_P7'.format(id))(P6_in)
else:
P3_in, P4_in, P5_in, P6_in, P7_in = features
P3_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P3'.format(id))(P3_in)
P4_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P4'.format(id))(P4_in)
P5_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P5'.format(id))(P5_in)
P6_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P6'.format(id))(P6_in)
P7_in = ConvBlock(num_channels,
kernel_size=1,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_P7'.format(id))(P7_in)
# upsample
P7_U = layers.UpSampling2D()(P7_in)
P6_td = wBiFPNAdd()([P7_U, P6_in])
P6_td = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_U_P6'.format(id))(P6_td)
P6_U = layers.UpSampling2D()(P6_td)
P5_td = wBiFPNAdd()([P6_U, P5_in])
P5_td = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_U_P5'.format(id))(P5_td)
P5_U = layers.UpSampling2D()(P5_td)
P4_td = wBiFPNAdd()([P5_U, P4_in])
P4_td = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_U_P4'.format(id))(P4_td)
P4_U = layers.UpSampling2D()(P4_td)
P3_out = wBiFPNAdd()([P4_U, P3_in])
P3_out = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_U_P3'.format(id))(P3_out)
# downsample
P3_D = layers.MaxPooling2D(strides=(2, 2))(P3_out)
P4_out = wBiFPNAdd()([P3_D, P4_td, P4_in])
P4_out = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_D_P4'.format(id))(P4_out)
P4_D = layers.MaxPooling2D(strides=(2, 2))(P4_out)
P5_out = wBiFPNAdd()([P4_D, P5_td, P5_in])
P5_out = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_D_P5'.format(id))(P5_out)
P5_D = layers.MaxPooling2D(strides=(2, 2))(P5_out)
P6_out = wBiFPNAdd()([P5_D, P6_td, P6_in])
P6_out = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_D_P6'.format(id))(P6_out)
P6_D = layers.MaxPooling2D(strides=(2, 2))(P6_out)
P7_out = wBiFPNAdd()([P6_D, P7_in])
P7_out = DepthwiseConvBlock(kernel_size=3,
strides=1,
freeze_bn=freeze_bn,
name='BiFPN_{}_D_P7'.format(id))(P7_out)
return P3_out, P4_out, P5_out, P6_out, P7_out
def build_wBiFPN(features, num_channels, id, freeze_bn=False):
if id == 0:
_, _, C3, C4, C5 = features
P3_in = C3
P4_in = C4
P5_in = C5
P6_in = layers.Conv2D(num_channels,
kernel_size=1,
padding='same',
name='resample_p6/conv2d')(C5)
P6_in = layers.BatchNormalization(momentum=MOMENTUM,
epsilon=EPSILON,
name='resample_p6/bn')(P6_in)
# P6_in = BatchNormalization(freeze=freeze_bn, name='resample_p6/bn')(P6_in)
P6_in = layers.MaxPooling2D(pool_size=3,
strides=2,
padding='same',
name='resample_p6/maxpool')(P6_in)
P7_in = layers.MaxPooling2D(pool_size=3,
strides=2,
padding='same',
name='resample_p7/maxpool')(P6_in)
P7_U = layers.UpSampling2D()(P7_in)
P6_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode0/add')(
[P6_in, P7_U])
P6_td = layers.Activation(lambda x: tf.nn.swish(x))(P6_td)
P6_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
P5_in_1 = layers.Conv2D(
num_channels,
kernel_size=1,
padding='same',
name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/conv2d')(P5_in)
P5_in_1 = layers.BatchNormalization(
momentum=MOMENTUM,
epsilon=EPSILON,
name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
# P5_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
P6_U = layers.UpSampling2D()(P6_td)
P5_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode1/add')(
[P5_in_1, P6_U])
P5_td = layers.Activation(lambda x: tf.nn.swish(x))(P5_td)
P5_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
P4_in_1 = layers.Conv2D(
num_channels,
kernel_size=1,
padding='same',
name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/conv2d')(P4_in)
P4_in_1 = layers.BatchNormalization(
momentum=MOMENTUM,
epsilon=EPSILON,
name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
# P4_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
P5_U = layers.UpSampling2D()(P5_td)
P4_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode2/add')(
[P4_in_1, P5_U])
P4_td = layers.Activation(lambda x: tf.nn.swish(x))(P4_td)
P4_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
P3_in = layers.Conv2D(
num_channels,
kernel_size=1,
padding='same',
name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/conv2d')(P3_in)
P3_in = layers.BatchNormalization(
momentum=MOMENTUM,
epsilon=EPSILON,
name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
# P3_in = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
P4_U = layers.UpSampling2D()(P4_td)
P3_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode3/add')(
[P3_in, P4_U])
P3_out = layers.Activation(lambda x: tf.nn.swish(x))(P3_out)
P3_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
P4_in_2 = layers.Conv2D(
num_channels,
kernel_size=1,
padding='same',
name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/conv2d')(P4_in)
P4_in_2 = layers.BatchNormalization(
momentum=MOMENTUM,
epsilon=EPSILON,
name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
# P4_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
P3_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P3_out)
P4_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode4/add')(
[P4_in_2, P4_td, P3_D])
P4_out = layers.Activation(lambda x: tf.nn.swish(x))(P4_out)
P4_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)
P5_in_2 = layers.Conv2D(
num_channels,
kernel_size=1,
padding='same',
name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/conv2d')(P5_in)
P5_in_2 = layers.BatchNormalization(
momentum=MOMENTUM,
epsilon=EPSILON,
name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
# P5_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
P4_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P4_out)
P5_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode5/add')(
[P5_in_2, P5_td, P4_D])
P5_out = layers.Activation(lambda x: tf.nn.swish(x))(P5_out)
P5_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)
P5_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P5_out)
P6_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode6/add')(
[P6_in, P6_td, P5_D])
P6_out = layers.Activation(lambda x: tf.nn.swish(x))(P6_out)
P6_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)
P6_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P6_out)
P7_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode7/add')(
[P7_in, P6_D])
P7_out = layers.Activation(lambda x: tf.nn.swish(x))(P7_out)
P7_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)
else:
P3_in, P4_in, P5_in, P6_in, P7_in = features
P7_U = layers.UpSampling2D()(P7_in)
P6_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode0/add')(
[P6_in, P7_U])
P6_td = layers.Activation(lambda x: tf.nn.swish(x))(P6_td)
P6_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
P6_U = layers.UpSampling2D()(P6_td)
P5_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode1/add')(
[P5_in, P6_U])
P5_td = layers.Activation(lambda x: tf.nn.swish(x))(P5_td)
P5_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
P5_U = layers.UpSampling2D()(P5_td)
P4_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode2/add')(
[P4_in, P5_U])
P4_td = layers.Activation(lambda x: tf.nn.swish(x))(P4_td)
P4_td = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
P4_U = layers.UpSampling2D()(P4_td)
P3_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode3/add')(
[P3_in, P4_U])
P3_out = layers.Activation(lambda x: tf.nn.swish(x))(P3_out)
P3_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
P3_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P3_out)
P4_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode4/add')(
[P4_in, P4_td, P3_D])
P4_out = layers.Activation(lambda x: tf.nn.swish(x))(P4_out)
P4_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)
P4_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P4_out)
P5_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode5/add')(
[P5_in, P5_td, P4_D])
P5_out = layers.Activation(lambda x: tf.nn.swish(x))(P5_out)
P5_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)
P5_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P5_out)
P6_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode6/add')(
[P6_in, P6_td, P5_D])
P6_out = layers.Activation(lambda x: tf.nn.swish(x))(P6_out)
P6_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)
P6_D = layers.MaxPooling2D(pool_size=3, strides=2,
padding='same')(P6_out)
P7_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode7/add')(
[P7_in, P6_D])
P7_out = layers.Activation(lambda x: tf.nn.swish(x))(P7_out)
P7_out = SeparableConvBlock(
num_channels=num_channels,
kernel_size=3,
strides=1,
name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)
return P3_out, P4_td, P5_td, P6_td, P7_out