import tensorflow.keras.backend as K
from tensorflow.keras.models import Model
from tensorflow.keras import Input
from tensorflow.keras.layers import Conv2D, PReLU, UpSampling2D, concatenate , Reshape, Dense, Permute, MaxPool2D
from tensorflow.keras.layers import GlobalAveragePooling2D, Activation, add, GaussianNoise, BatchNormalization, multiply
from tensorflow.keras.optimizers import SGD
from loss import custom_loss
K.set_image_data_format("channels_last")
def unet_model(input_shape, modified_unet=True, learning_rate=0.01, start_channel=64,
number_of_levels=3, inc_rate=2, output_channels=4, saved_model_dir=None):
"""
Builds UNet model
Parameters
----------
input_shape : tuple
Shape of the input data (height, width, channel)
modified_unet : bool
Whether to use modified UNet or the original UNet
learning_rate : float
Learning rate for the model. The default is 0.01.
start_channel : int
Number of channels of the first conv. The default is 64.
number_of_levels : int
The depth size of the U-structure. The default is 3.
inc_rate : int
Rate at which the conv channels will increase. The default is 2.
output_channels : int
The number of output layer channels. The default is 4
saved_model_dir : str
If spesified, the model weights will be loaded from this path. The default is None.
Returns
-------
model : keras.model
The created keras model with respect to the input parameters
"""
input_layer = Input(shape=input_shape, name='the_input_layer')
if modified_unet:
x = GaussianNoise(0.01, name='Gaussian_Noise')(input_layer)
x = Conv2D(64, 2, padding='same')(x)
x = level_block_modified(x, start_channel, number_of_levels, inc_rate)
x = BatchNormalization(axis = -1)(x)
x = PReLU(shared_axes=[1, 2])(x)
else:
x = level_block(input_layer, start_channel, number_of_levels, inc_rate)
x = Conv2D(output_channels, 1, padding='same')(x)
output_layer = Activation('softmax')(x)
model = Model(inputs = input_layer, outputs = output_layer)
if modified_unet:
print("The modified UNet was built!")
else:
print("The original UNet was built!")
if saved_model_dir:
model.load_weights(saved_model_dir)
print("the model weights were successfully loaded!")
sgd = SGD(lr=learning_rate, momentum=0.9, decay=0)
model.compile(optimizer=sgd, loss=custom_loss)
return model
def se_block(x, ratio=16):
"""
creates a squeeze and excitation block
https://arxiv.org/abs/1709.01507
Parameters
----------
x : tensor
Input keras tensor
ratio : int
The reduction ratio. The default is 16.
Returns
-------
x : tensor
A keras tensor
"""
channel_axis = 1 if K.image_data_format() == "channels_first" else -1
filters = x.shape[channel_axis]
se_shape = (1, 1, filters)
se = GlobalAveragePooling2D()(x)
se = Reshape(se_shape)(se)
se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)
se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)
if K.image_data_format() == 'channels_first':
se = Permute((3, 1, 2))(se)
x = multiply([x, se])
return x
def level_block(x, dim, level, inc):
if level > 0:
m = conv_layers(x, dim)
x = MaxPool2D(pool_size=(2, 2))(m)
x = level_block(x,int(inc*dim), level-1, inc)
x = UpSampling2D(size=(2, 2))(x)
x = Conv2D(dim, 2, padding='same')(x)
m = concatenate([m,x])
x = conv_layers(m, dim)
else:
x = conv_layers(x, dim)
return x
def level_block_modified(x, dim, level, inc):
if level > 0:
m = res_block(x, dim, encoder_path=True)##########
x = Conv2D(int(inc*dim), 2, strides=2, padding='same')(m)
x = level_block_modified(x, int(inc*dim), level-1, inc)
x = UpSampling2D(size=(2, 2))(x)
x = Conv2D(dim, 2, padding='same')(x)
m = concatenate([m,x])
m = se_block(m, 8)
x = res_block(m, dim, encoder_path=False)
else:
x = res_block(x, dim, encoder_path=True) #############
return x
def conv_layers(x, dim):
x = Conv2D(dim, 3, padding='same')(x)
x = Activation("relu")(x)
x = Conv2D(dim, 3, padding='same')(x)
x = Activation("relu")(x)
return x
def res_block(x, dim, encoder_path=True):
m = BatchNormalization(axis = -1)(x)
m = PReLU(shared_axes = [1, 2])(m)
m = Conv2D(dim, 3, padding='same')(m)
m = BatchNormalization(axis = -1)(m)
m = PReLU(shared_axes = [1, 2])(m)
m = Conv2D(dim, 3, padding='same')(m)
if encoder_path:
x = add([x, m])
else:
x = Conv2D(dim, 1, padding='same', use_bias=False)(x)
x = add([x,m])
return x
