model_functions.py

import keras as k
import numpy as np
import keras.backend as K
import tensorflow as tf

import segmentor_training
import loss_functions
import f_score_metrics


class ModelWrapper:
    def __init__(self,
                 segmentation_model,
                 denoise_model,
                 encoder_classification_model,
                 decoder_classification_model,
                 detailed_model,
                 encoder_layers,
                 intermediate_layers,
                 decoder_layers,
                 imagenet_classification_model,
                 splimgnet_classification_model):
        self.segmentation_model = segmentation_model
        self.denoise_model = denoise_model
        self.encoder_classification_model = encoder_classification_model
        self.decoder_classification_model = decoder_classification_model
        self.detailed_model = detailed_model

        self.encoder_layers = encoder_layers
        self.intermediate_layers = intermediate_layers
        self.decoder_layers = decoder_layers
        self.imagenet_classification_model = imagenet_classification_model
        self.splimgnet_classification_model = splimgnet_classification_model

def classifier(layers, filters, classes, input, model_input, name):
    x = input
    for i in range(layers):
        x = k.layers.SeparableConv2D(filters, 3, strides=2, padding="same")(x)
        x = k.layers.BatchNormalization()(x)
        x = k.layers.Activation("relu")(x)
    x = k.layers.AveragePooling2D(pool_size=x.shape[1])(x)
    x = k.layers.Flatten()(x)
    imgnet_classifier_output = k.layers.Dense(classes, activation="sigmoid")(x)

    # Define the model
    classification_model = k.Model(model_input, imgnet_classifier_output, name=name)
    return classification_model


def unet(img_size,
         num_classes=6,
         xcoder_layer_filters=[8, 16, 32], intermediate_layer_filters=[64, 64],
         encoder_depth=2, decoder_depth=2,
         name=None,
         activation="relu",
         verbose=False, add_to_layer_name="",
         clf_depth=2):

    if name is None:
        name = "unete" + str(encoder_depth) + "d" + str(decoder_depth) + "_" + ("-".join(str(x) for x in xcoder_layer_filters) + "_" + "-".join(
            str(x) for x in intermediate_layer_filters))

    inputs = k.Input(shape=img_size + (3,), dtype="uint8")
    inputs = k.layers.Rescaling(scale=1./127.5 - 1.)(inputs)
    # keeping a list of skip connections
    skip_connections = [inputs]
    # First half of the network: downsampling inputs

    # Blocks 1, 2, 3 are identical apart from the feature depth.
    for i, filters in enumerate(xcoder_layer_filters):
        if i == 0:
            x = FullConvolutionBA(filters, activation, name="Encoder_0" + add_to_layer_name)(inputs)
        else:
            x = DepthwiseSeparableBA(filters, activation, name="Encoder_" + str(encoder_depth*i) + add_to_layer_name)(x)

        for j in range(encoder_depth - 2):
            x = DepthwiseSeparableBA(filters, activation, name="Encoder_" + str(encoder_depth * i + j+1) + add_to_layer_name)(x)
        x = DepthwiseSeparableBA(filters, activation, strides=2, name="Encoder_" + str(encoder_depth*i +encoder_depth-1) + add_to_layer_name)(x)

        # Project residual
        skip_connections += [x]

    encoder_output = skip_connections.pop()

    # intermediate layers
    for i, filters in enumerate(intermediate_layer_filters):
        x = DepthwiseSeparableBA(filters, activation, name="Intermediate_"+str(i) + add_to_layer_name)(x)

    intermediate_output = x

    # Second half of the network: upsampling inputs
    for i, filters in enumerate(xcoder_layer_filters[::-1]):
        x = k.layers.UpSampling2D(2)(x)

        x = k.layers.Concatenate()([x, skip_connections.pop()])

        for j in range(decoder_depth):
            x = DepthwiseSeparableBA(filters, activation, name="Decoder_" + str(decoder_depth*i + j) + add_to_layer_name)(x)

    # Add a per-pixel classification layer
    outputs = k.layers.Conv2D(num_classes, 3, activation="softmax", padding="same", name="Output" + add_to_layer_name)(x)

    # Define the model
    model = k.Model(inputs, outputs, name=name)
    if verbose:
        model.summary(line_length=200)

    # denoiser model
    denoise_outputs = k.layers.Conv2D(3, 3, padding="same")(x)
    denoise_model = k.Model(inputs, denoise_outputs, name=name + "_denoiser")
    if verbose:
        denoise_model.summary(line_length=200)

    enc_classification_model = classifier(clf_depth, 512, num_classes, encoder_output, inputs, name + "_encclassifier")
    if verbose:
        enc_classification_model.summary(line_length=200)

    imgnet_classification_model = classifier(clf_depth, 512, 200, encoder_output, inputs, name + "_imgnetclassifier")
    if verbose:
        imgnet_classification_model.summary(line_length=200)

    splimgnet_classification_model = classifier(clf_depth, 512, 200 + num_classes, encoder_output, inputs, name + "_splimgnetclassifier")
    if verbose:
        splimgnet_classification_model.summary(line_length=200)

    dec_classification_output = k.layers.MaxPooling2D(pool_size=outputs.shape[1])(outputs)
    dec_classification_output = k.layers.Reshape((dec_classification_output.shape[-1],))(dec_classification_output)
    dec_classification_model = k.Model(inputs, dec_classification_output, name=name + "_decclassifier")
    if verbose:
        dec_classification_model.summary(line_length=200)

    detailed_model = k.Model(inputs, [encoder_output, intermediate_output, outputs], name=name + "_detailed")
    if verbose:
        detailed_model.summary(line_length=200)

    encoder_layers = []
    for layer in model.layers:
        if "Encoder" in layer.name:
            encoder_layers.append(layer)

    intermediate_layers = []
    for layer in model.layers:
        if "Intermediate" in layer.name:
            intermediate_layers.append(layer)

    decoder_layers = []
    for layer in model.layers:
        if "Decoder" in layer.name or "Output" in layer.name:
            decoder_layers.append(layer)

    wrapper = ModelWrapper(model,
                           denoise_model,
                           enc_classification_model,
                           dec_classification_model,
                           detailed_model,
                           encoder_layers,
                           intermediate_layers,
                           decoder_layers,
                           imgnet_classification_model,
                           splimgnet_classification_model)
    return wrapper


