Creating different Unet models in tensorflow

different type of unet model from tversky focal loss paper will be implemented here

print(os.environ.get('CUDA_PATH'))

/usr/lib/cuda

Start with normal unet

---

source

pooling

 pooling (inputs, max_pool_only=False, both=True, pool_size=2)

---

source

conv_block

 conv_block (inputs, filter_no, kernel_size, batch_nm=True, dropout=True,
             drp_rt=0.1, kernel_initializer='glorot_normal')

Create a conv block with batch norma and dropout

	Type	Default	Details
inputs
filter_no
kernel_size
batch_nm	bool	True
dropout	bool	True
drp_rt	float	0.1
kernel_initializer	str	glorot_normal	option he_normal

input_size = tf.random.normal((1, 1152, 1632, 1))
out_ = conv_block(input_size, 32, 3, batch_nm=False, dropout=False)
test_eq([1,1152, 1632, 32], out_.shape)

---

source

double_conv

 double_conv (inputs, filter_no, kernel_size, batch_nm=True, dropout=True,
              drp_rt=0.1, kernel_initializer='glorot_normal')

Create double conv block with batch norma and dropout

	Type	Default	Details
inputs
filter_no
kernel_size
batch_nm	bool	True
dropout	bool	True
drp_rt	float	0.1
kernel_initializer	str	glorot_normal	option he_normal

out_c = double_conv(inputs=input_size, filter_no=32, kernel_size=3, batch_nm=True, dropout=True, drp_rt=0.1)
out_p = pooling(out_c, both=True)

test_eq([1,1152, 1632, 32], out_c.shape)
test_eq([1,576, 816, 64], out_p.shape)

h, w, c = 1152, 1632, 1

---

source

encoder_block

 encoder_block (input_size:tensorflow.python.framework.tensor.Tensor,
                filter_size:List[int]=[16, 32, 64, 128, 256],
                kernel_initializer:str='glorot_normal')

Create Encoder block for Unet

	Type	Default	Details
input_size	Tensor
filter_size	typing.List[int]	[16, 32, 64, 128, 256]
kernel_initializer	str	glorot_normal	option he_normal

#h, w, c = 1052, 1632, 1
#h, w, c = 1152, 1632, 1
input_img = tf.random.normal((1, h, w, c))
inputs,enc_outputs, final_l = encoder_block(
    input_size=(h, w, c),
    filter_size=[16,32,64,128],
    kernel_initializer='glorot_normal' #option he_normal
    )

enc_outputs

[<KerasTensor: shape=(None, 1152, 1632, 16) dtype=float32 (created by layer 'dropout_16')>,
 <KerasTensor: shape=(None, 576, 816, 32) dtype=float32 (created by layer 'dropout_18')>,
 <KerasTensor: shape=(None, 288, 408, 64) dtype=float32 (created by layer 'dropout_20')>,
 <KerasTensor: shape=(None, 144, 204, 128) dtype=float32 (created by layer 'dropout_22')>]

final_l

<KerasTensor: shape=(None, 72, 102, 256) dtype=float32 (created by layer 'dropout_24')>

test_eq((None, 1152, 1632, 1), inputs.shape)

expected_shapes = [
    (None, 1152, 1632, 16),
    (None, 576, 816, 32),
    (None, 288, 408, 64),
    (None, 144, 204, 128),
]

for (expected, actual) in zip(expected_shapes, enc_outputs):
    test_eq(expected, actual.shape)

len(enc_outputs)

4

---

source

crop_and_concat

 crop_and_concat (upsampled:tensorflow.python.framework.tensor.Tensor,
                  skip_connection:tensorflow.python.framework.tensor.Tenso
                  r)

Crop or pad the skip connection and concatenate it with the upsampled tensor. Args: - upsampled (tf.Tensor): Tensor from the decoder path. - skip_connection (tf.Tensor): Corresponding tensor from the encoder path. Returns: - tf.Tensor: Tensor after cropping or padding and concatenation.

skip_connection = tf.random.normal([1, 16, 16, 128])  # Simulating a feature map from the encoder
upsampled = tf.random.normal([1, 14, 14, 128])
crop_concat_image = crop_and_concat(upsampled, skip_connection)
test_eq(crop_concat_image.shape, [1, 14, 14, 256])

 Cropping branch

---

source

decoder_block

 decoder_block (filter_list:List[int],
                encoder_outputs:List[tensorflow.python.framework.tensor.Te
                nsor],
                bottleneck:tensorflow.python.framework.tensor.Tensor)

Create a decoder block for unet

inputs,enc_outputs, final_l = encoder_block(
    input_size=(h, w, c),
    filter_size=[16,32,64,128],
    kernel_initializer='glorot_normal' #option he_normal
    )
