diff --git a/cpp/open3d/geometry/Octree.cpp b/cpp/open3d/geometry/Octree.cpp index 23efcb06e81..5266c5ee034 100644 --- a/cpp/open3d/geometry/Octree.cpp +++ b/cpp/open3d/geometry/Octree.cpp @@ -666,6 +666,43 @@ bool Octree::IsPointInBound(const Eigen::Vector3d& point, return rc; } +Octree::NodeSphereRelationValue Octree::NodeSphereRelation( + const Eigen::Vector3d& origin, + const double& half_size, + const Eigen::Vector3d& point, + const double& radius) { + Eigen::Vector3d center_of_cube = origin.array() + half_size; + bool flag = true; + const double radius2 = radius * radius; + for (int dx = -1; dx <= 1; dx += 2) { + for (int dy = -1; dy <= 1; dy += 2) { + for (int dz = -1; dz <= 1; dz += 2) { + Eigen::Vector3d corner = + center_of_cube + Eigen::Vector3d(dx * half_size, + dy * half_size, + dz * half_size); + if ((corner - point).squaredNorm() > radius2) { + flag = false; + break; + } + } + } + } + if (flag) return FULLY_INSIDE; + + // ref: https://www.zhihu.com/question/24251545/answer/27184960 + // Arvo, J. (1990). A simple method for box-sphere intersection + // testing. In Graphics gems (pp. 335-339). + // 1. Move octant space to origin, transform sphere to Quadrant I + // 2. Minimum length from sphere center to cube + double u2 = ((point - center_of_cube).array().abs() - half_size) + .max(0.0) + .square() + .sum(); + + return u2 <= radius2 ? OVERLAP : NO_OVERLAP; +} + void Octree::Traverse( const std::function&, const std::shared_ptr&)>& f) { @@ -745,6 +782,63 @@ Octree::LocateLeafNode(const Eigen::Vector3d& point) const { return std::make_pair(target_leaf_node, target_leaf_node_info); } +std::vector Octree::RadiusSearch( + const std::vector& point_cloud_points, + const Eigen::Vector3d& point, + const double& radius) { + std::vector node_indices; + const double radius2 = radius * radius; + + auto f_radius_search = + [&node_indices, &radius2, &point_cloud_points, &point, &radius]( + const std::shared_ptr& node, + const std::shared_ptr& node_info) -> bool { + bool skip_children = false; + switch (NodeSphereRelation(node_info->origin_, node_info->size_ / 2, + point, radius)) { + case FULLY_INSIDE: + skip_children = true; + if (auto leaf_node = + std::dynamic_pointer_cast( + node)) { + node_indices.reserve(node_indices.size() + + leaf_node->indices_.size()); + node_indices.insert(node_indices.end(), + leaf_node->indices_.begin(), + leaf_node->indices_.end()); + + } else if (auto internal_node = std::dynamic_pointer_cast< + OctreeInternalPointNode>(node)) { + node_indices.reserve(node_indices.size() + + internal_node->indices_.size()); + node_indices.insert(node_indices.end(), + internal_node->indices_.begin(), + internal_node->indices_.end()); + } + break; + case OVERLAP: + if (auto leaf_node = + std::dynamic_pointer_cast( + node)) { + for (auto& index : leaf_node->indices_) { + if ((point_cloud_points[index] - point).squaredNorm() <= + radius2) + node_indices.push_back(index); + } + } + break; + case NO_OVERLAP: + skip_children = true; + break; + } + + return skip_children; + }; + + Traverse(f_radius_search); + return node_indices; +} + std::shared_ptr Octree::ToVoxelGrid() const { auto voxel_grid = std::make_shared(); voxel_grid->CreateFromOctree(*this); diff --git a/cpp/open3d/geometry/Octree.h b/cpp/open3d/geometry/Octree.h index 5b3b6d90c52..ab84c3de70f 100644 --- a/cpp/open3d/geometry/Octree.h +++ b/cpp/open3d/geometry/Octree.h @@ -382,6 +382,17 @@ class Octree : public Geometry3D, public utility::IJsonConvertible { /// \param point Coordinates of the point. LocateLeafNode(const Eigen::Vector3d& point) const; + /// \brief Returns indices of points within search radius + /// + /// \param point_cloud_points Coordinates of pointcloud + /// used to construct the Octree + /// \param point Origin of the search sphere + /// \param radius Radius of the search sphere + std::vector RadiusSearch( + const std::vector& point_cloud_points, + const Eigen::Vector3d& point, + const double& radius); + /// \brief Return true if point within bound, that is, /// origin <= point < origin + size. /// @@ -392,6 +403,24 @@ class Octree : public Geometry3D, public utility::IJsonConvertible { const Eigen::Vector3d& origin, const double& size); + enum NodeSphereRelationValue { + FULLY_INSIDE = 0, // Octree node inside radius search ball + OVERLAP = 1, + NO_OVERLAP = 2, + }; + + /// \brief Check the spatial relationship between + /// Octree Node and search sphere + /// \param origin Coordinate of the of minimum point of the node cube + /// \param size Size of the octree node cube + /// \param point Origin of the search sphere + /// \param radius Radius of the search sphere + static NodeSphereRelationValue NodeSphereRelation( + const Eigen::Vector3d& origin, + const double& size, + const Eigen::Vector3d& point, + const double& radius); + /// Returns true if the Octree is completely the same, used for testing. bool operator==(const Octree& other) const; diff --git a/cpp/pybind/geometry/octree.cpp b/cpp/pybind/geometry/octree.cpp index 64001b7efab..2c9c2d43ff6 100644 --- a/cpp/pybind/geometry/octree.cpp +++ b/cpp/pybind/geometry/octree.cpp @@ -335,6 +335,9 @@ void pybind_octree_definitions(py::module &m) { .def("locate_leaf_node", &Octree::LocateLeafNode, "point"_a, "Returns leaf OctreeNode and OctreeNodeInfo where the query" "point should reside.") + .def("radius_search", &Octree::RadiusSearch, "point_cloud_points"_a, + "sphere_origin"_a, "radius"_a, + "Returns indices of points within search radius.") .def_static("is_point_in_bound", &Octree::IsPointInBound, "point"_a, "origin"_a, "size"_a, "Return true if point within bound, that is, origin<= " diff --git a/cpp/tests/geometry/Octree.cpp b/cpp/tests/geometry/Octree.cpp index 9ed6f4cf33f..9a468e82f58 100644 --- a/cpp/tests/geometry/Octree.cpp +++ b/cpp/tests/geometry/Octree.cpp @@ -357,6 +357,79 @@ TEST(Octree, FragmentPLYLocate) { } } +TEST(Octree, RadiusSearchCube) { + std::vector points{ + Eigen::Vector3d(0.5, 0.5, 0.5), Eigen::Vector3d(1.5, 0.5, 0.5), + Eigen::Vector3d(0.5, 1.5, 0.5), Eigen::Vector3d(1.5, 1.5, 0.5), + Eigen::Vector3d(0.5, 0.5, 1.5), Eigen::Vector3d(1.5, 0.5, 1.5), + Eigen::Vector3d(0.5, 1.5, 1.5), Eigen::Vector3d(1.5, 1.5, 1.5), + }; + + for (auto& depth : {1, 2, 3}) { + geometry::Octree octree(depth, Eigen::Vector3d(0, 0, 0), 2); + for (size_t i = 0; i < points.size(); ++i) { + octree.InsertPoint( + points[i], + geometry::OctreePointColorLeafNode::GetInitFunction(), + geometry::OctreePointColorLeafNode::GetUpdateFunction( + i, Eigen::Vector3d(0, 0, 0)), + geometry::OctreeInternalPointNode::GetInitFunction(), + geometry::OctreeInternalPointNode::GetUpdateFunction(i)); + } + + { + auto indices = octree.RadiusSearch( + points, Eigen::Vector3d(0.5, 0.5, 0.5), 0.1); + ASSERT_EQ(indices.size(), 1); + EXPECT_EQ(indices[0], 0); + } + + { + auto indices = octree.RadiusSearch( + points, Eigen::Vector3d(0.5, 0.5, 0.5), 1.0); + std::set expected = {0, 1, 2, 4}; + std::set actual(indices.begin(), indices.end()); + EXPECT_EQ(actual, expected); + } + + { + auto indices = octree.RadiusSearch( + points, Eigen::Vector3d(0.5, 0.5, 0.5), 0.999); + std::set expected = {0}; + std::set actual(indices.begin(), indices.end()); + EXPECT_EQ(actual, expected); + } + + { + auto indices = octree.RadiusSearch( + points, Eigen::Vector3d(2.0, 2.0, 2.0), 0.2); + EXPECT_TRUE(indices.empty()); + } + } +} + +TEST(Octree, RadiusSearchRandom) { + geometry::PointCloud pc; + int size = 1000; + Eigen::Vector3d dmin(0.0, 0.0, 0.0); + Eigen::Vector3d dmax(1.0, 1.0, 1.0); + pc.points_.resize(size); + Rand(pc.points_, dmin, dmax, 0); + + auto octree = geometry::Octree(4); + octree.ConvertFromPointCloud(pc, 0.01); + std::vector indices_ = + octree.RadiusSearch(pc.points_, pc.GetCenter(), 0.2); + + std::vector expected; + expected.reserve(size); + for (int i = 0; i < size; ++i) + if ((pc.points_[i] - pc.GetCenter()).norm() <= 0.2) + expected.push_back(i); + std::sort(indices_.begin(), indices_.end()); + EXPECT_EQ(indices_, expected); +} + TEST(Octree, ConvertFromPointCloudBoundSinglePoint) { geometry::Octree octree(10); geometry::PointCloud pcd; diff --git a/examples/python/geometry/octree_radius_search.py b/examples/python/geometry/octree_radius_search.py new file mode 100644 index 00000000000..ca36018f1f8 --- /dev/null +++ b/examples/python/geometry/octree_radius_search.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# - Open3D: www.open3d.org - +# ---------------------------------------------------------------------------- +# Copyright (c) 2018-2024 www.open3d.org +# SPDX-License-Identifier: MIT +# ---------------------------------------------------------------------------- + +import open3d as o3d +import numpy as np + +if __name__ == "__main__": + N = 2000 + armadillo_data = o3d.data.ArmadilloMesh() + pcd = o3d.io.read_triangle_mesh( + armadillo_data.path).sample_points_poisson_disk(N) + # Fit to unit cube. + pcd.scale(1 / np.max(pcd.get_max_bound() - pcd.get_min_bound()), + center=pcd.get_center()) + pcd.colors = o3d.utility.Vector3dVector(np.random.uniform(0, 1, + size=(N, 3))) + + octree = o3d.geometry.Octree(max_depth=4) + octree.convert_from_point_cloud(pcd, size_expand=0.01) + indices_ = octree.radius_search(pcd.points, pcd.get_center(), 0.2) + new_pcd = pcd.select_by_index(indices_) + + sphere = o3d.geometry.TriangleMesh.create_sphere(0.2) + sphere.translate(pcd.get_center()) + + o3d.visualization.draw([octree, new_pcd, sphere]) diff --git a/python/test/test_octree.py b/python/test/test_octree.py index 27dc6f28a03..5c388d76915 100644 --- a/python/test/test_octree.py +++ b/python/test/test_octree.py @@ -140,3 +140,18 @@ def test_locate_leaf_node(): assert node_info.depth == max_depth # Leaf node's size must match assert node_info.size == octree.size / np.power(2, max_depth) + + +def test_radius_search(): + pcd_data = o3d.data.PLYPointCloud() + pcd = o3d.io.read_point_cloud(pcd_data.path) + + expected = np.linalg.norm(np.array(pcd.points) - pcd.get_center(), axis=-1) + expected = np.where(expected <= 0.2) + + for max_depth in range(1, 5, 2): + octree = o3d.geometry.Octree(max_depth) + octree.convert_from_point_cloud(pcd, 0.01) + + indices_ = octree.radius_search(pcd.points, pcd.get_center(), 0.2) + assert np.all(sorted(indices_) == expected[0])