aix/models.py

#from aix.constants import *

import tensorflow as tf


def initialize_model(img_width, img_height, img_channels):
    #,
    #                 optimizer = 'adam', model_loss = 'binary_crossentropy', data_augm = False):

    inputs = tf.keras.layers.Input((img_width, img_height, img_channels), name='input')

    #Contraction path
    c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block1_conv2d_1')(inputs)
    c1 = tf.keras.layers.Dropout(0.1, name='block1_dropout')(c1)
    c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block1_conv2d_2')(c1)
    p1 = tf.keras.layers.MaxPooling2D((2, 2), name='block1_max_pooling')(c1)

    c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block2_conv2d_1')(p1)
    c2 = tf.keras.layers.Dropout(0.1, name='block2_dropout')(c2)
    c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block2_conv2d_2')(c2)
    p2 = tf.keras.layers.MaxPooling2D((2, 2), name='block2_max_pooling')(c2)

    c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block3_conv2d_1')(p2)
    c3 = tf.keras.layers.Dropout(0.2, name='block3_dropout')(c3)
    c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block3_conv2d_2')(c3)
    p3 = tf.keras.layers.MaxPooling2D((2, 2), name='block3_max_pooling')(c3)

    c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block4_conv2d_1')(p3)
    c4 = tf.keras.layers.Dropout(0.2, name='block4_dropout')(c4)
    c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block4_conv2d_2')(c4)
    p4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), name='block4_max_pooling')(c4)

    c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block5_conv2d_1')(p4)
    c5 = tf.keras.layers.Dropout(0.3, name='block5_dropout')(c5)
    c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block5_conv2d_2')(c5)

    #Expansive path
    u6 = tf.keras.layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same', name='block6_conv2d_transpose')(c5)
    u6 = tf.keras.layers.concatenate([u6, c4], name='block6_concatenate')
    c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block6_conv2d_1')(u6)
    c6 = tf.keras.layers.Dropout(0.2, name='block6_dropout')(c6)
    c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block6_conv2d_2')(c6)

    u7 = tf.keras.layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same', name='block7_conv2d_transpose')(c6)
    u7 = tf.keras.layers.concatenate([u7, c3], name='block7_concatenate')
    c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block7_conv2d_1')(u7)
    c7 = tf.keras.layers.Dropout(0.2, name='block7_dropout')(c7)
    c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block7_conv2d_2')(c7)

    u8 = tf.keras.layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same', name='block8_conv2d_transpose')(c7)
    u8 = tf.keras.layers.concatenate([u8, c2], name='block8_concatenate')
    c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block8_conv2d_1')(u8)
    c8 = tf.keras.layers.Dropout(0.1, name='block8_dropout')(c8)
    c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block8_conv2d_2')(c8)

    u9 = tf.keras.layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same', name='block9_conv2d_transpose')(c8)
    u9 = tf.keras.layers.concatenate([u9, c1], axis=3, name='block9_concatenate')
    c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block9_conv2d_1')(u9)
    c9 = tf.keras.layers.Dropout(0.1, name='block9_dropout')(c9)
    c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same', name='block9_conv2d_2')(c9)

    outputs = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid', name='output')(c9)

    model = tf.keras.Model(inputs = [inputs], outputs = [outputs])
    #model.compile(optimizer = optimizer, loss = model_loss, metrics = [dice_coef])
    #model.summary()

    return model


