Ahsen Khaliq commited on
Commit
6c5755e
1 Parent(s): cb326fe

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import torch
3
+ import urllib.request
4
+ import gradio as gr
5
+ import matplotlib.pyplot as plt
6
+ import numpy as np
7
+ from PIL import Image
8
+
9
+ url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
10
+ urllib.request.urlretrieve(url, filename)
11
+
12
+ model_type = "DPT_Large" # MiDaS v3 - Large (highest accuracy, slowest inference speed)
13
+ #model_type = "DPT_Hybrid" # MiDaS v3 - Hybrid (medium accuracy, medium inference speed)
14
+ #model_type = "MiDaS_small" # MiDaS v2.1 - Small (lowest accuracy, highest inference speed)
15
+
16
+ midas = torch.hub.load("intel-isl/MiDaS", model_type)
17
+
18
+ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
19
+ midas.to(device)
20
+ midas.eval()
21
+
22
+ midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
23
+
24
+ if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
25
+ transform = midas_transforms.dpt_transform
26
+ else:
27
+ transform = midas_transforms.small_transform
28
+
29
+ def inference(img):
30
+ img = cv2.imread(img.name)
31
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
32
+
33
+ input_batch = transform(img).to(device)
34
+
35
+ with torch.no_grad():
36
+ prediction = midas(input_batch)
37
+
38
+ prediction = torch.nn.functional.interpolate(
39
+ prediction.unsqueeze(1),
40
+ size=img.shape[:2],
41
+ mode="bicubic",
42
+ align_corners=False,
43
+ ).squeeze()
44
+
45
+ output = prediction.cpu().numpy()
46
+ formatted = (output * 255 / np.max(output)).astype('uint8')
47
+ img = Image.fromarray(formatted)
48
+ return img
49
+
50
+ inputs = gr.inputs.Image(type='file', label="Original Image")
51
+ outputs = gr.outputs.Image(type="pil",label="Output Image")
52
+
53
+ title = "DPT-Large"
54
+ description = "Gradio demo for DPT-Large:Vision Transformers for Dense Prediction.To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
55
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2103.13413'>Vision Transformers for Dense Prediction</a> | <a href='https://github.com/intel-isl/MiDaS'>Github Repo</a></p>"
56
+
57
+ examples=[['dog.jpg']]
58
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False,examples=examples).launch(debug=True)