from layers import BilinearUpSampling2D
from tensorflow.keras.models import load_model
from utils import load_images, predict
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr

custom_objects = {'BilinearUpSampling2D': BilinearUpSampling2D, 'depth_loss_function': None}
print('Loading model...')
model = load_model("model/model.h5", custom_objects=custom_objects, compile=False)
print('Successfully loaded model...')
examples = ['examples/377_image.png', 'examples/470_image.png', 'examples/499_image.png',
            'examples/626_image.png', 'examples/358_image.png']


def infer(image):
    inputs = load_images([image])
    outputs = predict(model, inputs)
    plasma = plt.get_cmap('plasma')
    rescaled = outputs[0][:, :, 0]
    rescaled = rescaled - np.min(rescaled)
    rescaled = rescaled / np.max(rescaled)
    image_out = plasma(rescaled)[:, :, :3]
    return image_out


iface = gr.Interface(
    fn=infer,
    title="Monocular Depth Estimation",
    description = "Unet architecture with Densenet201 backbone for estimating the depth of image 📏",
    inputs=[gr.inputs.Image(label="image", type="numpy", shape=(640, 480))],
    outputs="image",
    cache_examples=True,
    article = "Author: <a href=\"https://huggingface.co/vumichien\">Vu Minh Chien</a>.",
    examples=examples).launch(debug=True)