Skip to content

Commit

Permalink
Docstring updates for NearestNeighborSearch and compute_metrics (#7081)
Browse files Browse the repository at this point in the history 
* add docstrings for NearestNeighborSearch
* docstring improvements for tensor PointCloud and TriangleMesh
* update contribution recipes with example py docstring
* add sphinx-tabs and sphinx-copybutton
* add missing blank line in docstring for tb summary
  • Loading branch information
@benjaminum
benjaminum authored Dec 11, 2024
1 parent 7354eb1 commit f2461a3
Show file tree
Hide file tree
Showing 10 changed files with 401 additions and 92 deletions.
272 changes: 236 additions & 36 deletions cpp/pybind/core/nns/nearest_neighbor_search.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,35 @@ namespace nns {
void pybind_core_nns_declarations(py::module &m_nns) {
py::class_<NearestNeighborSearch, std::shared_ptr<NearestNeighborSearch>>
nns(m_nns, "NearestNeighborSearch",
"NearestNeighborSearch class for nearest neighbor search. "
"Construct a NearestNeighborSearch object with input "
"dataset_points of shape {n_dataset, d}.");
R"(NearestNeighborSearch class for nearest neighbor search.
This class holds multiple index types to accelerate various nearest neighbor
search operations for a dataset of points with shape {n,d} with `n` as the number
of points and `d` as the dimension of the points. The class supports knn search,
fixed-radius search, multi-radius search, and hybrid search.
Example:
The following example demonstrates how to perform knn search using the class::
import open3d as o3d
import numpy as np
dataset = np.random.rand(10,3)
query_points = np.random.rand(5,3)
nns = o3d.core.nns.NearestNeighborSearch(dataset)
# initialize the knn_index before we can use knn_search
nns.knn_index()
# perform knn search to get the 3 closest points with respect to the
# Euclidean distance. The returned distance is given as the squared
# distances
indices, squared_distances = nns.knn_search(query_points, knn=3)
)");
}

void pybind_core_nns_definitions(py::module &m_nns) {
auto nns = static_cast<py::class_<NearestNeighborSearch,
std::shared_ptr<NearestNeighborSearch>>>(
Expand All @@ -35,7 +60,14 @@ void pybind_core_nns_definitions(py::module &m_nns) {

// Index functions.
nns.def("knn_index", &NearestNeighborSearch::KnnIndex,
"Set index for knn search.");
R"(Initialize the index for knn search.
This function needs to be called once before performing search operations.
Returns:
True on success.
)");

nns.def(
"fixed_radius_index",
[](NearestNeighborSearch &self, utility::optional<double> radius) {
Expand All @@ -45,9 +77,28 @@ void pybind_core_nns_definitions(py::module &m_nns) {
return self.FixedRadiusIndex(radius.value());
}
},
py::arg("radius") = py::none());
py::arg("radius") = py::none(),
R"(Initialize the index for fixed-radius search.
This function needs to be called once before performing search operations.
Args:
radius (float, optional): Radius value for fixed-radius search. Required
for GPU fixed radius index.
Returns:
True on success.
)");

nns.def("multi_radius_index", &NearestNeighborSearch::MultiRadiusIndex,
"Set index for multi-radius search.");
R"(Initialize the index for multi-radius search.
This function needs to be called once before performing search operations.
Returns:
True on success.
)");

nns.def(
"hybrid_index",
[](NearestNeighborSearch &self, utility::optional<double> radius) {
Expand All @@ -57,11 +108,57 @@ void pybind_core_nns_definitions(py::module &m_nns) {
return self.HybridIndex(radius.value());
}
},
py::arg("radius") = py::none());
py::arg("radius") = py::none(),
R"(Initialize the index for hybrid search.
This function needs to be called once before performing search operations.
Args:
radius (float, optional): Radius value for hybrid search. Required
for GPU hybrid index.
Returns:
True on success.
)");

// Search functions.
nns.def("knn_search", &NearestNeighborSearch::KnnSearch, "query_points"_a,
"knn"_a, "Perform knn search.");
"knn"_a,
R"(Perform knn search.
Note:
To use knn_search initialize the index using knn_index before calling this function.
Args:
query_points (open3d.core.Tensor): Query points with shape {n, d}.
knn (int): Number of neighbors to search per query point.
Example:
The following searches the 3 nearest neighbors for random dataset and query points::
import open3d as o3d
import numpy as np
dataset = np.random.rand(10,3)
query_points = np.random.rand(5,3)
nns = o3d.core.nns.NearestNeighborSearch(dataset)
# initialize the knn_index before we can use knn_search
nns.knn_index()
# perform knn search to get the 3 closest points with respect to the
# Euclidean distance. The returned distance is given as the squared
# distances
indices, squared_distances = nns.knn_search(query_points, knn=3)
Returns:
Tuple of Tensors (indices, squared_distances).
- indices: Tensor of shape {n, knn}.
- squared_distances: Tensor of shape {n, knn}. The distances are squared L2 distances.
)");

