import random
from keras.models import load_model
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
def jaccard_coef(y_true, y_pred):
  y_true_flatten = K.flatten(y_true)
  y_pred_flatten = K.flatten(y_pred)
  intersection = K.sum(y_true_flatten * y_pred_flatten)
  final_coef_value = (intersection + 1.0) / (K.sum(y_true_flatten) + K.sum(y_pred_flatten) - intersection + 1.0)
  return final_coef_value
    
model = load_model('S_model.h5',custom_objects=({"jaccard_coef":jaccard_coef}))
image = gr.inputs.Image(shape = (256,256))
def  predict_image(image):
    test_image_input = np.expand_dims(image, 0)
    prediction = model.predict(test_image_input)
    predicted_image = np.argmax(prediction, axis=3)
    predicted_image = predicted_image[0,:,:]
    fig = plt.figure()
    plt.imshow(predicted_image)
   
    return fig

gr.Interface(fn = predict_image,inputs = image, outputs=['plot'],interpretation = 'default').launch()