import os
import cv2
from PIL import Image 
import numpy as np 
import segmentation_models as sm
from matplotlib import pyplot as plt
import random

from keras import backend as K
from keras.models import load_model

import gradio as gr

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    

weights = [0.2, 0.2, 0.2, 0.2, 0.2]
dice_loss = sm.losses.DiceLoss(class_weights=weights)
focal_loss = sm.losses.CategoricalFocalLoss()
total_loss = dice_loss + (1 * focal_loss)

satellite_model = load_model('model/satellite_segmentation_full.h5', custom_objects={'dice_loss_plus_1focal_loss': total_loss, 'jaccard_coef': jaccard_coef})

def process_input_image(image_source):
    image = np.expand_dims(image_source, 0)

    prediction = satellite_model.predict(image)
    predicted_colored = np.argmax(prediction, axis=3)

    predicted_colored = predicted_colored[0,:,:]
    predicted_colored = predicted_colored * 50
    return 'Predicted Masked Image', predicted_colored

my_app = gr.Interface(fn=process_input_image, 
                      inputs=gr.inputs.Image(label="Please select the source image", shape=(256, 256)), 
                      outputs="image",
                      title="Satellite Image Segmentation Application UI with Gradio")

my_app.launch()