fix occgrid func

valgur · Oct 6, 2020 · 5314a7f · 5314a7f
1 parent ac5ca36
commit 5314a7f
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 2 deletions.
diff --git a/src/cupoch/geometry/occupancygrid.cu b/src/cupoch/geometry/occupancygrid.cu
@@ -271,15 +271,15 @@ OccupancyGrid& OccupancyGrid::Clear() {
 Eigen::Vector3f OccupancyGrid::GetMinBound() const {
     return (min_bound_.cast<int>() - Eigen::Vector3i::Constant(resolution_ / 2))
                            .cast<float>() *
-                   voxel_size_ -
+                   voxel_size_ +
            origin_;
 }
 
 Eigen::Vector3f OccupancyGrid::GetMaxBound() const {
     return (max_bound_.cast<int>() -
             Eigen::Vector3i::Constant(resolution_ / 2 - 1))
                            .cast<float>() *
-                   voxel_size_ -
+                   voxel_size_ +
            origin_;
 }
 
@@ -363,6 +363,27 @@ OccupancyGrid& OccupancyGrid::Reconstruct(float voxel_size, int resolution) {
     return *this;
 }
 
+OccupancyGrid& OccupancyGrid::SetFreeArea(const Eigen::Vector3f& min_bound,
+                                          const Eigen::Vector3f& max_bound) {
+    const Eigen::Vector3i half_res = Eigen::Vector3i::Constant(resolution_ / 2);
+    Eigen::Vector3i imin_bound = ((min_bound - origin_) / voxel_size_).array().floor().matrix().cast<int>() + half_res;
+    Eigen::Vector3i imax_bound = ((max_bound - origin_) / voxel_size_).array().floor().matrix().cast<int>() + half_res;
+    Eigen::Vector3i diff = imax_bound - imin_bound + Eigen::Vector3i::Ones();
+    min_bound_ = imin_bound.cast<unsigned short>();
+    max_bound_ = imax_bound.cast<unsigned short>();
+    extract_range_voxels_functor func(diff, resolution_,
+                                      imin_bound);
+    thrust::for_each(thrust::make_permutation_iterator(voxels_.begin(),
+                            thrust::make_transform_iterator(thrust::make_counting_iterator<size_t>(0), func)),
+                     thrust::make_permutation_iterator(voxels_.begin(),
+                            thrust::make_transform_iterator(thrust::make_counting_iterator<size_t>(diff[0] * diff[1] * diff[2]), func)),
+                     [pml = prob_miss_log_] __device__ (geometry::OccupancyVoxel& v) {
+                         v.prob_log_ = (isnan(v.prob_log_)) ? 0 : v.prob_log_;
+                         v.prob_log_ += pml;
+                     });
+    return *this;
+}
+
 OccupancyGrid& OccupancyGrid::Insert(
         const utility::device_vector<Eigen::Vector3f>& points,
         const Eigen::Vector3f& viewpoint,

diff --git a/src/cupoch/geometry/occupancygrid.h b/src/cupoch/geometry/occupancygrid.h
@@ -99,6 +99,8 @@ class OccupancyGrid : public DenseGrid<OccupancyVoxel> {
     ExtractOccupiedVoxels() const;
 
     OccupancyGrid& Reconstruct(float voxel_size, int resolution);
+    OccupancyGrid& SetFreeArea(const Eigen::Vector3f& min_bound,
+                               const Eigen::Vector3f& max_bound);
 
     OccupancyGrid& Insert(const utility::device_vector<Eigen::Vector3f>& points,
                           const Eigen::Vector3f& viewpoint,

diff --git a/src/python/cupoch_pybind/geometry/occupancygrid.cpp b/src/python/cupoch_pybind/geometry/occupancygrid.cpp
@@ -106,6 +106,7 @@ void pybind_occupanygrid(py::module &m) {
                          &geometry::OccupancyGrid::Insert),
                  "Function to insert occupancy grid from pointcloud.",
                  "pointcloud"_a, "viewpoint"_a, "max_range"_a = -1.0)
+            .def("set_free_area", &geometry::OccupancyGrid::SetFreeArea)
             .def_readwrite("voxel_size", &geometry::OccupancyGrid::voxel_size_)
             .def_readwrite("resolution", &geometry::OccupancyGrid::resolution_)
             .def_readwrite("origin", &geometry::OccupancyGrid::origin_)

diff --git a/src/tests/geometry/occupancygrid.cpp b/src/tests/geometry/occupancygrid.cpp
@@ -82,4 +82,11 @@ TEST(OccupancyGrid, Insert) {
     EXPECT_TRUE(thrust::get<0>(res4));
     auto res5 = occupancy_grid->GetVoxel(Eigen::Vector3f(0.0, 0.0, 4.5));
     EXPECT_FALSE(thrust::get<0>(res5));
+}
+
+TEST(OccupancyGrid, SetFreeArea) {
+    auto occupancy_grid = std::make_shared<geometry::OccupancyGrid>();
+    occupancy_grid->SetFreeArea(Eigen::Vector3f(0, 0, 0), Eigen::Vector3f(0.1, 0.1, 0.1));
+    auto res = occupancy_grid->ExtractFreeVoxels();
+    EXPECT_EQ(res->size(), 27);
 }