def initialize_model_v2(img_width, img_height, img_channels,
                     optimizer = 'adam', model_loss = 'binary_crossentropy', data_augm = False):
    
    n = 1
    inputs = tf.keras.layers.Input((img_width, img_height, img_channels))
    #Contraction path
    c1 = tf.keras.layers.Conv2D(16 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(inputs)
    #c1 = tf.keras.layers.Dropout(0.1)(c1)
    c1 = tf.keras.layers.Conv2D(16 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c1)
    p1 = tf.keras.layers.MaxPooling2D((2, 2))(c1)

    c2 = tf.keras.layers.Conv2D(32 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p1)
    #c2 = tf.keras.layers.Dropout(0.1)(c2)
    c2 = tf.keras.layers.Conv2D(32 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c2)
    p2 = tf.keras.layers.MaxPooling2D((2, 2))(c2)

    c3 = tf.keras.layers.Conv2D(64 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p2)
    #c3 = tf.keras.layers.Dropout(0.2)(c3)
    c3 = tf.keras.layers.Conv2D(64 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c3)
    p3 = tf.keras.layers.MaxPooling2D((2, 2))(c3)

    c4 = tf.keras.layers.Conv2D(128 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p3)
    #c4 = tf.keras.layers.Dropout(0.2)(c4)
    c4 = tf.keras.layers.Conv2D(128 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c4)
    p4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(c4)

    c5 = tf.keras.layers.Conv2D(256 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p4)
    #c5 = tf.keras.layers.Dropout(0.5)(c5)
    c5 = tf.keras.layers.Conv2D(256 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c5)
    c5 = tf.keras.layers.Dropout(0.5)(c5)

    #Expansive path
    u6 = tf.keras.layers.Conv2DTranspose(128 * n, (2, 2), strides=(2, 2), padding='same')(c5)
    u6 = tf.keras.layers.concatenate([u6, c4])
    c6 = tf.keras.layers.Conv2D(128 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u6)
    #c6 = tf.keras.layers.Dropout(0.2)(c6)
    c6 = tf.keras.layers.Conv2D(128 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c6)

    u7 = tf.keras.layers.Conv2DTranspose(64 * n, (2, 2), strides=(2, 2), padding='same')(c6)
    u7 = tf.keras.layers.concatenate([u7, c3])
    c7 = tf.keras.layers.Conv2D(64 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u7)
    #c7 = tf.keras.layers.Dropout(0.2)(c7)
    c7 = tf.keras.layers.Conv2D(64 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c7)

    u8 = tf.keras.layers.Conv2DTranspose(32 * n, (2, 2), strides=(2, 2), padding='same')(c7)
    u8 = tf.keras.layers.concatenate([u8, c2])
    c8 = tf.keras.layers.Conv2D(32 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u8)
    #c8 = tf.keras.layers.Dropout(0.1)(c8)
    c8 = tf.keras.layers.Conv2D(32 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c8)

    u9 = tf.keras.layers.Conv2DTranspose(16 * n, (2, 2), strides=(2, 2), padding='same')(c8)
    u9 = tf.keras.layers.concatenate([u9, c1], axis=3)
    c9 = tf.keras.layers.Conv2D(16 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u9)
    #c9 = tf.keras.layers.Dropout(0.1)(c9)
    c9 = tf.keras.layers.Conv2D(16 * n, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c9)

    outputs = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid')(c9)

    model = tf.keras.Model(inputs = [inputs], outputs = [outputs])
    #model.compile(optimizer = optimizer, loss = model_loss, metrics = [dice_coef])
    #model.summary()

    return model

# From pix2pix (https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py)
#
def pix2pix_upsample(filters, size, norm_type='batchnorm', apply_dropout=False):
    """Upsamples an input.
        Conv2DTranspose => Batchnorm => Dropout => Relu
        Args:
            filters: number of filters
            size: filter size
            norm_type: Normalization type; either 'batchnorm' or 'instancenorm'.
            apply_dropout: If True, adds the dropout layer
        Returns:
            Upsample Sequential Model
    """

    initializer = tf.random_normal_initializer(0., 0.02)

    result = tf.keras.Sequential()
    result.add(
        tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                      padding='same',
                                      kernel_initializer=initializer,
                                      use_bias=False))

    if norm_type.lower() == 'batchnorm':
        result.add(tf.keras.layers.BatchNormalization())
    elif norm_type.lower() == 'instancenorm':
        result.add(InstanceNormalization())

    if apply_dropout:
        result.add(tf.keras.layers.Dropout(0.5))

    result.add(tf.keras.layers.ReLU())

    return result

# Adapted From: https://www.tensorflow.org/tutorials/images/segmentation

def unet_model(output_channels:int, input_shape=[128, 128, 3],
               optimizer = 'adam', model_loss = 'binary_crossentropy'):
    base_model = tf.keras.applications.MobileNetV2(input_shape=input_shape, include_top=False)

    # Use the activations of these layers
    layer_names = [
        'block_1_expand_relu',   # 64x64
        'block_3_expand_relu',   # 32x32
        'block_6_expand_relu',   # 16x16
        'block_13_expand_relu',  # 8x8
        'block_16_project',      # 4x4
    ]
    base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

    # Create the feature extraction model
    down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)

    down_stack.trainable = False

    up_stack = [
        pix2pix_upsample(512, 3),  # 4x4 -> 8x8
        pix2pix_upsample(256, 3),  # 8x8 -> 16x16
        pix2pix_upsample(128, 3),  # 16x16 -> 32x32
        pix2pix_upsample(64, 3),   # 32x32 -> 64x64
    ]

    inputs = tf.keras.layers.Input(shape=input_shape)

    # Downsampling through the model
    skips = down_stack(inputs)
    x = skips[-1]
    skips = reversed(skips[:-1])

    # Upsampling and establishing the skip connections
    for up, skip in zip(up_stack, skips):
        x = up(x)
        concat = tf.keras.layers.Concatenate()
        x = concat([x, skip])

    # This is the last layer of the model
    last = tf.keras.layers.Conv2DTranspose(
        filters=output_channels, kernel_size=3, strides=2,
        padding='same')  #64x64 -> 128x128

    x = last(x)

    model = tf.keras.Model(inputs=inputs, outputs=x)

    model.compile(optimizer=optimizer,
                 #loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  loss=model_loss,
                  metrics=['acc', dice_coef]
                  )


    return model