import gradio as gr
import tensorflow as tf
import numpy as np
from keras.utils import normalize

def dice_coef(y_true, y_pred):
    smooth = 1e-5
    intersection = K.sum(y_true * y_pred, axis=[1, 2, 3])
    union = K.sum(y_true, axis=[1, 2, 3]) + K.sum(y_pred, axis=[1, 2, 3])
    return K.mean((2.0 * intersection + smooth) / (union + smooth), axis=0)

def predict_segmentation(image):
    SIZE_X = 128
    SIZE_Y = 128

    img = cv2.resize(image, (SIZE_Y, SIZE_X))
    img = np.expand_dims(img, axis=2)
    img = normalize(img, axis=1)

    # Prepare image for prediction
    img = np.expand_dims(img, axis=0)

    # Predict
    prediction = model.predict(img)
    predicted_img = np.argmax(prediction, axis=3)[0, :, :]

    return predicted_img

# Load the model
model = tf.keras.models.load_model("path_to_your_model_directory", custom_objects={'dice_coef': dice_coef})

# Gradio Interface
iface = gr.Interface(
    fn=predict_segmentation,
    inputs="image",
    outputs="image",
    live=False
)

iface.launch()