[Feature] Find neighbours points within a given radius through Octree

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<bool(const std::shared_ptr<OctreeNode>&,
                                  const std::shared_ptr<OctreeNodeInfo>&)>& 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<size_t> Octree::RadiusSearch(
+        const std::vector<Eigen::Vector3d>& point_cloud_points,
+        const Eigen::Vector3d& point,
+        const double& radius) {
+    std::vector<size_t> node_indices;
+    const double radius2 = radius * radius;
+
+    auto f_radius_search =
+            [&node_indices, &radius2, &point_cloud_points, &point, &radius](
+                    const std::shared_ptr<OctreeNode>& node,
+                    const std::shared_ptr<OctreeNodeInfo>& 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<OctreePointColorLeafNode>(
+                                    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<OctreePointColorLeafNode>(
+                                    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<geometry::VoxelGrid> Octree::ToVoxelGrid() const {
     auto voxel_grid = std::make_shared<geometry::VoxelGrid>();
     voxel_grid->CreateFromOctree(*this);

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<size_t> RadiusSearch(
+            const std::vector<Eigen::Vector3d>& 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;
 

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<= "

cpp/tests/geometry/Octree.cpp

@@ -357,6 +357,79 @@ TEST(Octree, FragmentPLYLocate) {
     }
 }
 
+TEST(Octree, RadiusSearchCube) {
+    std::vector<Eigen::Vector3d> 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<size_t> expected = {0, 1, 2, 4};
+            std::set<size_t> 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<size_t> expected = {0};
+            std::set<size_t> 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<size_t> indices_ =
+            octree.RadiusSearch(pc.points_, pc.GetCenter(), 0.2);
+
+    std::vector<size_t> 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;

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])

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])