def register_FCGF(self, xyz0, xyz1, inlier_thr=0.00): ''' Main algorithm of DeepGlobalRegistration ''' self.reg_timer.tic() with torch.no_grad(): # Step 0: voxelize and generate sparse input xyz0, coords0, feats0 = self.preprocess(xyz0) xyz1, coords1, feats1 = self.preprocess(xyz1) # Step 1: Feature extraction self.feat_timer.tic() fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0) fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1) self.feat_timer.toc() # Step 2: Coarse correspondences corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1) # > Case 1: Safeguard RANSAC + (Optional) ICP pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) T = self.safeguard_registration(pcd0, pcd1, corres_idx0, corres_idx1, feats0, feats1, 2 * self.voxel_size, num_iterations=80000) safeguard_time = self.reg_timer.toc() print(f'=> Safeguard takes {safeguard_time:.2} s') return T
def do_single_pair_evaluation(feature_path, set_name, traj, voxel_size, tau_1=0.1, tau_2=0.05, num_rand_keypoints=-1): trans_gth = np.linalg.inv(traj.pose) i = traj.metadata[0] j = traj.metadata[1] name_i = "%s_%03d" % (set_name, i) name_j = "%s_%03d" % (set_name, j) # coord and feat form a sparse tensor. data_i = np.load(os.path.join(feature_path, name_i + ".npz")) coord_i, points_i, feat_i = data_i['xyz'], data_i['points'], data_i[ 'feature'] data_j = np.load(os.path.join(feature_path, name_j + ".npz")) coord_j, points_j, feat_j = data_j['xyz'], data_j['points'], data_j[ 'feature'] # use the keypoints in 3DMatch if num_rand_keypoints > 0: # Randomly subsample N points Ni, Nj = len(points_i), len(points_j) inds_i = np.random.choice(Ni, min(Ni, num_rand_keypoints), replace=False) inds_j = np.random.choice(Nj, min(Nj, num_rand_keypoints), replace=False) sample_i, sample_j = points_i[inds_i], points_j[inds_j] key_points_i = ME.utils.fnv_hash_vec(np.floor(sample_i / voxel_size)) key_points_j = ME.utils.fnv_hash_vec(np.floor(sample_j / voxel_size)) key_coords_i = ME.utils.fnv_hash_vec(np.floor(coord_i / voxel_size)) key_coords_j = ME.utils.fnv_hash_vec(np.floor(coord_j / voxel_size)) inds_i = np.where(np.isin(key_coords_i, key_points_i))[0] inds_j = np.where(np.isin(key_coords_j, key_points_j))[0] coord_i, feat_i = coord_i[inds_i], feat_i[inds_i] coord_j, feat_j = coord_j[inds_j], feat_j[inds_j] coord_i = make_open3d_point_cloud(coord_i) coord_j = make_open3d_point_cloud(coord_j) hit_ratio = evaluate_feature_3dmatch(coord_i, coord_j, feat_i, feat_j, trans_gth, tau_1) # logging.info(f"Hit ratio of {name_i}, {name_j}: {hit_ratio}, {hit_ratio >= tau_2}") if hit_ratio >= tau_2: return True else: return False
def register(self, xyz0, xyz1, inlier_thr=0.00): ''' Main algorithm of DeepGlobalRegistration ''' self.reg_timer.tic() T = o3d.registration.registration_icp( make_open3d_point_cloud(xyz0), make_open3d_point_cloud(xyz1), self.voxel_size * 2, np.identity(4), o3d.registration.TransformationEstimationPointToPoint()).transformation safeguard_time = self.reg_timer.toc() print(f'=> ICP takes {safeguard_time:.2} s') return T
def read_data(feature_path, name): data = np.load(os.path.join(feature_path, name + ".npz")) xyz = make_open3d_point_cloud(data['xyz']) feat = make_open3d_feature_from_numpy(data['feature']) print("In benchmark_util.py - data['feature'] : ", data['feature']) print("In benchmark_util.py - data['feature'] shape : ", data['feature'].shape) print("In benchmark_util.py - feat : ", feat) print('-'*20) return data['points'], xyz, feat
def register(self, xyz0, xyz1, T_gt=None): ''' Main algorithm of DeepGlobalRegistration ''' self.reg_timer.tic() with torch.no_grad(): # Step 0: voxelize and generate sparse input xyz0, coords0, feats0 = self.preprocess(xyz0) xyz1, coords1, feats1 = self.preprocess(xyz1) # Step 1: Feature extraction self.feat_timer.tic() ## generate fpfh features pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) pcd0.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*2, max_nn=30)) fpfh0 = o3d.registration.compute_fpfh_feature(pcd0, o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*5, max_nn=100)) fpfh_np0 = np.array(fpfh0.data).T pcd1.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*2, max_nn=30)) fpfh1 = o3d.registration.compute_fpfh_feature(pcd1, o3d.geometry.KDTreeSearchParamHybrid(radius=0.05*5, max_nn=100)) fpfh_np1 = np.array(fpfh1.data).T fcgf_feats0 = torch.from_numpy(fpfh_np0).cuda() fcgf_feats1 = torch.from_numpy(fpfh_np1).cuda() import torch.nn.functional as F fcgf_feats0 = F.normalize(fcgf_feats0, dim=-1) fcgf_feats1 = F.normalize(fcgf_feats1, dim=-1) # fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0) # fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1) self.feat_timer.toc() # np.savez_compressed( # os.path.join('KITTI_trainval/', f'kitti_pair_{i}'), # xyz0=xyz0.detach().cpu().numpy(), # features0=fcgf_feats0.detach().cpu().numpy(), # xyz1=xyz1.detach().cpu().numpy(), # features1=fcgf_feats1.detach().cpu().numpy(), # gt_trans=T_gt # ) # Step 2: Coarse correspondences corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1) # _, corres_idx0 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1) # _, corres_idx1 = self.fcgf_feature_matching(fcgf_feats1, fcgf_feats0) # mutual_nearest = (corres_idx1[corres_idx0] == torch.arange(corres_idx0.shape[0]).cuda() ) # # corr = [] # # for i in range(len(corres_idx0)): # # if corres_idx1[corres_idx0[i]] == i: # # corr.append([i, corres_idx0[i]]) # mutual_corres_idx0 = torch.where(mutual_nearest == 1)[0] # mutual_corres_idx1 = corres_idx0[mutual_nearest] # corres_idx0, corres_idx1 = mutual_corres_idx0, mutual_corres_idx1 # Step 3: Inlier feature generation # coords[corres_idx0]: 1D temporal + 3D spatial coord # coords[corres_idx1, 1:]: 3D spatial coord # => 1D temporal + 6D spatial coord inlier_coords = torch.cat((coords0[corres_idx0], coords1[corres_idx1, 1:]), dim=1).int() inlier_feats = self.inlier_feature_generation(xyz0, xyz1, coords0, coords1, fcgf_feats0, fcgf_feats1, corres_idx0, corres_idx1) # Step 4: Inlier likelihood estimation and truncation import time s_time = time.time() logit = self.inlier_prediction(inlier_feats.contiguous(), coords=inlier_coords) weights = logit.sigmoid() if self.clip_weight_thresh > 0: weights[weights < self.clip_weight_thresh] = 0 wsum = weights.sum().item() # Step 5: Registration. Note: torch's gradient may be required at this stage # > Case 0: Weighted Procrustes + Robust Refinement wsum_threshold = max(200, len(weights) * 0.05) # wsum_threshold = -1 sign = '>=' if wsum >= wsum_threshold else '<' print(f'=> Weighted sum {wsum:.2f} {sign} threshold {wsum_threshold}') T = np.identity(4) safe_guard = 0 if wsum >= wsum_threshold: try: rot, trans, opt_output = GlobalRegistration(xyz0[corres_idx0], xyz1[corres_idx1], weights=weights.detach().cpu(), break_threshold_ratio=1e-4, quantization_size=2 * self.voxel_size, verbose=False) T[0:3, 0:3] = rot.detach().cpu().numpy() T[0:3, 3] = trans.detach().cpu().numpy() dgr_time = self.reg_timer.toc() print(f'=> DGR takes {dgr_time:.2} s') except RuntimeError: # Will directly go to Safeguard print('###############################################') print('# WARNING: SVD failed, weights sum: ', wsum) print('# Falling back to Safeguard') print('###############################################') else: safe_guard = 1 # > Case 1: Safeguard RANSAC + (Optional) ICP pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) T = self.safeguard_registration(pcd0, pcd1, corres_idx0, corres_idx1, feats0, feats1, 2 * self.voxel_size, num_iterations=80000) safeguard_time = self.reg_timer.toc() print(f'=> Safeguard takes {safeguard_time:.2} s') if self.use_icp: T = o3d.registration.registration_icp( make_open3d_point_cloud(xyz0), make_open3d_point_cloud(xyz1), self.voxel_size * 2, T, o3d.registration.TransformationEstimationPointToPoint()).transformation e_time = time.time() eps = np.finfo(float).eps inlier_xyz0 = np.array(xyz0[corres_idx0]) inlier_xyz1 = np.array(xyz1[corres_idx1]) pcd0 = o3d.geometry.PointCloud() pcd0.points = o3d.utility.Vector3dVector(inlier_xyz0) pcd0.transform(T_gt) inlier_xyz0_warped = np.array(pcd0.points) dis = np.sqrt( np.sum( (inlier_xyz0_warped- inlier_xyz1)**2, -1) + eps) is_correct = dis < 2 * self.voxel_size target = torch.from_numpy(is_correct).squeeze() neg_target = (~target).to(torch.bool) pred = weights > self.clip_weight_thresh pred_on_pos, pred_on_neg = pred[target], pred[neg_target] tp = pred_on_pos.sum().item() fp = pred_on_neg.sum().item() tn = (~pred_on_neg).sum().item() fn = (~pred_on_pos).sum().item() precision = tp / (tp + fp + eps) recall = tp / (tp + fn + eps) return T, precision, recall, e_time - s_time, safe_guard
def __getitem__(self, idx): drive = self.files[idx][0] t0, t1 = self.files[idx][1], self.files[idx][2] all_odometry = self.get_video_odometry(drive, [t0, t1]) positions = [self.odometry_to_positions(odometry) for odometry in all_odometry] fname0 = self._get_velodyne_fn(drive, t0) fname1 = self._get_velodyne_fn(drive, t1) # XYZ and reflectance xyzr0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4) xyzr1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4) xyz0 = xyzr0[:, :3] xyz1 = xyzr1[:, :3] key = '%d_%d_%d' % (drive, t0, t1) filename = self.icp_path + '/' + key + '.npy' if key not in kitti_icp_cache: if not os.path.exists(filename): # work on the downsampled xyzs, 0.05m == 5cm _, sel0 = ME.utils.sparse_quantize(xyz0 / 0.05, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / 0.05, return_index=True) M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T) @ np.linalg.inv(self.velo2cam)).T xyz0_t = self.apply_transform(xyz0[sel0], M) pcd0 = make_open3d_point_cloud(xyz0_t) pcd1 = make_open3d_point_cloud(xyz1[sel1]) reg = o3d.registration.registration_icp( pcd0, pcd1, 0.2, np.eye(4), o3d.registration.TransformationEstimationPointToPoint(), o3d.registration.ICPConvergenceCriteria(max_iteration=200)) pcd0.transform(reg.transformation) # pcd0.transform(M2) or self.apply_transform(xyz0, M2) M2 = M @ reg.transformation # o3d.draw_geometries([pcd0, pcd1]) # write to a file np.save(filename, M2) else: M2 = np.load(filename) kitti_icp_cache[key] = M2 else: M2 = kitti_icp_cache[key] if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, np.pi / 4) T1 = sample_random_trans(xyz1, self.randg, np.pi / 4) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) if len(matches) < 1000: raise ValueError(f"{drive}, {t0}, {t1}") # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans)
def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) xyz0 = data0["pcd"] xyz1 = data1["pcd"] color0 = data0["color"] color1 = data1["color"] matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization _, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(pcd0.colors) npts1 = len(pcd1.colors) feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((npts0, 1))) feats_train1.append(np.ones((npts1, 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def read_data(feature_path, name): data = np.load(os.path.join(feature_path, name + ".npz")) xyz = make_open3d_point_cloud(data['xyz']) feat = make_open3d_feature_from_numpy(data['feature']) return data['points'], xyz, feat
def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) # 这俩点云读进来是可以拼在一起的,只不过都只是完整点云的一部分 xyz0 = data0["pcd"]#50万左右个点 xyz1 = data1["pcd"] color0 = data0["color"]#每个点都有它自己的颜色,且不一样 color1 = data1["color"] matching_search_voxel_size = self.matching_search_voxel_size #0.0375 if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) #rotation_range:360 且旋转平移之后的点云中心坐标为0 T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) # xyz0到xyz1的相对位姿 trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization voxel_size 0.025 _, sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size, return_index=True) # 用稀疏点云生成pcd pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches # matches包含了点云1的每一点在点云二的(一定范围内的 #0.0375)匹配点 matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features,特征都为1 npts0 = len(pcd0.colors) npts1 = len(pcd1.colors) feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((npts0, 1))) feats_train1.append(np.ones((npts1, 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) # 返回体素化之后的np点云,稀疏点云坐标、特征,匹配关系,真实旋转平移 return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def main(config): test_loader = make_data_loader( config, config.test_phase, 1, num_threads=config.test_num_workers, shuffle=True) num_feats = 1 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') Model = load_model(config.model) model = Model( num_feats, config.model_n_out, bn_momentum=config.bn_momentum, conv1_kernel_size=config.conv1_kernel_size, normalize_feature=config.normalize_feature) checkpoint = torch.load(config.save_dir + '/checkpoint.pth') model.load_state_dict(checkpoint['state_dict']) model = model.to(device) model.eval() success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter(), AverageMeter() data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer() test_iter = test_loader.__iter__() N = len(test_iter) n_gpu_failures = 0 # downsample_voxel_size = 2 * config.voxel_size for i in range(len(test_iter)): data_timer.tic() try: data_dict = test_iter.next() except ValueError: n_gpu_failures += 1 logging.info(f"# Erroneous GPU Pair {n_gpu_failures}") continue data_timer.toc() xyz0, xyz1 = data_dict['pcd0'], data_dict['pcd1'] T_gth = data_dict['T_gt'] xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy() pcd0 = make_open3d_point_cloud(xyz0np) pcd1 = make_open3d_point_cloud(xyz1np) with torch.no_grad(): feat_timer.tic() sinput0 = ME.SparseTensor( data_dict['sinput0_F'].to(device), coordinates=data_dict['sinput0_C'].to(device)) F0 = model(sinput0).F.detach() sinput1 = ME.SparseTensor( data_dict['sinput1_F'].to(device), coordinates=data_dict['sinput1_C'].to(device)) F1 = model(sinput1).F.detach() feat_timer.toc() feat0 = make_open3d_feature(F0, 32, F0.shape[0]) feat1 = make_open3d_feature(F1, 32, F1.shape[0]) reg_timer.tic() distance_threshold = config.voxel_size * 1.0 ransac_result = o3d.registration.registration_ransac_based_on_feature_matching( pcd0, pcd1, feat0, feat1, distance_threshold, o3d.registration.TransformationEstimationPointToPoint(False), 4, [ o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), o3d.registration.CorrespondenceCheckerBasedOnDistance( distance_threshold) ], o3d.registration.RANSACConvergenceCriteria(4000000, 10000)) T_ransac = torch.from_numpy( ransac_result.transformation.astype(np.float32)) reg_timer.toc() # Translation error rte = np.linalg.norm(T_ransac[:3, 3] - T_gth[:3, 3]) rre = np.arccos((np.trace(T_ransac[:3, :3].t() @ T_gth[:3, :3]) - 1) / 2) # Check if the ransac was successful. successful if rte < 2m and rre < 5◦ # http://openaccess.thecvf.com/content_ECCV_2018/papers/Zi_Jian_Yew_3DFeat-Net_Weakly_Supervised_ECCV_2018_paper.pdf if rte < 2: rte_meter.update(rte) if not np.isnan(rre) and rre < np.pi / 180 * 5: rre_meter.update(rre) if rte < 2 and not np.isnan(rre) and rre < np.pi / 180 * 5: success_meter.update(1) else: success_meter.update(0) logging.info(f"Failed with RTE: {rte}, RRE: {rre}") if i % 10 == 0: logging.info( f"{i} / {N}: Data time: {data_timer.avg}, Feat time: {feat_timer.avg}," + f" Reg time: {reg_timer.avg}, RTE: {rte_meter.avg}," + f" RRE: {rre_meter.avg}, Success: {success_meter.sum} / {success_meter.count}" + f" ({success_meter.avg * 100} %)") data_timer.reset() feat_timer.reset() reg_timer.reset() logging.info( f"RTE: {rte_meter.avg}, var: {rte_meter.var}," + f" RRE: {rre_meter.avg}, var: {rre_meter.var}, Success: {success_meter.sum} " + f"/ {success_meter.count} ({success_meter.avg * 100} %)")
def __getitem__(self, idx): drive = self.files[idx][0] t0, t1 = self.files[idx][1], self.files[idx][2] all_odometry = self.get_video_odometry(drive, [t0, t1]) positions = [self.odometry_to_positions(odometry) for odometry in all_odometry] fname0 = self._get_velodyne_fn(drive, t0) fname1 = self._get_velodyne_fn(drive, t1) # XYZ and reflectance xyzr0 = np.fromfile(fname0, dtype=np.float32).reshape(-1, 4) xyzr1 = np.fromfile(fname1, dtype=np.float32).reshape(-1, 4) xyz0 = xyzr0[:, :3] xyz1 = xyzr1[:, :3] coords0 = (xyz0 - xyz0.min(0)) / 0.05 coords1 = (xyz1 - xyz1.min(0)) / 0.05 sel0 = ME.utils.sparse_quantize(coords0, return_index=True) sel1 = ME.utils.sparse_quantize(coords1, return_index=True) xyz0 = xyz0[sel0] xyz1 = xyz1[sel1] # r0 = xyzr0[:, -1].reshape(-1, 1) # r1 = xyzr1[:, -1].reshape(-1, 1) # pcd0 = make_open3d_point_cloud(xyz0_t, 0.7 * np.ones((len(xyz0), 3))) # pcd1 = make_open3d_point_cloud(xyz1, 0.3 * np.ones((len(xyz1), 3))) key = '%d_%d_%d' % (drive, t0, t1) filename = self.icp_path + '/' + key + '.npy' if key not in kitti_icp_cache: if not os.path.exists(filename): if self.IS_ODOMETRY: M = (self.velo2cam @ positions[0].T @ np.linalg.inv(positions[1].T) @ np.linalg.inv(self.velo2cam)).T else: M = self.get_position_transform(positions[0], positions[1], invert=True).T xyz0_t = self.apply_transform(xyz0, M) pcd0 = make_open3d_point_cloud(xyz0_t) pcd1 = make_open3d_point_cloud(xyz1) reg = o3d.registration_icp(pcd0, pcd1, 0.2, np.eye(4), o3d.TransformationEstimationPointToPoint(), o3d.ICPConvergenceCriteria(max_iteration=200)) pcd0.transform(reg.transformation) # pcd0.transform(M2) or self.apply_transform(xyz0, M2) M2 = M @ reg.transformation # o3d.draw_geometries([pcd0, pcd1]) # write to a file np.save(filename, M2) else: M2 = np.load(filename) kitti_icp_cache[key] = M2 else: M2 = kitti_icp_cache[key] if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, np.pi / 4) T1 = sample_random_trans(xyz1, self.randg, np.pi / 4) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 # Voxelization coords0 = np.floor(xyz0 / self.voxel_size) coords1 = np.floor(xyz1 / self.voxel_size) sel0 = ME.utils.sparse_quantize(coords0, return_index=True) sel1 = ME.utils.sparse_quantize(coords1, return_index=True) coords0, coords1 = coords0[sel0], coords1[sel1] # r0, r1 = r0[sel0], r1[sel1] pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) # Select features and points using the returned voxelized indices pcd0.points = o3d.utility.Vector3dVector(np.array(pcd0.points)[sel0]) pcd1.points = o3d.utility.Vector3dVector(np.array(pcd1.points)[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) if len(matches) < 1000: raise ValueError(f"{drive}, {t0}, {t1}") feats_train0, feats_train1 = [], [] feats_train0.append(np.ones((len(sel0), 1))) feats_train1.append(np.ones((len(sel1), 1))) feats0 = np.hstack(feats_train0) feats1 = np.hstack(feats_train1) # Get coords xyz0 = np.array(pcd0.points) xyz1 = np.array(pcd1.points) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) return (xyz0, xyz1, coords0, coords1, feats0, feats1, matches, trans)
def __getitem__(self, idx): file0 = os.path.join(self.root, self.files[idx][0]) file1 = os.path.join(self.root, self.files[idx][1]) data0 = np.load(file0) data1 = np.load(file1) xyz0 = data0["pcd"] xyz1 = data1["pcd"] matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 if self.random_rotation: T0 = sample_random_trans(xyz0, self.randg, self.rotation_range) T1 = sample_random_trans(xyz1, self.randg, self.rotation_range) trans = T1 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = np.identity(4) # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Select features and points using the returned voxelized indices # 3DMatch color is not helpful # pcd0.colors = o3d.utility.Vector3dVector(color0[sel0]) # pcd1.colors = o3d.utility.Vector3dVector(color1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] # xyz as feats if self.use_xyz_feature: feats_train0.append(unique_xyz0_th - unique_xyz0_th.mean(0)) feats_train1.append(unique_xyz1_th - unique_xyz1_th.mean(0)) else: feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) extra_package = {'idx': idx, 'file0': file0, 'file1': file1} return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans, extra_package)
def register(self, xyz0, xyz1, inlier_thr=0.00): ''' Main algorithm of DeepGlobalRegistration ''' self.reg_timer.tic() with torch.no_grad(): # Step 0: voxelize and generate sparse input xyz0, coords0, feats0 = self.preprocess(xyz0) xyz1, coords1, feats1 = self.preprocess(xyz1) # Step 1: Feature extraction self.feat_timer.tic() fcgf_feats0 = self.fcgf_feature_extraction(feats0, coords0) fcgf_feats1 = self.fcgf_feature_extraction(feats1, coords1) self.feat_timer.toc() # Step 2: Coarse correspondences corres_idx0, corres_idx1 = self.fcgf_feature_matching(fcgf_feats0, fcgf_feats1) # Step 3: Inlier feature generation # coords[corres_idx0]: 1D temporal + 3D spatial coord # coords[corres_idx1, 1:]: 3D spatial coord # => 1D temporal + 6D spatial coord inlier_coords = torch.cat((coords0[corres_idx0], coords1[corres_idx1, 1:]), dim=1).int() inlier_feats = self.inlier_feature_generation(xyz0, xyz1, coords0, coords1, fcgf_feats0, fcgf_feats1, corres_idx0, corres_idx1) # Step 4: Inlier likelihood estimation and truncation logit = self.inlier_prediction(inlier_feats.contiguous(), coords=inlier_coords) weights = logit.sigmoid() if self.clip_weight_thresh > 0: weights[weights < self.clip_weight_thresh] = 0 wsum = weights.sum().item() # Step 5: Registration. Note: torch's gradient may be required at this stage # > Case 0: Weighted Procrustes + Robust Refinement wsum_threshold = max(200, len(weights) * 0.05) sign = '>=' if wsum >= wsum_threshold else '<' print(f'=> Weighted sum {wsum:.2f} {sign} threshold {wsum_threshold}') T = np.identity(4) if wsum >= wsum_threshold: try: rot, trans, opt_output = GlobalRegistration(xyz0[corres_idx0], xyz1[corres_idx1], weights=weights.detach().cpu(), break_threshold_ratio=1e-4, quantization_size=2 * self.voxel_size, verbose=False) T[0:3, 0:3] = rot.detach().cpu().numpy() T[0:3, 3] = trans.detach().cpu().numpy() dgr_time = self.reg_timer.toc() print(f'=> DGR takes {dgr_time:.2} s') except RuntimeError: # Will directly go to Safeguard print('###############################################') print('# WARNING: SVD failed, weights sum: ', wsum) print('# Falling back to Safeguard') print('###############################################') else: # > Case 1: Safeguard RANSAC + (Optional) ICP pcd0 = make_open3d_point_cloud(xyz0) pcd1 = make_open3d_point_cloud(xyz1) T = self.safeguard_registration(pcd0, pcd1, corres_idx0, corres_idx1, feats0, feats1, 2 * self.voxel_size, num_iterations=80000) safeguard_time = self.reg_timer.toc() print(f'=> Safeguard takes {safeguard_time:.2} s') if self.use_icp: T = o3d.registration.registration_icp( make_open3d_point_cloud(xyz0), make_open3d_point_cloud(xyz1), self.voxel_size * 2, T, o3d.registration.TransformationEstimationPointToPoint()).transformation return T
def __getitem__(self, idx): drive = self.U.pairs[idx, 0] t0 = self.U.pairs[idx, 1] t1 = self.U.pairs[idx, 2] M2, xyz0, xyz1 = self.U.get_pair(idx) if self.random_rotation: T0 = sample_almost_planar_rotation(self.randg) T1 = sample_almost_planar_rotation(self.randg) trans = T1 @ M2 @ np.linalg.inv(T0) xyz0 = self.apply_transform(xyz0, T0) xyz1 = self.apply_transform(xyz1, T1) else: trans = M2 matching_search_voxel_size = self.matching_search_voxel_size if self.random_scale and random.random() < 0.95: scale = self.min_scale + \ (self.max_scale - self.min_scale) * random.random() matching_search_voxel_size *= scale xyz0 = scale * xyz0 xyz1 = scale * xyz1 M2[:3, 3] *= scale # Voxelization xyz0_th = torch.from_numpy(xyz0) xyz1_th = torch.from_numpy(xyz1) _, sel0 = ME.utils.sparse_quantize(xyz0_th / self.voxel_size, return_index=True) _, sel1 = ME.utils.sparse_quantize(xyz1_th / self.voxel_size, return_index=True) # Make point clouds using voxelized points pcd0 = make_open3d_point_cloud(xyz0[sel0]) pcd1 = make_open3d_point_cloud(xyz1[sel1]) # Get matches matches = get_matching_indices(pcd0, pcd1, trans, matching_search_voxel_size) # Get features npts0 = len(sel0) npts1 = len(sel1) feats_train0, feats_train1 = [], [] unique_xyz0_th = xyz0_th[sel0] unique_xyz1_th = xyz1_th[sel1] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = torch.floor(unique_xyz0_th / self.voxel_size) coords1 = torch.floor(unique_xyz1_th / self.voxel_size) if self.transform: coords0, feats0 = self.transform(coords0, feats0) coords1, feats1 = self.transform(coords1, feats1) extra_package = {'drive': drive, 't0': t0, 't1': t1} return (unique_xyz0_th.float(), unique_xyz1_th.float(), coords0.int(), coords1.int(), feats0.float(), feats1.float(), matches, trans, extra_package)
def main(config): test_loader = make_data_loader(config, config.test_phase, 1, num_threads=config.test_num_thread, shuffle=False) num_feats = 1 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') Model = load_model(config.model) model = Model(num_feats, config.model_n_out, bn_momentum=config.bn_momentum, conv1_kernel_size=config.conv1_kernel_size, normalize_feature=config.normalize_feature) checkpoint = torch.load(config.save_dir + '/checkpoint.pth') model.load_state_dict(checkpoint['state_dict']) model = model.to(device) model.eval() success_meter, rte_meter, rre_meter = AverageMeter(), AverageMeter( ), AverageMeter() data_timer, feat_timer, reg_timer = Timer(), Timer(), Timer() test_iter = test_loader.__iter__() N = len(test_iter) n_gpu_failures = 0 # downsample_voxel_size = 2 * config.voxel_size list_results_to_save = [] for i in range(len(test_iter)): data_timer.tic() try: data_dict = test_iter.next() except ValueError: n_gpu_failures += 1 logging.info(f"# Erroneous GPU Pair {n_gpu_failures}") continue data_timer.toc() xyz0, xyz1 = data_dict['pcd0'], data_dict['pcd1'] T_gth = data_dict['T_gt'] xyz0np, xyz1np = xyz0.numpy(), xyz1.numpy() #import pdb # pdb.set_trace() pcd0 = make_open3d_point_cloud(xyz0np) pcd1 = make_open3d_point_cloud(xyz1np) with torch.no_grad(): feat_timer.tic() sinput0 = ME.SparseTensor( data_dict['sinput0_F'].to(device), coordinates=data_dict['sinput0_C'].to(device)) F0 = model(sinput0).F.detach() sinput1 = ME.SparseTensor( data_dict['sinput1_F'].to(device), coordinates=data_dict['sinput1_C'].to(device)) F1 = model(sinput1).F.detach() feat_timer.toc() # saving files to pkl print(i) dict_sample = { "pts_source": xyz0np, "feat_source": F0.cpu().detach().numpy(), "pts_target": xyz1np, "feat_target": F1.cpu().detach().numpy() } list_results_to_save.append(dict_sample) import pickle path_results_to_save = "fcgf.results.pkl" print('Saving results to ', path_results_to_save) pickle.dump(list_results_to_save, open(path_results_to_save, 'wb')) print('Saved!') import pdb pdb.set_trace()
def __getitem__(self, idx): # split, idx): if self.split == "test": sample = self.list_test_sample[idx] else: T_noise = self.generate_noise_T() sample = self.get_sample(idx, T_noise) xyz0 = sample["source"].points xyz1 = sample["target"].points trans = torch.Tensor(self.data["T_map"][idx]) matching_search_voxel_size = self.matching_search_voxel_size import copy # sel0 = ME.utils.sparse_quantize( xyz0.contiguous() / self.voxel_size, return_index=True)[1] sel1 = ME.utils.sparse_quantize( xyz1.contiguous() / self.voxel_size, return_index=True)[1] unique_xyz0_th = xyz0[sel0] # [ind_0] unique_xyz1_th = xyz1[sel1] # [ind_1] pcd0 = make_open3d_point_cloud(unique_xyz0_th) pcd1 = make_open3d_point_cloud(unique_xyz1_th) # Get matches matches = get_matching_indices( pcd0, pcd1, trans, matching_search_voxel_size) # Get features feats_train0, feats_train1 = [], [] npts0 = unique_xyz0_th.shape[0] npts1 = unique_xyz1_th.shape[0] feats_train0.append(torch.ones((npts0, 1))) feats_train1.append(torch.ones((npts1, 1))) feats0 = torch.cat(feats_train0, 1) feats1 = torch.cat(feats_train1, 1) coords0 = np.floor(unique_xyz0_th / self.voxel_size) coords1 = np.floor(unique_xyz1_th / self.voxel_size) #coords0_mean = coords0.min(axis=0).int() #coords0 -=coords0_mean if False:#len(matches) < 300: # idx == 113:#len(matches) < print(self.split, "num matches = ", len(matches)) print(coords0) pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(coords0) o3d.io.write_point_cloud("coords0_before_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(coords1) o3d.io.write_point_cloud("coords1_before_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz1_th) o3d.io.write_point_cloud("unique_xyz1_th_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th) o3d.io.write_point_cloud("unique_xyz0_th_%d.ply" % idx, pcd_target) # pcd_target = o3d.geometry.PointCloud() pcd_target.points = o3d.utility.Vector3dVector(unique_xyz0_th) pcd_target.transform(trans) o3d.io.write_point_cloud( "unique_xyz0_th_trans_%d.ply" % idx, pcd_target) # import pdb pdb.set_trace() #if self.transform: # add noises to the point clouds # coords0, feats0 = self.transform(coords0, feats0) # coords1, feats1 = self.transform(coords1, feats1) return (unique_xyz0_th, unique_xyz1_th, coords0, coords1, feats0.float(), feats1.float(), matches, trans)