def swift_net(img_size, num_classes=6,
              encoder_layer_filters=[8, 16, 32], intermediate_layer_filters=[64, 64], decoder_layer_filters=32,
              name=None,
              encoder_depth = 2, decoder_depth=1,
              activation="relu",
              sum_lateral=True,
              verbose=False, add_to_layer_name=""):
    if name is None:
        name = f"swiftfe{encoder_depth}d{decoder_depth}_" + ("-".join(str(x) for x in encoder_layer_filters) + "_" + "-".join(
            str(x) for x in intermediate_layer_filters)) + "_" + str(decoder_layer_filters) + str(sum_lateral)
    inputs = k.Input(shape=img_size + (3,), dtype="uint8")
    inputs = k.layers.Rescaling(scale=1./127.5 - 1.)(inputs)
    # keeping a list of skip connections
    skip_connections = [inputs]
    encoder_layers = []
    # First half of the network: downsampling inputs

    # Blocks 1, 2, 3 are identical apart from the feature depth.
    for i, filters in enumerate(encoder_layer_filters):
        if i == 0:
            x = FullConvolutionBA(filters, activation, name="Encoder_0" + add_to_layer_name)(inputs)
        else:
            x = DepthwiseSeparableBA(filters, activation, name="Encoder_" + str(encoder_depth*i) + add_to_layer_name)(x)
        encoder_layers.append(x)

        for j in range(encoder_depth - 2):
            x = DepthwiseSeparableBA(filters, activation, name="Encoder_" + str(encoder_depth * i + j+1) + add_to_layer_name)(x)
            encoder_layers.append(x)
        x = DepthwiseSeparableBA(filters, activation, strides=2, name="Encoder_" + str(encoder_depth*i + encoder_depth-1) + add_to_layer_name)(x)
        encoder_layers.append(x)

        # Project residual
        skip_connections += [DepthwiseSeparableBA(decoder_layer_filters, activation, name="skip_" + str(i + 1) + add_to_layer_name, size=1)(x)]

    # intermediate layers
    intermediate_layers = []
    for i, filters in enumerate(intermediate_layer_filters):
        x = DepthwiseSeparableBA(filters, activation, name="Intermediate_" + str(i) + add_to_layer_name)(x)
        intermediate_layers.append(x)

    skip_connections.pop()

    # Second half of the network: upsampling inputs
    decoder_layers = []
    for i in range(len(encoder_layer_filters)):
        x = DepthwiseSeparableBA(decoder_layer_filters, activation, name="Decoder_" + str(decoder_depth * i) + add_to_layer_name)(x)
        decoder_layers.append(x)

        x = k.layers.UpSampling2D(2)(x)
        decoder_layers.append(x)

        # add skip connection. For the last block, we need to concatenate
        merge = k.layers.Add() if sum_lateral and i != len(encoder_layer_filters)-1 else k.layers.Concatenate()
        x = merge([x, skip_connections.pop()])

        for j in range(decoder_depth-1):
            x = DepthwiseSeparableBA(decoder_layer_filters, activation, name="Decoder_" + str(decoder_depth * i + j + 1) + add_to_layer_name)(x)
            decoder_layers.append(x)

        decoder_layers.append(x)

    # Add a per-pixel classification layer
    outputs = k.layers.Conv2D(num_classes, 3, activation="softmax", padding="same", name="Output" + add_to_layer_name)(x)
    decoder_layers.append(outputs)

    # Define the model
    model = k.Model(inputs, outputs, name=name)
    if verbose:
        model.summary(line_length=200)

    # denoiser model
    denoise_outputs = k.layers.Conv2D(3, 3, padding="same")(x)

    # Define the model
    denoise_model = k.Model(inputs, denoise_outputs, name=name + "_denoiser")
    if verbose:
        denoise_model.summary(line_length=200)

    # encoder classification model
    shape = np.prod(encoder_layers[-1].shape[1:])
    enc_clf_x = k.layers.Reshape((shape,))(encoder_layers[-1])
    enc_classifier_output = k.layers.Dense(num_classes, activation="sigmoid")(enc_clf_x)

    # Define the model
    enc_classification_model = k.Model(inputs, enc_classifier_output, name=name + "_encclassifier")
    if verbose:
        enc_classification_model.summary(line_length=200)

    # classification model
    shape = np.prod(x.shape[1:])
    clf_x = k.layers.Reshape((shape,))(x)
    dec_classifier_output = k.layers.Dense(num_classes, activation="sigmoid")(clf_x)

    # Define the model
    dec_classification_model = k.Model(inputs, dec_classifier_output, name=name + "_decclassifier")
    if verbose:
        dec_classification_model.summary(line_length=200)

    detailed_model = k.Model(inputs, [encoder_layers[-1], intermediate_layers[-1], outputs], name=name + "_detailed")
    if verbose:
        detailed_model.summary(line_length=200)
    wrapper = ModelWrapper(model,
                           denoise_model,
                           enc_classification_model,
                           dec_classification_model,
                           detailed_model,
                           encoder_layers,
                           intermediate_layers,
                           decoder_layers)
    return wrapper


