Sentdex commited on
Commit
8286d35
1 Parent(s): 1b0fa2d

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +68 -0
app.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import torch
3
+ import gradio as gr
4
+ import numpy as np
5
+ from PIL import Image
6
+ import time
7
+
8
+ midas = torch.hub.load("intel-isl/MiDaS", "MiDaS")
9
+
10
+ use_large_model = True
11
+
12
+ if use_large_model:
13
+ midas = torch.hub.load("intel-isl/MiDaS", "MiDaS")
14
+ else:
15
+ midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
16
+
17
+ device = "cpu"
18
+ midas.to(device)
19
+
20
+ midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
21
+
22
+ if use_large_model:
23
+ transform = midas_transforms.default_transform
24
+ else:
25
+ transform = midas_transforms.small_transform
26
+
27
+
28
+ def depth(img):
29
+ original_image = img
30
+ cv_image = np.array(img)
31
+ img = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
32
+
33
+ input_batch = transform(img).to(device)
34
+ with torch.no_grad():
35
+ prediction = midas(input_batch)
36
+
37
+ prediction = torch.nn.functional.interpolate(
38
+ prediction.unsqueeze(1),
39
+ size=img.shape[:2],
40
+ mode="bicubic",
41
+ align_corners=False,
42
+ ).squeeze()
43
+
44
+ output = prediction.cpu().numpy()
45
+ formatted = (output * 255 / np.max(output)).astype('uint8')
46
+ img = Image.fromarray(formatted)
47
+
48
+ # create new image with with original_image and img side by side
49
+ new_im = Image.new('RGB', (original_image.width * 2, original_image.height))
50
+ new_im.paste(original_image, (0,0))
51
+ new_im.paste(img, (original_image.width,0))
52
+
53
+ # save the image to a file: (removed for hosting on HF)
54
+ #new_im.save(f'RGBDs/{int(time.time())}_RGBD.png')
55
+
56
+
57
+ return new_im
58
+
59
+
60
+ inputs = gr.inputs.Image(type='pil', label="Original Image")
61
+ outputs = gr.outputs.Image(type="pil",label="Output Image")
62
+
63
+ title = "RGB to RGBD for Looking Glass (using MiDaS)"
64
+ description = "Takes an RGB image and creates the depth + combines to the RGB image. Depth is predicted by MiDaS. This is a demo of the Looking Glass. For more information, visit https://lookingglassfactory.com"
65
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1907.01341v3'>Towards Robust Monocular Depth Estimation: Mixing Datasets for Zero-shot Cross-dataset Transfer</a> | <a href='https://github.com/intel-isl/MiDaS'>Github Repo</a></p>"
66
+
67
+
68
+ gr.Interface(depth, inputs, outputs, title=title, description=description, article=article).launch()