Merge pull request #9 from nicolas-chaulet/ball_query_inverse

'reverse search'

Author: nicolas-chaulet

cuda/src/ball_query.cpp

@@ -28,8 +28,7 @@ at::Tensor ball_query_dense(at::Tensor new_xyz, at::Tensor xyz, const float radi
     CHECK_CUDA(xyz);
   }
 
-  at::Tensor idx =
-      torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
+  at::Tensor idx = torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
                    at::device(new_xyz.device()).dtype(at::ScalarType::Int));
 
   if (new_xyz.type().is_cuda()) {
@@ -67,11 +66,11 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x,
 		CHECK_CUDA(batch_y);
 	}
 
-	at::Tensor idx = torch::full({x.size(0), nsample}, y.size(0),
-				      at::device(x.device()).dtype(at::ScalarType::Long));
+	at::Tensor idx = torch::full({y.size(0), nsample}, x.size(0),
+				at::device(y.device()).dtype(at::ScalarType::Long));
 
-	at::Tensor dist = torch::full({x.size(0), nsample}, -1,
-				      at::device(x.device()).dtype(at::ScalarType::Float));
+	at::Tensor dist = torch::full({y.size(0), nsample}, -1,
+			    at::device(y.device()).dtype(at::ScalarType::Float));
 
 	cudaSetDevice(x.get_device());
 	auto batch_sizes = (int64_t *)malloc(sizeof(int64_t));

cuda/src/ball_query_gpu.cu

@@ -75,8 +75,8 @@ __global__ void query_ball_point_kernel_partial_dense(int size_x,
 					(x[n_x * 3 + d] - y[n_y * 3 + d]);
 			}
 			if(dist <= radius2){
-				idx_out[n_x * nsample + count] = n_y;
-				dist_out[n_x * nsample + count] = dist;
+				idx_out[n_y * nsample + count] = n_x;
+				dist_out[n_y * nsample + count] = dist;
 				count++;
 		       }
 			if(count >= nsample){

test/test_ballquerry_partial.py

@@ -20,8 +20,8 @@ def test_simple_gpu(self):
         idx = idx.detach().cpu().numpy()
         dist2 = dist2.detach().cpu().numpy()
 
-        idx_answer = np.asarray([[1, 1], [0, 1], [1, 1], [1, 1]])
-        dist2_answer = np.asarray([[-1, -1], [0.01, -1], [-1, -1], [-1, -1]]).astype(np.float32)
+        idx_answer = np.asarray([[1, 4]])
+        dist2_answer = np.asarray([[ 0.0100, -1.0000]]).astype(np.float32)
 
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)

torch_points/torchpoints.py

@@ -9,7 +9,6 @@
     import torch_points.points_cuda as tpcuda
 
 
-
 class FurthestPointSampling(Function):
     @staticmethod
     def forward(ctx, xyz, npoint):
@@ -289,15 +288,16 @@ def ball_query_dense(radius, nsample, xyz, new_xyz):
     """
     return BallQueryDense.apply(radius, nsample, xyz, new_xyz)
 
+
 class BallQueryPartialDense(Function):
     @staticmethod
     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
 
@@ -315,7 +315,7 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y):
     nsample : int
         maximum number of features in the balls
     x : torch.Tensor
-        (M, 3) xyz coordinates of the features
+        (M, 3) xyz coordinates of the features (The neighbours are going to be looked for there)
     y : torch.Tensor
         (N, npoint, 3) centers of the ball query
     batch_x : torch.Tensor
@@ -326,11 +326,12 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y):
     Returns
     -------
     torch.Tensor
-        idx: (M, nsample) Default value: N. It contains the indexes of the element within y at radius distance to x
-        dist2: (M, nsample) Default value: -1. It contains the square distances of the element within y at radius distance to x
+        idx: (N, nsample) Default value: N. It contains the indexes of the element within y at radius distance to x
+        dist2: (N, nsample) Default value: -1. It contains the square distances of the element within y at radius distance to x
     """
     return BallQueryPartialDense.apply(radius, nsample, x, y, batch_x, batch_y)
 
+
 def ball_query(radius: float, nsample: int, x, y, batch_x=None, batch_y=None, mode=None):
     if mode is None:
         raise Exception('The mode should be defined within ["PARTIAL_DENSE | DENSE"]')
@@ -347,4 +348,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))