from skimage.util import montage as montage2d
from utils import load_model, preprocess_image, attempt_download_from_hub
import matplotlib.pyplot as plt

import gradio as gr

model_path = 'deprem-ml/deprem-keras-satellite-semantic-mapping'

def keras_inference(img_data, model_path):
    model_path = attempt_download_from_hub(model_path)
    seg_model = load_model(model_path)
    out_img = preprocess_image(img_data)
    pred_y = seg_model.predict(out_img)

    plt.imshow(montage2d(pred_y[:, :, :, 0]), cmap = 'bone_r')
    plt.savefig('output.png')
    return 'output.png'

inputs = [
    gr.Image(type='filepath', label='Image'),
    gr.Dropdown([model_path], value=model_path, label='Model Path')
]

outputs = gr.Image(label='Segmentation')

examples = [
    ['data/testv1.jpg', model_path],
    ['data/testv2.jpg', model_path],
    ['data/testv3.jpg', model_path],
]

title = 'Segmenting Buildings in Satellite Images with Keras'

demo_app = gr.Interface(
    keras_inference, 
    inputs, 
    outputs, 
    title=title, 
    examples=examples,
    cache_examples=True,
)

demo_app.launch(debug=True, enable_queue=True)