import gradio as gr

from layers import BilinearUpSampling2D
import matplotlib.pyplot as plt
import numpy as np

from huggingface_hub import from_pretrained_keras

custom_objects = {'BilinearUpSampling2D': BilinearUpSampling2D, 'depth_loss_function': None}
print('Loading model...')
model = from_pretrained_keras("keras-io/monocular-depth-estimation", custom_objects=custom_objects, compile=False)
print('Successfully loaded model...')

import importlib
import utils
importlib.reload(utils)

def infer(image, min_th, max_th):
    print('_'*20)
    inputs = utils.load_images([image])
    outputs = utils.predict(model, inputs)

    plasma = plt.get_cmap('plasma')
    rescaled = outputs[0][:, :, 0]
    print("Min Max Bef", np.min(rescaled), np.max(rescaled))
    rescaled = rescaled - np.min(rescaled)
    rescaled = rescaled / np.max(rescaled)

    image_out = plasma(rescaled)[:, :, :3]

    print("Min Max Aft", np.min(rescaled), np.max(rescaled))

    print("Shape Scaled:",rescaled.shape)
    filtered = rescaled
    # filtered[filtered[:, :, 0] < min_th/100, 0] = 0
    # filtered[filtered[:, :, 0] < min_th/100, 1] = 0
    # filtered[filtered[:, :, 0] < min_th/100, 2] = 0
    # filt_arr = filtered[((filtered[:,0] > min_th/100) & (filtered[:,0] < max_th/100))]
    filt_arr = (filtered > min_th/100) * filtered * (filtered < max_th/100)


    print("Shape Image:",image.shape)
    print("Shape Image filt:",im_filt.shape)
    print("Shape Image Heat:",image_out.shape)
    im_filt = plasma(filt_arr)[:, :, :3]
    return image_out, im_filt, image

# def detr(im):
#     return im

gr_input = [
    gr.inputs.Image(label="image", type="numpy", shape=(640, 480))
    ,gr.inputs.Slider(minimum=0, maximum=100, step=5, default=0, label="Minimum Threshold")
    ,gr.inputs.Slider(minimum=0, maximum=100, step=5, default=100, label="Maximum Threshold")
]

gr_output = [
    gr.outputs.Image(type="pil",label="HeatMap Image"),
    gr.outputs.Image(type="pil",label="Filtered Image"),
    gr.outputs.Image(type="pil",label="Output Image")
]

iface = gr.Interface(
    fn=infer,
    title="Space Title Here",
    description = "Description Here",
    inputs  = gr_input,
    outputs = gr_output
    )
iface.launch()