def disp_(cfg, ori_result, data):
from dmb.data.datasets.evaluation.stereo.eval import remove_padding
from dmb.data.datasets.evaluation.stereo.eval_hooks import disp_evaluation
disps = ori_result['disps']
# remove the padding when data augmentation
ori_size = data['original_size']
disps = remove_padding(disps, ori_size)
# evaluation
whole_error_dict, data = disp_evaluation(cfg.copy(), disps, data)
result = {
'Disparity': disps,
'GroundTruth': data['leftDisp'],
'Error': whole_error_dict,
}
if hasattr(cfg.model, 'cmn'):
# confidence measurement network
ori_size = data['original_size']
confs = ori_result['confs']
confs = remove_padding(confs, ori_size)
result.update(Confidence=confs)
return result
def flow_(cfg, ori_result, data):
from dmb.data.datasets.evaluation.flow.eval_hooks import flow_evaluation
from dmb.data.datasets.evaluation.flow.eval_hooks import remove_padding
flows = ori_result['flows']
ori_size = data['original_size']
flows = remove_padding(flows, ori_size)
# evaluation
whole_error_dict, data = flow_evaluation(cfg.copy(), flows, data)
result = {
'Flow': flows,
'GroundTruth': data['flow'],
'Error': whole_error_dict,
}
return result
def _inference_single(model, batchDict, img_transform, device):
cfg = model.cfg.copy()
procData, oriData = _prepare_data(batchDict, img_transform, cfg, device)
with torch.no_grad():
result, _ = model(procData)
result = to_cpu(result)
assert isinstance(result, dict)
for k, v in result.items():
assert isinstance(v, (tuple, list))
for i in range(len(v)):
vv = v[i]
if torch.is_tensor(vv):
# inverse up/down sample
vv = F.interpolate(vv * 1.0 / cfg.scale_factor,
scale_factor=1.0 / cfg.scale_factor,
mode='bilinear',
align_corners=False)
ori_size = procData['original_size']
if cfg.pad_to_shape is not None:
vv = remove_padding(vv, ori_size)
v[i] = vv
result[k] = v
logData = {
'Result': result,
'OriginalData': oriData,
}
save_root = osp.join(
cfg.log_dir, batchDict['left_image_path'].split('/')[-1].split('.')[0])
mkdir_or_exist(save_root)
save_path = osp.join(save_root, 'result.pkl')
print('Result of {} will be saved to {}!'.format(
batchDict['left_image_path'].split('/')[-1], save_path))
mmcv.dump(logData, save_path)
return logData
def after_train_epoch(self, runner):
if not self.every_n_epochs(runner, self.interval):
return
runner.logger.info(
"Start Visualizing on {} dataset({} images).".format(self.dataset.name, len(self.dataset))
)
# get prog bar
if runner.rank == 0:
prog_bar = mmcv.ProgressBar(len(self.dataset))
else:
prog_bar = None
runner.model.eval()
results = [None for _ in range(len(self.dataset))]
for idx in range(runner.rank, len(self.dataset), runner.world_size):
data = self.dataset[idx]
data_gpu = scatter(
collate([data], samples_per_gpu=1),
[torch.cuda.current_device()]
)[0]
# compute output
with torch.no_grad():
ori_result, _ = runner.model(data_gpu)
# remove the padding when data augmentation
disps = ori_result['disps']
costs = ori_result['costs']
ori_size = data_gpu['original_size']
disps = remove_padding(disps, ori_size)
target = data_gpu['leftDisp'] if 'leftDisp' in data_gpu else None
if target is not None:
target = remove_padding(target, ori_size)
result = {
'Disparity': disps,
'GroundTruth': target,
}
if hasattr(self.cfg.model, 'cmn'):
# confidence measurement network
confs = ori_result['confs']
confs = remove_padding(confs, ori_size)
result.update(Confidence=confs)
if self.cfg.model.eval.is_cost_return:
if self.cfg.model.eval.is_cost_to_cpu:
costs = [cost.cpu() for cost in costs]
result['Cost'] = costs
# convert result to suitable visualization image
item = self.dataset.data_list[idx]
result['leftImage'] = imread(
osp.join(self.cfg.data.vis.data_root, item['left_image_path'])
).astype(np.float32)
result['rightImage'] = imread(
osp.join(self.cfg.data.vis.data_root, item['right_image_path'])
).astype(np.float32)
image_name = item['left_image_path'].split('/')[-1]
result = prepare_visualize(result, runner.epoch + 1, self.cfg.work_dir, image_name)
results[idx] = result
batch_size = runner.world_size
if runner.rank == 0:
for _ in range(batch_size):
prog_bar.update()
if runner.rank == 0:
print('\n')
dist.barrier()
for i in range(1, min(runner.world_size, len(self.dataset))):
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(i))
tmp_results = mmcv.load(tmp_file)
for idx in range(i, len(results), runner.world_size):
results[idx] = tmp_results[idx]
os.remove(tmp_file)
self.visualize(runner, results)
else:
tmp_file = osp.join(runner.work_dir, 'temp_{}.pkl'.format(runner.rank))
mmcv.dump(results, tmp_file)
dist.barrier()
dist.barrier()
torch.cuda.empty_cache()
def multi_gpu_test(model, dataset, cfg, show=False, tmpdir=None):
model.eval()
results = []
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for idx in range(rank, len(dataset), world_size):
data = dataset[idx]
# None type data cannot be scatter, here we pick out the not None type data
notNoneData = {}
for k, v in zip(data.keys(), data.values()):
if v is not None:
notNoneData[k] = v
notNoneData = scatter(collate([notNoneData], samples_per_gpu=1),
[torch.cuda.current_device()])[0]
data.update(notNoneData)
# TODO: evaluate after generate all predictions!
with torch.no_grad():
result, _ = model(data)
disps = result['disps']
ori_size = data['original_size']
disps = remove_padding(disps, ori_size)
target = None
error_dict = dict()
if 'leftDisp' in data:
target = data['leftDisp']
target = remove_padding(target, ori_size)
error_dict = do_evaluation(disps[0], target,
cfg.model.eval.lower_bound,
cfg.model.eval.upper_bound)
if cfg.model.eval.eval_occlusion and 'rightDisp' in data:
data['leftDisp'] = remove_padding(data['leftDisp'],
ori_size)
data['rightDisp'] = remove_padding(data['rightDisp'],
ori_size)
occ_error_dict = do_occlusion_evaluation(
disps[0], data['leftDisp'], data['rightDisp'],
cfg.model.eval.lower_bound, cfg.model.eval.upper_bound)
error_dict.update(occ_error_dict)
result = {
'Disparity': disps,
'GroundTruth': target,
'Error': error_dict,
}
filter_result = {}
filter_result.update(Error=result['Error'])
if show:
item = dataset.data_list[idx]
result['leftImage'] = imread(
osp.join(cfg.data.test.data_root,
item['left_image_path'])).astype(np.float32)
result['rightImage'] = imread(
osp.join(cfg.data.test.data_root,
item['right_image_path'])).astype(np.float32)
image_name = item['left_image_path'].split('/')[-1]
save_result(result, cfg.out_dir, image_name)
if hasattr(cfg, 'sparsification_plot'):
filter_result['Error'].update(sparsification_eval(result, cfg))
results.append(filter_result)
if rank == 0:
batch_size = world_size
for _ in range(batch_size):
prog_bar.update()
# collect results from all ranks
results = collect_results(results, len(dataset), tmpdir)
return results
def forward(self, batch):
ref_image, target_image = batch['leftImage'], batch['rightImage']
target = batch['leftDisp'] if 'leftDisp' in batch else None
ref_fm, target_fm = self.backbone(ref_image, target_image)
raw_cost = self.cost_computation(ref_fm, target_fm,
int(self.max_disp // 4))
costs = self.cost_aggregation(raw_cost)
confidence_costs = [cen(c) for c, cen in zip(costs, self.conf_est_net)]
confidences = [torch.sigmoid(c) for c in confidence_costs]
variances = [
self.confidence_coefficient * (1 - conf) +
self.confidence_init_value for conf in confidences
]
disps = [
self.disp_predictor(cost, temperature=self.sa_temperature)
for cost in costs
]
if self.training:
assert target is not None, "Ground truth disparity map should be given"
losses = {}
focal_losses = self.focal_loss_evaluator(costs, target, variances)
losses.update(focal_losses)
l1_losses = self.l1_loss_evaluator(disps, target)
l1_losses = {
k: v * self.l1_loss_weight
for k, v in zip(l1_losses.keys(), l1_losses.values())
}
losses.update(l1_losses)
nll_losses = self.conf_loss_evaluator(confidence_costs, target)
nll_losses = {
k: v * self.conf_loss_weight
for k, v in zip(nll_losses.keys(), nll_losses.values())
}
losses.update(nll_losses)
return losses
else:
confidences = remove_padding(confidences, batch['original_size'])
disps = remove_padding(disps, batch['original_size'])
error_dict = {}
if target is not None:
target = remove_padding(target, batch['original_size'])
error_dict = do_evaluation(disps[0], target,
self.cfg.model.eval.lower_bound,
self.cfg.model.eval.upper_bound)
if self.cfg.model.eval.eval_occlusion and 'leftDisp' in batch and 'rightDisp' in batch:
batch['leftDisp'] = remove_padding(batch['leftDisp'],
batch['original_size'])
batch['rightDisp'] = remove_padding(batch['rightDisp'],
batch['original_size'])
occ_error_dict = do_occlusion_evaluation(
disps[0], batch['leftDisp'], batch['rightDisp'],
self.cfg.model.eval.lower_bound,
self.cfg.model.eval.upper_bound)
error_dict.update(occ_error_dict)
result = {
'Disparity': disps,
'GroundTruth': target,
'Confidence': confidences,
'Error': error_dict,
}
if self.cfg.model.eval.is_cost_return:
if self.cfg.model.eval.is_cost_to_cpu:
costs = [cost.cpu() for cost in costs]
result['Cost'] = costs
return result
def multi_gpu_test(model, dataset, cfg, show=False, tmpdir=None):
model.eval()
results = []
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for idx in range(rank, len(dataset), world_size):
data = dataset[idx]
# None type data cannot be scatter, here we pick out the not None type data
notNoneData = {}
for k, v in zip(data.keys(), data.values()):
if v is not None:
notNoneData[k] = v
notNoneData = scatter(
collate([notNoneData], samples_per_gpu=1),
[torch.cuda.current_device()]
)[0]
data.update(notNoneData)
# TODO: evaluate after generate all predictions!
with torch.no_grad():
ori_result, _ = model(data)
# remove the padding when data augmentation
disps = ori_result['disps']
ori_size = data['original_size']
disps = remove_padding(disps, ori_size)
# process the ground truth disparity map
data['leftDisp'] = data['leftDisp'] if 'leftDisp' in data else None
if data['leftDisp'] is not None:
data['leftDisp'] = remove_padding(data['leftDisp'], ori_size)
data['rightDisp'] = data['rightDisp'] if 'rightDisp' in data else None
if data['rightDisp'] is not None:
data['rightDisp'] = remove_padding(data['rightDisp'], ori_size)
# evaluation
whole_error_dict = disp_evaluation(cfg.copy(), disps, data['leftDisp'], data['rightDisp'])
result = {
'Disparity': disps,
'GroundTruth': data['leftDisp'],
'Error': whole_error_dict,
}
if hasattr(cfg.model, 'cmn'):
# confidence measurement network
confs = ori_result['confs']
confs = remove_padding(confs, ori_size)
result.update(Confidence=confs)
filter_result = {}
filter_result.update(Error=result['Error'])
if show:
item = dataset.data_list[idx]
result['leftImage'] = imread(
osp.join(cfg.data.test.data_root, item['left_image_path'])
).astype(np.float32)
result['rightImage'] = imread(
osp.join(cfg.data.test.data_root, item['right_image_path'])
).astype(np.float32)
image_name = item['left_image_path'].split('/')[-1]
save_result(result, cfg.out_dir, image_name)
if hasattr(cfg, 'sparsification_plot'):
eval_disparity_id = cfg.get('eval_disparity_id', [0])
whole_error_dict = {}
for id in eval_disparity_id:
sparsification_plot_dict = sparsification_eval(result, cfg, id=id)
for key, val in sparsification_plot_dict.items():
whole_error_dict['metric_confidence_{}/'.format(id) + key] = val
filter_result['Error'].update(whole_error_dict)
results.append(filter_result)
if rank == 0:
batch_size = world_size
for _ in range(batch_size):
prog_bar.update()
# collect results from all ranks
results = collect_results(results, len(dataset), tmpdir)
return results