| 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.convert('RGB').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() |