import numpy as np
import cv2
import matplotlib.pyplot as plt
import tensorflow as tf
smooth=1.
#-----------------------------------------------------------------------------------------------------------------------------------------------------------#
'''Function for returning dice coefficient'''
def DICE_COEFF(y_true, y_pred):
    y_true = K.flatten(y_true)
    y_pred = K.flatten(y_pred)
    intersection = K.sum(y_true * y_pred)
    union = K.sum(y_true) + K.sum(y_pred)
    return (2.0 * intersection + smooth) / (union + smooth)
'''Dice Coefficient Loss'''
def dice_coef_loss(y_true, y_pred):
    return 1 - DICE_COEFF(y_true, y_pred)
#---------------------------------------------------------#
'''Function for combining binary cross entropy with dice coeffcients for loss function'''
def bce_dice_loss(y_true, y_pred):
    bce = tf.keras.losses.BinaryCrossentropy()
    return dice_coef_loss(y_true, y_pred) + bce(y_true, y_pred)

#----------------------------------------------------------------------------------------------------------------------------------------------------------#

'''Function for Jacards coefficient'''
def IOU_JACARD(y_true, y_pred):
    y_true=K.flatten(y_true)
    y_pred=K.flatten(y_pred)
    intersection = K.sum(y_true * y_pred)
    sum_jac = K.sum(y_true + y_pred)
    jac = (intersection + smooth) / (sum_jac - intersection + smooth)
    return jac
#----------------------------------------------------------------------------------------------------------------------------------------------------------#

model = tf.keras.models.load_model('MRI_Attention_UNet_ResNet.hdf5',custom_objects={'bce_dice_loss':bce_dice_loss,'IOU_JACARD': IOU_JACARD,'DICE_COEFF':DICE_COEFF})



def segment(image):

 

    images = np.array(image) 

    # Converting it to 'float32'
    images = images.astype('float32')

    # Normalize the Numpy array (if desired)
    images = images / 255.0

    # Convert the Numpy array to a TensorFlow tensor
    images = tf.convert_to_tensor(images)
    images=tf.image.resize(images,[256,256])
    images=np.array(images)
    images=tf.expand_dims(images,axis=0)
    images=model.predict(images) 
    images=np.array(images)
    images=images.reshape((256,256))
    print(images.shape)
    return images


import gradio as gr

# Write 1 line of Python to create a simple GUI
gr.Interface(fn=segment, inputs="image", outputs="image").launch();