def test(self): print('Start to evaluate on test datasets...') _logger, _log_path = prepare_logger(self.config, log_path=os.path.join( self.snapshot_dir, 'results')) pred_transforms = [] total_rotation = [] all_inlier_ratios = [] num_iter = int( len(self.loader['test'].dataset) // self.loader['test'].batch_size) c_loader_iter = self.loader['test'].__iter__() self.model.eval() with torch.no_grad(): for idx in tqdm(range(num_iter)): # loop through this epoch inputs = c_loader_iter.next() ################################## # load inputs to device. for k, v in inputs.items(): if type(v) == list: inputs[k] = [item.to(self.device) for item in v] elif type(v) == dict: pass else: inputs[k] = v.to(self.device) rot_trace = inputs['sample']['transform_gt'][:, 0, 0] + inputs['sample']['transform_gt'][:, 1, 1] + \ inputs['sample']['transform_gt'][:, 2, 2] rotdeg = torch.acos( torch.clamp(0.5 * (rot_trace - 1), min=-1.0, max=1.0)) * 180.0 / np.pi total_rotation.append(np.abs(to_array(rotdeg))) ################################### # forward pass feats, scores_overlap, scores_saliency = self.model( inputs) #[N1, C1], [N2, C2] scores_overlap = scores_overlap.detach().cpu() scores_saliency = scores_saliency.detach().cpu() len_src = inputs['stack_lengths'][0][0] src_feats, tgt_feats = feats[:len_src], feats[len_src:] src_pcd, tgt_pcd = inputs['src_pcd_raw'], inputs['tgt_pcd_raw'] src_overlap, tgt_overlap = scores_overlap[: len_src], scores_overlap[ len_src:] src_saliency, tgt_saliency = scores_saliency[: len_src], scores_saliency[ len_src:] ######################################## # run probabilistic sampling n_points = 450 src_scores = src_overlap * src_saliency tgt_scores = tgt_overlap * tgt_saliency if (src_pcd.size(0) > n_points): idx = np.arange(src_pcd.size(0)) probs = (src_scores / src_scores.sum()).numpy().flatten() idx = np.random.choice(idx, size=n_points, replace=False, p=probs) src_pcd, src_feats = src_pcd[idx], src_feats[idx] if (tgt_pcd.size(0) > n_points): idx = np.arange(tgt_pcd.size(0)) probs = (tgt_scores / tgt_scores.sum()).numpy().flatten() idx = np.random.choice(idx, size=n_points, replace=False, p=probs) tgt_pcd, tgt_feats = tgt_pcd[idx], tgt_feats[idx] ######################################## # run ransac distance_threshold = 0.02 ts_est = ransac_pose_estimation( src_pcd, tgt_pcd, src_feats, tgt_feats, mutual=False, distance_threshold=distance_threshold, ransac_n=3) pred_transforms.append(ts_est) total_rotation = np.concatenate(total_rotation, axis=0) _logger.info(('Rotation range in data: {}(avg), {}(max)'.format( np.mean(total_rotation), np.max(total_rotation)))) pred_transforms = torch.from_numpy( np.array(pred_transforms)).float()[:, None, :, :] c_loader_iter = self.loader['test'].__iter__() num_processed, num_total = 0, len(pred_transforms) metrics_for_iter = [ defaultdict(list) for _ in range(pred_transforms.shape[1]) ] with torch.no_grad(): for idx in tqdm(range(num_iter)): # loop through this epoch inputs = c_loader_iter.next() batch_size = 1 for i_iter in range(pred_transforms.shape[1]): cur_pred_transforms = pred_transforms[ num_processed:num_processed + batch_size, i_iter, :, :] metrics = compute_metrics(inputs['sample'], cur_pred_transforms) for k in metrics: metrics_for_iter[i_iter][k].append(metrics[k]) num_processed += batch_size for i_iter in range(len(metrics_for_iter)): metrics_for_iter[i_iter] = { k: np.concatenate(metrics_for_iter[i_iter][k], axis=0) for k in metrics_for_iter[i_iter] } summary_metrics = summarize_metrics(metrics_for_iter[i_iter]) print_metrics(_logger, summary_metrics, title='Evaluation result (iter {})'.format(i_iter))
else: raise NotImplementedError return model def main(): # Load data_loader test_dataset = get_test_datasets(_args) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=_args.val_batch_size, shuffle=False) model = get_model() model.eval() evaluate(test_loader, model) _logger.info('Finished') if __name__ == '__main__': # Arguments and logging parser = rpmnet_eval_arguments() _args = parser.parse_args() _logger, _log_path = prepare_logger(_args, log_path=_args.eval_save_path) os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu) if _args.gpu >= 0 and (_args.method == 'rpm' or _args.method == 'rpmnet'): os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu) _device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') else: _device = torch.device('cpu') main()
import torch.utils.data from tqdm import tqdm from arguments import rpmnet_train_arguments from common.colors import BLUE, ORANGE from common.misc import prepare_logger from common.torch import dict_all_to_device, CheckPointManager, TorchDebugger from common.math_torch import se3 from data_loader.datasets import get_train_datasets from eval import compute_metrics, summarize_metrics, print_metrics from models.rpmnet import get_model # Set up arguments and logging parser = rpmnet_train_arguments() _args = parser.parse_args() _logger, _log_path = prepare_logger(_args) if _args.gpu >= 0: os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu) _device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') else: _device = torch.device('cpu') def main(): train_set, val_set = get_train_datasets(_args) run(train_set, val_set) def compute_losses(data: Dict, pred_transforms: List, endpoints: Dict, loss_type: str = 'mae', reduction: str = 'mean') -> Dict: """Compute losses