def generate_discriminator(input, discriminator_depth):
    discr_input = k.Input(shape=input.shape[1:])
    x = k.layers.SeparableConv2D(32, 4, padding="same", strides=2)(discr_input)
    x = k.layers.LeakyReLU(0.2)(x)
    x = k.layers.BatchNormalization()(x)

    for i in range(discriminator_depth-1):
        x = k.layers.SeparableConv2D((i+1) * 64, 4, padding="same", strides=2)(x)
        x = k.layers.LeakyReLU(0.2)(x)
        x = k.layers.BatchNormalization()(x)

    x = k.layers.SeparableConv2D(1, 4, padding="same", strides=1, activation="sigmoid")(x)
    x = k.layers.UpSampling2D(2**discriminator_depth,)(x)
    discriminator_pixel_output = k.layers.Reshape(x.shape[1:3])(x)
    pixel_discriminator = k.Model(inputs=discr_input, outputs=discriminator_pixel_output, name="pixel_discriminator")
    return pixel_discriminator


class DepthwiseSeparableBA(k.layers.Layer):

    def __init__(self, filters, activation, strides=1, name=None, size=3, **kwargs):
        super().__init__(name=name, **kwargs)
        self.filters = filters
        self.activation = activation
        self.strides = strides
        self.convolution = k.layers.SeparableConv2D(filters, size, padding="same", strides=strides)
        self.batch_norm = k.layers.BatchNormalization()
        self.activation = k.layers.Activation(activation)

    def call(self, inputs):
        x = self.convolution(inputs)
        x = self.batch_norm(x)
        x = self.activation(x)
        return x

    def get_config(self):
        return {"filters": self.filters,
                "activation": self.activation,
                "strides": self.strides}