nns.def(
"fixed_radius_search",
[](NearestNeighborSearch &self, Tensor query_points, double radius,
Expand All @@ -74,38 +171,141 @@ void pybind_core_nns_definitions(py::module &m_nns) {
}
},
py::arg("query_points"), py::arg("radius"),
py::arg("sort") = py::none());
py::arg("sort") = py::none(),
R"(Perform fixed-radius search.
Note:
To use fixed_radius_search initialize the index using fixed_radius_index before calling this function.
Args:
query_points (open3d.core.Tensor): Query points with shape {n, d}.
radius (float): Radius value for fixed-radius search. Note that this
parameter can differ from the radius used to initialize the index
for convenience, which may cause the index to be rebuilt for GPU
devices.
sort (bool, optional): Sort the results by distance. Default is True.
Returns:
Tuple of Tensors (indices, splits, distances).
- indices: The indices of the neighbors.
- distances: The squared L2 distances.
- splits: The splits of the indices and distances defining the start
and exclusive end of each query point's neighbors. The shape is {num_queries+1}
Example:
The following searches the neighbors within a radius of 1.0 a set of data and query points::
# define data and query points
points = np.array([
[0.1,0.1,0.1],
[0.9,0.9,0.9],
[0.5,0.5,0.5],
[1.7,1.7,1.7],
[1.8,1.8,1.8],
[0.3,2.4,1.4]], dtype=np.float32)
queries = np.array([
[1.0,1.0,1.0],
[0.5,2.0,2.0],
[0.5,2.1,2.1],
[100,100,100],
], dtype=np.float32)
nns = o3d.core.nns.NearestNeighborSearch(points)
nns.fixed_radius_index(radius=1.0)
neighbors_index, neighbors_distance, neighbors_splits = nns.fixed_radius_search(queries, radius=1.0, sort=True)
# returns neighbors_index = [1, 2, 5, 5]
# neighbors_distance = [0.03 0.75 0.56000006 0.62]
# neighbors_splits = [0, 2, 3, 4, 4]
for i, (start_i, end_i) in enumerate(zip(neighbors_splits, neighbors_splits[1:])):
start_i = start_i.item()
end_i = end_i.item()
print(f"query_point {i} has the neighbors {neighbors_index[start_i:end_i].numpy()} "
f"with squared distances {neighbors_distance[start_i:end_i].numpy()}")
)");

nns.def("multi_radius_search", &NearestNeighborSearch::MultiRadiusSearch,
"query_points"_a, "radii"_a,
"Perform multi-radius search. Each query point has an independent "
"radius.");
R"(Perform multi-radius search. Each query point has an independent radius.
Note:
To use multi_radius_search initialize the index using multi_radius_index before calling this function.
Args:
query_points (open3d.core.Tensor): Query points with shape {n, d}.
radii (open3d.core.Tensor): Radii of query points. Each query point has one radius.
Returns:
Tuple of Tensors (indices, splits, distances).
- indices: The indices of the neighbors.
- distances: The squared L2 distances.
- splits: The splits of the indices and distances defining the start
and exclusive end of each query point's neighbors. The shape is {num_queries+1}
Example:
The following searches the neighbors with an individual radius for each query point::
# define data, query points and radii
points = np.array([
[0.1,0.1,0.1],
[0.9,0.9,0.9],
[0.5,0.5,0.5],
[1.7,1.7,1.7],
[1.8,1.8,1.8],
[0.3,2.4,1.4]], dtype=np.float32)
queries = np.array([
[1.0,1.0,1.0],
[0.5,2.0,2.0],
[0.5,2.1,2.1],
[100,100,100],
], dtype=np.float32)
radii = np.array([0.5, 1.0, 1.5, 2], dtype=np.float32)
nns = o3d.core.nns.NearestNeighborSearch(points)
nns.multi_radius_index()
neighbors_index, neighbors_distance, neighbors_splits = nns.multi_radius_search(queries, radii)
# returns neighbors_index = [1 5 5 3 4]
# neighbors_distance = [0.03 0.56 0.62 1.76 1.87]
# neighbors_splits = [0 1 2 5 5]
for i, (start_i, end_i) in enumerate(zip(neighbors_splits, neighbors_splits[1:])):
start_i = start_i.item()
end_i = end_i.item()
print(f"query_point {i} has the neighbors {neighbors_index[start_i:end_i].numpy()} "
f"with squared distances {neighbors_distance[start_i:end_i].numpy()}")
)");

nns.def("hybrid_search", &NearestNeighborSearch::HybridSearch,
"query_points"_a, "radius"_a, "max_knn"_a,
"Perform hybrid search.");

