kdtree cuda

valgur · Nov 22, 2019 · 2dfd5cc · 2dfd5cc
1 parent 9fd8171
commit 2dfd5cc
Show file tree

Hide file tree

Showing 22 changed files with 2,403 additions and 220 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -106,7 +106,13 @@ set(CUDA_ARCH -gencode arch=compute_30,code=sm_30
               -gencode arch=compute_60,code=sm_60
               -gencode arch=compute_61,code=sm_61
               -gencode arch=compute_70,code=sm_70)
-set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --expt-relaxed-constexpr --expt-extended-lambda ${CUDA_ARCH})
+set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS}
+    --expt-relaxed-constexpr
+    --expt-extended-lambda
+    -Xcudafe "--diag_suppress=integer_sign_change"
+    -Xcudafe "--diag_suppress=partial_override"
+    -Xcudafe "--diag_suppress=virtual_function_decl_hidden"
+    ${CUDA_ARCH})
 
 add_subdirectory(third_party)
 

diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 neka-nat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/src/cupoc/geometry/CMakeLists.txt b/src/cupoc/geometry/CMakeLists.txt
@@ -1,2 +1,6 @@
-cuda_add_library(cupoc_geometry pointcloud.cu estimate_normals.cu geometry3d.cu kdtree_flann.cu)
-target_link_libraries(cupoc_geometry cupoc_utility)
+cuda_add_library(cupoc_geometry pointcloud.cu
+                 estimate_normals.cu
+                 down_sample.cu
+                 geometry3d.cu
+                 kdtree_flann.cu)
+target_link_libraries(cupoc_geometry cupoc_utility flann_cuda)
diff --git a/src/cupoc/geometry/down_sample.cu b/src/cupoc/geometry/down_sample.cu
@@ -0,0 +1,136 @@
+#include "cupoc/geometry/pointcloud.h"
+#include "cupoc/geometry/geometry3d.h"
+#include "cupoc/utility/console.h"
+#include "cupoc/utility/helper.h"
+#include <thrust/gather.h>
+
+
+using namespace cupoc;
+using namespace cupoc::geometry;
+
+namespace {
+
+struct compute_key_functor {
+    compute_key_functor(const Eigen::Vector3f& voxel_min_bound, float voxel_size)
+        : voxel_min_bound_(voxel_min_bound), voxel_size_(voxel_size) {};
+    const Eigen::Vector3f voxel_min_bound_;
+    const float voxel_size_;
+    __device__
+    Eigen::Vector3i operator()(const Eigen::Vector3f_u& pt) {
+        auto ref_coord = (pt - voxel_min_bound_) / voxel_size_;
+        return Eigen::Vector3i(int(floor(ref_coord(0))), int(floor(ref_coord(1))), int(floor(ref_coord(2))));
+    }
+};
+
+template<typename OutputIterator, class... Args>
+__host__
+int CalcAverageByKey(thrust::device_vector<Eigen::Vector3i>& keys,
+                     OutputIterator buf_begins, OutputIterator output_begins) {
+    const size_t n = keys.size();
+    thrust::sort_by_key(keys.begin(), keys.end(), buf_begins);
+
+    thrust::device_vector<Eigen::Vector3i> keys_out(n);
+    thrust::device_vector<int> counts(n);
+    auto end1 = thrust::reduce_by_key(keys.begin(), keys.end(),
+                                      thrust::make_constant_iterator(1),
+                                      keys_out.begin(), counts.begin());
+    int n_out = static_cast<int>(end1.second - counts.begin());
+    counts.resize(n_out);
+
+    thrust::equal_to<Eigen::Vector3i> binary_pred;
+    add_tuple_functor<Args...> add_func;
+    auto end2 = thrust::reduce_by_key(keys.begin(), keys.end(), buf_begins,
+                                      keys_out.begin(), output_begins,
+                                      binary_pred, add_func);
+
+    devided_tuple_functor<Args...> dv_func;
+    thrust::transform(output_begins, output_begins + n_out,
+                      counts.begin(), output_begins,
+                      dv_func);
+    return n_out;
+}
+
+}
+
+utility::shared_ptr<PointCloud> PointCloud::SelectDownSample(const thrust::device_vector<size_t> &indices, bool invert) const {
+    auto output = utility::shared_ptr<PointCloud>(new PointCloud());
+    const bool has_normals = HasNormals();
+    const bool has_colors = HasColors();
+
+    output->points_.resize(indices.size());
+    thrust::gather(indices.begin(), indices.end(), points_.begin(), output->points_.begin());
+    return output;
+}
+
+utility::shared_ptr<PointCloud> PointCloud::VoxelDownSample(float voxel_size) const {
+    auto output = utility::shared_ptr<PointCloud>(new PointCloud());
+    if (voxel_size <= 0.0) {
+        utility::LogWarning("[VoxelDownSample] voxel_size <= 0.\n");
+        return output;
+    }
+
+    const Eigen::Vector3f voxel_size3 = Eigen::Vector3f(voxel_size, voxel_size, voxel_size);
+    const Eigen::Vector3f voxel_min_bound = GetMinBound() - voxel_size3 * 0.5;
+    const Eigen::Vector3f voxel_max_bound = GetMaxBound() + voxel_size3 * 0.5;
+
+    if (voxel_size * std::numeric_limits<int>::max() < (voxel_max_bound - voxel_min_bound).maxCoeff()) {
+        utility::LogWarning("[VoxelDownSample] voxel_size is too small.\n");
+        return output;
+    }
+
+    const int n = points_.size();
+    const bool has_normals = HasNormals();
+    const bool has_colors = HasColors();
+    compute_key_functor ck_func(voxel_min_bound, voxel_size);
+    thrust::device_vector<Eigen::Vector3i> keys(n);
+    thrust::transform(points_.begin(), points_.end(), keys.begin(), ck_func);
+
+    thrust::device_vector<Eigen::Vector3f_u> sorted_points = points_;
+    output->points_.resize(n);
+    if (!has_normals && !has_colors) {
+        typedef thrust::tuple<thrust::device_vector<Eigen::Vector3f_u>::iterator> IteratorTuple;
+        typedef thrust::zip_iterator<IteratorTuple> ZipIterator;
+        auto n_out = CalcAverageByKey<ZipIterator, Eigen::Vector3f_u>(keys,
+                    thrust::make_zip_iterator(thrust::make_tuple(sorted_points.begin())),
+                    thrust::make_zip_iterator(thrust::make_tuple(output->points_.begin())));
+        output->points_.resize(n_out);
+    } else if (has_normals && !has_colors) {
+        thrust::device_vector<Eigen::Vector3f_u> sorted_normals = normals_;
+        output->normals_.resize(n);
+        typedef thrust::tuple<thrust::device_vector<Eigen::Vector3f_u>::iterator, thrust::device_vector<Eigen::Vector3f_u>::iterator> IteratorTuple;
+        typedef thrust::zip_iterator<IteratorTuple> ZipIterator;
+        auto n_out = CalcAverageByKey<ZipIterator, Eigen::Vector3f_u, Eigen::Vector3f_u>(keys,
+                    thrust::make_zip_iterator(thrust::make_tuple(sorted_points.begin(), sorted_normals.begin())),
+                    thrust::make_zip_iterator(thrust::make_tuple(output->points_.begin(), output->normals_.begin())));
+        output->points_.resize(n_out);
+        output->normals_.resize(n_out);
+    } else if (!has_normals && has_colors) {
+        thrust::device_vector<Eigen::Vector3f_u> sorted_colors = colors_;
+        output->colors_.resize(n);
+        typedef thrust::tuple<thrust::device_vector<Eigen::Vector3f_u>::iterator, thrust::device_vector<Eigen::Vector3f_u>::iterator> IteratorTuple;
+        typedef thrust::zip_iterator<IteratorTuple> ZipIterator;
+        auto n_out = CalcAverageByKey<ZipIterator, Eigen::Vector3f_u, Eigen::Vector3f_u>(keys,
+                    thrust::make_zip_iterator(thrust::make_tuple(sorted_points.begin(), sorted_colors.begin())),
+                    thrust::make_zip_iterator(thrust::make_tuple(output->points_.begin(), output->colors_.begin())));
+        output->points_.resize(n_out);
+        output->colors_.resize(n_out);
+    } else {
+        thrust::device_vector<Eigen::Vector3f_u> sorted_normals = normals_;
+        thrust::device_vector<Eigen::Vector3f_u> sorted_colors = colors_;
+        output->normals_.resize(n);
+        output->colors_.resize(n);
+        typedef thrust::tuple<thrust::device_vector<Eigen::Vector3f_u>::iterator, thrust::device_vector<Eigen::Vector3f_u>::iterator, thrust::device_vector<Eigen::Vector3f_u>::iterator> IteratorTuple;
+        typedef thrust::zip_iterator<IteratorTuple> ZipIterator;
+        auto n_out = CalcAverageByKey<ZipIterator, Eigen::Vector3f_u, Eigen::Vector3f_u, Eigen::Vector3f_u>(keys,
+                    thrust::make_zip_iterator(thrust::make_tuple(sorted_points.begin(), sorted_normals.begin(), sorted_colors.begin())),
+                    thrust::make_zip_iterator(thrust::make_tuple(output->points_.begin(), output->normals_.begin(), output->colors_.begin())));
+        output->points_.resize(n_out);
+        output->normals_.resize(n_out);
+        output->colors_.resize(n_out);
+    }
+
+    utility::LogDebug(
+            "Pointcloud down sampled from {:d} points to {:d} points.\n",
+            (int)points_.size(), (int)output->points_.size());
+    return output;
+}