def inference_kitti(self, data):
#crop kitti images into 3 parts
with torch.no_grad():
self.a_l_real = data['A_l'].cuda()
_, b_l_classes = self.depth_model(self.a_l_real)
self.b_l_fake = bins_to_depth(b_l_classes)
self.a_m_real = data['A_m'].cuda()
_, b_m_classes = self.depth_model(self.a_m_real)
self.b_m_fake = bins_to_depth(b_m_classes)
self.a_r_real = data['A_r'].cuda()
_, b_r_classes = self.depth_model(self.a_r_real)
self.b_r_fake = bins_to_depth(b_r_classes)
out = kitti_merge_imgs(self.b_l_fake, self.b_m_fake, self.b_r_fake,
torch.squeeze(data['B_raw']).shape,
data['crop_lmr'])
return {'b_fake': out}
def criterion(self, pred_softmax, pred_logit, data, epoch):
pred_depth = bins_to_depth(pred_softmax)
loss_metric = self.weight_cross_entropy_loss(pred_logit,
data['B_bins'],
data['B'].cuda())
loss_normal = self.virtual_normal_loss(data['B'].cuda(), pred_depth)
loss = {}
loss['metric_loss'] = loss_metric
loss['virtual_normal_loss'] = cfg.MODEL.DIFF_LOSS_WEIGHT * loss_normal
loss['total_loss'] = loss['metric_loss'] + loss['virtual_normal_loss']
return loss
def forward(self, data):
# Input data is a_real, predicted data is b_fake, groundtruth is b_real
self.a_real = data['A'].cuda()
self.b_fake_logit, self.b_fake_softmax = self.depth_model(self.a_real)
pred_depth = bins_to_depth(self.b_fake_softmax)
refined_depth = self.refine_model(pred_depth.detach(), self.a_real)
return {
'b_fake_logit': self.b_fake_logit,
'b_fake_softmax': self.b_fake_softmax,
"refined_depth": refined_depth
}
def inference(self, data):
with torch.no_grad():
out = self.forward(data)
if cfg.MODEL.PREDICTION_METHOD == 'classification':
pred_depth = bins_to_depth(out['b_fake_softmax'])
elif cfg.MODEL.PREDICTION_METHOD == 'regression':
# for regression methods
pred_depth = torch.nn.functional.sigmoid(out['b_fake_logit'])
else:
raise ValueError("Unknown prediction methods")
out = pred_depth
return {'b_fake': out}
def inference(self, data):
with torch.no_grad():
out = self.forward(data)
pred_depth = bins_to_depth(out['b_fake_softmax'])
b_fake = pred_depth
return {'b_fake': b_fake}
model.eval()
# load checkpoint
if test_args.load_ckpt:
load_ckpt(test_args, model)
model.cuda()
model = torch.nn.DataParallel(model)
path = os.path.join(cfg.ROOT_DIR,
'./test_any_imgs_examples') # the dir of imgs
imgs_list = os.listdir(path)
for i in imgs_list:
print(i)
with torch.no_grad():
img = cv2.imread(os.path.join(path, i))
img_resize = cv2.resize(img,
(int(img.shape[1]), int(img.shape[0])),
interpolation=cv2.INTER_LINEAR)
img_torch = scale_torch(img_resize, 255)
img_torch = img_torch[None, :, :, :].cuda()
_, pred_depth_softmax = model.module.depth_model(img_torch)
pred_depth = bins_to_depth(pred_depth_softmax)
pred_depth = pred_depth.cpu().numpy().squeeze()
pred_depth_scale = (pred_depth / pred_depth.max() * 60000).astype(
np.uint16) # scale 60000 for visualization
cv2.imwrite(os.path.join(path,
i.split('.')[0] + '-raw.png'),
pred_depth_scale)