out_, dec_out_list = decoder_block(
    filter_list=[16,32,64,128],
    encoder_outputs=enc_outputs,
    bottleneck=final_l,
    )

expected_decoder_shapes = [
    (None, 144, 204, 256),
    (None, 288, 408, 128),
    (None, 576, 816, 64),
    (None, 1152, 1632, 32),
]

for (expected, actual) in zip(expected_decoder_shapes, dec_out_list):
    test_eq(expected, actual.shape)
test_eq((None, 1152, 1632, 16), out_.shape)

---

source

unet_model

 unet_model (input_size:Tuple[int]=(1152, 1632, 1),
             filter_list:List[int]=[16, 32, 64, 128],
             kernel_initializer:str='glorot_normal', n_classes:int=1)

Create a unet model

	Type	Default	Details
input_size	typing.Tuple[int]	(1152, 1632, 1)
filter_list	typing.List[int]	[16, 32, 64, 128]
kernel_initializer	str	glorot_normal	option he_normal
n_classes	int	1

model_ = unet_model(
    input_size=(1152, 1632, 1),
    filter_list=[16,32,64,128],
    n_classes=1,
    )

tf.keras.utils.plot_model(model_, show_shapes=True, show_layer_activations=True)

Attention Unet

---

source

unet_gating_signal

 unet_gating_signal (input:tensorflow.python.framework.tensor.Tensor,
                     out_size:int, is_batchnorm:bool, name:str,
                     kinit:str='glorot_normal')

this is simply 1x1 convolution, bn, activation

---

source

expend_as

 expend_as (tensor, rep, name)

---

source

attention_blocks

 attention_blocks (x, g, inter_shape, name)

take g which is the spatially smaller signal, do a conv to get the same number of feature channels as x (bigger spatially) do a conv on x to also get same feature channels (theta_x) then, upsample g to be same size as x add x and g (concat_xg) relu, 1x1 conv, then sigmoid then upsample the final - this gives us attn coefficients

---

source

decoder_block_attention_gates

 decoder_block_attention_gates (filter_list:List[int],
                                encoder_outputs:List[tensorflow.python.fra
                                mework.tensor.Tensor], bottleneck:tensorfl
                                ow.python.framework.tensor.Tensor)

Create a decoder block for unet with attention gates

inputs, encoder_output, bottle_neck_enc = encoder_block(
    input_size=(256, 256,1),
    filter_size=[64,128,256,512],
    )
dec_final, decoder_outputs = decoder_block_attention_gates(
    filter_list=[64,128,256,512],
    encoder_outputs=encoder_output,
    bottleneck=bottle_neck_enc,
    )

testing decoder output shape

expected_attn_deocder_shapes = [
    (None, 32, 32, 512),
    (None, 64, 64, 256),
    (None, 128, 128, 128),
    (None, 256, 256, 64),
]
for (expected, actual) in zip(expected_attn_deocder_shapes, decoder_outputs):
    test_eq(expected, actual.shape)

---

source

unet_model_attention_gates

 unet_model_attention_gates (input_size:Tuple[int]=(256, 256, 1),
                             filter_list:List[int]=[64, 128, 256, 512],
                             kernel_initializer:str='glorot_normal',
                             n_classes:int=1)

Create a unet model

	Type	Default	Details
input_size	typing.Tuple[int]	(256, 256, 1)
filter_list	typing.List[int]	[64, 128, 256, 512]
kernel_initializer	str	glorot_normal	option he_normal
n_classes	int	1

attn_unet_model_ =unet_model_attention_gates(
    input_size=(256, 256, 1),
    filter_list=[64,128,256,512],
    n_classes=1,
    )

tf.keras.utils.plot_model(attn_unet_model_, show_shapes=True, show_layer_activations=True)

Residual Attention Unet

---

source

res_double_conv

 res_double_conv (inputs, filter_no, kernel_size, batch_nm=True,
                  dropout=True, drp_rt=0.1,
                  kernel_initializer='glorot_normal')

Create double conv block with batch norma and dropout

	Type	Default	Details
inputs
filter_no
kernel_size
batch_nm	bool	True
dropout	bool	True
drp_rt	float	0.1
kernel_initializer	str	glorot_normal	option he_normal

---

source

residual_encoder_block

 residual_encoder_block
                         (input_size:tensorflow.python.framework.tensor.Te
                         nsor, filter_size:List[int]=[16, 32, 64, 128,
                         256], kernel_initializer:str='glorot_normal')

Create Encoder block for residual Unet

	Type	Default	Details
input_size	Tensor
filter_size	typing.List[int]	[16, 32, 64, 128, 256]
kernel_initializer	str	glorot_normal	option he_normal

