mirror of https://github.com/koide3/small_gicp.git
improve knn performance of voxelmaps
This commit is contained in:
parent
fccb6195a8
commit
b1d84b0829
|
|
@ -7,6 +7,7 @@
|
|||
#include <Eigen/Geometry>
|
||||
#include <small_gicp/ann/traits.hpp>
|
||||
#include <small_gicp/points/traits.hpp>
|
||||
#include <small_gicp/ann/knn_result.hpp>
|
||||
|
||||
namespace small_gicp {
|
||||
|
||||
|
|
@ -67,21 +68,9 @@ public:
|
|||
return 0;
|
||||
}
|
||||
|
||||
size_t min_index = -1;
|
||||
double min_sq_dist = std::numeric_limits<double>::max();
|
||||
|
||||
for (size_t i = 0; i < points.size(); i++) {
|
||||
const double sq_dist = (points[i] - pt).squaredNorm();
|
||||
if (sq_dist < min_sq_dist) {
|
||||
min_index = i;
|
||||
min_sq_dist = sq_dist;
|
||||
}
|
||||
}
|
||||
|
||||
*k_index = min_index;
|
||||
*k_sq_dist = min_sq_dist;
|
||||
|
||||
return 1;
|
||||
KnnResult<1> result(k_index, k_sq_dist);
|
||||
knn_search(pt, result);
|
||||
return result.num_found();
|
||||
}
|
||||
|
||||
/// @brief Find k nearest neighbors.
|
||||
|
|
@ -90,28 +79,32 @@ public:
|
|||
/// @param k_index Indices of nearest neighbors
|
||||
/// @param k_sq_dist Squared distances to nearest neighbors
|
||||
/// @return Number of found points
|
||||
size_t knn_search(const Eigen::Vector4d& pt, int k, size_t* k_index, double* k_sq_dist) const {
|
||||
size_t knn_search(const Eigen::Vector4d& pt, int k, size_t* k_indices, double* k_sq_dists) const {
|
||||
if (points.empty()) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::priority_queue<std::pair<size_t, double>> queue;
|
||||
KnnResult<-1> result(k_indices, k_sq_dists, k);
|
||||
knn_search(pt, result);
|
||||
return result.num_found();
|
||||
}
|
||||
|
||||
/// @brief Find k nearest neighbors.
|
||||
/// @param pt Query point
|
||||
/// @param k Number of neighbors
|
||||
/// @param k_index Indices of nearest neighbors
|
||||
/// @param k_sq_dist Squared distances to nearest neighbors
|
||||
/// @return Number of found points
|
||||
template <typename Result>
|
||||
void knn_search(const Eigen::Vector4d& pt, Result& result) const {
|
||||
if (points.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < points.size(); i++) {
|
||||
const double sq_dist = (points[i] - pt).squaredNorm();
|
||||
queue.push({i, sq_dist});
|
||||
if (queue.size() > k) {
|
||||
queue.pop();
|
||||
}
|
||||
result.push(i, sq_dist);
|
||||
}
|
||||
|
||||
const size_t n = queue.size();
|
||||
while (!queue.empty()) {
|
||||
k_index[queue.size() - 1] = queue.top().first;
|
||||
k_sq_dist[queue.size() - 1] = queue.top().second;
|
||||
queue.pop();
|
||||
}
|
||||
|
||||
return n;
|
||||
}
|
||||
|
||||
public:
|
||||
|
|
@ -151,6 +144,11 @@ struct Traits<FlatContainer<HasNormals, HasCovs>> {
|
|||
static size_t knn_search(const FlatContainer<HasNormals, HasCovs>& container, const Eigen::Vector4d& pt, size_t k, size_t* k_index, double* k_sq_dist) {
|
||||
return container.knn_search(pt, k, k_index, k_sq_dist);
|
||||
}
|
||||
|
||||
template <typename Result>
|
||||
static void knn_search(const FlatContainer<HasNormals, HasCovs>& container, const Eigen::Vector4d& pt, Result& result) {
|
||||
container.knn_search(pt, result);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace traits
|
||||
|
|
|
|||
|
|
@ -79,6 +79,11 @@ struct Traits<GaussianVoxel> {
|
|||
static size_t knn_search(const GaussianVoxel& voxel, const Eigen::Vector4d& pt, size_t k, size_t* k_index, double* k_sq_dist) {
|
||||
return nearest_neighbor_search(voxel, pt, k_index, k_sq_dist);
|
||||
}
|
||||
|
||||
template <typename Result>
|
||||
static void knn_search(const GaussianVoxel& voxel, const Eigen::Vector4d& pt, Result& result) {
|
||||
result.push(0, (voxel.mean - pt).squaredNorm());
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace traits
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@
|
|||
#include <Eigen/Geometry>
|
||||
|
||||
#include <small_gicp/ann/traits.hpp>
|
||||
#include <small_gicp/ann/knn_result.hpp>
|
||||
#include <small_gicp/ann/flat_container.hpp>
|
||||
#include <small_gicp/points/traits.hpp>
|
||||
#include <small_gicp/util/fast_floor.hpp>
|
||||
|
|
@ -103,13 +104,10 @@ public:
|
|||
const size_t voxel_index = found->second;
|
||||
const auto& voxel = flat_voxels[voxel_index]->second;
|
||||
|
||||
size_t point_index;
|
||||
if (traits::nearest_neighbor_search(voxel, pt, &point_index, sq_dist) == 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
*index = calc_index(voxel_index, point_index);
|
||||
return 1;
|
||||
const auto index_transform = [=](size_t i) { return calc_index(voxel_index, i); };
|
||||
KnnResult<1, decltype(index_transform)> result(index, sq_dist, -1, index_transform);
|
||||
traits::Traits<VoxelContents>::knn_search(voxel, pt, result);
|
||||
return result.num_found();
|
||||
}
|
||||
|
||||
/// @brief Find k nearest neighbors
|
||||
|
|
@ -128,15 +126,10 @@ public:
|
|||
const size_t voxel_index = found->second;
|
||||
const auto& voxel = flat_voxels[voxel_index]->second;
|
||||
|
||||
std::vector<size_t> point_indices(k);
|
||||
std::vector<double> sq_dists(k);
|
||||
const size_t num_found = traits::knn_search(voxel, pt, k, point_indices.data(), sq_dists.data());
|
||||
|
||||
for (size_t i = 0; i < num_found; i++) {
|
||||
k_indices[i] = calc_index(voxel_index, point_indices[i]);
|
||||
k_sq_dists[i] = sq_dists[i];
|
||||
}
|
||||
return num_found;
|
||||
const auto index_transform = [=](size_t i) { return calc_index(voxel_index, i); };
|
||||
KnnResult<-1, decltype(index_transform)> result(k_indices, k_sq_dists, k, index_transform);
|
||||
traits::Traits<VoxelContents>::knn_search(voxel, pt, result);
|
||||
return result.num_found();
|
||||
}
|
||||
|
||||
/// @brief Calculate the global point index from the voxel index and the point index.
|
||||
|
|
|
|||
|
|
@ -22,9 +22,14 @@ public:
|
|||
double epsilon = 0.0; ///< Early termination threshold
|
||||
};
|
||||
|
||||
/// @brief Identity transform (alternative to std::identity in C++20).
|
||||
struct identity_transform {
|
||||
size_t operator()(size_t i) const { return i; }
|
||||
};
|
||||
|
||||
/// @brief K-nearest neighbor search result container.
|
||||
/// @tparam N Number of neighbors to search. If N == -1, the number of neighbors is dynamicaly determined.
|
||||
template <int N>
|
||||
template <int N, typename IndexTransform = identity_transform>
|
||||
struct KnnResult {
|
||||
public:
|
||||
static constexpr size_t INVALID = std::numeric_limits<size_t>::max();
|
||||
|
|
@ -33,7 +38,12 @@ public:
|
|||
/// @param indices Buffer to store indices (must be larger than k=max(N, num_neighbors))
|
||||
/// @param distances Buffer to store distances (must be larger than k=max(N, num_neighbors))
|
||||
/// @param num_neighbors Number of neighbors to search (must be -1 for static case N > 0)
|
||||
explicit KnnResult(size_t* indices, double* distances, int num_neighbors = -1) : capacity(num_neighbors), num_found_neighbors(0), indices(indices), distances(distances) {
|
||||
explicit KnnResult(size_t* indices, double* distances, int num_neighbors = -1, const IndexTransform& index_transform = identity_transform())
|
||||
: index_transform(index_transform),
|
||||
capacity(num_neighbors),
|
||||
num_found_neighbors(0),
|
||||
indices(indices),
|
||||
distances(distances) {
|
||||
if constexpr (N > 0) {
|
||||
if (num_neighbors >= 0) {
|
||||
std::cerr << "warning: Specifying dynamic num_neighbors=" << num_neighbors << " for a static KNN result container (N=" << N << ")" << std::endl;
|
||||
|
|
@ -72,7 +82,7 @@ public:
|
|||
}
|
||||
|
||||
if constexpr (N == 1) {
|
||||
indices[0] = index;
|
||||
indices[0] = index_transform(index);
|
||||
distances[0] = distance;
|
||||
} else {
|
||||
int insert_loc = std::min<int>(num_found_neighbors, buffer_size() - 1);
|
||||
|
|
@ -81,7 +91,7 @@ public:
|
|||
distances[insert_loc] = distances[insert_loc - 1];
|
||||
}
|
||||
|
||||
indices[insert_loc] = index;
|
||||
indices[insert_loc] = index_transform(index);
|
||||
distances[insert_loc] = distance;
|
||||
}
|
||||
|
||||
|
|
@ -89,10 +99,11 @@ public:
|
|||
}
|
||||
|
||||
public:
|
||||
const int capacity; ///< Maximum number of neighbors to search
|
||||
int num_found_neighbors; ///< Number of found neighbors
|
||||
size_t* indices; ///< Indices of neighbors
|
||||
double* distances; ///< Distances to neighbors
|
||||
const IndexTransform index_transform; ///< Point index transformation (e.g., local point index to global point/voxel index)
|
||||
const int capacity; ///< Maximum number of neighbors to search
|
||||
int num_found_neighbors; ///< Number of found neighbors
|
||||
size_t* indices; ///< Indices of neighbors
|
||||
double* distances; ///< Distances to neighbors
|
||||
};
|
||||
|
||||
} // namespace small_gicp
|
||||
|
|
|
|||
Loading…
Reference in New Issue