model

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
from tensorflow.keras.applications import *
import numpy as np
import cv2
layers = tf.keras.layers

def squeeze_excite_block(inputs, ratio=8):
    init = inputs
    channel_axis = -1
    filters = init.shape[channel_axis]
    se_shape = (1, 1, filters)

    se = GlobalAveragePooling2D()(init)
    se = Reshape(se_shape)(se)
    se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)
    se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)

    x = Multiply()([init, se])
    return x

def conv_block(inputs, filters):
    x = inputs

    x = Conv2D(filters, (3, 3), padding="same")(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    x = Conv2D(filters, (3, 3), padding="same")(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)

    #x = squeeze_excite_block(x)

    return x

def encoder1(inputs):
    skip_connections = []

    model = VGG19(include_top=False, weights='imagenet', input_tensor=inputs)
    names = ["block1_conv2", "block2_conv2", "block3_conv4", "block4_conv4"]
    for name in names:
        skip_connections.append(model.get_layer(name).output)

    output = model.get_layer("block5_conv4").output
    return output, skip_connections

def decoder1(inputs, skip_connections):
    num_filters = [256, 128, 64, 32]
    skip_connections.reverse()
    x = inputs

    for i, f in enumerate(num_filters):
        x = UpSampling2D((2, 2), interpolation='bilinear')(x)
        x = Concatenate()([x, skip_connections[i]])
        x = conv_block(x, f)

    return x



def encoder2(inputs):
    num_filters = [32, 64, 128, 256]
    skip_connections = []
    x = inputs

    for i, f in enumerate(num_filters):
        x = conv_block(x, f)
        skip_connections.append(x)
        x = MaxPool2D((2, 2))(x)

    return x, skip_connections

def decoder2(inputs, skip_1, skip_2):
    num_filters = [256, 128, 64, 32]
    skip_2.reverse()
    x = inputs

    for i, f in enumerate(num_filters):
        x = UpSampling2D((2, 2), interpolation='bilinear')(x)
        x = Concatenate()([x, skip_1[i], skip_2[i]])
        x = conv_block(x, f)

    return x

def output_block(inputs):
    x = Conv2D(1, (1, 1), padding="same")(inputs)
    x = Activation('sigmoid')(x)
    return x
def output_block1(inputs):
    x = Conv2D(1, (1, 1), padding="same")(inputs)
    return x

def Upsample(tensor, size):
    """Bilinear upsampling"""
    def _upsample(x, size):
        return tf.image.resize(images=x, size=size)
    return Lambda(lambda x: _upsample(x, size), output_shape=size)(tensor)

def DDSPP(x, filter):
    shape = x.shape

    y1 = AveragePooling2D(pool_size=(shape[1], shape[2]))(x)
    y1 = Conv2D(filter, 1, padding="same")(y1)
    y1 = BatchNormalization()(y1)
    y1 = Activation("relu")(y1)
    y1 = UpSampling2D((shape[1], shape[2]), interpolation='bilinear')(y1)
    
    y1 = Concatenate()([x, y1])
    y2 = Conv2D(filter, 1, dilation_rate=2, padding="same", use_bias=False)(x)
    y2 = BatchNormalization()(y2)
    y2 = Activation("relu")(y2)
    
    y2 = Concatenate()([x, y1, y2])
    y3 = Conv2D(filter, 3, dilation_rate=4, padding="same", use_bias=False)(x)
    y3 = BatchNormalization()(y3)
    y3 = Activation("relu")(y3)
    
    y3 = Concatenate()([x, y1, y2, y3])
    y4 = Conv2D(filter, 3, dilation_rate=8, padding="same", use_bias=False)(x)
    y4 = BatchNormalization()(y4)
    y4 = Activation("relu")(y4)
    
    y4 = Concatenate()([x, y1, y2, y3, y4])
    y5 = Conv2D(filter, 3, dilation_rate=12, padding="same", use_bias=False)(x)
    y5 = BatchNormalization()(y5)
    y5 = Activation("relu")(y5)

    y = Concatenate()([x, y1, y2, y3, y4, y5])

    y = Conv2D(filter, 1, dilation_rate=1, padding="same", use_bias=False)(y)
    y = BatchNormalization()(y)
    y = Activation("relu")(y)

    return y

# vertical edge detection
sobel_x = np.array([[-1, 0, 1],
                   [-2, 0, 2],
                   [-1, 0, 1]])

def build_model1(shape):
    inputs = Input(shape)
    x, skip_1 = encoder1(inputs)
    x = DDSPP(x, 64)
    x = decoder1(x, skip_1)
    outputs1 = output_block(x)

  

    model = Model(inputs, outputs1)
    #model.compile(optimizer = Adam(lr = learning_rate), loss = [], metrics = ['accuracy'])

    return model

def build_model2(shape):
    inputs = Input(shape)
    x, skip_1 = encoder1(inputs)
    x = DDSPP(x, 64)
    x = decoder1(x, skip_1)
    outputs1 = output_block1(x)

    x = inputs * outputs1

    x, skip_2 = encoder2(outputs1)
    x = DDSPP(x, 64)
    x = decoder2(x, skip_1, skip_2)
    outputs2 = output_block(x)
    #filtered_image1 = tf.image.sobel_edges(outputs2)

    outputs = Concatenate()([outputs1, outputs2])

    model = Model(inputs, outputs)
    #model.compile(optimizer = Adam(lr = learning_rate), loss = [], metrics = ['accuracy'])

    return model

model1 = build_model1((512, 512, 3))
model2 = build_model2((512, 512, 3))

model2.summary()

# train_steps = (len(train_x)//batch_size)
# valid_steps = (len(valid_x)//batch_size)

# if len(train_x) % batch_size != 0:
#     train_steps += 1

# if len(valid_x) % batch_size != 0:
#     valid_steps += 1

# model.fit(genearte,
#         epochs=epochs,
#         validation_data=generate1,
#         steps_per_epoch=train_steps,
#         validation_steps=valid_steps,
#         shuffle=False)

# steps = len(test_img)//BS
# preds_test = model.predict_generator(test_img, steps, verbose=1)
#next(preds_test)[0].shape
#count =0
#for img in preds_test:
#    print ('image', img.shape)
#    filtered_image1 = tf.image.sobel_edges(img[0])
#    filtered_image2 = cv2.Sobel(src=filtered_image1, ddepth=cv2.CV_64F, dx=1, dy=1, ksize=5) 
#    plt.imshow(filtered_image1[0])
#    count+=1
#    if count>1:
#        break
#