---

source

residual_decoder_block

 residual_decoder_block (filter_list:List[int],
                         encoder_outputs:List[tensorflow.python.framework.
                         tensor.Tensor], bottleneck:tensorflow.python.fram
                         ework.tensor.Tensor)

Create a decoder block for unet with attention gates

---

source

residual_attn_unet

 residual_attn_unet (input_size:Tuple[int]=(256, 256, 1),
                     filter_list:List[int]=[64, 128, 256, 512],
                     kernel_initializer:str='glorot_normal',
                     n_classes:int=1)

Create a residual attention unet model

	Type	Default	Details
input_size	typing.Tuple[int]	(256, 256, 1)
filter_list	typing.List[int]	[64, 128, 256, 512]
kernel_initializer	str	glorot_normal	option he_normal
n_classes	int	1

resdual_attn_unet = residual_attn_unet(
    input_size=(256, 256, 1),
    filter_list=[64,128,256,512],
    n_classes=1,
    )

tf.keras.utils.plot_model(resdual_attn_unet, show_shapes=True, show_layer_activations=True)

Checking purpose

from tensorflow.keras import layers

def res_conv_block(x, filter_size, size, dropout, batch_norm=False):
    '''
    Residual convolutional layer.
    Two variants....
    Either put activation function before the addition with shortcut
    or after the addition (which would be as proposed in the original resNet).
    
    1. conv - BN - Activation - conv - BN - Activation 
                                          - shortcut  - BN - shortcut+BN
                                          
    2. conv - BN - Activation - conv - BN   
                                     - shortcut  - BN - shortcut+BN - Activation                                     
    
    Check fig 4 in https://arxiv.org/ftp/arxiv/papers/1802/1802.06955.pdf
    '''

    conv = layers.Conv2D(size, (filter_size, filter_size), padding='same')(x)
    if batch_norm is True:
        conv = layers.BatchNormalization(axis=3)(conv)
    conv = layers.Activation('relu')(conv)
    
    conv = layers.Conv2D(size, (filter_size, filter_size), padding='same')(conv)
    if batch_norm is True:
        conv = layers.BatchNormalization(axis=3)(conv)
    #conv = layers.Activation('relu')(conv)    #Activation before addition with shortcut
    if dropout > 0:
        conv = layers.Dropout(dropout)(conv)

    shortcut = layers.Conv2D(size, kernel_size=(1, 1), padding='same')(x)
    if batch_norm is True:
        shortcut = layers.BatchNormalization(axis=3)(shortcut)

    res_path = layers.add([shortcut, conv])
    res_path = layers.Activation('relu')(res_path)    #Activation after addition with shortcut (Original residual block)
    return res_path


def gating_signal(input, out_size, batch_norm=False):
    """
    resize the down layer feature map into the same dimension as the up layer feature map
    using 1x1 conv
    :return: the gating feature map with the same dimension of the up layer feature map
    """
    x = layers.Conv2D(out_size, (1, 1), padding='same')(input)
    if batch_norm:
        x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    return x

def attention_block(x, gating, inter_shape):
    shape_x = K.int_shape(x)
    shape_g = K.int_shape(gating)

# Getting the x signal to the same shape as the gating signal
    theta_x = layers.Conv2D(inter_shape, (2, 2), strides=(2, 2), padding='same')(x)  # 16
    shape_theta_x = K.int_shape(theta_x)

