import os
import gradio as gr
import cv2
import torch
import urllib.request
from PIL import Image

import matplotlib.pyplot as plt

def update(slider, img):

    if not os.path.exists('temp'):
      os.system('mkdir temp')

    filename = "temp/image.jpg"

    img.save(filename, "JPEG")

    model_type = "DPT_Hybrid" 
    midas = torch.hub.load("intel-isl/MiDaS", model_type)
    
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    midas.to(device)
    midas.eval()

    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

    if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
        transform = midas_transforms.dpt_transform
    else:
        transform = midas_transforms.small_transform

    img = cv2.imread(filename)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    input_batch = transform(img).to(device)

    with torch.no_grad():
        prediction = midas(input_batch)

        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=img.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()

    output = prediction.cpu().numpy()

    out_im = Image.fromarray(output)
    out_im.save("temp/image_depth.jpeg", "JPEG")

    return f'temp/image_depth.jpeg'



with gr.Blocks() as demo:
    gr.Markdown("Start typing below and then click **Run** to see the output.")
    inp = [gr.Slider(1,15, default = 2, label='StepCycles',step= 1)]

    with gr.Row():
        inp.append(gr.Image(type="pil", label="Input"))
        out = gr.Image(type="file", label="Output")
    btn = gr.Button("Run")
    btn.click(fn=update, inputs=inp, outputs=out)

demo.launch()