fix problem of cpu

Author: humanpose1

cpu/include/ball_query.h

@@ -9,3 +9,7 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 						   at::Tensor query_batch,
 						   at::Tensor support_batch,
 						   float radius, int max_num, int mode);
+
+std::pair<at::Tensor, at::Tensor> dense_ball_query(at::Tensor query,
+						   at::Tensor support,
+						   float radius, int max_num, int mode);

cpu/src/bindings.cpp

@@ -35,4 +35,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 		"mode=1 means a matrix of edges of size Num_edge x 2"
 	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
 	      "query"_a, "support"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
+    m.def("dense_ball_query", &dense_ball_query,
+	  "compute the radius search of a batch of point cloud using nanoflann"
+	      "-query : a pytorch tensor of size B x N1 x 3,. used to query the nearest neighbors"
+	      "- support : a pytorch tensor of size B x N2 x 3. used to build the tree"
+	      "-  radius : float number, size of the ball for the radius search."
+	      "- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors will be computed). "
+	      " - mode : int number that indicate which format for the neighborhood"
+	      " mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
+	      "mode=1 means a matrix of edges of size Num_edge x 2"
+	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2 and return a tensor containing the squared distance of the neighbors",
+	      "query"_a, "support"_a, "radius"_a, "max_num"_a=-1, "mode"_a=0);
 }

cpu/src/neighbors.cpp

@@ -82,16 +82,20 @@ int nanoflann_neighbors(vector<scalar_t>& queries,
 
 		i0 = 0;
 
+		int token = 0;
 		for (auto& inds : list_matches){
+			token = inds[0].first;
 			for (int j = 0; j < max_count; j++){
 				if (j < inds.size()){
 					neighbors_indices[i0 * max_count + j] = inds[j].first;
 					dists[i0 * max_count + j] = (float) inds[j].second;
+
+
 				}
 
 				else {
-					neighbors_indices[i0 * max_count + j] = -1;
-					dists[i0 * max_count + j] = radius * radius;
+					neighbors_indices[i0 * max_count + j] = token;
+					dists[i0 * max_count + j] = -1;
 				}
 			}
 			i0++;
@@ -239,8 +243,8 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
 					dists[i0 * max_count + j] = (float) inds_dists[j].second;
 				}
 				else {
-					neighbors_indices[i0 * max_count + j] = supports.size();
-					dists[i0 * max_count + j] = radius * radius;
+					neighbors_indices[i0 * max_count + j] = supports.size()/3;
+					dists[i0 * max_count + j] = -1;
 				}
 
 			}

cpu/src/torch_nearest_neighbors.cpp

@@ -135,3 +135,23 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 	}
 	return std::make_pair(out.clone(), out_dists.clone());
 }
+
+
+std::pair<at::Tensor, at::Tensor> dense_ball_query(at::Tensor query,
+						   at::Tensor support,
+						   float radius, int max_num, int mode){
+
+	int b = query.size(0);
+	vector<at::Tensor> batch_idx;
+	vector<at::Tensor> batch_dist;
+	for (int i=0; i < b; i++){
+
+		auto out_pair = ball_query(query[i], support[i], radius, max_num, mode);
+		batch_idx.push_back(out_pair.first);
+		batch_dist.push_back(out_pair.second);
+	}
+	auto out_idx = torch::stack(batch_idx);
+	auto out_dist = torch::stack(batch_dist);
+	return std::make_pair(out_idx, out_dist);
+
+}

test/test_ballquerry.py

@@ -19,8 +19,12 @@ def test_simple_cpu(self):
 
     def test_cpu_gpu_equality(self):
         a = torch.randn(5, 1000, 3)
-        npt.assert_array_equal(ball_query_dense(0.1, 17, a, a).detach().numpy(),
-                               ball_query_dense(0.1, 17, a.cuda(), a.cuda()).cpu().detach().numpy())
+        res_cpu = ball_query_dense(0.1, 17, a, a).detach().numpy()
+        res_cuda = ball_query_dense(0.1, 17, a.cuda(), a.cuda()).cpu().detach().numpy()
+        for i in range(a.shape[0]):
+            for j in range(a.shape[1]):
+                # Because it is not necessary the same order
+                assert set(res_cpu[i][j]) == set(res_cuda[i][j])
 
 
 if __name__ == "__main__":

torch_points/torchpoints.py

@@ -252,17 +252,10 @@ def forward(ctx, radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xyz=No
         if new_xyz.is_cuda:
             return tpcuda.ball_query_dense(new_xyz, xyz, radius, nsample)
         else:
-            b = xyz.size(0)
-            npoints = new_xyz.size(1)
-            n = xyz.size(1)
-            batch_new_xyz = torch.arange(0, b, dtype=torch.long).repeat(npoints, 1).T.reshape(-1)
-            batch_xyz = torch.arange(0, b, dtype=torch.long).repeat(n, 1).T.reshape(-1)
-            ind, dist = tpcpu.batch_ball_query(new_xyz.view(-1, 3),
-                                               xyz.view(-1, 3),
-                                               batch_new_xyz,
-                                               batch_xyz,
-                                               radius, nsample)
-            return ind.view(b, npoints, nsample)
+            ind, dist = tpcpu.dense_ball_query(new_xyz,
+                                               xyz,
+                                               radius, nsample, mode=0)
+            return ind
 
     @staticmethod
     def backward(ctx, a=None):
@@ -294,12 +287,16 @@ class BallQueryPartialDense(Function):
     def forward(ctx, radius, nsample, x, y, batch_x, batch_y):
         # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
         if x.is_cuda:
-                return tpcuda.ball_query_partial_dense(x, y,
-                                                       batch_x,
-                                                       batch_y,
-                                                       radius, nsample)
+            return tpcuda.ball_query_partial_dense(x, y,
+                                                   batch_x,
+                                                   batch_y,
+                                                   radius, nsample)
         else:
-            raise NotImplementedError
+            ind, dist = tpcpu.dense_ball_query(x, y,
+                                               batch_x,
+                                               batch_y,
+                                               radius, nsample, mode=0)
+            return ind, dist
 
     @staticmethod
     def backward(ctx, a=None):
@@ -319,9 +316,9 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y):
     y : torch.Tensor
         (N, npoint, 3) centers of the ball query
     batch_x : torch.Tensor
-        (M, ) Contains indexes to indicate within batch it belongs to. 
+        (M, ) Contains indexes to indicate within batch it belongs to.
     batch_y : torch.Tensor
-        (N, ) Contains indexes to indicate within batch it belongs to  
+        (N, ) Contains indexes to indicate within batch it belongs to
 
     Returns
     -------
@@ -347,4 +344,4 @@ def ball_query(radius: float, nsample: int, x, y, batch_x=None, batch_y=None, mo
             raise Exception('batch_x and batch_y should not be provided')
         return ball_query_dense(radius, nsample, x, y)
     else:
-        raise Exception('unrecognized mode {}'.format(mode))
+        raise Exception('unrecognized mode {}'.format(mode))