ARTeLab/DTM_Estimation

1from collections import OrderedDict2import torch3from models.model import GLPDepth4from PIL import Image5from torchvision import transforms6import matplotlib.pyplot as plt7import numpy as np8 9DEVICE='cpu'10 11def load_mde_model(path):12    model = GLPDepth(max_depth=700.0, is_train=False).to(DEVICE)13    model_weight = torch.load(path, map_location=torch.device('cpu'))14    model_weight = model_weight['model_state_dict']15    if 'module' in next(iter(model_weight.items()))[0]:16        model_weight = OrderedDict((k[7:], v) for k, v in model_weight.items())17    model.load_state_dict(model_weight)18    model.eval()19    return model20 21model = load_mde_model('best_model.ckpt')22preprocess = transforms.Compose([23    transforms.Resize((512, 512)),24    transforms.ToTensor()25]) 26 27input_img = Image.open('demo_imgs/fake.jpg')28torch_img = preprocess(input_img).to(DEVICE).unsqueeze(0)29with torch.no_grad():30    output_patch = model(torch_img)31output_patch = output_patch['pred_d'].squeeze().cpu().detach().numpy()32print(output_patch.shape)33 34plt.imshow(output_patch, cmap='jet', vmin=0, vmax=np.max(output_patch))35plt.colorbar()36plt.savefig('test.png')