def MiDaS(pretrained=True, **kwargs):
""" # This docstring shows up in hub.help()
MiDaS model for monocular depth estimation
pretrained (bool): load pretrained weights into model
"""
model = MidasNet()
if pretrained:
checkpoint = "https://github.com/intel-isl/MiDaS/releases/download/v2/model.pt"
state_dict = torch.hub.load_state_dict_from_url(checkpoint,
progress=True)
model.load_state_dict(state_dict)
return model
def main(image_path, model_path='model.pt', output_path=None):
print("Loading model...")
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = MidasNet(model_path, non_negative=True)
model.to(device)
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()
print("Creating depth maps...")
rgb_path = os.path.abspath(image_path)
if os.path.isdir(rgb_path):
for file in os.listdir(rgb_path):
test(model, os.path.join(rgb_path, file), output_path)
else:
test(model, rgb_path, output_path)
print("Done.")