import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add
from tensorflow.keras.initializers import Zeros


def build_fcn32s(nb_classes, target_size=(None, None)):
    inputs = Input(shape=(*target_size, 3))
    vgg = VGG16(weights='imagenet', include_top=False, input_tensor=inputs, input_shape=(*target_size, 3))
    x = Conv2D(4096, (7, 7), activation='relu', padding='same')(vgg.output)
    x = Dropout(0.5)(x)
    x = Conv2D(4096, (1, 1), activation='relu', padding='same')(x)
    x = Dropout(0.5)(x)
    x = Conv2D(nb_classes, (1, 1), padding='same', kernel_initializer='he_normal')(x)
    x = Conv2DTranspose(nb_classes, (64, 64), strides=(32, 32), use_bias=False, padding='same', activation='softmax', name='fcn32s-transpose')(x)

    model = Model(inputs=inputs, outputs=x)
    return model


def build_fcn16s(nb_classes, target_size=(None, None)):
    inputs = Input(shape=(*target_size, 3))
    vgg = VGG16(weights='imagenet', include_top=False, input_tensor=inputs, input_shape=(*target_size, 3))
    x = Conv2D(4096, (7, 7), activation='relu', padding='same')(vgg.output)
    x = Dropout(0.5)(x)
    x = Conv2D(4096, (1, 1), activation='relu', padding='same')(x)
    x = Dropout(0.5)(x)
    x = Conv2D(nb_classes, (1, 1), padding='same', kernel_initializer='he_normal')(x)

    x = Conv2DTranspose(nb_classes, (4, 4), strides=(2, 2), use_bias=False, padding='same', activation='relu', name='fcn16s-transpose-first')(x)

    skip_con = Conv2D(nb_classes, (1, 1), strides=(1, 1), padding='same', bias_initializer=Zeros(), kernel_initializer=Zeros(), name='fcn16s-skip-con')(vgg.get_layer(name="block4_pool").output)
    x = Add()([x, skip_con])
    x = Conv2DTranspose(nb_classes, (32, 32), strides=(16, 16), use_bias=False, padding='same', activation='softmax', name='fcn16s-transpose-second')(x)

    model = Model(inputs=inputs, outputs=x)
    return model


tensorflow.keras.models.Model tensorflow.keras.layers.Conv2DTranspose tensorflow.keras.applications.vgg16.VGG16 tensorflow.keras.layers.Conv2D tensorflow.keras.layers.Add tensorflow.keras.initializers.Zeros tensorflow.keras.layers.Dropout tensorflow.keras.layers.Input
[(10, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': '(*target_size, 3)'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (11, 'tensorflow.keras.applications.vgg16.VGG16', 'VGG16', ([], {'weights': '"""imagenet"""', 'include_top': '(False)', 'input_tensor': 'inputs', 'input_shape': '(*target_size, 3)'}), False, 'from tensorflow.keras.applications.vgg16 import VGG16\n'), (19, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'inputs', 'outputs': 'x'}), False, 'from tensorflow.keras.models import Model\n'), (24, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': '(*target_size, 3)'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (25, 'tensorflow.keras.applications.vgg16.VGG16', 'VGG16', ([], {'weights': '"""imagenet"""', 'include_top': '(False)', 'input_tensor': 'inputs', 'input_shape': '(*target_size, 3)'}), False, 'from tensorflow.keras.applications.vgg16 import VGG16\n'), (38, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'inputs', 'outputs': 'x'}), False, 'from tensorflow.keras.models import Model\n'), (12, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(4096)', '(7, 7)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (13, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (14, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(4096)', '(1, 1)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (15, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (16, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['nb_classes', '(1, 1)'], {'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (17, 'tensorflow.keras.layers.Conv2DTranspose', 'Conv2DTranspose', (['nb_classes', '(64, 64)'], {'strides': '(32, 32)', 'use_bias': '(False)', 'padding': '"""same"""', 'activation': '"""softmax"""', 'name': '"""fcn32s-transpose"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (26, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(4096)', '(7, 7)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (27, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (28, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(4096)', '(1, 1)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (29, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (30, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['nb_classes', '(1, 1)'], {'padding': '"""same"""', 'kernel_initializer': '"""he_normal"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (32, 'tensorflow.keras.layers.Conv2DTranspose', 'Conv2DTranspose', (['nb_classes', '(4, 4)'], {'strides': '(2, 2)', 'use_bias': '(False)', 'padding': '"""same"""', 'activation': '"""relu"""', 'name': '"""fcn16s-transpose-first"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (35, 'tensorflow.keras.layers.Add', 'Add', ([], {}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (36, 'tensorflow.keras.layers.Conv2DTranspose', 'Conv2DTranspose', (['nb_classes', '(32, 32)'], {'strides': '(16, 16)', 'use_bias': '(False)', 'padding': '"""same"""', 'activation': '"""softmax"""', 'name': '"""fcn16s-transpose-second"""'}), False, 'from tensorflow.keras.layers import Input, Conv2D, Dropout, Conv2DTranspose, Add\n'), (34, 'tensorflow.keras.initializers.Zeros', 'Zeros', ([], {}), False, 'from tensorflow.keras.initializers import Zeros\n'), (34, 'tensorflow.keras.initializers.Zeros', 'Zeros', ([], {}), False, 'from tensorflow.keras.initializers import Zeros\n')]