# Data Handling
from huggingface_hub import from_pretrained_keras
import numpy as np
import cv2
import imutils
import tensorflow as tf
from tensorflow import keras
import gradio as gr
from tensorflow.keras.models import load_model

model = load_model('./augmented_unet_pretrained.h5', compile=False)

def segmentation(inp):
                    
    #inp = cv2.cvtColor(inp, cv2.COLOR_BGR2RGB)      # Input image
    inp = cv2.resize(inp, (256, 256))               # Resize
    inp = (inp.astype('float32')) / 255. 
    test_input = inp
    # (Must Add cropping for real time images)

    # Predictions
    prediction_on_test = np.expand_dims(test_input, 0)
    prediction_on_test = model.predict(prediction_on_test)
    prediction_on_test = prediction_on_test > 0.5
    predicted_img = prediction_on_test[0,:,:,0]   

     # EXTRACTING CONTOURS
   
    predicted = predicted_img.astype(np.uint8)
    cnts = cv2.findContours(image=predicted, mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_NONE)
    contours = imutils.grab_contours(cnts)
    contoured = test_input.copy()
    contoured = (contoured * 255).astype(np.uint8)
    cv2.drawContours(image=contoured, contours=contours, contourIdx=-1, color=(255, 0, 0), thickness=1, lineType=cv2.LINE_AA)


     # Circumference of detected Mask
    if contours :
        a = "Polynya Detected"
        for i in range(len(contours)):
            circum = cv2.arcLength(contours[i], True)
            circum = round(circum,2)
            b = str(circum) + '\t' + "px"
    else:
        a = "No Polynya Detected"
        b = "0.0 px"

    return(contoured, a, b)
    
image = gr.Image(label = 'Input Image')
out1 = gr.Image(label = 'Result')
out2 = gr.Textbox(label = 'Label')
out3 = gr.Textbox(label = 'Circumference in Pixel Unit')

interface = gr.Interface(fn = segmentation, inputs = image, outputs = [out1, out2, out3], 
                         title= 'Polynya Detection',
                         description= 'Let the system detect if there is a Polynya in your image or not.', 
                         share=True)
                 
interface.launch()