def ball_search_np(pc, kpt, knn, search_radius, subsample_ratio=1): if subsample_ratio > 1: pc_sub = subsample_pc(pc, pc.shape[0] // subsample_ratio) else: pc_sub = pc # knn-ball search nn = min(10000, pc_sub.shape[0]) nbrs = nnbrs(n_neighbors=nn, algorithm='ball_tree').fit(pc_sub) dists, indices = nbrs.kneighbors(pc[kpt]) true_indices = [] maxcount = 0 for i in range(len(dists)): if dists[i].max() > search_radius: lidx = np.where(dists[i] > search_radius)[0][0] # print(f'{lidx} vs {knn}') if lidx >= knn: true_indices.append(np.random.choice(indices[i][:lidx], knn)) else: choice = np.random.choice(range(lidx - 1), knn - lidx) true_indices.append( np.append(indices[i][:lidx], indices[i][choice])) else: true_indices.append(np.random.choice(indices[i], knn)) maxcount += 1 print("inclusion ratio: ", 1 - float(maxcount) / float(len(dists))) return np.array(true_indices, dtype=np.int32), pc_sub
def find_scenes_overlap(pc1, pc2, T, k=5000, margin=1e-2, subsample=False): if subsample: pc1 = subsample_pc(pc1, pc1.shape[0]//10) pc2 = subsample_pc(pc2, pc2.shape[0]//10) if T is not None: pc2_t = transform_np(pc2, T) else: pc2_t = pc2 nbrs = nnbrs(n_neighbors=1, algorithm='ball_tree').fit(pc2_t) dists, indices = nbrs.kneighbors(pc1) pc1idx = np.argwhere(dists<=margin)[:,0] pc2idx = indices[pc1idx].reshape(-1) print("Matched points: ", pc1idx.shape[0]) if pc1idx.shape[0] > 10 * k: choice = np.random.choice(pc1idx.shape[0], k, replace=False) pc1idx = pc1idx[choice] pc2idx = pc2idx[choice] return pc1idx, pc2idx, pc2_t else: # save_ply("example_src.ply",pc1) # save_ply("example_tgt.ply", pc2_t) raise ValueError("Not enough overlapping points between this pair of scenes.")
def generate_kp(src, src_sub, src_down, src_fpfh, j, sdir, save_path, cfg): from_o3d = lambda x: np.asarray(x.points) tgt_pcd = o3d.io.read_point_cloud( os.path.join(sdir, 'cloud_bin_%d.ply' % j)) tgt_pose = np.loadtxt(os.path.join(sdir, 'cloud_bin_%d_pose.txt' % j), dtype=np.float32) tgt_pcd = tgt_pcd.transform(tgt_pose) tgt = from_o3d(tgt_pcd) # tgt = load_ply(os.path.join(sdir, 'cloud_bin_%d.ply'%j)) tgt_sub = subsample_pc(tgt, min(cfg.subsample_maxpoints, tgt.shape[0])) if test_scenes_overlap(src_sub, tgt_sub, cfg.overlap_ratio, cfg.dist_margin, cfg.verbose): # extract keypoints through fpfh matching tgt_down, tgt_fpfh = downsample_and_compute_fpfh(tgt_pcd, cfg) tgt_fpfh = np.asarray(tgt_fpfh.data).T kpti, kptj = cross_filtering_via_fpfh(src_down, from_o3d(tgt_down), src_fpfh, tgt_fpfh, cfg) if kpti is not None: kpts = [] for pcd, kpt in zip([src, tgt], [kpti, kptj]): nbrs = nnbrs(n_neighbors=1, algorithm='ball_tree').fit(pcd) dists, indices = nbrs.kneighbors(kpt) # sanity check? if np.sum(dists.squeeze() > 0.03) > 0: print( "WARNING: SOME SEARCHED POINTS MAY BE TOO FAR FROM SELECTED KEYPOINTS!!!" ) kpts.append(indices) kpts = np.concatenate(kpts, axis=1) print("Keypoint saved to", save_path) np.save(save_path, kpts.astype(np.int32))
def crop_point_cloud(data, k=0.05): N, C = data.shape crop_center = data[np.random.randint(N)] nbrs = nnbrs(n_neighbors=int(k*N)).fit(data) _, indices = nbrs.kneighbors(crop_center[None]) cropped = np.delete(data, indices.flatten(), axis=0) return cropped
def nnbrsearch(feat1, feat2, knn=1, dist=False): nbrs = nnbrs(n_neighbors=knn) nbrs.fit(feat2) rst = nbrs.kneighbors(feat1, return_distance=dist) if dist: return (rst[0].squeeze(), rst[1].squeeze()) else: return rst.squeeze()
def find_scenes_overlap(pc1, pc2, T, k=5000, margin=5e-3): pc2_t = pctk.transform_np(pc2, T) nbrs = nnbrs(n_neighbors=1, algorithm='ball_tree').fit(pc2_t) dists, indices = nbrs.kneighbors(pc1) pc1idx = np.argwhere(dists<=margin)[:,0] pc2idx = indices[pc1idx].reshape(-1) print("Matched points: ", pc1idx.shape[0]) if pc1idx.shape[0] > k: choice = np.random.choice(pc1idx.shape[0], k, replace=False) pc1idx = pc1idx[choice] pc2idx = pc2idx[choice] return pc1idx, pc2idx, pc2_t else: pctk.save_ply("example_src.ply",pc1) pctk.save_ply("example_tgt.ply", pc2_t) import ipdb; ipdb.set_trace() raise ValueError("Not enough overlapping points between this pair of scenes.")
def test_scenes_overlap(pc1, pc2, overlap_ratio=0.3, margin=1e-2, verbose=True): nbrs = nnbrs(n_neighbors=1, algorithm='ball_tree').fit(pc2) dists, indices = nbrs.kneighbors(pc1) pc1idx = np.argwhere(dists <= margin)[:, 0] pc2idx = indices[pc1idx].reshape(-1) if verbose: print("Matched points: ", pc1idx.shape[0]) n_overlap = pc1idx.shape[0] n_pts = max(pc1.shape[0], pc2.shape[0]) if verbose: print("Overlap ratio is %f" % (n_overlap / n_pts)) return n_overlap >= overlap_ratio * n_pts