unet.py

def unet(x, ch, depth, cut_path=False):
    conv_kwargs = {'kernel_size': (3, 3), 'strides'=(1, 1), 'dilation_rate'=(1, 1),
                  'padding'='same', 'use_bias'=True, 'bias_initializer'='zeros', 'kernel_initializer'='he_normal'}
    for _ in range(2):
        x = tf.keras.layers.Convolution2D(ch, **conv_kwargs)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation('relu')(x)
    
    if(depth != 0):
        path = x
        x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
        x = unet(x, ch=ch*2, depth=depth-1, cut_path=cut_path)
        x = tf.keras.layers.UpSampling2D((2, 2))(x)
        if(cut_path==False):
            x = tf.keras.layers.Concatenate(axis=-1)([x, path])
                    
    for _ in range(2):
        x = tf.keras.layers.Convolution2D(ch, **conv_kwargs)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation('relu')(x)
    return x