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/00015_colors.png', 'examples/00084_colors.png', 'examples/00033_colors.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,
    inputs=[gr.inputs.Image(label="image", type="numpy", shape=(640, 480))],
    outputs="image",
    examples=examples).launch(debug=True)