"""
Authors: Elena Vasileva, Zoran Ivanovski
E-mail: elenavasileva95@gmail.com, mars@feit.ukim.edu.mk
Course: Mashinski vid, FEEIT, Spring 2021
Date: 09.03.2021

Description: function library
             model operations: construction, loading, saving
Python version: 3.6
"""

# python imports
from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate
from keras.models import Model, model_from_json


def load_model(model_path, weights_path):
    """
    loads a pre-trained model configuration and calculated weights
    :param model_path: path of the serialized model configuration file (.json) [string]
    :param weights_path: path of the serialized model weights file (.h5) [string]
    :return: model - keras model object
    """

    # --- load model configuration ---
    json_file = open(model_path, 'r')
    model_json = json_file.read()
    json_file.close()
    model = model_from_json(model_json)     # load model architecture

    model.load_weights(weights_path)     # load weights

    return model


def construct_model_unet_orig(input_shape):
    """
    construct semantic segmentation model architecture (encoder-decoder)
    :param input_shape: list of input dimensions (height, width, depth) [tuple]
    :return:  model - Keras model object
    """

    input = Input(shape=input_shape)

    # --- encoder ---

    conv1 = Conv2D(filters=64, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(input)
    conv11 = Conv2D(filters=64, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    pool1 = MaxPool2D(pool_size=(2, 2))(conv11)

    conv2 = Conv2D(filters=128, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
    conv22 = Conv2D(filters=128, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
    pool2 = MaxPool2D(pool_size=(2, 2))(conv22)

    conv3 = Conv2D(filters=256, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
    conv33 = Conv2D(filters=256, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
    pool3 = MaxPool2D(pool_size=(2, 2))(conv33)

    conv4 = Conv2D(filters=512, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
    conv44 = Conv2D(filters=512, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
    pool4 = MaxPool2D(pool_size=(2, 2))(conv44)


    # --- decoder ---
    conv5 = Conv2D(filters=1024, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4)
    conv55 = Conv2D(filters=512, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)

    up1 = UpSampling2D(size=(2, 2))(conv55)
    merge1 = Concatenate(axis=3)([conv44, up1])
    deconv1 = Conv2DTranspose(filters=512, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(merge1)
    deconv11 = Conv2DTranspose(filters=256, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(deconv1)

    up2 = UpSampling2D(size=(2, 2))(deconv11)
    merge2 = Concatenate(axis=3)([conv33, up2])
    deconv2 = Conv2DTranspose(filters=256, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(merge2)
    deconv22 = Conv2DTranspose(filters=128, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(deconv2)

    up3 = UpSampling2D(size=(2, 2))(deconv22)
    merge3 = Concatenate(axis=3)([conv22, up3])
    deconv3 = Conv2DTranspose(filters=128, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(merge3)
    deconv33 = Conv2DTranspose(filters=64, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(deconv3)

    up4 = UpSampling2D(size=(2, 2))(deconv33)
    merge4 = Concatenate(axis=3)([conv11, up4])
    deconv4 = Conv2DTranspose(filters=64, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(merge4)
    deconv44 = Conv2DTranspose(filters=64, kernel_size=3, activation='relu', padding='same', kernel_initializer='he_normal')(deconv4)

    output = Conv2DTranspose(filters=input_shape[2], kernel_size=1, padding='same', activation='sigmoid')(deconv44)

    model = Model(input=input, output=output)

    return model

[((920, 947), 'keras.models.model_from_json', 'model_from_json', (['model_json'], {}), '(model_json)\n', (935, 947), False, 'from keras.models import Model, model_from_json\n'), ((1322, 1346), 'keras.layers.Input', 'Input', ([], {'shape': 'input_shape'}), '(shape=input_shape)\n', (1327, 1346), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((4381, 4414), 'keras.models.Model', 'Model', ([], {'input': 'input', 'output': 'output'}), '(input=input, output=output)\n', (4386, 4414), False, 'from keras.models import Model, model_from_json\n'), ((1383, 1487), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(64)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=64, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (1389, 1487), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1504, 1608), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(64)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=64, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (1510, 1608), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1624, 1651), 'keras.layers.MaxPool2D', 'MaxPool2D', ([], {'pool_size': '(2, 2)'}), '(pool_size=(2, 2))\n', (1633, 1651), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1673, 1778), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(128)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=128, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (1679, 1778), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1795, 1900), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(128)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=128, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (1801, 1900), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1916, 1943), 'keras.layers.MaxPool2D', 'MaxPool2D', ([], {'pool_size': '(2, 2)'}), '(pool_size=(2, 2))\n', (1925, 1943), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((1965, 2070), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(256)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=256, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (1971, 2070), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2087, 2192), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(256)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=256, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (2093, 2192), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2208, 2235), 'keras.layers.MaxPool2D', 'MaxPool2D', ([], {'pool_size': '(2, 2)'}), '(pool_size=(2, 2))\n', (2217, 2235), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2257, 2362), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(512)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=512, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (2263, 2362), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2379, 2484), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(512)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=512, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (2385, 2484), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2500, 2527), 'keras.layers.MaxPool2D', 'MaxPool2D', ([], {'pool_size': '(2, 2)'}), '(pool_size=(2, 2))\n', (2509, 2527), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2572, 2678), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(1024)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=1024, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (2578, 2678), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2695, 2800), 'keras.layers.Conv2D', 'Conv2D', ([], {'filters': '(512)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=512, kernel_size=3, activation='relu', padding='same',\n    kernel_initializer='he_normal')\n", (2701, 2800), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2815, 2840), 'keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': '(2, 2)'}), '(size=(2, 2))\n', (2827, 2840), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2862, 2881), 'keras.layers.Concatenate', 'Concatenate', ([], {'axis': '(3)'}), '(axis=3)\n', (2873, 2881), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((2911, 3026), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(512)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=512, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (2926, 3026), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3045, 3160), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(256)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=256, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (3060, 3160), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3176, 3201), 'keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': '(2, 2)'}), '(size=(2, 2))\n', (3188, 3201), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3225, 3244), 'keras.layers.Concatenate', 'Concatenate', ([], {'axis': '(3)'}), '(axis=3)\n', (3236, 3244), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3274, 3389), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(256)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=256, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (3289, 3389), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3408, 3523), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(128)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=128, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (3423, 3523), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3539, 3564), 'keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': '(2, 2)'}), '(size=(2, 2))\n', (3551, 3564), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3588, 3607), 'keras.layers.Concatenate', 'Concatenate', ([], {'axis': '(3)'}), '(axis=3)\n', (3599, 3607), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3637, 3752), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(128)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=128, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (3652, 3752), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3771, 3885), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(64)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=64, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (3786, 3885), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3901, 3926), 'keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': '(2, 2)'}), '(size=(2, 2))\n', (3913, 3926), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3950, 3969), 'keras.layers.Concatenate', 'Concatenate', ([], {'axis': '(3)'}), '(axis=3)\n', (3961, 3969), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((3999, 4113), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(64)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=64, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (4014, 4113), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((4132, 4246), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': '(64)', 'kernel_size': '(3)', 'activation': '"""relu"""', 'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), "(filters=64, kernel_size=3, activation='relu', padding=\n    'same', kernel_initializer='he_normal')\n", (4147, 4246), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n'), ((4265, 4361), 'keras.layers.Conv2DTranspose', 'Conv2DTranspose', ([], {'filters': 'input_shape[2]', 'kernel_size': '(1)', 'padding': '"""same"""', 'activation': '"""sigmoid"""'}), "(filters=input_shape[2], kernel_size=1, padding='same',\n    activation='sigmoid')\n", (4280, 4361), False, 'from keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, UpSampling2D, Input, Concatenate\n')]