// Docstrings.
static const std::unordered_map<std::string, std::string>
map_nearest_neighbor_search_method_docs = {
{"query_points", "The query tensor of shape {n_query, d}."},
{"radii",
"Tensor of shape {n_query,} containing multiple radii, "
"one for each query point."},
{"radius", "Radius value for radius search."},
{"max_knn",
"Maximum number of neighbors to search per query point."},
{"knn", "Number of neighbors to search per query point."}};
docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
"knn_search",
map_nearest_neighbor_search_method_docs);
docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
"multi_radius_search",
map_nearest_neighbor_search_method_docs);
docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
"fixed_radius_search",
map_nearest_neighbor_search_method_docs);
docstring::ClassMethodDocInject(m_nns, "NearestNeighborSearch",
"hybrid_search",
map_nearest_neighbor_search_method_docs);
R"(Perform hybrid search.
Hybrid search behaves similarly to fixed-radius search, but with a maximum number of neighbors to search per query point.
Note:
To use hybrid_search initialize the index using hybrid_index before calling this function.
Args:
query_points (open3d.core.Tensor): Query points with shape {n, d}.
radius (float): Radius value for hybrid search.
max_knn (int): Maximum number of neighbor to search per query.
Returns:
Tuple of Tensors (indices, distances, counts)
- indices: The indices of the neighbors with shape {n, max_knn}.
If there are less than max_knn neighbors within the radius then the
last entries are padded with -1.
- distances: The squared L2 distances with shape {n, max_knn}.
If there are less than max_knn neighbors within the radius then the
last entries are padded with 0.
- counts: Counts of neighbour for each query points with shape {n}.
)");
}

} // namespace nns
Expand Down
34 changes: 23 additions & 11 deletions cpp/pybind/t/geometry/pointcloud.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -774,17 +774,29 @@ the partition id for each point.
)");

pointcloud.def(
"compute_metrics", &PointCloud::ComputeMetrics, "pcd2"_a,
"metrics"_a, "params"_a,
R"(Compute various metrics between two point clouds. Currently, Chamfer distance, Hausdorff distance and F-Score <a href="../tutorial/reference.html#Knapitsch2017">[[Knapitsch2017]]</a> are supported. The Chamfer distance is the sum of the mean distance to the nearest neighbor from the points of the first point cloud to the second point cloud. The F-Score at a fixed threshold radius is the harmonic mean of the Precision and Recall. Recall is the percentage of surface points from the first point cloud that have the second point cloud points within the threshold radius, while Precision is the percentage of points from the second point cloud that have the first point cloud points within the threhold radius.
.. math:
\text{Chamfer Distance: } d_{CD}(X,Y) = \frac{1}{|X|}\sum_{i \in X} || x_i - n(x_i, Y) || + \frac{1}{|Y|}\sum_{i \in Y} || y_i - n(y_i, X) ||\\
\text{Hausdorff distance: } d_H(X,Y) = \max \left{ \max_{i \in X} || x_i - n(x_i, Y) ||, \max_{i \in Y} || y_i - n(y_i, X) || \right}\\
\text{Precision: } P(X,Y|d) = \frac{100}{|X|} \sum_{i \in X} || x_i - n(x_i, Y) || < d \\
\text{Recall: } R(X,Y|d) = \frac{100}{|Y|} \sum_{i \in Y} || y_i - n(y_i, X) || < d \\
\text{F-Score: } F(X,Y|d) = \frac{2 P(X,Y|d) R(X,Y|d)}{P(X,Y|d) + R(X,Y|d)} \\
pointcloud.def("compute_metrics", &PointCloud::ComputeMetrics, "pcd2"_a,
"metrics"_a, "params"_a,
R"(Compute various metrics between two point clouds.
Currently, Chamfer distance, Hausdorff distance and F-Score `[Knapitsch2017] <../tutorial/reference.html#Knapitsch2017>`_ are supported.
The Chamfer distance is the sum of the mean distance to the nearest neighbor
from the points of the first point cloud to the second point cloud. The F-Score
at a fixed threshold radius is the harmonic mean of the Precision and Recall.
Recall is the percentage of surface points from the first point cloud that have
the second point cloud points within the threshold radius, while Precision is
the percentage of points from the second point cloud that have the first point
cloud points within the threhold radius.
.. math::
:nowrap:
\begin{align}
\text{Chamfer Distance: } d_{CD}(X,Y) &= \frac{1}{|X|}\sum_{i \in X} || x_i - n(x_i, Y) || + \frac{1}{|Y|}\sum_{i \in Y} || y_i - n(y_i, X) ||\\
\text{Hausdorff distance: } d_H(X,Y) &= \max \left\{ \max_{i \in X} || x_i - n(x_i, Y) ||, \max_{i \in Y} || y_i - n(y_i, X) || \right\}\\
\text{Precision: } P(X,Y|d) &= \frac{100}{|X|} \sum_{i \in X} || x_i - n(x_i, Y) || < d \\
\text{Recall: } R(X,Y|d) &= \frac{100}{|Y|} \sum_{i \in Y} || y_i - n(y_i, X) || < d \\
\text{F-Score: } F(X,Y|d) &= \frac{2 P(X,Y|d) R(X,Y|d)}{P(X,Y|d) + R(X,Y|d)} \\
\end{align}
Args:
pcd2 (t.geometry.PointCloud): Other point cloud to compare with.
Expand Down
Loading

0 comments on commit f2461a3

Please sign in to comment.