cleaning up according to flake8; c++ api now working (except options)

Author: liaopeiyuan

csrc/cpu/radius_cpu.cpp

@@ -2,25 +2,26 @@
 
 #include "utils.h"
 
-torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support, 
+torch::Tensor radius_cpu(torch::Tensor q, torch::Tensor s, 
              torch::Tensor ptr_x, torch::Tensor ptr_y, 
 			 float radius, int max_num){
 
-	CHECK_CPU(query);
-	CHECK_CPU(support);
+	CHECK_CPU(q);
+	CHECK_CPU(s);
 
 	/*
 	x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
     y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
-	auto batch_x = ptr_x.clone();
-	auto batch_y = ptr_y.clone();
+	*/
 
-	batch_x._mul(2*radius);
-	batch_y._mul(2*radius);
+	auto batch_x = ptr_x.clone().reshape({-1, 1});
+	auto batch_y = ptr_y.clone().reshape({-1, 1});
 
-	auto query = torch::cat({query,batch_x},-1);
-	auto support = torch::cat({support,batch_y},-1);
-	*/
+	batch_x.mul_(2*radius);
+	batch_y.mul_(2*radius);
+
+	auto query = torch::cat({q,batch_x},-1);
+	auto support = torch::cat({s,batch_y},-1);
 
 	torch::Tensor out;
 	std::vector<long> neighbors_indices;
@@ -43,9 +44,18 @@ torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
 	});
 
 	long* neighbors_indices_ptr = neighbors_indices.data();
-	out = torch::from_blob(neighbors_indices_ptr, {neighbors_indices.size()/2, 2}, options=options);
 
-	return out.t().clone();
+	const long long tsize = static_cast<long long>(neighbors_indices.size()/2);
+	out = torch::from_blob(neighbors_indices_ptr, {tsize, 2}, options=options);
+	out = out.t();
+
+	auto result = torch::zeros_like(out);
+
+	auto index = torch::tensor({0,1});
+
+	result.index_copy_(0, index, out);
+
+	return result;
 }
 
 void get_size_batch(const vector<long>& batch, vector<long>& res){

csrc/cpu/radius_cpu.h

@@ -1,46 +1,11 @@
 #pragma once
 
 #include <torch/extension.h>
-#include <ATen/ATen.h>
 #include "utils/neighbors.h"
 #include "utils/neighbors.cpp"
 #include <iostream>
 #include "compat.h"
 
-torch::Tensor radius_cpu(torch::Tensor query,
-			 torch::Tensor support,torch::Tensor ptr_x,
-                     torch::Tensor ptr_y, 
-			 float radius, int max_num);
-/*
-using namespace pybind11::literals;
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-	m.def("radius_search",
-	      &radius_search,
-	      "compute the radius search of a point cloud using nanoflann"
-	      "-query : a pytorch tensor of size N1 x d,. used to query the nearest neighbors"
-	      "- support : a pytorch tensor of size N2 x d. 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.",
-	      "query"_a, "support"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
-	m.def("batch_radius_search",
-	      &batch_radius_search,
-		"compute the radius search of a point cloud for each batch using nanoflann"
-		"-query : a pytorch tensor (float) of size N1 x d,. used to query the nearest neighbors"
-		"- support : a pytorch tensor(float) of size N2 x d. used to build the tree"
-		"- query_batch : a pytorch tensor(long) contains indices of the batch of the query size N1"
-	      "NB : the batch must be sorted"
-		"- support_batch: a pytorch tensor(long) contains indices of the batch of the support size N2"
-	      "NB: the batch must be sorted"
-		"-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 wrt the radius will be computed)."
-		"- mode : int number that indicate which format for the neighborhood"
-		"mode=0 mean a matrix of neighbors(N2 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.",
-	      "query"_a, "support"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
-}
-*/
+torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
+                         torch::Tensor ptr_x, torch::Tensor ptr_y, 
+			 			 float radius, int max_num);

csrc/cpu/utils/cloud.h

@@ -26,14 +26,14 @@ struct PointCloud
 
 		// pts = std::vector<Point>((Point*)new_pts, (Point*)new_pts+new_pts.size()/3);
 		std::vector<std::vector<scalar_t>> temp(new_pts.size()/dim);
