import gradio as gr
from utils import colorize
import cv2
import math
import torchvision.transforms as transforms
import torch
import spaces

resolution = 4
base_area = resolution * 480 * 640
flip_test = True

DEVICE = "cuda"

EXAMPLES = {
    "indoor" : [],
    "outdoor" : []
}

@spaces.GPU
def predict_depth(model, image):
    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    frame_height, frame_width = frame.shape[:2]
    
    curr_area = frame_width * frame_height
    scale = math.sqrt(base_area / curr_area)
    new_w, new_h = int(scale * frame_width), int(scale * frame_height)
    frame = cv2.resize(frame, (new_w, new_h))
    
    frame = transforms.ToTensor()(frame).unsqueeze(0)
    if flip_test:
        frame = torch.cat([frame, frame.flip(-1)])
    frame = frame.to(DEVICE)
    model.to(DEVICE)
    with torch.no_grad():
        depth = model(frame)
    if flip_test:
        depth = ((depth[0] + depth[1].flip(-1))/2).unsqueeze(0)
    return depth.detach().cpu().numpy()[0, 0]

def create_demo(model, scene):
    gr.Markdown("### Depth Prediction demo")
    with gr.Row():
        input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
        depth_image = gr.Image(label="Depth Map", type='numpy', elem_id='img-display-output')
    submit = gr.Button("Submit")

    def on_submit(image):
        depth = predict_depth(model, image)
        colored_depth = colorize(depth, cmap='Spectral_r')
        return colored_depth
    
    submit.click(on_submit, inputs=[input_image], outputs=depth_image)
    gr.Examples(examples=EXAMPLES[scene], inputs=[input_image])