Spaces:
Build error
Build error
File size: 2,677 Bytes
dfcd969 a90faa2 dfcd969 a327756 dfcd969 b240372 dfcd969 a327756 dfcd969 b240372 dfcd969 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import cv2, torch
import urllib.request
import numpy as np
from PIL import Image
MODEL_DICT = {
"DPT_Large": "MiDaS v3 - Large (highest accuracy, slowest inference speed)",
"DPT_Hybrid": "MiDaS v3 - Hybrid (medium accuracy, medium inference speed)",
"MiDaS_small": "MiDaS v2.1 - Small (lowest accuracy, highest inference speed)"
}
def load_model(model_type = 'DPT_Large'):
assert model_type in MODEL_DICT.keys(), f'{model_type} is not a valid model_type: {MODEL_DICT.keys()}'
midas = torch.hub.load("intel-isl/MiDaS", model_type, force_reload=True)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(f'---DPT will use device: {device}')
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
return {
'midas': midas, 'device': device, 'transform': transform
}
def inference(img_array_rgb, model_obj, as_pil = False):
'''run DPT model and returns a PIL image'''
# img = cv2.imread(img.name)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
midas = model_obj['midas']
transform = model_obj['transform']
device = model_obj['device']
input_batch = transform(img_array_rgb).to(device)
with torch.no_grad():
prediction = midas(input_batch)
prediction = torch.nn.functional.interpolate(
prediction.unsqueeze(1),
size=img_array_rgb.shape[:2],
mode="bicubic",
align_corners=False,
).squeeze()
output = prediction.cpu().numpy()
formatted = (output * 255 / np.max(output)).astype('uint8')
img = Image.fromarray(formatted)
return img if as_pil else formatted
# inputs = gr.inputs.Image(type='file', label="Original Image")
# outputs = gr.outputs.Image(type="pil",label="Output Image")
# title = "DPT-Large"
# 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."
# article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2103.13413' target='_blank'>Vision Transformers for Dense Prediction</a> | <a href='https://github.com/intel-isl/MiDaS' target='_blank'>Github Repo</a></p>"
#
# examples=[['dog.jpg']]
# gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False,examples=examples, enable_queue=True).launch(debug=True)
|