def knn_query(k, support_pts, query_pts): """ :param support_pts: points you have, B*N1*3 :param query_pts: points you want to know the neighbour index, B*N2*3 :param k: Number of neighbours in knn search :return: neighbor_idx: neighboring points indexes, B*N2*k """ neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) return neighbor_idx.astype(np.int32)
def indices_deconv(self, input_pts, next_pts, K): indices = nearest_neighbors.knn_batch(input_pts.cpu().detach().numpy(), next_pts.cpu().detach().numpy(), K, omp=True) indices = torch.from_numpy(indices).long() if input_pts.is_cuda: indices = indices.cuda() return indices, next_pts