def forward(self, x, pos, batch): # FPS sampling id_clusters = fps(pos, ratio=self.ratio, batch=batch) # compute for each cluster the k nearest points sub_batch = batch[id_clusters] if batch is not None else None # beware of self loop id_k_neighbor = knn(pos, pos[id_clusters], k=self.k, batch_x=batch, batch_y=sub_batch) # transformation of features through a simple MLP x = self.mlp(x) # Max pool onto each cluster the features from knn in points x_out, _ = scatter_max(x[id_k_neighbor[1]], id_k_neighbor[0], dim_size=id_clusters.size(0), dim=0) # keep only the clusters and their max-pooled features sub_pos, out = pos[id_clusters], x_out return out, sub_pos, sub_batch
def fps(x, batch=None, ratio=0.5, random_start=True): r"""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space" <https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the most distant point with regard to the rest points. Args: x (Tensor): Node feature matrix :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`. batch (LongTensor, optional): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each node to a specific example. (default: :obj:`None`) ratio (float, optional): Sampling ratio. (default: :obj:`0.5`) random_start (bool, optional): If set to :obj:`False`, use the first node in :math:`\mathbf{X}` as starting node. (default: obj:`True`) :rtype: :class:`LongTensor` .. code-block:: python import torch from torch_geometric.nn import fps x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) batch = torch.tensor([0, 0, 0, 0]) index = fps(x, batch, ratio=0.5) """ if torch_cluster is None: raise ImportError('`fps` requires `torch-cluster`.') return torch_cluster.fps(x, batch, ratio, random_start)
def test_fps(dtype, device): x = tensor([ [-1, -1], [-1, +1], [+1, +1], [+1, -1], [-2, -2], [-2, +2], [+2, +2], [+2, -2], ], dtype, device) batch = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device) out = fps(x, batch, ratio=0.5, random_start=False) assert out.tolist() == [0, 2, 4, 6] out = fps(x, ratio=0.5, random_start=False) assert out.sort()[0].tolist() == [0, 5, 6, 7]
def test_random_fps(device): N = 1024 for _ in range(5): pos = torch.randn((2 * N, 3), device=device) batch_1 = torch.zeros(N, dtype=torch.long, device=device) batch_2 = torch.ones(N, dtype=torch.long, device=device) batch = torch.cat([batch_1, batch_2]) idx = fps(pos, batch, ratio=0.5) assert idx.min() >= 0 and idx.max() < 2 * N
def test_fps_speed(dtype, device): return batch_size, num_nodes = 100, 10000 x = torch.randn((batch_size * num_nodes, 3), dtype=dtype, device=device) batch = torch.arange(batch_size, dtype=torch.long, device=device) batch = batch.view(-1, 1).repeat(1, num_nodes).view(-1) out = fps(x, batch, ratio=0.5, random_start=True) assert out.size(0) == batch_size * num_nodes * 0.5 assert out.min().item() >= 0 and out.max().item() < batch_size * num_nodes batch_size, num_nodes, dim = 100, 300, 128 x = torch.randn((batch_size * num_nodes, dim), dtype=dtype, device=device) batch = torch.arange(batch_size, dtype=torch.long, device=device) batch = batch.view(-1, 1).repeat(1, num_nodes).view(-1) out = fps(x, batch, ratio=0.5, random_start=True) assert out.size(0) == batch_size * num_nodes * 0.5 assert out.min().item() >= 0 and out.max().item() < batch_size * num_nodes
def getPatchDescriptors(mesh, radius=15.0, n_max=20, l_max=10, feature_name='zd', sample_ratio=1.0): patches = getRadialGeodesicPatches(mesh, radius, add_self_loops=True, to_csr=True) M = MeshDescriptor(mesh, patches, n_max=n_max, l_max=l_max, scale_patches=True, center_patches=True) if sample_ratio < 1.0: # use fps centroids as patch loci from torch_cluster import fps import torch idx = fps(torch.tensor(mesh.vertices), batch=None, ratio=sample_ratio, random_start=False) idx = idx.numpy() descriptors = [] for i in idx: descriptors.append(M.getPatchDescriptors(i)) descriptors = np.array(descriptors) # map to all vertices descriptors = mapPointFeaturesToMesh(mesh, mesh.vertices[idx], descriptors, map_to='nearest') else: # loop over vertices, compute moments for every patch descriptors = [] for i in range(M.Nv): descriptors.append(M.getPatchDescriptors(i)) descriptors = np.array(descriptors) # add features to mesh feature_names = [] for i in range(descriptors.shape[1]): key = "{}{}".format(feature_name, i) mesh.vertex_attributes[key] = descriptors[:, i] feature_names.append(key) return feature_names
def forward(self, batch): indices = fps(batch.pos, batch.batch, ratio=self.ratio, random_start=False) new_batch = torch_geometric.data.Batch(batch=batch.batch[indices], x=batch.x[indices], pos=batch.pos[indices]) return new_batch
def furthest_point_sample(xyz: torch.tensor, num_samples: int): # How to use fps: examples/PointNet2ASIS/tests/fps_test.py xyz = xyz.transpose(1, 2).contiguous() device = xyz.device B, N, C = xyz.shape xyz = xyz.view(B * N, C) batch = torch.arange(0, B, dtype=torch.long, device=device) batch = batch.view(-1, 1).repeat(1, N).view(B * N) fps_idx = fps(xyz, batch, num_samples / N, True) fps_idx = fps_idx.view(B, num_samples) idx_base = torch.arange(0, B, device=device).view(-1, 1) * N fps_idx = fps_idx - idx_base return fps_idx
def __call__(self, sample): keys = sorted([x for x in dir(sample) if 'edge_index' in x]) num_vertices = [sample.num_nodes] + sample.num_vertices # sample.edge_index = knn_graph(sample.pos, self._k[0]) # knn_edges = knn_graph(sample.pos, self._k[0] * self._d) # dilated_idx = [index for index in range(knn_edges.shape[1])[0::self._d]] # sample.edge_index = knn_edges[:, dilated_idx] for level, key in enumerate(keys): if level == len(keys) - 1: break pos_key = key.replace('edge_index', 'pos').replace('hierarchy_', '') subset_points_idx = fps(sample[pos_key], ratio=sample.num_vertices[level] / num_vertices[level]) # if level == 0: # sample.y = sample.y[subset_points_idx] num_vertices[level + 1] = subset_points_idx.shape[0] sample.num_vertices[level] = num_vertices[level + 1] sample['pos_' + str(level + 1)] = sample[pos_key][subset_points_idx] sample[f"hierarchy_trace_index_{level+1}"] = knn( sample['pos_' + str(level + 1)], sample[pos_key], 1)[1, :] # sample[f"hierarchy_edge_index_{level+1}"] = knn_graph(sample['pos_' + str(level+1)], self._k[level+1]) # knn_edges = knn_graph(sample['pos_' + str(level+1)], self._k[level+1] * self._d) # dilated_idx = [index for index in range(knn_edges.shape[1])[0::self._d]] # sample[f"hierarchy_edge_index_{level+1}"] = knn_edges[:, dilated_idx] keys = sorted([x for x in dir(sample) if x.startswith('x_')]) for key in keys: delattr(sample, key) # keys = sorted([x for x in dir(sample) if 'pos_' in x]) # for key in keys: # delattr(sample, key) return sample
def fps_sampling(pair_ind, pos, num_pos_pairs, ind=0): """ perform fps sampling to choose the positive pairs Parameters: pair_ind: torch tensor which represents index of pair size N x 2 pos: torch tensor which represents the point cloud of size M x 3 num_pos_pairs: int which number of pairs we want() ind: must be 0 or 1 """ small_pos_source = pos[pair_ind[:, ind]] batch = torch.zeros(small_pos_source.shape[0]).long() ratio = float(num_pos_pairs) / len(pair_ind) if (ratio <= 0 or ratio >= 1): raise ValueError("ratio cannot have this value: {}".format(ratio)) index = fps(small_pos_source, batch, ratio=ratio, random_start=False) return index
def farthest_point_sample_fast(xyz, npoint): """ Input: xyz: pointcloud data, [B, N, C] npoint: number of samples Return: centroids: sampled pointcloud index, [B, npoint] """ device = xyz.device B, N, C = xyz.shape centroids = torch.zeros(B, npoint, dtype=torch.long).to(device) r = npoint / N * 0.9999999 + 0.00000001 for i in range(B): centroids[i] = fps(xyz[i], ratio=r) return centroids
def create_patch_pair(depth_path, mask_path, im_cam, gt, save_name, md_pcd_pts): raw_depth = io.load_depth(depth_path) mask = io.load_im(mask_path) img_pcd = PointCloud.create_from_depth_image( depth=Image(masked_where(mask == 0.0, raw_depth).filled(0.0)), intrinsic=PHCamIntrinsic(*IM_SIZE, *[im_cam['cam_K'][i] for i in K]), depth_scale=im_cam['depth_scale'], depth_trunc=150000) img_pcd.voxel_down_sample(VOXEL_SIZE) if np.asarray(img_pcd.points).shape[0] in PCD_PTS_RANGE or IS_TARGET: cam_R, cam_t = gt['cam_R_m2c'], gt['cam_t_m2c'] # Select reference points on image using farthest point sampling img_pcd_pts_fps = torch.as_tensor(img_pcd.points).to(DEVICE) img_ref_idxs = fps(img_pcd_pts_fps, ratio=FPS_RATIO).to('cpu').numpy() # Calculate model reference points img_ref_pts = np.asarray(img_pcd.points)[img_ref_idxs] md_ref_pts = (img_ref_pts - cam_t.T) @ np.linalg.inv(cam_R).T # Recreate model point cloud md_ref_idxs = np.arange(md_ref_pts.shape[0]) md_pcd_pts = np.concatenate([md_ref_pts, md_pcd_pts], axis=0) md_pcd = PointCloud() md_pcd.points = Vector3dVector(md_pcd_pts) # Calculate and save PPFs img_save_path = f'image/{save_name}' create_local_patches(img_pcd, img_ref_idxs, img_save_path) md_save_path = f'model/{save_name}' create_local_patches(md_pcd, md_ref_idxs, md_save_path) entry = [save_name, img_ref_idxs.shape[0]] else: entry = [] return entry
def getMutiplePrototypes(self, feat, k): """ Extract multiple prototypes by points separation and assembly Args: feat: input point features, shape:(n_points, feat_dim) Return: prototypes: output prototypes, shape: (n_prototypes, feat_dim) """ # sample k seeds as initial centers with Farthest Point Sampling (FPS) n = feat.shape[0] assert n > 0 ratio = k / n if ratio < 1: fps_index = fps(feat, None, ratio=ratio, random_start=False).unique() num_prototypes = len(fps_index) farthest_seeds = feat[fps_index] # compute the point-to-seed distance distances = F.pairwise_distance(feat[..., None], farthest_seeds.transpose(0, 1)[None, ...], p=2) # (n_points, n_prototypes) # hard assignment for each point assignments = torch.argmin(distances, dim=1) # (n_points,) # aggregating each cluster to form prototype prototypes = torch.zeros((num_prototypes, self.feat_dim)).cuda() for i in range(num_prototypes): selected = torch.nonzero(assignments == i).squeeze(1) selected = feat[selected, :] prototypes[i] = selected.mean(0) return prototypes else: return feat
def down_sample(self, data): ratio = (float(self.point_num_th) / data.shape[0]) if data.shape[0] != 0 else 10086 if ratio >= 1.0: return None, None # mid_x = data[:, 0].mean() # mid_y = data[:, 1].mean() # inds = [ # (data[:, 0] >= mid_x) & (data[:, 1] >= mid_y), # (data[:, 0] >= mid_x) & (data[:, 1] < mid_y), # (data[:, 0] < mid_x) & (data[:, 1] >= mid_y), # (data[:, 0] < mid_x) & (data[:, 1] < mid_y), # ] # for i in inds: # cur_i = fps(data[i], torch.zeros(data[i].shape[0]).cuda().long(), ratio=ratio, random_start=False) # cur_index_torch = torch.zeros(data[i].shape[0]).cuda().long() # cur_index_torch[cur_i] = 1 # index_torch[i] = cur_index_torch # step = data.shape[0] // 4 # split = [i * step for i in range(5)] # split[-1] = data.shape[0] # index_torch = torch.zeros(data.shape[0]).cuda().long() # for i in range(4): # cur_i = fps(data[split[i]:split[i+1]], torch.zeros(data[split[i]:split[i+1]].shape[0]).cuda().long(), ratio=ratio, random_start=False) # index_torch[cur_i + i*step] = 1 index_torch = fps(data, torch.zeros(data.shape[0]).cuda().long(), ratio=ratio, random_start=False) # index_torch = index_torch == 1 index = index_torch.detach().cpu().numpy() # index = np.random.choice(data.shape[0], int(ratio * data.shape[0])) # index_torch = torch.from_numpy(index).cuda() return index, index_torch
def farthest_sampling(point_clouds: PointClouds3D, ratio: float) -> PointClouds3D: """ Args: point_clouds: Pointclouds object """ points_packed = point_clouds.points_packed() from torch_cluster import fps packed_to_cloud_idx = num_points_2_packed_to_cloud_idx( point_clouds.num_points_per_cloud()) fps_idx = fps(points_packed, packed_to_cloud_idx, ratio) sampled_points = points_packed[fps_idx] # back to pointclouds object point_lst = [ sampled_points[packed_to_cloud_idx[fps_idx] == b] for b in range(len(point_clouds)) ] sampled_point_clouds = point_clouds.__class__(point_lst) if (normals_packed := point_clouds.normals_packed()) is not None: normals_packed = normals_packed[fps_idx] sampled_point_clouds.update_normals_(normals_packed)
# select only corresponding points pcd1_corr = pcd1.select_down_sample(corrs[:, 0]) pcd2_corr = pcd2.select_down_sample(corrs[:, 1]) pcd1 = pcd1.voxel_down_sample(voxel_size) pcd2 = pcd2.voxel_down_sample(voxel_size) # apply ground truth transformation to bring them in the same reference frame pcd1_corr.transform(T1) pcd2_corr.transform(T2) # FPS tensor_pcd1_frag = torch.Tensor(np.asarray(pcd1_corr.points)).to(device) fps_pcd1_idx = fps(tensor_pcd1_frag, ratio=batch_size / tensor_pcd1_frag.shape[0], random_start=True) _pcd2_frag_tree = o3d.geometry.KDTreeFlann(pcd2_corr) fps_pcd1_pts = np.asarray(pcd1_corr.points)[fps_pcd1_idx.cpu()] fps_pcd2_idx = torch.empty(fps_pcd1_idx.shape, dtype=int) # find nearest neighbors on the other point cloud for i, pt in enumerate(fps_pcd1_pts): _, patch_idx, _ = _pcd2_frag_tree.search_knn_vector_xd(pt, 1) fps_pcd2_idx[i] = patch_idx[0] # visualise point clouds with FPS + NN result overlaid # to_viz = []
def main(): parser = argparse.ArgumentParser() parser.add_argument('--folder', '-f', help='Path to data folder') parser.add_argument('--benchmark', '-b', help='Path to benchmark folder') parser.add_argument('--outpath', '-o', help='Path to output folder') parser.add_argument('--saveply', '-s', action='store_true', help='Save color ply or not') args = parser.parse_args() print(args) #label_tsv = args.benchmark + "/scannet-labels.combined.tsv" train_list_file = args.benchmark + "/train.txt" test_list_file = args.benchmark + "/test.txt" val_list_file = args.benchmark + "/val.txt" #label_shapenetcore55 = args.benchmark + "/classes_ObjClassification-ShapeNetCore55.txt" ##########################################################Read Source########################################################## print("read scene dir:", args.folder) scenedir = dir(args.folder, 'd') print("read trainval list:", train_list_file) train_scene_list = [] with open(train_list_file, 'r') as train_f: for line in train_f.readlines(): sceneid = line.strip().split("scene")[1] spaceid = sceneid.split("_")[0] scanid = sceneid.split("_")[1] train_scene_list.append(spaceid + scanid) print("read test list:", test_list_file) test_scene_list = [] with open(test_list_file, 'r') as train_f: for line in train_f.readlines(): sceneid = line.strip().split("scene")[1] spaceid = sceneid.split("_")[0] scanid = sceneid.split("_")[1] test_scene_list.append(spaceid + scanid) print("read val list:", val_list_file) val_scene_list = [] with open(val_list_file, 'r') as train_f: for line in train_f.readlines(): sceneid = line.strip().split("scene")[1] spaceid = sceneid.split("_")[0] scanid = sceneid.split("_")[1] val_scene_list.append(spaceid + scanid) # split scene to train and test process_train_list = [] process_test_list = [] for scene in scenedir: sceneid = scene.strip().split("scene")[1] spaceid = sceneid.split("_")[0] scanid = sceneid.split("_")[1] scenename = spaceid + scanid if scenename in train_scene_list: process_train_list.append(scene) elif scenename in test_scene_list: process_test_list.append(scene) print("Train all:", len(train_scene_list), "Test all:", len(test_scene_list), "Dir all:", len(scenedir)) print("Process Train:", len(process_train_list), "Process Test:", len(process_test_list)) ##########################################################Process Data########################################################## print("Process Train Scene:") for scene in process_train_list: ply_file = scene + "/scene" + sceneid + "_vh_clean_2.ply" # Read ply file print("\nRead ply file:", ply_file) plydata = PlyData.read(ply_file).elements[0].data pts_num = len(plydata) print("points num:", pts_num) data = torch.tensor(plydata) batch = torch.zeros(pts_num) index = fps(data, batch, ratio=0.5, random_start=False) data = data[index] torch.save(data) print("Process Test Scene:") for scene in process_test_list: scene2instances(scene, args.outpath + "/test/", [label_map, label_info], label_shapenetcore55_map, args.saveply)
def fps2(x: Tensor, ratio: Tensor) -> Tensor: return fps(x, None, ratio, False)
def all_getter(self,idx): save_id,cloud_ind_0,cloud_ind_1,common_voxel = self.all_valid_combs[idx]['combination'] ground_height = self.save_dict[save_id]['ground_height'] clouds = self.save_dict[save_id]['clouds'] center = self.all_valid_combs[idx]['voxel_center'] are_same = (cloud_ind_1 == cloud_ind_0) cloud_0,cloud_1 = clouds[cloud_ind_0],clouds[cloud_ind_1] voxel_1_small = get_voxel(cloud_1,center,self.final_voxel_size) voxel_0_large = get_voxel(cloud_0,center,self.context_voxel_size) voxel_1_small = voxel_1_small[fps(voxel_1_small, torch.zeros(voxel_1_small.shape[0]).long( ), ratio=self.n_samples/voxel_1_small.shape[0], random_start=False), :] voxel_1_small = voxel_1_small[:self.n_samples, :] voxel_0_large = voxel_0_large[fps(voxel_0_large, torch.zeros(voxel_0_large.shape[0]).long( ), ratio=self.n_samples_context/voxel_0_large.shape[0], random_start=False), :] voxel_0_large = voxel_0_large[:self.n_samples_context,:] if self.include_all: voxel_0_small = get_voxel(cloud_0,center,self.final_voxel_size) voxel_1_small_original = voxel_1_small.clone() voxel_0_small = voxel_0_small[fps(voxel_0_small, torch.zeros(voxel_0_small.shape[0]).long( ), ratio=self.n_samples/voxel_0_small.shape[0], random_start=False), :] voxel_0_small = voxel_0_small[:self.n_samples,:] voxel_0_small_original = voxel_0_small.clone() voxel_0_small_self, voxel_0_large_self,inverse = self.last_processing(voxel_0_small,voxel_0_large) voxel_1_large = get_voxel(cloud_1,center,self.context_voxel_size) voxel_1_large = voxel_1_large[fps(voxel_1_large, torch.zeros(voxel_1_large.shape[0]).long( ), ratio=self.n_samples/voxel_1_large.shape[0], random_start=False), :] voxel_1_large = voxel_1_large[:self.n_samples,:] voxel_1_large_self, voxel_1_small_self,inverse = self.last_processing(voxel_1_large,voxel_1_small) voxel_1_large_self, voxel_1_small_self,inverse = self.last_processing(voxel_1_large,voxel_1_small) voxel_opposite_small , voxel_opposite_large, inverse = self.last_processing(voxel_0_small,voxel_1_large) #Only augment in train if are_same: voxel_1_small = voxel_1_small.clone() if self.mode == 'train': voxel_0_large[:, :3] += torch.rand_like(voxel_0_large[:, :3])*0.01 voxel_0_large, voxel_1_small,inverse = self.last_processing(voxel_0_large, voxel_1_small) if self.mode == 'train': rads = torch.rand((1))*math.pi*2 if self.rotation_augment: rot_mat = rotate_xy(rads) voxel_0_large[:, :2] = torch.matmul(voxel_0_large[:, :2], rot_mat) voxel_1_small[:, :2] = torch.matmul(voxel_1_small[:, :2], rot_mat) # Distance from ground as extra context extra_context = inverse['mean'][2] - ground_height extra_context = extra_context.unsqueeze(-1) if self.include_all: return voxel_0_large, voxel_1_small,extra_context, voxel_1_large_self, voxel_1_small_self, voxel_opposite_small , voxel_opposite_large, voxel_0_small_self, voxel_0_large_self, voxel_0_small_original ,voxel_1_small_original else: return voxel_0_large, voxel_1_small,extra_context