-		for(unsigned int i=0; i < new_pts.size(); i++){
+		for(size_t i=0; i < new_pts.size(); i++){
 			if(i%dim == 0){
 
 				//Point point;
 				std::vector<scalar_t> point(dim);
 				//std::vector<scalar_t> vect(sizeof(scalar_t)*dim, 0)
 				//point.pt = temp;
-				for (unsigned int j = 0; j < dim; j++) {
+				for (size_t j = 0; j < (size_t)dim; j++) {
 					point[j]=new_pts[i+j];
 					//point.pt[j] = new_pts[i+j];
 				}
@@ -47,9 +47,8 @@ struct PointCloud
 		std::vector<std::vector<scalar_t>> temp(size);
 		for(int i=0; i < size; i++){
 			//std::vector<scalar_t> temp(sizeof(scalar_t)*dim, 0);
-			//point.pt = temp;
 			std::vector<scalar_t> point(dim);
-			for (unsigned int j = 0; j < dim; j++) {
+			for (size_t j = 0; j < (size_t)dim; j++) {
 					point[j] = new_pts[dim*(begin+i)+j];
 			}
 

csrc/cpu/utils/neighbors.cpp

@@ -1,5 +1,5 @@
 
-// Taken from https://github.com/HuguesTHOMAS/KPConv
+// 3D Version https://github.com/HuguesTHOMAS/KPConv
 
 #include "neighbors.h"
 
@@ -10,16 +10,10 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 	// Initiate variables
 	// ******************
 
-	// square radius
-
 	const scalar_t search_radius = static_cast<scalar_t>(radius*radius);
 
-	// indices
-	int i0 = 0;
-
 	// Counting vector
-	int max_count = 1;
-	float d2;
+	size_t max_count = 1;
 
 	// Nanoflann related variables
 	// ***************************
@@ -32,7 +26,7 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 	pcd_query.set(queries, dim);
 
 	// Tree parameters
-	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
+	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(15 /* max leaf */);
 
 	// KDTree type definition
 	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
@@ -51,6 +45,9 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 
 	float eps = 0.00001;
 
+	// indices
+	size_t i0 = 0;
+
 	for (auto& p0 : pcd_query.pts){
 
 		// Find neighbors
@@ -62,7 +59,6 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 
 		list_matches[i0].reserve(max_count);
 		std::vector<std::pair<size_t, scalar_t> > ret_matches;
-		//nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
 
 		const size_t nMatches = index->radiusSearch(&query_pt[0], search_radius+eps, ret_matches, search_params);
 		list_matches[i0] = ret_matches;
@@ -84,19 +80,19 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 		max_count = max_num;
 	}
 
-	int size = 0; // total number of edges
+	size_t size = 0; // total number of edges
 	for (auto& inds : list_matches){
 		if(inds.size() <= max_count)
 			size += inds.size();
 		else
 			size += max_count;
 	}
-	
+
 	neighbors_indices.resize(size*2);
-	int i1 = 0; // index of the query points
-	int u = 0; // curent index of the neighbors_indices
+	size_t i1 = 0; // index of the query points
+	size_t u = 0; // curent index of the neighbors_indices
 	for (auto& inds : list_matches){
-		for (int j = 0; j < max_count; j++){
+		for (size_t j = 0; j < max_count; j++){
 			if(j < inds.size()){
 				neighbors_indices[u] = inds[j].first;
 				neighbors_indices[u + 1] = i1;
@@ -111,165 +107,4 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
 
 
 
-}
-template<typename scalar_t>
-int batch_nanoflann_neighbors (vector<scalar_t>& queries,
-                               vector<scalar_t>& supports,
-                               vector<long>& q_batches,
-                               vector<long>& s_batches,
-                               vector<long>& neighbors_indices,
-                               float radius, int dim, int max_num,
-			       int mode){
-
-
-// Initiate variables
-// ******************
-// indices
-	int i0 = 0;
-
-// Square radius
-	float r2 = radius * radius;
-
-	// Counting vector
-	int max_count = 0;
-	float d2;
-
-
-	// batch index
-	long b = 0;
-	long sum_qb = 0;
-	long sum_sb = 0;
-
-	// Nanoflann related variables
-	// ***************************
-
-	// CLoud variable
-	PointCloud<scalar_t> current_cloud;
-	PointCloud<scalar_t> query_pcd;
-	query_pcd.set(queries, dim);
-	vector<vector<pair<size_t, scalar_t> > > all_inds_dists(query_pcd.pts.size());
-
-	// Tree parameters
-	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
-
-	// KDTree type definition
-	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
-
-// Pointer to trees
-	my_kd_tree_t* index;
-    // Build KDTree for the first batch element
-	current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
-	index = new my_kd_tree_t(dim, current_cloud, tree_params);
-	index->buildIndex();
-// Search neigbors indices
-// ***********************
-// Search params
-	nanoflann::SearchParams search_params;
-	search_params.sorted = true;
-
-	for (auto& p0 : query_pcd.pts){
-// Check if we changed batch
-
-		scalar_t query_pt[dim];
-		std::copy(p0.begin(), p0.end(), query_pt); 
-
-		/*
-		std::cout << "\n ========== \n";
-		for(int i=0; i < dim; i++)
-			std::cout << query_pt[i] << '\n';
-		std::cout << "\n ========== \n";
-		*/
-	
-		if (i0 == sum_qb + q_batches[b]){
-			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], dim);
-// Build KDTree of the current element of the batch
-			delete index;
-			index = new my_kd_tree_t(dim, current_cloud, tree_params);
-			index->buildIndex();
-		}
-// Initial guess of neighbors size
-		all_inds_dists[i0].reserve(max_count);
-// Find neighbors
-		size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params);
-// Update max count
-
-		std::vector<std::pair<size_t, float> > indices_dists;
-		nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
-
-		index->findNeighbors(resultSet, query_pt, search_params);
-
-		if (nMatches > max_count)
-			max_count = nMatches;
-// Increment query idx
-		i0++;
-	}
-	// how many neighbors do we keep
-	if(max_num > 0) {
-		max_count = max_num;
-	}
-// Reserve the memory
-	if(mode == 0){
-		neighbors_indices.resize(query_pcd.pts.size() * max_count);
-		i0 = 0;
-		sum_sb = 0;
-		sum_qb = 0;
-		b = 0;
-
-		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];
-				b++;
-			}
-
-			for (int j = 0; j < max_count; j++){
-				if (j < inds_dists.size())
-					neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb;
-				else
-					neighbors_indices[i0 * max_count + j] = supports.size();
-
-			}
-
-			i0++;
-		}
-		delete index;
-	}
-	else if(mode == 1){
-		int size = 0; // total number of edges
-		for (auto& inds_dists : all_inds_dists){
-			if(inds_dists.size() <= max_count)
-				size += inds_dists.size();
-			else
-				size += max_count;
-		}
-		neighbors_indices.resize(size * 2);
-		i0 = 0;
-		sum_sb = 0;
-		sum_qb = 0;
-		b = 0;
-		int u = 0;
-		for (auto& inds_dists : all_inds_dists){
-			if (i0 == sum_qb + q_batches[b]){
-				sum_qb += q_batches[b];
-				sum_sb += s_batches[b];
-				b++;
-			}
-			for (int j = 0; j < max_count; j++){
-				if (j < inds_dists.size()){
-					neighbors_indices[u] = inds_dists[j].first + sum_sb;
-					neighbors_indices[u + 1] = i0;
-					u += 2;
-				}
-			}
-			i0++;
-		}
-	}
-	return max_count;
-}
+}

csrc/cpu/utils/neighbors.h

@@ -10,13 +10,4 @@ using namespace std;
 
 template<typename scalar_t>
 int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
-			vector<long>& neighbors_indices, float radius, int dim, int max_num, int mode);
-
-
-template<typename scalar_t>
-int batch_nanoflann_neighbors(vector<scalar_t>& queries,
-                                vector<scalar_t>& supports,
-                                vector<long>& q_batches,
-                                vector<long>& s_batches,
-                                vector<long>& neighbors_indices,
-			        float radius, int dim, int max_num, int mode);
+			vector<long>& neighbors_indices, float radius, int dim, int max_num, int mode);