# Getting the gating signal to the same number of filters as the inter_shape
    phi_g = layers.Conv2D(inter_shape, (1, 1), padding='same')(gating)
    upsample_g = layers.Conv2DTranspose(inter_shape, (3, 3),
                                 strides=(shape_theta_x[1] // shape_g[1], shape_theta_x[2] // shape_g[2]),
                                 padding='same')(phi_g)  # 16

    concat_xg = layers.add([upsample_g, theta_x])
    act_xg = layers.Activation('relu')(concat_xg)
    psi = layers.Conv2D(1, (1, 1), padding='same')(act_xg)
    sigmoid_xg = layers.Activation('sigmoid')(psi)
    shape_sigmoid = K.int_shape(sigmoid_xg)
    upsample_psi = layers.UpSampling2D(size=(shape_x[1] // shape_sigmoid[1], shape_x[2] // shape_sigmoid[2]))(sigmoid_xg)  # 32

    upsample_psi = repeat_elem(upsample_psi, shape_x[3])

    y = layers.multiply([upsample_psi, x])

    result = layers.Conv2D(shape_x[3], (1, 1), padding='same')(y)
    result_bn = layers.BatchNormalization()(result)
    return result_bn


def repeat_elem(tensor, rep):
    # lambda function to repeat Repeats the elements of a tensor along an axis
    #by a factor of rep.
    # If tensor has shape (None, 256,256,3), lambda will return a tensor of shape 
    #(None, 256,256,6), if specified axis=3 and rep=2.

     return layers.Lambda(lambda x, repnum: K.repeat_elements(x, repnum, axis=3),
                          arguments={'repnum': rep})(tensor)

def Attention_ResUNet(input_shape, NUM_CLASSES=1, dropout_rate=0.0, batch_norm=True):
    '''
    Rsidual UNet, with attention 
    
    '''
    # network structure
    FILTER_NUM = 64 # number of basic filters for the first layer
    FILTER_SIZE = 3 # size of the convolutional filter
    UP_SAMP_SIZE = 2 # size of upsampling filters
    # input data
    # dimension of the image depth
    inputs = layers.Input(input_shape, dtype=tf.float32)
    axis = 3

    # Downsampling layers
    # DownRes 1, double residual convolution + pooling
    conv_128 = res_conv_block(inputs, FILTER_SIZE, FILTER_NUM, dropout_rate, batch_norm)
    pool_64 = layers.MaxPooling2D(pool_size=(2,2))(conv_128)
    # DownRes 2
    conv_64 = res_conv_block(pool_64, FILTER_SIZE, 2*FILTER_NUM, dropout_rate, batch_norm)
    pool_32 = layers.MaxPooling2D(pool_size=(2,2))(conv_64)
    # DownRes 3
    conv_32 = res_conv_block(pool_32, FILTER_SIZE, 4*FILTER_NUM, dropout_rate, batch_norm)
    pool_16 = layers.MaxPooling2D(pool_size=(2,2))(conv_32)
    # DownRes 4
    conv_16 = res_conv_block(pool_16, FILTER_SIZE, 8*FILTER_NUM, dropout_rate, batch_norm)
    pool_8 = layers.MaxPooling2D(pool_size=(2,2))(conv_16)
    # DownRes 5, convolution only
    conv_8 = res_conv_block(pool_8, FILTER_SIZE, 16*FILTER_NUM, dropout_rate, batch_norm)

    # Upsampling layers
    # UpRes 6, attention gated concatenation + upsampling + double residual convolution
    gating_16 = gating_signal(conv_8, 8*FILTER_NUM, batch_norm)
    att_16 = attention_block(conv_16, gating_16, 8*FILTER_NUM)
    up_16 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(conv_8)
    up_16 = layers.concatenate([up_16, att_16], axis=axis)
    up_conv_16 = res_conv_block(up_16, FILTER_SIZE, 8*FILTER_NUM, dropout_rate, batch_norm)
    # UpRes 7
    gating_32 = gating_signal(up_conv_16, 4*FILTER_NUM, batch_norm)
    att_32 = attention_block(conv_32, gating_32, 4*FILTER_NUM)
    up_32 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_16)
    up_32 = layers.concatenate([up_32, att_32], axis=axis)
    up_conv_32 = res_conv_block(up_32, FILTER_SIZE, 4*FILTER_NUM, dropout_rate, batch_norm)
    # UpRes 8
    gating_64 = gating_signal(up_conv_32, 2*FILTER_NUM, batch_norm)
    att_64 = attention_block(conv_64, gating_64, 2*FILTER_NUM)
    up_64 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_32)
    up_64 = layers.concatenate([up_64, att_64], axis=axis)
    up_conv_64 = res_conv_block(up_64, FILTER_SIZE, 2*FILTER_NUM, dropout_rate, batch_norm)
    # UpRes 9
    gating_128 = gating_signal(up_conv_64, FILTER_NUM, batch_norm)
    att_128 = attention_block(conv_128, gating_128, FILTER_NUM)
    up_128 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_64)
    up_128 = layers.concatenate([up_128, att_128], axis=axis)
    up_conv_128 = res_conv_block(up_128, FILTER_SIZE, FILTER_NUM, dropout_rate, batch_norm)

    # 1*1 convolutional layers
    
    conv_final = layers.Conv2D(NUM_CLASSES, kernel_size=(1,1))(up_conv_128)
    conv_final = layers.BatchNormalization(axis=axis)(conv_final)
    conv_final = layers.Activation('sigmoid')(conv_final)  #Change to softmax for multichannel

    # Model integration
    model = tf.keras.models.Model(inputs, conv_final, name="AttentionResUNet")
    return model

input_shape = (256, 256, 1)
attn_res_unet_ = Attention_ResUNet(input_shape, NUM_CLASSES=1, dropout_rate=0.0, batch_norm=True)
tf.keras.utils.plot_model(attn_res_unet_, show_shapes=True, show_layer_activations=True)