fix segfault cpu

Author: humanpose1

cpu/include/cloud.h

@@ -57,6 +57,7 @@ struct PointCloud
 		pts = temp;
 	}
 	void set_batch(std::vector<scalar_t> new_pts, int begin, int size){
+
 		std::vector<PointXYZ> temp(size);
 		for(int i=0; i < size; i++){
 			PointXYZ point;

cpu/src/neighbors.cpp

@@ -190,24 +190,30 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
 	search_params.sorted = true;
 	for (auto& p0 : query_pcd.pts){
 // Check if we changed batch
-
-		if (i0 == sum_qb + q_batches[b]){
+		if (i0 == sum_qb + q_batches[b] && b < s_batches.size()){
 			sum_qb += q_batches[b];
 			sum_sb += s_batches[b];
+
 			b++;
 
 // Change the points
 			current_cloud.pts.clear();
 			current_cloud.set_batch(supports, sum_sb, s_batches[b]);
 // Build KDTree of the current element of the batch
 			delete index;
+
 			index = new my_kd_tree_t(3, current_cloud, tree_params);
 			index->buildIndex();
 		}
 // Initial guess of neighbors size
+
+
 		all_inds_dists[i0].reserve(max_count);
 // Find neighbors
+		//std::cerr << p0.x << p0.y << p0.z<<std::endl;
 		scalar_t query_pt[3] = { p0.x, p0.y, p0.z};
+
+
 		size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params);
 // Update max count
 
@@ -221,8 +227,10 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
 		max_count = max_num;
 	}
 // Reserve the memory
+
 	if(mode == 0){
 		neighbors_indices.resize(query_pcd.pts.size() * max_count);
+
 		dists.resize(query_pcd.pts.size() * max_count);
 		i0 = 0;
 		sum_sb = 0;
@@ -231,6 +239,7 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
 
 		for (auto& inds_dists : all_inds_dists){// Check if we changed batch
 
+
 			if (i0 == sum_qb + q_batches[b]){
 				sum_qb += q_batches[b];
 				sum_sb += s_batches[b];

cpu/src/torch_nearest_neighbors.cpp

@@ -61,20 +61,22 @@ std::pair<at::Tensor, at::Tensor> ball_query(at::Tensor query,
 
 void cumsum(const vector<long>& batch, vector<long>& res){
 
-	res.resize(batch[batch.size()-1]-batch[0]+1, 0);
+	res.resize(batch[batch.size()-1]-batch[0]+2, 0);
 	long ind = batch[0];
 	long incr = 1;
-	for(int i=1; i < batch.size(); i++){
-
-		if(batch[i] == ind)
-			incr++;
-		else{
-			res[ind-batch[0]] = incr;
-			incr =1;
-			ind = batch[i];
+	if(res.size() > 1){
+		for(int i=1; i < batch.size(); i++){
+			if(batch[i] == ind)
+				incr++;
+			else{
+				res[ind-batch[0]+1] = incr;
+				incr =1;
+				ind = batch[i];
+			}
 		}
+
 	}
-	res[ind-batch[0]] = incr;
+	res[ind-batch[0]+1] = incr;
 }
 
 std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
@@ -89,9 +91,11 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 	std::vector<long> query_batch_stl = std::vector<long>(data_qb, data_qb+query_batch.size(0));
 	std::vector<long> cumsum_query_batch_stl;
 	cumsum(query_batch_stl, cumsum_query_batch_stl);
+
 	std::vector<long> support_batch_stl = std::vector<long>(data_sb, data_sb+support_batch.size(0));
 	std::vector<long> cumsum_support_batch_stl;
 	cumsum(support_batch_stl, cumsum_support_batch_stl);
+
 	std::vector<long> neighbors_indices;
 
 	auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
@@ -107,6 +111,7 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 	std::vector<scalar_t> supports_stl = std::vector<scalar_t>(data_s,
 								   data_s + support.size(0)*support.size(1));
 
+
 	max_count = batch_nanoflann_neighbors<scalar_t>(queries_stl,
 							supports_stl,
 							cumsum_query_batch_stl,
@@ -117,6 +122,7 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor query,
 							max_num,
 							mode);
 	});
+
 	long* neighbors_indices_ptr = neighbors_indices.data();
 	auto neighbors_dists_ptr = neighbors_dists.data();
 

test/test_ballquerry_partial.py

@@ -1,6 +1,7 @@
 import unittest
 import torch
 from torch_points import ball_query
+from torch_points.points_cpu import ball_query as cpu_ball_query
 from torch_cluster import radius_cuda
 import numpy.testing as npt
 import numpy as np
@@ -26,6 +27,7 @@ def test_simple_gpu(self):
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)
 
+
     def test_simple_cpu(self):
         x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0], [0.1, 0, 0]]).to(torch.float)
         y = torch.tensor([[0, 0, 0]]).to(torch.float)
@@ -46,5 +48,22 @@ def test_simple_cpu(self):
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)
 
+    def test_random_cpu(self):
+        a = torch.randn(1000, 3).to(torch.float)
+        b = torch.randn(1500, 3).to(torch.float)
+        batch_a = torch.randint(1, (1000,)).sort(0)[0].long()
+        batch_b = torch.randint(1, (1500,)).sort(0)[0].long()
+        idx, dist2 = ball_query(1.0, 12, a, b, batch_a, batch_b, mode="PARTIAL_DENSE")
+        idx, dist2 = ball_query(1.0, 12, b, a, batch_b, batch_a, mode="PARTIAL_DENSE")
+        idx = idx.detach().cpu().numpy()
+        dist2 = dist2.detach().cpu().numpy()
+        idx2, _ = cpu_ball_query(a, b, 1.0, 12)
+        print(idx[5], print(idx2[5]))
+
+
+
+
+
+
 if __name__ == "__main__":
     unittest.main()