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