class FullConvolutionBA(k.layers.Layer):

    def __init__(self, filters, activation, strides=1, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        self.filters = filters
        self.activation = activation
        self.strides = strides
        self.convolution = k.layers.Conv2D(filters, 3, padding="same", strides=strides)
        self.batch_norm = k.layers.BatchNormalization()
        self.activation = k.layers.Activation(activation)

    def call(self, inputs):
        x = self.convolution(inputs)
        x = self.batch_norm(x)
        x = self.activation(x)
        return x

    def get_config(self):
        return {"filters": self.filters,
                "activation": self.activation,
                "strides":self.strides}


def freeze(layers):
    for layer in layers:
        layer.trainable = False


def microwave(layers):
    for layer in layers:
        layer.trainable = True


class GAN(k.Model):
    def __init__(self, generator, detailed_model, discriminator_depth, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.generator = generator
        self.detailed_model = detailed_model
        self.discriminator = generate_discriminator(generator.output, discriminator_depth)
        self.discriminator_weight = 0.

    def train_step(self, data):
        x, (y_gen, y_dis) = data

        res = dict()
        # train gen
        self.discriminator.trainable = False
        for l in self.discriminator.layers:
            l.trainable = False
        self.generator.trainable = True
        for l in self.generator.layers:
            l.trainable = True
        gen_hist = self.generator.train_step((x, y_gen))
        for key, val in gen_hist.items():
            res["seg_" + key] = val

        # train discr
        self.discriminator.trainable = True
        for l in self.discriminator.layers:
            l.trainable = True
        self.generator.trainable= False
        for l in self.generator.layers:
            l.trainable = False
        y_pred = self.generator(x, training=False)
        dis_hist = self.discriminator.train_step((y_pred, y_dis))
        for key, val in dis_hist.items():
            res["dis_" + key] = val

        return res

    def test_step(self, data):
        x, (y_gen, y_dis) = data
        res = dict()
        gen_hist = self.generator.test_step((x, y_gen))
        for key, val in gen_hist.items():
            res["seg_" + key] = val

        y_pred = self.generator(x, training=False)
        dis_hist = self.discriminator.test_step((y_pred, y_dis))
        for key, val in dis_hist.items():
            res["dis_" + key] = val

        return res

    def compile(self,
                generator_optimizer=None,
                discriminator_optimizer=None,
                generator_loss=tf.keras.losses.sparse_categorical_crossentropy,
                discriminator_weight=0.001,
                positive_weight=1.0, negative_weight=1.0,
                generator_metrics=[loss_functions.FBetaMetric(beta=0.5, average="macro")],
                discriminator_metrics=["accuracy", k.metrics.Precision(), k.metrics.Recall()]):
        super().compile()
        if generator_optimizer is None:
            generator_optimizer = k.optimizers.RMSprop(0.001)
        if discriminator_optimizer is None:
            discriminator_optimizer = k.optimizers.Adam(0.001)
        self.discriminator_weight = discriminator_weight
        discriminator_loss = loss_functions.WeightedBinaryCrossentropy(weight_positive=positive_weight, weight_negative=negative_weight)
        self.discriminator.compile(optimizer=discriminator_optimizer,
                                   loss=discriminator_loss,
                                   metrics=discriminator_metrics)
        self.generator.compile(optimizer=generator_optimizer,
                               loss=loss_functions.DiscriminatorLoss(base_loss=tf.keras.losses.sparse_categorical_crossentropy,
                                                                     discriminator=self.discriminator,
                                                                     discriminator_weight=discriminator_weight),
                               metrics=generator_metrics)