From f647014f54cf3506f9b042a41972cc1540d0d3a4 Mon Sep 17 00:00:00 2001 From: Muhammad Date: Sun, 16 Aug 2026 19:51:39 +0300 Subject: [PATCH] Add FlannIndex abstract interface and implement NanoFlannIndex subclass (#1744) * Add FlannIndex abstract interface and implement NanoFlannIndex subclass * Refactored: made NanoFlann a new NN type instead of inheriting FlannIndex. Added tests. Vendoring nanoflann.h directly in the repo. RegistrationVis now use NANOFLANN_INDEX_KDTREE_SINGLE (instead of FLANN_INDEX_KDTREE_SINGLE) flann index for 2d points matching. * cleanup comments, added FlannIndex doxygen * Fixing windows tests * updating flaky test * Simplified interface, added flann kdtree single approach selectable by parameters. * RegVis: symmetry of nanoflann for two branches of guess feature matching * cv::BFMatcher baseline * Small cmake optimization FLANN_KDTREE_MEM_OPT only defined for FlannIndex * Refactored where FLANN_KDTREE_MEM_OPT is defined * fixed file name already exist * cleanup * fixup build --------- Co-authored-by: matlabbe --- .github/workflows/cmake-linux.yml | 2 +- .github/workflows/cmake-macos.yml | 2 +- .github/workflows/cmake-windows.yml | 2 +- .github/workflows/coverage.yml | 2 +- CMakeLists.txt | 12 +- corelib/include/rtabmap/core/FlannIndex.h | 149 +- corelib/include/rtabmap/core/Parameters.h | 8 +- .../include/rtabmap/core/RegistrationVis.h | 7 + corelib/include/rtabmap/core/VWDictionary.h | 13 + corelib/src/CMakeLists.txt | 11 +- corelib/src/FlannIndex.cpp | 342 +- corelib/src/RegistrationVis.cpp | 127 +- corelib/src/VWDictionary.cpp | 97 +- corelib/src/nanoflann/NanoFlannIndex.cpp | 535 ++ corelib/src/nanoflann/NanoFlannIndex.h | 159 + corelib/src/nanoflann/nanoflann.h | 4469 +++++++++++++++++ corelib/src/nanoflann/readme.txt | 7 + corelib/src/rtflann/{config.h => config.h.in} | 7 + corelib/src/rtflann/defines.h | 2 +- corelib/test/CMakeLists.txt | 19 + corelib/test/FlannIndexBackends.h | 313 ++ corelib/test/perf_flann_index.cpp | 629 +++ corelib/test/test_flann_index.cpp | 548 ++ corelib/test/test_registrationvis.cpp | 49 +- corelib/test/test_rtabmap_integration.cpp | 3 +- corelib/test/test_vwdictionary.cpp | 153 +- guilib/src/ui/preferencesDialog.ui | 24 +- tools/Matcher/main.cpp | 25 +- 28 files changed, 7537 insertions(+), 179 deletions(-) create mode 100644 corelib/src/nanoflann/NanoFlannIndex.cpp create mode 100644 corelib/src/nanoflann/NanoFlannIndex.h create mode 100644 corelib/src/nanoflann/nanoflann.h create mode 100644 corelib/src/nanoflann/readme.txt rename corelib/src/rtflann/{config.h => config.h.in} (82%) create mode 100644 corelib/test/FlannIndexBackends.h create mode 100644 corelib/test/perf_flann_index.cpp create mode 100644 corelib/test/test_flann_index.cpp diff --git a/.github/workflows/cmake-linux.yml b/.github/workflows/cmake-linux.yml index c9aab0d8..dd4a9dc7 100644 --- a/.github/workflows/cmake-linux.yml +++ b/.github/workflows/cmake-linux.yml @@ -101,4 +101,4 @@ jobs: env: PYTHONNOUSERSITE: 1 run: | - ctest -C ${{env.BUILD_TYPE}} -V \ No newline at end of file + ctest -C ${{env.BUILD_TYPE}} -V -LE performance \ No newline at end of file diff --git a/.github/workflows/cmake-macos.yml b/.github/workflows/cmake-macos.yml index e386e7b9..1f8caac0 100644 --- a/.github/workflows/cmake-macos.yml +++ b/.github/workflows/cmake-macos.yml @@ -124,7 +124,7 @@ jobs: env: PYTHONNOUSERSITE: 1 run: | - ctest -C ${{env.BUILD_TYPE}} -V + ctest -C ${{env.BUILD_TYPE}} -V -LE performance # A ctest SEGFAULT is reported as just "SEGFAULT" with no backtrace, which # makes a crash inside a long integration test invisible -- the log simply diff --git a/.github/workflows/cmake-windows.yml b/.github/workflows/cmake-windows.yml index 2cc18648..d40df3b3 100644 --- a/.github/workflows/cmake-windows.yml +++ b/.github/workflows/cmake-windows.yml @@ -125,7 +125,7 @@ jobs: PYTHONHOME: ${{env.VCPKG_EXPORT_PATH}}/installed/x64-windows-release/tools/python3 PYTHONNOUSERSITE: 1 run: | - ctest -C ${{env.BUILD_TYPE}} -V --timeout 300 + ctest -C ${{env.BUILD_TYPE}} -V --timeout 300 -LE performance - name: Info # Skipped for CUDA: the binary links ZED (sl_zed64.dll -> nvcuvid/nvEncodeAPI64), diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 2a146678..20cd2154 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -128,7 +128,7 @@ jobs: # # The replays still run (and gate) in the cmake-linux / macos / # windows jobs; they are just not measured here. - ctest -V -LE long + ctest -V -LE "long|performance" - name: Generate LCOV report run: | diff --git a/CMakeLists.txt b/CMakeLists.txt index 8b219d42..86019736 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -63,10 +63,6 @@ ELSE() option(FLANN_KDTREE_MEM_OPT "Disable multi-threaded FLANN kd-tree to minimize memory allocations" ON) ENDIF() -IF(FLANN_KDTREE_MEM_OPT) - ADD_DEFINITIONS("-DFLANN_KDTREE_MEM_OPT") -ENDIF(FLANN_KDTREE_MEM_OPT) - IF(WIN32 AND NOT MINGW) ADD_DEFINITIONS("-DNOMINMAX") ADD_DEFINITIONS("-wd4100 -wd4127 -wd4150 -wd4191 -wd4242 -wd4244 -wd4251 -wd4305 -wd4365 -wd4512 -wd4514 -wd4548 -wd4571 -wd4619 -wd4625 -wd4626 -wd4628 -wd4668 -wd4710 -wd4711 -wd4738 -wd4820 -wd4946 -wd4986") @@ -182,6 +178,13 @@ if(ANDROID OR IOS OR CMAKE_SYSTEM_NAME STREQUAL "iOS") set(BUILD_TESTING OFF CACHE BOOL "Build the testing tree." FORCE) endif() +# Performance tests are benchmarks: they report times, memory and recall +# instead of asserting on them, as those depend on the machine. They are built +# by default so they cannot rot, but are labelled "performance" so that a run +# opts in with `ctest -L performance` and CI opts out with `ctest -LE +# performance`. They can also be run directly, e.g. bin/test_flann_index_perf. +OPTION(BUILD_PERF_TESTS "Build the performance tests (see ctest -L performance)" ON) + include(CTest) if(BUILD_TESTING) # Add GTest using FetchContent @@ -1565,6 +1568,7 @@ MESSAGE(STATUS " BUILD_APP = ${BUILD_APP}") MESSAGE(STATUS " BUILD_TOOLS = ${BUILD_TOOLS}") MESSAGE(STATUS " BUILD_EXAMPLES = ${BUILD_EXAMPLES}") MESSAGE(STATUS " BUILD_TESTING = ${BUILD_TESTING}") +MESSAGE(STATUS " BUILD_PERF_TESTS = ${BUILD_PERF_TESTS}") MESSAGE(STATUS " ENABLE_COVERAGE = ${ENABLE_COVERAGE}") MESSAGE(STATUS " ENABLE_FORMAT_ERRORS = ${ENABLE_FORMAT_ERRORS}") IF(NOT WIN32) diff --git a/corelib/include/rtabmap/core/FlannIndex.h b/corelib/include/rtabmap/core/FlannIndex.h index 8e387400..4838a7ba 100644 --- a/corelib/include/rtabmap/core/FlannIndex.h +++ b/corelib/include/rtabmap/core/FlannIndex.h @@ -34,36 +34,117 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. namespace rtabmap { +class NanoFlannIndex; + +/** + * @class FlannIndex + * @brief Nearest neighbor index over a set of features + * + * Wraps the search structures of the vendored rtflann and nanoflann libraries + * behind one interface, the structure being chosen with flann_algorithm_t at + * build time. Used for the visual word dictionary (VWDictionary) and for the + * 2D point searches of visual registration (RegistrationVis). + * + * The features are not copied: the index refers to the matrices it is given and + * keeps them alive, cv::Mat data being reference counted, so they must not be + * modified in place while it is in use. Every point it holds is + * designated by an index, assigned in the order the points were added and + * stable for the lifetime of the index: removePoint() leaves a hole rather + * than renumbering the points after it. + */ class RTABMAP_CORE_EXPORT FlannIndex { public: - // A forward of the internal enum, indexes should match. See src/rtflann/defines.h + /** + * @enum flann_algorithm_t + * @brief The index structure built by buildIndex() + * + * The values under 8 are forwarded from rtflann's own enum and have to + * match it (see src/rtflann/defines.h); the nanoflann ones are + * rtabmap-specific and kept outside its range (0-7, 254, 255). A value is + * written in the serialized index header and checked back on load, so none + * of them may be renumbered. + * + * The nanoflann structures take float features only (nanoflann has no + * Hamming metric) and search exactly, ignoring "checks". That makes them + * the fastest ones for 2D and 3D points, and the wrong ones for + * descriptors: an exact search visits more and more of the tree as the + * dimension grows, down to being as slow as an exhaustive search. Prefer + * the approximate rtflann kd-trees for those. + */ enum flann_algorithm_t { - FLANN_INDEX_LINEAR = 0, - FLANN_INDEX_KDTREE = 1, - FLANN_INDEX_KDTREE_SINGLE = 4, - FLANN_INDEX_LSH = 6, + FLANN_INDEX_LINEAR = 0, ///< Exhaustive search + FLANN_INDEX_KDTREE = 1, ///< 4 randomized kd-trees, searched approximately + FLANN_INDEX_KDTREE_SINGLE = 4, ///< Single kd-tree, searched exactly + FLANN_INDEX_LSH = 6, ///< Locality-Sensitive Hashing (binary descriptors) + + /// nanoflann kd-tree. With a rebalancing factor of 1 it is built once, + /// which is the cheapest to build and to search; over 1 it is the + /// weight-balanced tree accepting addPoints()/removePoint(), which + /// cannot be serialized while some of its points are removed. + NANOFLANN_INDEX_KDTREE_SINGLE = 100, }; FlannIndex(); virtual ~FlannIndex(); + /** @brief Drop the index and everything it holds, back to the state of a new one. */ void release(); + + /** + * @brief Serialize the index, to be given back to loadIndex() + * @param computeChecksum Add a checksum of the indexed features to the + * data, which loadIndex() compares against the features it is given + * @return The serialized index, empty when there is nothing to serialize or + * when the structure in use cannot be + * + * The format depends on the architecture and on the versions of the + * vendored libraries: loadIndex() refuses an index it cannot read, leaving + * it to be rebuilt. + */ std::vector serializeIndex(bool computeChecksum = true) const; + /** @return Number of indexed features, the removed ones excluded. */ size_t indexedFeatures() const; - // return Bytes + /** + * @return Bytes used by the index, the features themselves excluded as + * they are only referred to. + */ size_t memoryUsed() const; - // Note that useDistanceL1 doesn't have any effect if LSH is used + /** + * @brief Build the index over the given features, releasing any previous one + * @param algorithm The structure to build + * @param features One feature per row, CV_32FC1 or, for the rtflann + * structures only, CV_8UC1 for binary descriptors (Hamming distance) + * @param useDistanceL1 Search with the L1 distance instead of L2, ignored + * by LSH and by the binary descriptors + * @param rebalancingFactor Fraction (factor-1)/factor of the index that can + * be left removed before it is rebuilt, e.g. half of it for 2. Set + * to 1 to never rebuild it. + */ void buildIndex( flann_algorithm_t algorithm, const cv::Mat & features, bool useDistanceL1 = false, float rebalancingFactor = 2.0f); - // Return false if the indexData doesn't correspond to expected features used and parameters. + + /** + * @brief Load an index serialized by serializeIndex(), releasing any previous one + * @param indexData The serialized index + * @param algorithm The structure it was built with + * @param features The very same features it was built with, in the same + * order: the index refers to them by their row + * @param useDistanceL1 The distance it was built with + * @param rebalancingFactor See buildIndex(). The serialized data carries the + * one the index was built with, which is deprecated and ignored: + * this one is used instead. + * @param errorMsg Filled with what didn't match when the index is refused + * @return False if the data doesn't correspond to the given features and + * parameters, in which case the index is left released + */ bool loadIndex( const std::vector & indexData, flann_algorithm_t algorithm, @@ -71,6 +152,7 @@ public: bool useDistanceL1 = false, float rebalancingFactor = 2.0f, std::string * errorMsg = NULL); + /** @brief Load an index from a raw buffer, see the overload above. */ bool loadIndex( const unsigned char * indexData, size_t indexDataSize, @@ -80,16 +162,46 @@ public: float rebalancingFactor = 2.0f, std::string * errorMsg = NULL); + /** @return Whether an index has been built or loaded. */ bool isBuilt(); + /** @return Type of the indexed features (CV_32FC1 or CV_8UC1). */ int featuresType() const {return featuresType_;} + /** @return Dimension of the indexed features. */ int featuresDim() const {return featuresDim_;} + /** + * @brief Add features to the index + * @param features One feature per row, of the type and dimension the index + * was built with + * @return The index assigned to each of them, empty when the structure + * doesn't accept points after it is built + */ std::vector addPoints(const cv::Mat & features); + /** + * @brief Remove an indexed feature, by the index addPoints() gave for it + * + * The feature is only marked as removed: it is skipped by the searches, but + * keeps taking memory until the index is rebuilt (see the rebalancing + * factor of buildIndex()). Not supported by every structure. + */ void removePoint(unsigned int index); - // return squared distances (indices should be casted in size_t) + /** + * @brief Search the k nearest neighbors of each query + * @param query One feature per row, of the type and dimension the index was + * built with + * @param indices Neighbors found, one query per row, CV_32SC1. The + * neighbors that couldn't be found are set to -1. + * @param dists Their squared distances, CV_32FC1, or CV_32SC1 for the + * Hamming distances of binary descriptors + * @param knn Number of neighbors to search for + * @param checks Number of leaves an approximate search visits, the exact + * structures ignoring it + * @param eps Search for eps-approximate neighbors + * @param sorted Give the neighbors back by increasing distance + */ void knnSearch( const cv::Mat & query, cv::Mat & indices, @@ -99,7 +211,21 @@ public: float eps = 0.0, bool sorted = true) const; - // return squared distances + /** + * @brief Search the neighbors of each query within a radius + * @param query One feature per row, of the type and dimension the index was + * built with + * @param indices Neighbors found, one vector per query + * @param dists Their squared distances, one vector per query + * @param radius Search radius, squared internally: it is a distance, not a + * squared one + * @param maxNeighbors Maximum number of neighbors per query, the nearest + * ones being kept. 0 for all of them. + * @param checks Number of leaves an approximate search visits, the exact + * structures ignoring it + * @param eps Search for eps-approximate neighbors + * @param sorted Give the neighbors back by increasing distance + */ void radiusSearch( const cv::Mat & query, std::vector > & indices, @@ -111,7 +237,8 @@ public: bool sorted = true) const; private: - void * index_; + void * index_; // rtflann backend + NanoFlannIndex * nanoIndex_; // nanoflann backend, only one of the two is set unsigned int nextIndex_; int featuresType_; int featuresDim_; diff --git a/corelib/include/rtabmap/core/Parameters.h b/corelib/include/rtabmap/core/Parameters.h index 202624c2..1a848faf 100644 --- a/corelib/include/rtabmap/core/Parameters.h +++ b/corelib/include/rtabmap/core/Parameters.h @@ -252,10 +252,10 @@ class RTABMAP_CORE_EXPORT Parameters RTABMAP_PARAM(Mem, RotateImagesUpsideUp, bool, false, "Rotate images so that upside is up if they are not already. This can be useful in case the robots don't have all same camera orientation but are using the same map, so that not rotation-invariant visual features can still be used across the fleet."); // KeypointMemory (Keypoint-based) - RTABMAP_PARAM(Kp, NNStrategy, int, 1, "kNNFlannNaive=0, kNNFlannKdTree=1, kNNFlannLSH=2, kNNBruteForce=3, kNNBruteForceGPU=4"); + RTABMAP_PARAM(Kp, NNStrategy, int, 1, "FLANN Linear=0, FLANN KdTree=1, FLANN LSH=2, Brute Force=3, Brute Force GPU=4, FLANN KdTree Single=5, NanoFLANN KdTree=6"); RTABMAP_PARAM(Kp, IncrementalDictionary, bool, true, ""); - RTABMAP_PARAM(Kp, IncrementalFlann, bool, true, uFormat("When using FLANN based strategy, add/remove points to its index without always rebuilding the index (the index is built only when the dictionary increases of the factor \"%s\" in size).", kKpFlannRebalancingFactor().c_str())); - RTABMAP_PARAM(Kp, FlannRebalancingFactor, float, 2.0, uFormat("Factor used when rebuilding the incremental FLANN index (see \"%s\"). Set <=1 to disable.", kKpIncrementalFlann().c_str())); + RTABMAP_PARAM(Kp, IncrementalFlann, bool, true, uFormat("When using FLANN based strategy, add/remove points to its index without always rebuilding the index (the index is only rebuilt when too many of its features have been removed, see \"%s\").", kKpFlannRebalancingFactor().c_str())); + RTABMAP_PARAM(Kp, FlannRebalancingFactor, float, 2.0, uFormat("Rebuild the incremental FLANN index (see \"%s\") once the ratio (factor-1)/factor of its features has been removed, e.g. half of them for a factor of 2. Rebuilding frees the memory of the removed features and speeds up the searches. Features are mostly removed when memory management is enabled (\"%s\" or \"%s\"). Set to 1 to never rebuild, which also uses less memory as the features don't have to be referenced one by one.", kKpIncrementalFlann().c_str(), kRtabmapTimeThr().c_str(), kRtabmapMemoryThr().c_str())); RTABMAP_PARAM(Kp, ByteToFloat, bool, false, uFormat("For %s=1, binary descriptors are converted to float by converting each byte to float instead of converting each bit to float. When converting bytes instead of bits, less memory is used and search is faster at the cost of slightly less accurate matching.", kKpNNStrategy().c_str())); RTABMAP_PARAM(Kp, MaxDepth, float, 0, "Filter extracted keypoints by depth (0=inf)."); RTABMAP_PARAM(Kp, MinDepth, float, 0, "Filter extracted keypoints by depth."); @@ -780,7 +780,7 @@ class RTABMAP_CORE_EXPORT Parameters RTABMAP_PARAM(Vis, GridRows, int, 1, uFormat("Number of rows of the grid used to extract uniformly \"%s / grid cells\" features from each cell.", kVisMaxFeatures().c_str())); RTABMAP_PARAM(Vis, GridCols, int, 1, uFormat("Number of columns of the grid used to extract uniformly \"%s / grid cells\" features from each cell.", kVisMaxFeatures().c_str())); RTABMAP_PARAM(Vis, CorType, int, 0, "Correspondences computation approach: 0=Features Matching, 1=Optical Flow"); - RTABMAP_PARAM(Vis, CorNNType, int, 1, uFormat("[%s=0] kNNFlannNaive=0, kNNFlannKdTree=1, kNNFlannLSH=2, kNNBruteForce=3, kNNBruteForceGPU=4, BruteForceCrossCheck=5, SuperGlue=6, GMS=7. Used for features matching approach.", kVisCorType().c_str())); + RTABMAP_PARAM(Vis, CorNNType, int, 1, uFormat("[%s=0] FLANN Linear=0, FLANN KdTree=1, FLANN LSH=2, Brute Force=3, Brute Force GPU=4, Brute Force Cross Check=5, SuperGlue=6, GMS=7, FLANN KdTree Single=8, NanoFLANN KdTree=9. Used for features matching approach.", kVisCorType().c_str())); RTABMAP_PARAM(Vis, CorNNDR, float, 0.8, uFormat("[%s=0] NNDR: nearest neighbor distance ratio. Used for knn features matching approach.", kVisCorType().c_str())); RTABMAP_PARAM(Vis, CorGuessWinSize, int, 40, uFormat("[%s=0] Matching window size (pixels) around projected points when a guess transform is provided to find correspondences. 0 means disabled.", kVisCorType().c_str())); RTABMAP_PARAM(Vis, CorGuessMatchToProjection, bool, false, uFormat("[%s=0] Match frame's corners to source's projected points (when guess transform is provided) instead of projected points to frame's corners.", kVisCorType().c_str())); diff --git a/corelib/include/rtabmap/core/RegistrationVis.h b/corelib/include/rtabmap/core/RegistrationVis.h index 86dbeb03..b288d0f2 100644 --- a/corelib/include/rtabmap/core/RegistrationVis.h +++ b/corelib/include/rtabmap/core/RegistrationVis.h @@ -75,6 +75,13 @@ public: int getMinInliers() const {return _minInliers;} /** @return **Vis/CorNNType** nearest-neighbor strategy. */ int getNNType() const {return _nnType;} + /** @return Name of the **Vis/CorNNType** nearest-neighbor strategy in use. */ + std::string getNNTypeName() const {return getNNTypeName(_nnType);} + + /** + * @brief Name of a Vis/CorNNType value + */ + static std::string getNNTypeName(int nnType); /** @return **Vis/CorNNDR** ratio test threshold. */ float getNNDR() const {return _nndr;} /** @return **Vis/EstimationType** (0: 3D→3D, 1: PnP, 2: epipolar). */ diff --git a/corelib/include/rtabmap/core/VWDictionary.h b/corelib/include/rtabmap/core/VWDictionary.h index ab6b7b8f..c7963bc2 100644 --- a/corelib/include/rtabmap/core/VWDictionary.h +++ b/corelib/include/rtabmap/core/VWDictionary.h @@ -68,6 +68,11 @@ public: /** * @enum NNStrategy * @brief Nearest neighbor search strategies for descriptor matching + * + * The values are those of the Kp/NNStrategy parameter, saved in user + * configurations and databases: append, never renumber. The Vis/CorNNType + * parameter has strategies of its own, its values are mapped to these ones + * by RegistrationVis::nnStrategyFromCorNNType(). */ enum NNStrategy{ kNNFlannNaive, ///< FLANN naive search (exhaustive) @@ -75,6 +80,8 @@ public: kNNFlannLSH, ///< FLANN Locality-Sensitive Hashing (ideal for binary descriptors) kNNBruteForce, ///< Brute force CPU search kNNBruteForceGPU, ///< Brute force GPU-accelerated search (requires CUDA) + kNNFlannKdTreeSingle, ///< FLANN single exact kd-tree index (rebuilt whenever a word is added, for an index built once) + kNNNanoFlannKdTree, ///< nanoflann kd-tree index (float descriptors only, incremental) kNNUndef ///< Undefined strategy }; @@ -106,11 +113,17 @@ public: return "BRUTE FORCE"; case kNNBruteForceGPU: return "BRUTE FORCE GPU"; + case kNNNanoFlannKdTree: + return "NANOFLANN KD-TREE"; + case kNNFlannKdTreeSingle: + return "FLANN KD-TREE SINGLE"; default: return "Unknown"; } } + + public: /** * @brief Constructor diff --git a/corelib/src/CMakeLists.txt b/corelib/src/CMakeLists.txt index 781c27d8..24bdb8f5 100644 --- a/corelib/src/CMakeLists.txt +++ b/corelib/src/CMakeLists.txt @@ -131,7 +131,8 @@ SET(SRC_FILES rtflann/ext/lz4.c rtflann/ext/lz4hc.c FlannIndex.cpp - + nanoflann/NanoFlannIndex.cpp + #clams stuff clams/discrete_depth_distortion_model_helpers.cpp clams/discrete_depth_distortion_model.cpp @@ -872,6 +873,14 @@ add_definitions(${PCL_DEFINITIONS}) # Add binary that is built from the source file "main.cpp". # The extension is automatically found. +# rtflann's config.h is generated, as it is upstream, so that the +# FLANN_KDTREE_MEM_OPT option travels in a header instead of a compile +# definition: toggling it then recompiles the sources including rtflann rather +# than every object file of the library (see the header for the details). +configure_file( + ${CMAKE_CURRENT_SOURCE_DIR}/rtflann/config.h.in + ${CMAKE_CURRENT_BINARY_DIR}/rtflann/config.h) + ADD_LIBRARY(rtabmap_core ${SRC_FILES} ${RESOURCES_HEADERS}) ADD_LIBRARY(rtabmap::core ALIAS rtabmap_core) diff --git a/corelib/src/FlannIndex.cpp b/corelib/src/FlannIndex.cpp index 2f0317ba..08b89be0 100644 --- a/corelib/src/FlannIndex.cpp +++ b/corelib/src/FlannIndex.cpp @@ -34,12 +34,14 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include #include "rtflann/flann.hpp" +#include "nanoflann/NanoFlannIndex.h" #include namespace rtabmap { FlannIndex::FlannIndex(): index_(0), + nanoIndex_(0), nextIndex_(0), featuresType_(0), featuresDim_(0), @@ -54,6 +56,13 @@ FlannIndex::~FlannIndex() void FlannIndex::release() { + if(nanoIndex_) + { + UDEBUG("Clearing nanoflann index..."); + delete nanoIndex_; + nanoIndex_ = 0; + UDEBUG("Clearing nanoflann index... done!"); + } if(index_) { UDEBUG("Clearing flann index..."); @@ -86,7 +95,112 @@ void FlannIndex::release() #define FLANN_INDEX_HEADER_SIZE 12 +// The rebalancing factor is turned into the fraction of removed features an +// index is allowed to hold before being rebuilt. A factor of 2 used to mean +// "rebuild once the index has doubled in size", it now means "rebuild once half +// of it has been removed": growing an index doesn't degrade it enough to be +// worth a rebuild, removing from it does, as removed features are only marked +// as such and stay in the index until it is rebuilt (see +// corelib/test/test_flann_index.cpp). This is nanoflann's alpha_deleted, which +// both backends now share. +static float removedRatioThreshold(float rebalancingFactor) +{ + if(rebalancingFactor <= 1.0f) + { + return 1.0f; // never rebuilt, a ratio of 1 is never reached + } + return (rebalancingFactor-1.0f)/rebalancingFactor; +} + +template +static bool needsRebuild(const T * index, float removedRatio) +{ + const size_t total = index->size() + index->removedCount(); + return removedRatio < 1.0f && + total > 0 && + float(index->removedCount()) > removedRatio * float(total); +} + +static bool isNanoFlannAlgorithm(FlannIndex::flann_algorithm_t algorithm) +{ + return algorithm == FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE; +} + +static unsigned int computeCrc(const cv::Mat & data) +{ + boost::crc_32_type result; + result.process_bytes(data.data, data.total()*data.elemSize()); + return result.checksum(); +} + +// Fill the header prefixing a serialized index. Shared by both backends: the +// same fields are checked back by loadIndex() whichever one produced the index. +static void fillIndexHeader( + int * header, + int algorithm, + int featuresDim, + bool useDistanceL1, + float rebalancingFactor, // Deprecated + int dataRows, + int dataCols, + int dataType, + unsigned int crcValue, + int indexSize) +{ + int rebalancingFactorAsInt; // Deprecated + memcpy(&rebalancingFactorAsInt, &rebalancingFactor, sizeof(rebalancingFactor)); // Deprecated + int crcValueAsInt; + memcpy(&crcValueAsInt, &crcValue, sizeof(crcValue)); + // Not checked on load: kept so that a later change of the format, adding or + // removing a field, can tell which one it is reading. + header[0] = RTABMAP_VERSION_MAJOR; + header[1] = RTABMAP_VERSION_MINOR; + header[2] = RTABMAP_VERSION_PATCH; + header[3] = algorithm; + header[4] = featuresDim; + header[5] = useDistanceL1?1:0; + header[6] = rebalancingFactorAsInt; // Deprecated + header[7] = dataRows; + header[8] = dataCols; + header[9] = dataType; + header[10] = crcValueAsInt; + header[11] = indexSize; + UDEBUG("Header: \"%d.%d.%d\" alg=%d dim=%d L1=%d factor=%f data(%dx%d type=%d, crc=%X) %d", + header[0],header[1],header[2], + header[3], + header[4], + header[5], + rebalancingFactor, // Deprecated + header[7], header[8], header[9], crcValue, + header[11]); +} + std::vector FlannIndex::serializeIndex(bool computeChecksum) const { + if(nanoIndex_) + { + std::vector nanoIndexData = nanoIndex_->serializeIndex(); + if(!nanoIndexData.empty()) + { + const size_t headerSizeBytes = sizeof(int)*FLANN_INDEX_HEADER_SIZE; + const cv::Mat dataset = nanoIndex_->indexedPoints(); + std::vector indexData(headerSizeBytes + nanoIndexData.size()); + int header[FLANN_INDEX_HEADER_SIZE]; + fillIndexHeader(header, + algorithm_, + featuresDim_, + useDistanceL1_, + rebalancingFactor_, + dataset.rows, + dataset.cols, + dataset.type(), + computeChecksum?computeCrc(dataset):0, + (int)nanoIndexData.size()); + memcpy(indexData.data(), header, headerSizeBytes); + memcpy(indexData.data()+headerSizeBytes, nanoIndexData.data(), nanoIndexData.size()); + return indexData; + } + return std::vector(); + } if(index_ && !addedDescriptors_.empty()) { #ifdef WIN32 @@ -147,6 +261,10 @@ std::vector FlannIndex::serializeIndex(bool computeChecksum) cons if(computeChecksum){ removedDescriptors.insert(removedIndexes_.begin(), removedIndexes_.end()); } + // A descriptor header can cover more than one point (see the end of + // buildIndex() and addPoints()), the index of its row r being + // iter.first+r. addedDescriptors_ is sorted by index, so walking it + // gives the points back in the order they were added. for(const auto & iter: addedDescriptors_) { UASSERT(!iter.second.empty()); @@ -163,59 +281,43 @@ std::vector FlannIndex::serializeIndex(bool computeChecksum) cons else { UASSERT(dataType == iter.second.type()); } - if(computeChecksum){ - if(removedDescriptors.find(iter.first) == removedDescriptors.end()) { - if(dataset.empty()) { - dataset = iter.second.clone(); - } - else { - dataset.push_back(iter.second); - } - } - else { - dataRows -= iter.second.rows; - } - } } - if(!computeChecksum) { - for(const auto & index: removedIndexes_) + // Each removed index is one point, whatever the headers cover. + dataRows -= (int)removedIndexes_.size(); + + if(computeChecksum && dataRows > 0) { + // The checksum is compared against the one of the features + // given back to loadIndex(), so it is computed over the points + // still indexed, in the same order. + dataset.create(dataRows, dataCols, dataType); + int row = 0; + for(const auto & iter: addedDescriptors_) { - dataRows -= addedDescriptors_.at(index).rows; + for(int r=0; r FlannIndex::serializeIndex(bool computeChecksum) cons size_t FlannIndex::indexedFeatures() const { + if(nanoIndex_) + { + return nanoIndex_->indexedFeatures(); + } if(!index_) { return 0; @@ -258,6 +364,10 @@ size_t FlannIndex::indexedFeatures() const // return Bytes size_t FlannIndex::memoryUsed() const { + if(nanoIndex_) + { + return nanoIndex_->memoryUsed(); + } if(!index_) { return 0; @@ -303,6 +413,22 @@ void FlannIndex::buildIndex( rebalancingFactor_ = rebalancingFactor; algorithm_ = algorithm; + if(isNanoFlannAlgorithm(algorithm)) + { + // The tree keeps its own copy of the points and rebuilds itself, so + // addedDescriptors_ is not used here. + nanoIndex_ = new NanoFlannIndex(); + // Nothing to rebuild for a factor of 1: the tree that cannot be added + // to is the cheapest one, and it upgrades itself if points are added + // after all. + nanoIndex_->buildIndex( + features, + useDistanceL1_, + rebalancingFactor_ > 1.0f, + removedRatioThreshold(rebalancingFactor_)); + return; + } + rtflann::IndexParams params; switch (algorithm) @@ -384,8 +510,8 @@ bool FlannIndex::loadIndex( algorithm, features, useDistanceL1, - rebalancingFactor), - error; + rebalancingFactor, + error); } bool FlannIndex::loadIndex( const unsigned char * indexData, @@ -396,16 +522,23 @@ bool FlannIndex::loadIndex( float rebalancingFactor, std::string * error) { - UASSERT(indexData!=NULL); if(indexDataSize == 0) { UWARN("Trying to load empty index...."); + if(error) { + *error = "Trying to load an empty index."; + } return false; } - + UASSERT(indexData!=NULL); #ifdef WIN32 - UERROR("FLANN index deserialization is not yet implemented on Windows. Index cannot be loaded from memory buffer."); - return false; -#else + if(!isNanoFlannAlgorithm(algorithm)) { + UERROR("FLANN index deserialization is not yet implemented on Windows. Index cannot be loaded from memory buffer."); + if(error) { + *error = "FLANN index deserialization is not yet implemented on Windows."; + } + return false; + } +#endif // Check if the features match the expected data from the index size_t headerSizeBytes = sizeof(int)*FLANN_INDEX_HEADER_SIZE; @@ -451,6 +584,15 @@ bool FlannIndex::loadIndex( } return false; } + if(isNanoFlannAlgorithm(algorithm) && (savedRebalancingFactor > 1.0f) != (rebalancingFactor > 1.0f)) { + // The factor is what tells the nanoflann structures apart, and they + // don't serialize to the same thing. + if(error) { + *error = uFormat("Serialized index was built with a rebalancing factor of %f, which doesn't select the same structure as %f.", + savedRebalancingFactor, rebalancingFactor); + } + return false; + } if(savedDistanceL1 != useDistanceL1) { if(error) { *error = uFormat("Serialized \"use distance L1\" (%s) doesn't match the expected one (%s).", savedDistanceL1?"true":"false", useDistanceL1?"true":"false"); @@ -514,6 +656,28 @@ bool FlannIndex::loadIndex( UDEBUG("algorithm=%d", (int)algorithm); + if(isNanoFlannAlgorithm(algorithm)) + { + nanoIndex_ = new NanoFlannIndex(); + if(!nanoIndex_->loadIndex( + features, + useDistanceL1_, + rebalancingFactor_ > 1.0f, + indexData+headerSizeBytes, + indexDataSize-headerSizeBytes, + removedRatioThreshold(rebalancingFactor_), + 10, + error)) + { + this->release(); + return false; + } + return true; + } + +#ifdef WIN32 + return false; // rtflann deserialization is not implemented on Windows, rejected above +#else rtflann::IndexParams params; switch (algorithm) @@ -587,11 +751,15 @@ bool FlannIndex::loadIndex( bool FlannIndex::isBuilt() { - return index_!=0; + return index_!=0 || nanoIndex_!=0; } std::vector FlannIndex::addPoints(const cv::Mat & features) { + if(nanoIndex_) + { + return nanoIndex_->addPoints(features); + } if(!index_) { UERROR("Flann index not yet created!"); @@ -601,16 +769,16 @@ std::vector FlannIndex::addPoints(const cv::Mat & features) UASSERT(features.cols == featuresDim_); bool indexRebuilt = false; size_t removedPts = 0; + const float removedRatio = removedRatioThreshold(rebalancingFactor_); if(featuresType_ == CV_8UC1) { rtflann::Matrix points(features.data, features.rows, features.cols); rtflann::Index > * index = (rtflann::Index >*)index_; removedPts = index->removedCount(); index->addPoints(points, 0); - // Rebuild index if it is now X times in size - if(rebalancingFactor_ > 1.0f && size_t(float(index->sizeAtBuild()) * rebalancingFactor_) < index->size()+index->removedCount()) + if(needsRebuild(index, removedRatio)) { - UDEBUG("Rebuilding FLANN index: %d -> %d", (int)index->sizeAtBuild(), (int)(index->size()+index->removedCount())); + UDEBUG("Rebuilding FLANN index: %d removed of %d", (int)index->removedCount(), (int)(index->size()+index->removedCount())); index->buildIndex(); } // if no more removed points, the index has been rebuilt @@ -624,10 +792,9 @@ std::vector FlannIndex::addPoints(const cv::Mat & features) rtflann::Index > * index = (rtflann::Index >*)index_; removedPts = index->removedCount(); index->addPoints(points, 0); - // Rebuild index if it doubles in size - if(rebalancingFactor_ > 1.0f && size_t(float(index->sizeAtBuild()) * rebalancingFactor_) < index->size()+index->removedCount()) + if(needsRebuild(index, removedRatio)) { - UDEBUG("Rebuilding FLANN index: %d -> %d", (int)index->sizeAtBuild(), (int)(index->size()+index->removedCount())); + UDEBUG("Rebuilding FLANN index: %d removed of %d", (int)index->removedCount(), (int)(index->size()+index->removedCount())); index->buildIndex(); } // if no more removed points, the index has been rebuilt @@ -638,10 +805,9 @@ std::vector FlannIndex::addPoints(const cv::Mat & features) rtflann::Index > * index = (rtflann::Index >*)index_; removedPts = index->removedCount(); index->addPoints(points, 0); - // Rebuild index if it doubles in size - if(rebalancingFactor_ > 1.0f && size_t(float(index->sizeAtBuild()) * rebalancingFactor_) < index->size()+index->removedCount()) + if(needsRebuild(index, removedRatio)) { - UDEBUG("Rebuilding FLANN index: %d -> %d", (int)index->sizeAtBuild(), (int)(index->size()+index->removedCount())); + UDEBUG("Rebuilding FLANN index: %d removed of %d", (int)index->removedCount(), (int)(index->size()+index->removedCount())); index->buildIndex(); } // if no more removed points, the index has been rebuilt @@ -652,10 +818,9 @@ std::vector FlannIndex::addPoints(const cv::Mat & features) rtflann::Index > * index = (rtflann::Index >*)index_; removedPts = index->removedCount(); index->addPoints(points, 0); - // Rebuild index if it doubles in size - if(rebalancingFactor_ > 1.0f && size_t(float(index->sizeAtBuild()) * rebalancingFactor_) < index->size()+index->removedCount()) + if(needsRebuild(index, removedRatio)) { - UDEBUG("Rebuilding FLANN index: %d -> %d", (int)index->sizeAtBuild(), (int)(index->size()+index->removedCount())); + UDEBUG("Rebuilding FLANN index: %d removed of %d", (int)index->removedCount(), (int)(index->size()+index->removedCount())); index->buildIndex(); } // if no more removed points, the index has been rebuilt @@ -674,13 +839,27 @@ std::vector FlannIndex::addPoints(const cv::Mat & features) removedIndexes_.clear(); } - // incremental FLANN: we should add all headers separately in case we remove - // some indexes (to keep underlying matrix data allocated) std::vector indexes; + indexes.reserve(features.rows); for(int i=0; i 1.0f) + { + for(int i=0; i FlannIndex::addPoints(const cv::Mat & features) void FlannIndex::removePoint(unsigned int index) { + if(nanoIndex_) + { + nanoIndex_->removePoint(index); + return; + } if(!index_) { UERROR("Flann index not yet created!"); @@ -728,6 +912,12 @@ void FlannIndex::knnSearch( float eps, bool sorted) const { + if(nanoIndex_) + { + // exact search, "checks", "eps" and "sorted" don't apply + nanoIndex_->knnSearch(query, indices, dists, knn); + return; + } if(!index_) { UERROR("Flann index not yet created!"); @@ -767,10 +957,12 @@ void FlannIndex::knnSearch( indices.create(query.rows, knn, CV_32S); int * ptr = indices.ptr(); - for(size_t i=0 ; i::max()?-1:(int)indicesBuffer[i]; - ptr[i+1] = indicesBuffer[i+1] == std::numeric_limits::max()?-1:(int)indicesBuffer[i+1]; } } @@ -784,6 +976,12 @@ void FlannIndex::radiusSearch( float eps, bool sorted) const { + if(nanoIndex_) + { + // "checks" doesn't apply + nanoIndex_->radiusSearch(query, indices, dists, radius, maxNeighbors, eps, sorted); + return; + } if(!index_) { UERROR("Flann index not yet created!"); diff --git a/corelib/src/RegistrationVis.cpp b/corelib/src/RegistrationVis.cpp index 0b6323a6..b7597c0c 100644 --- a/corelib/src/RegistrationVis.cpp +++ b/corelib/src/RegistrationVis.cpp @@ -38,6 +38,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include #include #include +#include #include #include #include @@ -56,7 +57,6 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include #endif -#include #ifdef RTABMAP_PYTHON @@ -65,6 +65,52 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. namespace rtabmap { +// The dictionary strategy a Vis/CorNNType value stands for. Vis/CorNNType +// shares the values of Kp/NNStrategy for the strategies the dictionary +// implements, and extends them with matching approaches of its own, hence the +// mapping. Return VWDictionary::kNNUndef for the values RegistrationVis handles +// itself (BruteForceCrossCheck, SuperGlue, GMS). +static VWDictionary::NNStrategy nnStrategyFromCorNNType(int nnType) +{ + // 0 to 4 are the dictionary strategies themselves, 5, 6 and 7 are the + // approaches RegistrationVis implements (BruteForceCrossCheck, SuperGlue + // and GMS), and the ones after them are dictionary strategies again, at an + // offset of the three above. + if(nnType >= 0 && nnType <= VWDictionary::kNNBruteForceGPU) + { + return (VWDictionary::NNStrategy)nnType; + } + if(nnType > 7) + { + const int strategy = nnType - 3; + if(strategy < VWDictionary::kNNUndef) + { + return (VWDictionary::NNStrategy)strategy; + } + } + return VWDictionary::kNNUndef; +} + +std::string RegistrationVis::getNNTypeName(int nnType) +{ + const VWDictionary::NNStrategy strategy = nnStrategyFromCorNNType(nnType); + if(strategy != VWDictionary::kNNUndef) + { + return VWDictionary::nnStrategyName(strategy); + } + switch(nnType) + { + case 5: + return "BRUTE FORCE CROSS CHECK"; + case 6: + return "PY MATCHER"; + case 7: + return "GMS"; + default: + return "Unknown"; + } +} + RegistrationVis::RegistrationVis(const ParametersMap & parameters, Registration * child) : Registration(parameters, child), _minInliers(Parameters::defaultVisMinInliers()), @@ -123,6 +169,12 @@ RegistrationVis::RegistrationVis(const ParametersMap & parameters, Registration uInsert(_featureParameters, ParametersPair(Parameters::kKpGridRows(), _featureParameters.at(Parameters::kVisGridRows()))); uInsert(_featureParameters, ParametersPair(Parameters::kKpGridCols(), _featureParameters.at(Parameters::kVisGridCols()))); uInsert(_featureParameters, ParametersPair(Parameters::kKpNewWordsComparedTogether(), "false")); + // The dictionary used to match descriptors (see computeTransformationImpl()) + // is built once and searched once, then thrown away: the words added while + // searching it are never indexed. Nothing is gained by keeping its index + // ready to be added to, and the bookkeeping that needs costs a descriptor + // reference per feature on every registration. + uInsert(_featureParameters, ParametersPair(Parameters::kKpIncrementalFlann(), "false")); this->parseParameters(parameters); } @@ -237,9 +289,10 @@ void RegistrationVis::parseParameters(const ParametersMap & parameters) if(uContains(parameters, Parameters::kVisCorNNType())) { - if(_nnType cornersProjectedMat((float*)cornersProjected.data(), cornersProjected.size(), 2); - rtflann::Index > index(cornersProjectedMat, rtflann::KDTreeIndexParams()); - index.buildIndex(); + // Index the projected keypoints. A rebalancing factor of 1: + // the index is thrown away with the frame, nothing is ever + // added to or removed from it. cv::Point2f being two floats, + // the points are indexed where they are. + cv::Mat cornersProjectedMat((int)cornersProjected.size(), 2, CV_32FC1, (void*)cornersProjected.data()); + FlannIndex flannIndex; + flannIndex.buildIndex(FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, cornersProjectedMat, false, 1.0f); std::vector< std::vector > indices; std::vector > dists; float radius = (float)_guessWinSize; // pixels std::vector pointsTo; cv::KeyPoint::convert(kptsTo, pointsTo); - rtflann::Matrix pointsToMat((float*)pointsTo.data(), pointsTo.size(), 2); - index.radiusSearch(pointsToMat, indices, dists, radius*radius, rtflann::SearchParams()); + cv::Mat pointsToMat((int)pointsTo.size(), 2, CV_32FC1, (void*)pointsTo.data()); + flannIndex.radiusSearch(pointsToMat, indices, dists, radius); - UASSERT(indices.size() == pointsToMat.rows); + UASSERT(indices.size() == (size_t)pointsToMat.rows); UASSERT(descriptorsFrom.cols == descriptorsTo.cols); UASSERT(descriptorsFrom.rows == (int)kptsFrom.size()); UASSERT((int)pointsToMat.rows == descriptorsTo.rows); - UASSERT(pointsToMat.rows == kptsTo.size()); + UASSERT(pointsToMat.rows == (int)kptsTo.size()); UDEBUG("radius search done for guess"); // Process results (Nearest Neighbor Distance Ratio) @@ -1107,9 +1163,21 @@ Transform RegistrationVis::computeTransformationImpl( std::map addedWordsFrom; // std::map duplicates; // int newWords = 0; + // The projected words that a keypoint of the frame was found + // near, as the other branch collects them: several keypoints + // can be near the same one, hence the set. OdometryF2M uses + // them to know which words of its map are still seen. + std::set projectedIDs; cv::Mat descriptors(10, descriptorsTo.cols, descriptorsTo.type()); - for(unsigned int i = 0; i < pointsToMat.rows; ++i) + for(int i = 0; i < pointsToMat.rows; ++i) { + for(unsigned int j=0; j= 2) { @@ -1200,9 +1268,10 @@ Transform RegistrationVis::computeTransformationImpl( ++newWords; } } - UDEBUG("addedWordsFrom=%d/%d (duplicates=%d, newWords=%d), kptsTo=%d, wordsTo=%d, words3From=%d", + info.projectedIDs = std::vector(projectedIDs.begin(), projectedIDs.end()); + UDEBUG("addedWordsFrom=%d/%d (duplicates=%d, newWords=%d), kptsTo=%d, wordsTo=%d, words3From=%d, projectedIDs=%d", (int)addedWordsFrom.size(), (int)cornersProjected.size(), (int)duplicates.size(), newWords, - (int)kptsTo.size(), (int)wordsTo.size(), (int)words3From.size()); + (int)kptsTo.size(), (int)wordsTo.size(), (int)words3From.size(), (int)info.projectedIDs.size()); // create fake ids for not matched words from "from" int addWordsFromNotMatched = 0; @@ -1224,25 +1293,29 @@ Transform RegistrationVis::computeTransformationImpl( else { UDEBUG("match projected to frame"); + // Index the frame's keypoints. A rebalancing factor of 1: + // the index is thrown away with the frame, nothing is ever + // added to or removed from it. cv::Point2f being two floats, + // the points are indexed where they are. std::vector pointsTo; cv::KeyPoint::convert(kptsTo, pointsTo); - rtflann::Matrix pointsToMat((float*)pointsTo.data(), pointsTo.size(), 2); - rtflann::Index > index(pointsToMat, rtflann::KDTreeIndexParams()); - index.buildIndex(); + cv::Mat pointsToMat((int)pointsTo.size(), 2, CV_32FC1, (void*)pointsTo.data()); + FlannIndex flannIndex; + flannIndex.buildIndex(FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, pointsToMat, false, 1.0f); - std::vector< std::vector > indices; - std::vector > dists; + cv::Mat queryMat((int)cornersProjected.size(), 2, CV_32FC1, (void*)cornersProjected.data()); + + std::vector> indices; + std::vector> dists; float radius = (float)_guessWinSize; // pixels - rtflann::Matrix cornersProjectedMat((float*)cornersProjected.data(), cornersProjected.size(), 2); - index.radiusSearch(cornersProjectedMat, indices, dists, radius*radius, rtflann::SearchParams(32, 0, false)); - UASSERT(indices.size() == cornersProjectedMat.rows); - UASSERT(descriptorsFrom.cols == descriptorsTo.cols); - UASSERT(descriptorsFrom.rows == (int)kptsFrom.size()); + flannIndex.radiusSearch(queryMat, indices, dists, radius, 0, 32, 0.0, false); + + UASSERT(indices.size() == cornersProjected.size()); UASSERT((int)pointsToMat.rows == descriptorsTo.rows); - UASSERT(pointsToMat.rows == kptsTo.size()); + UASSERT(pointsToMat.rows == (int)kptsTo.size()); UDEBUG("radius search done for guess"); - + // Process results (Nearest Neighbor Distance Ratio) std::set addedWordsTo; std::set addedWordsFrom; @@ -1250,7 +1323,7 @@ Transform RegistrationVis::computeTransformationImpl( double bruteForceDescCopy = 0.0; UTimer bruteForceTimer; cv::Mat descriptors(10, descriptorsTo.cols, descriptorsTo.type()); - for(unsigned int i = 0; i < cornersProjectedMat.rows; ++i) + for(unsigned int i = 0; i < cornersProjected.size(); ++i) { int matchedIndexFrom = projectedIndexToDescIndex[i]; diff --git a/corelib/src/VWDictionary.cpp b/corelib/src/VWDictionary.cpp index 4a1bee12..d7acf014 100644 --- a/corelib/src/VWDictionary.cpp +++ b/corelib/src/VWDictionary.cpp @@ -56,6 +56,43 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. namespace rtabmap { +// Whether the strategy searches with a FlannIndex, as opposed to the brute +// force ones matching against the _dataTree matrix. +static bool isFlannStrategy(VWDictionary::NNStrategy strategy) +{ + return strategy == VWDictionary::kNNFlannNaive || + strategy == VWDictionary::kNNFlannKdTree || + strategy == VWDictionary::kNNFlannLSH || + strategy == VWDictionary::kNNNanoFlannKdTree || + strategy == VWDictionary::kNNFlannKdTreeSingle; +} + +// Whether the strategy indexes float descriptors in a kd-tree, in which case +// binary descriptors have to be converted first. +static bool isKdTreeStrategy(VWDictionary::NNStrategy strategy) +{ + return strategy == VWDictionary::kNNFlannKdTree || + strategy == VWDictionary::kNNNanoFlannKdTree || + strategy == VWDictionary::kNNFlannKdTreeSingle; +} + +static FlannIndex::flann_algorithm_t flannAlgorithm(VWDictionary::NNStrategy strategy) +{ + switch(strategy) + { + case VWDictionary::kNNFlannNaive: + return FlannIndex::FLANN_INDEX_LINEAR; + case VWDictionary::kNNFlannLSH: + return FlannIndex::FLANN_INDEX_LSH; + case VWDictionary::kNNNanoFlannKdTree: + return FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE; + case VWDictionary::kNNFlannKdTreeSingle: + return FlannIndex::FLANN_INDEX_KDTREE_SINGLE; + default: + return FlannIndex::FLANN_INDEX_KDTREE; // kNNFlannKdTree + } +} + const int VWDictionary::ID_START = 1; const int VWDictionary::ID_INVALID = 0; @@ -115,7 +152,17 @@ void VWDictionary::parseParameters(const ParametersMap & parameters) NNStrategy nnStrategy = (NNStrategy)std::atoi((*iter).second.c_str()); treeUpdated = this->setNNStrategy(nnStrategy); } - if(!treeUpdated && byteToFloat!=_byteToFloat && _strategy == kNNFlannKdTree) + if(_strategy == kNNFlannKdTreeSingle && _incrementalDictionary && _incrementalFlann) + { + UWARN("%s=%d (%s) rebuilds its whole index every time a word is added, which " + "is very slow with %s=true. It is meant for an index built once and " + "searched once, like the one matching the features of two frames.", + Parameters::kKpNNStrategy().c_str(), (int)_strategy, + nnStrategyName(_strategy).c_str(), + Parameters::kKpIncrementalFlann().c_str()); + } + + if(!treeUpdated && byteToFloat!=_byteToFloat && isKdTreeStrategy(_strategy)) { UINFO("KDTree: Binary to Float conversion approach has changed, re-initialize kd-tree."); this->rebuildIndex(); @@ -392,7 +439,7 @@ unsigned long VWDictionary::getMemoryUsed() const memoryUsage += _visualWords.size()*(sizeof(int) + _visualWords.rbegin()->second->getMemoryUsed() + sizeof(std::map::iterator)) + sizeof(std::map); if(_dataTree.empty() && _visualWords.begin()->second->getDescriptor().type() == CV_8U && - _strategy == kNNFlannKdTree) + isKdTreeStrategy(_strategy)) { // Binary descriptors were converted to float, and not included in _dataTree memoryUsage += _visualWords.size() * _visualWords.begin()->second->getDescriptor().total() * sizeof(float) * (_byteToFloat?1:8); @@ -507,7 +554,7 @@ void VWDictionary::update() if(!firstUpdate && _incrementalFlann && - _strategy < kNNBruteForce && + isFlannStrategy(_strategy) && _visualWords.size()) { ULOGGER_DEBUG("Incremental FLANN: Removing %d words...", (int)_removedIndexedWords.size()); @@ -535,7 +582,7 @@ void VWDictionary::update() if(w->getDescriptor().type() == CV_8U) { useDistanceL1_ = true; - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { descriptor = convertBinTo32F(w->getDescriptor(), _byteToFloat); } @@ -555,9 +602,7 @@ void VWDictionary::update() UDEBUG("Building FLANN index... (strategy=%s, byteToFloat=%s, useDistanceL1=%s, rebalancingFactor=%f)", nnStrategyName(_strategy).c_str(), _byteToFloat?"true":"false", useDistanceL1_?"true":"false", _rebalancingFactor); _flannIndex->buildIndex( - _strategy == kNNFlannNaive ? FlannIndex::FLANN_INDEX_LINEAR: - _strategy == kNNFlannLSH ? FlannIndex::FLANN_INDEX_LSH: - FlannIndex::FLANN_INDEX_KDTREE, // kNNFlannKdTree + flannAlgorithm(_strategy), descriptor, useDistanceL1_, _rebalancingFactor); UDEBUG("Building FLANN index... done!"); } @@ -577,7 +622,7 @@ void VWDictionary::update() ULOGGER_DEBUG("Incremental FLANN: Inserting %d words... done! (in %f s)", (int)_notIndexedWords.size(), timer.ticks()); } } - else if(_strategy >= kNNBruteForce && + else if(!isFlannStrategy(_strategy) && _notIndexedWords.size() && _removedIndexedWords.size() == 0 && _visualWords.size()) @@ -587,8 +632,8 @@ void VWDictionary::update() if(_dataTree.rows >= IMGIDX_ONE) { UWARN("%s=%d is not a FLANN strategy and the number of words in the vocabulary (%d) is over %d (IMGIDX_ONE), so opencv may " - "assert on an IMGIDX_ONE check when adding new words. Use a FLANN strategy instead (%s<%d).", - Parameters::kKpNNStrategy().c_str(), _strategy, _dataTree.rows, IMGIDX_ONE, Parameters::kKpNNStrategy().c_str(), kNNBruteForce); + "assert on an IMGIDX_ONE check when adding new words. Use a FLANN strategy instead (e.g. %s=%d).", + Parameters::kKpNNStrategy().c_str(), _strategy, _dataTree.rows, IMGIDX_ONE, Parameters::kKpNNStrategy().c_str(), kNNFlannKdTree); } //just add not indexed words @@ -633,7 +678,7 @@ void VWDictionary::update() if(_visualWords.begin()->second->getDescriptor().type() == CV_8U) { useDistanceL1_ = true; - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { type = CV_32F; if(!_byteToFloat) @@ -662,7 +707,7 @@ void VWDictionary::update() cv::Mat descriptor; if(iter->second->getDescriptor().type() == CV_8U) { - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { descriptor = convertBinTo32F(iter->second->getDescriptor(), _byteToFloat); } @@ -687,12 +732,10 @@ void VWDictionary::update() ULOGGER_DEBUG("_mapIndexId.size() = %d, words.size()=%d, _dim=%d",(int)_mapIndexId.size(), (int)_visualWords.size(), dim); ULOGGER_DEBUG("copying data = %f s", timer.ticks()); - if(_strategy < kNNBruteForce) + if(isFlannStrategy(_strategy)) { _flannIndex->buildIndex( - _strategy == kNNFlannNaive ? FlannIndex::FLANN_INDEX_LINEAR: - _strategy == kNNFlannLSH ? FlannIndex::FLANN_INDEX_LSH: - FlannIndex::FLANN_INDEX_KDTREE, // kNNFlannKdTree + flannAlgorithm(_strategy), _dataTree, useDistanceL1_, _incrementalDictionary&&_incrementalFlann?_rebalancingFactor:1); @@ -714,7 +757,7 @@ void VWDictionary::update() std::vector VWDictionary::serializeIndex() const { - if(_strategy >= kNNBruteForce) { + if(!isFlannStrategy(_strategy)) { UINFO("Not flann strategy, ignoring serialization..."); return std::vector(); } @@ -739,7 +782,7 @@ bool VWDictionary::deserializeIndex(const unsigned char * data, size_t size) return false; } UDEBUG("Loading flann index... (data size=%ld bytes)", size); - if(_strategy >= kNNBruteForce) { + if(!isFlannStrategy(_strategy)) { //ignore return false; } @@ -772,7 +815,7 @@ bool VWDictionary::deserializeIndex(const unsigned char * data, size_t size) if(_visualWords.begin()->second->getDescriptor().type() == CV_8U) { useDistanceL1_ = true; - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { type = CV_32F; if(!_byteToFloat) @@ -801,7 +844,7 @@ bool VWDictionary::deserializeIndex(const unsigned char * data, size_t size) cv::Mat descriptor; if(iter->second->getDescriptor().type() == CV_8U) { - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { descriptor = convertBinTo32F(iter->second->getDescriptor(), _byteToFloat); } @@ -830,9 +873,7 @@ bool VWDictionary::deserializeIndex(const unsigned char * data, size_t size) if(_flannIndex->loadIndex( data, size, - _strategy == kNNFlannNaive ? FlannIndex::FLANN_INDEX_LINEAR: - _strategy == kNNFlannLSH ? FlannIndex::FLANN_INDEX_LSH: - FlannIndex::FLANN_INDEX_KDTREE, + flannAlgorithm(_strategy), dataTree, useDistanceL1_, _incrementalDictionary && _incrementalFlann ? _rebalancingFactor:1, @@ -975,7 +1016,7 @@ std::list VWDictionary::addNewWords( if(descriptorsIn.type() == CV_8U) { useDistanceL1_ = true; - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { descriptors = convertBinTo32F(descriptorsIn, _byteToFloat); } @@ -1031,7 +1072,7 @@ std::list VWDictionary::addNewWords( //Find nearest neighbors UDEBUG("newPts.total()=%d _strategy=%d", descriptors.rows, _strategy); - if(_strategy == kNNFlannNaive || _strategy == kNNFlannKdTree || _strategy == kNNFlannLSH) + if(isFlannStrategy(_strategy)) { _flannIndex->knnSearch(descriptors, results, dists, k, KNN_CHECKS); } @@ -1308,7 +1349,7 @@ std::vector VWDictionary::findNN(const cv::Mat & queryIn) const cv::Mat query; if(queryIn.type() == CV_8U) { - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { query = convertBinTo32F(queryIn, _byteToFloat); } @@ -1353,7 +1394,7 @@ std::vector VWDictionary::findNN(const cv::Mat & queryIn) const //Find nearest neighbors UDEBUG("query.rows=%d ", query.rows); - if(_strategy == kNNFlannNaive || _strategy == kNNFlannKdTree || _strategy == kNNFlannLSH) + if(isFlannStrategy(_strategy)) { _flannIndex->knnSearch(query, results, dists, k, KNN_CHECKS); } @@ -1436,7 +1477,7 @@ std::vector VWDictionary::findNN(const cv::Mat & queryIn) const cv::Mat descriptor; if(vw->getDescriptor().type() == CV_8U) { - if(_strategy == kNNFlannKdTree) + if(isKdTreeStrategy(_strategy)) { descriptor = convertBinTo32F(vw->getDescriptor(), _byteToFloat); } diff --git a/corelib/src/nanoflann/NanoFlannIndex.cpp b/corelib/src/nanoflann/NanoFlannIndex.cpp new file mode 100644 index 00000000..027f151d --- /dev/null +++ b/corelib/src/nanoflann/NanoFlannIndex.cpp @@ -0,0 +1,535 @@ +/* +Copyright (c) 2010-2016, Mathieu Labbe - IntRoLab - Universite de Sherbrooke +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the Universite de Sherbrooke nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*/ + +#include "nanoflann/NanoFlannIndex.h" + +#include +#include + +#include + +#include "nanoflann/nanoflann.h" + +#include + +namespace rtabmap { + +namespace { + +// Points indexed by the tree. nanoflann is zero-copy: it only stores indexes in +// this container, which holds a pointer to the first coordinate of each point +// rather than a copy of it, the way rtflann's NNIndex::points_ does. The +// features they point into are kept alive by "blocks" below. The container is +// append-only so that the indexes handed out by addPoints() stay valid (removed +// points leave a hole behind). +struct PointCloud +{ + std::vector pts; + std::vector blocks; // owners of the rows pts points into + int dim = 0; + + inline size_t kdtree_get_point_count() const {return pts.size();} + inline float kdtree_get_pt(const size_t idx, const size_t d) const {return pts[idx][d];} + template bool kdtree_get_bbox(BBOX & /* bb */) const {return false;} +}; + +} + +// Type-erases the metric and the compile-time dimension of the tree, so that +// nanoflann's templates stay in this file. +class NanoFlannIndexImpl +{ +public: + virtual ~NanoFlannIndexImpl() {} + + // index every point currently in cloud + virtual void buildIndex() = 0; + virtual bool isIncremental() const = 0; + + // [start, end] indexes of points already appended to cloud + virtual void addPoints(size_t start, size_t end) = 0; + virtual void removePoint(size_t index) = 0; + virtual size_t size() const = 0; + virtual size_t usedMemory() const = 0; + virtual size_t knnSearch(const float * query, size_t knn, unsigned int * indices, float * dists) const = 0; + virtual size_t radiusSearch( + const float * query, + float radiusSqr, + std::vector > & matches, + const nanoflann::SearchParameters & params) const = 0; + virtual void saveIndex(std::ostream & stream) const = 0; + // throws std::runtime_error if the stream doesn't match this instantiation + virtual void loadIndex(std::istream & stream) = 0; + + PointCloud cloud; +}; + +namespace { + +template +class NanoFlannTree : public NanoFlannIndexImpl +{ +public: + // cloud is a base class member, so it is already constructed here. The + // incremental tree always starts empty, whatever the dataset holds. + NanoFlannTree(int dim, float alphaDeleted) : + tree_(dim, cloud, nanoflann::KDTreeIncrementalIndexParams(0.75f, alphaDeleted)) {} + + virtual void buildIndex() override + { + const size_t count = cloud.kdtree_get_point_count(); + if(count) + { + tree_.addPoints(0, (unsigned int)(count-1)); + } + } + virtual bool isIncremental() const override {return true;} + virtual void addPoints(size_t start, size_t end) override {tree_.addPoints((unsigned int)start, (unsigned int)end);} + virtual void removePoint(size_t index) override {tree_.removePoint((unsigned int)index);} + virtual size_t size() const override {return tree_.size();} + virtual size_t usedMemory() const override {return tree_.usedMemory();} + virtual size_t knnSearch(const float * query, size_t knn, unsigned int * indices, float * dists) const override + { + return tree_.knnSearch(query, knn, indices, dists); + } + virtual size_t radiusSearch( + const float * query, + float radiusSqr, + std::vector > & matches, + const nanoflann::SearchParameters & params) const override + { + return tree_.radiusSearch(query, radiusSqr, matches, params); + } + virtual void saveIndex(std::ostream & stream) const override {tree_.saveIndex(stream);} + virtual void loadIndex(std::istream & stream) override {tree_.loadIndex(stream);} + +private: + nanoflann::KDTreeSingleIndexIncrementalAdaptor tree_; +}; + +template +class NanoFlannStaticTree : public NanoFlannIndexImpl +{ +public: + // The static tree indexes the dataset as it is when it is built, and cloud + // is still empty here: the initial build is skipped, buildIndex() or + // loadIndex() is called once the points are in. + NanoFlannStaticTree(int dim, size_t leafMaxSize) : + tree_(dim, cloud, nanoflann::KDTreeSingleIndexAdaptorParams( + leafMaxSize, nanoflann::KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex)) {} + + virtual void buildIndex() override {tree_.buildIndex();} + virtual bool isIncremental() const override {return false;} + virtual void addPoints(size_t, size_t) override {UFATAL("Not supported by the static nanoflann index.");} + virtual void removePoint(size_t) override {UFATAL("Not supported by the static nanoflann index.");} + // no removed points to exclude + virtual size_t size() const override {return cloud.kdtree_get_point_count();} + virtual size_t usedMemory() const override {return tree_.usedMemory(tree_);} + virtual size_t knnSearch(const float * query, size_t knn, unsigned int * indices, float * dists) const override + { + return tree_.knnSearch(query, knn, indices, dists); + } + virtual size_t radiusSearch( + const float * query, + float radiusSqr, + std::vector > & matches, + const nanoflann::SearchParameters & params) const override + { + return tree_.radiusSearch(query, radiusSqr, matches, params); + } + virtual void saveIndex(std::ostream & stream) const override {tree_.saveIndex(stream);} + virtual void loadIndex(std::istream & stream) override {tree_.loadIndex(stream);} + +private: + nanoflann::KDTreeSingleIndexAdaptor tree_; +}; + +// L2_Simple is the metric recommended by nanoflann for 2D and 3D point clouds, +// L2 (with its partial distance early exit) for the higher dimensions of the +// descriptors. A compile-time dimension additionally keeps the per-node +// bounding boxes on the stack, so the two point cloud cases are instantiated +// with theirs. +template class Tree, class ... Args> +NanoFlannIndexImpl * createTree(int dim, bool useDistanceL1, Args ... args) +{ + if(useDistanceL1) + { + return new Tree, -1>(dim, args...); + } + if(dim == 2) + { + return new Tree, 2>(dim, args...); + } + if(dim == 3) + { + return new Tree, 3>(dim, args...); + } + return new Tree, -1>(dim, args...); +} + +NanoFlannIndexImpl * createImpl(int dim, bool useDistanceL1, bool incremental, float removedRatio, int leafMaxSize) +{ + if(incremental) + { + // nanoflann's alpha_deleted: the fraction of removed points above which + // a subtree is rebuilt, dropping them. + return createTree(dim, useDistanceL1, removedRatio); + } + UASSERT(leafMaxSize > 0); + return createTree(dim, useDistanceL1, (size_t)leafMaxSize); +} + +} + +NanoFlannIndex::NanoFlannIndex() : + index_(0), + featuresDim_(0), + useDistanceL1_(false), + removedRatio_(0.5f) +{ +} + +NanoFlannIndex::~NanoFlannIndex() +{ + this->release(); +} + +void NanoFlannIndex::release() +{ + delete index_; + index_ = 0; + featuresDim_ = 0; +} + +void NanoFlannIndex::buildIndex( + const cv::Mat & features, + bool useDistanceL1, + bool incremental, + float removedRatio, + int leafMaxSize) +{ + this->release(); + + UASSERT_MSG(features.type() == CV_32FC1, "Only 32F features are supported by the nanoflann index."); + UASSERT(features.cols > 0); + + featuresDim_ = features.cols; + useDistanceL1_ = useDistanceL1; + removedRatio_ = removedRatio; + index_ = createImpl(featuresDim_, useDistanceL1, incremental, removedRatio, leafMaxSize); + index_->cloud.dim = featuresDim_; + + this->appendPoints(features); + index_->buildIndex(); +} + +size_t NanoFlannIndex::indexedFeatures() const +{ + return index_?index_->size():0; +} + +// return Bytes +size_t NanoFlannIndex::memoryUsed() const +{ + if(!index_) + { + return 0; + } + // Like the rtflann backend, the features themselves are not counted: they + // are owned by the caller, only referenced here. + return sizeof(NanoFlannIndex) + + index_->cloud.pts.capacity() * sizeof(const float*) + + index_->cloud.blocks.capacity() * sizeof(cv::Mat) + + index_->usedMemory(); +} + +// Reference the points at the end of the storage, without indexing them, and +// return the index of the first one added. +size_t NanoFlannIndex::appendPoints(const cv::Mat & features) +{ + PointCloud & cloud = index_->cloud; + const size_t start = cloud.pts.size(); + if(cloud.pts.capacity() < start + (size_t)features.rows) + { + // Grow geometrically: reserving exactly what is needed would make every + // single point insertion reallocate and copy the whole storage. + cloud.pts.reserve(std::max(start + (size_t)features.rows, cloud.pts.capacity()*2)); + } + // Keeping the header alive is what keeps the rows valid, cv::Mat data being + // reference counted. One header covers the whole batch. + cloud.blocks.push_back(features); + for(int i=0; i(i)); + } + return start; +} + +// Swap the tree that is built once for the one that accepts points, keeping +// the points already indexed and the indexes they were given. +void NanoFlannIndex::makeIncremental() +{ + UDEBUG("Rebuilding the nanoflann index as an incremental one (%d points)", + (int)index_->cloud.pts.size()); + + const cv::Mat points = this->indexedPoints(); + delete index_; + index_ = createImpl(featuresDim_, useDistanceL1_, true, removedRatio_, 10); + index_->cloud.dim = featuresDim_; + if(!points.empty()) + { + this->appendPoints(points); + index_->buildIndex(); + } +} + +std::vector NanoFlannIndex::addPoints(const cv::Mat & features) +{ + if(!index_) + { + UERROR("Nanoflann index not yet created!"); + return std::vector(); + } + if(!index_->isIncremental()) + { + // Built as the tree that cannot be added to, but points are added after + // all: rebuild it as the one that can. + this->makeIncremental(); + } + UASSERT(features.type() == CV_32FC1); + UASSERT(features.cols == featuresDim_); + + std::vector indexes; + if(features.rows == 0) + { + return indexes; + } + + const size_t start = this->appendPoints(features); + index_->addPoints(start, start + (size_t)features.rows - 1); + + indexes.resize(features.rows); + for(size_t i=0; i NanoFlannIndex::serializeIndex() const +{ + if(!index_) + { + return std::vector(); + } + if(index_->size() != index_->cloud.kdtree_get_point_count()) + { + // The tree indexes holes in the point storage, which the features + // matrix given back to loadIndex() cannot reproduce. + UWARN("Points have been removed from the nanoflann index (%d indexed of %d points), " + "it cannot be serialized before being rebuilt.", + (int)index_->size(), (int)index_->cloud.kdtree_get_point_count()); + return std::vector(); + } + + std::ostringstream stream(std::ios_base::out | std::ios_base::binary); + index_->saveIndex(stream); + const std::string data = stream.str(); + return std::vector(data.begin(), data.end()); +} + +bool NanoFlannIndex::loadIndex( + const cv::Mat & features, + bool useDistanceL1, + bool incremental, + const unsigned char * indexData, + size_t indexDataSize, + float removedRatio, + int leafMaxSize, + std::string * errorMsg) +{ + this->release(); + + UASSERT_MSG(features.type() == CV_32FC1, "Only 32F features are supported by the nanoflann index."); + UASSERT(features.cols > 0); + + if(indexData == 0 || indexDataSize == 0) + { + if(errorMsg) + { + *errorMsg = "Trying to load an empty nanoflann index."; + } + return false; + } + + featuresDim_ = features.cols; + useDistanceL1_ = useDistanceL1; + removedRatio_ = removedRatio; + index_ = createImpl(featuresDim_, useDistanceL1, incremental, removedRatio, leafMaxSize); + index_->cloud.dim = featuresDim_; + this->appendPoints(features); + + // nanoflann checks its own magic number, version and type sizes, and + // throws when the stream wasn't written by the same instantiation. + try + { + std::istringstream stream( + std::string((const char *)indexData, indexDataSize), + std::ios_base::in | std::ios_base::binary); + index_->loadIndex(stream); + } + catch(const std::exception & e) + { + if(errorMsg) + { + *errorMsg = uFormat("Nanoflann index cannot be loaded: %s", e.what()); + } + this->release(); + return false; + } + + if(index_->size() != (size_t)features.rows) + { + if(errorMsg) + { + *errorMsg = uFormat("Serialized nanoflann index has %d points, but %d features were given.", + (int)index_->size(), features.rows); + } + this->release(); + return false; + } + return true; +} + +cv::Mat NanoFlannIndex::indexedPoints() const +{ + if(!index_ || index_->cloud.pts.empty()) + { + return cv::Mat(); + } + // The points are referenced row by row, so a continuous matrix of them has + // to be materialized. Only used to serialize the index. + cv::Mat points((int)index_->cloud.pts.size(), featuresDim_, CV_32FC1); + for(int i=0; i(i), index_->cloud.pts[i], featuresDim_*sizeof(float)); + } + return points; +} + +void NanoFlannIndex::removePoint(unsigned int index) +{ + if(!index_) + { + UERROR("Nanoflann index not yet created!"); + return; + } + if(!index_->isIncremental()) + { + // Same as addPoints(): a tree built without the intention of changing + // it can still be changed. + this->makeIncremental(); + } + // The point stays in cloud so that the indexes of the other points don't + // move, only the tree drops it. + index_->removePoint(index); +} + +void NanoFlannIndex::knnSearch( + const cv::Mat & query, + cv::Mat & indices, + cv::Mat & dists, + int knn) const +{ + if(!index_) + { + UERROR("Nanoflann index not yet created!"); + return; + } + UASSERT(query.type() == CV_32FC1 && query.cols == featuresDim_); + UASSERT(knn > 0); + + indices = cv::Mat(query.rows, knn, CV_32SC1, cv::Scalar(-1)); + dists = cv::Mat(query.rows, knn, CV_32FC1, cv::Scalar(-1.0f)); + + std::vector resultIndices(knn); + std::vector resultDists(knn); + for(int i=0; iknnSearch(query.ptr(i), knn, resultIndices.data(), resultDists.data()); + for(size_t j=0; j(i, j) = (int)resultIndices[j]; + dists.at(i, j) = resultDists[j]; + } + } +} + +void NanoFlannIndex::radiusSearch( + const cv::Mat & query, + std::vector > & indices, + std::vector > & dists, + float radius, + int maxNeighbors, + float eps, + bool sorted) const +{ + if(!index_) + { + UERROR("Nanoflann index not yet created!"); + return; + } + UASSERT(query.type() == CV_32FC1 && query.cols == featuresDim_); + + indices.resize(query.rows); + dists.resize(query.rows); + + // nanoflann compares squared distances, and sorting is required to know + // which neighbors are the closest ones when maxNeighbors is set. + const float radiusSqr = radius * radius; + nanoflann::SearchParameters params(eps, sorted || maxNeighbors>0); + + std::vector > matches; + for(int i=0; iradiusSearch(query.ptr(i), radiusSqr, matches, params); + if(maxNeighbors > 0 && found > (size_t)maxNeighbors) + { + found = (size_t)maxNeighbors; + } + + indices[i].resize(found); + dists[i].resize(found); + for(size_t j=0; j + +namespace rtabmap { + +class NanoFlannIndexImpl; + +/** + * kd-tree backed by nanoflann, held by FlannIndex when its + * NANOFLANN_INDEX_KDTREE_SINGLE algorithm is selected. Two trees are available, + * buildIndex() picking one with its "incremental" argument: + * + * Incremental (nanoflann's KDTreeSingleIndexIncrementalAdaptor), a single + * weight-balanced tree accepting points after it is built: + * - addPoints() inserts incrementally, and bulk-rebuilds when the batch is + * large relative to the tree. + * - removePoint() is lazy; a subtree is rebuilt, dropping its tombstones, once + * its removed fraction gets over "removedRatio". + * Indexes returned by addPoints() stay valid for the lifetime of the index: + * points are only appended and removals never renumber the remaining ones. + * + * Static (nanoflann's KDTreeSingleIndexAdaptor), built once from all the points + * given to buildIndex(): cheaper to build and to search. Use it for a dataset + * known to be fixed, like the image points searched during registration. Adding + * or removing points is still possible, it rebuilds itself as the incremental + * tree when that happens. + * + * Only float features are supported (nanoflann has no Hamming metric, binary + * descriptors have to be converted first), with the L2 or L1 metric. Searches + * are exact, there is no equivalent of rtflann's "checks" budget. + */ +class NanoFlannIndex +{ +public: + NanoFlannIndex(); + ~NanoFlannIndex(); + + NanoFlannIndex(const NanoFlannIndex &) = delete; + NanoFlannIndex & operator=(const NanoFlannIndex &) = delete; + + void release(); + + // features must be a CV_32FC1 matrix, one point per row. "incremental" + // selects the tree accepting addPoints()/removePoint(), for which + // "removedRatio" is the fraction of it that can be left removed before a + // rebuild (1 to never rebuild). "leafMaxSize" is the number of points under + // which the static tree stops splitting, it doesn't apply to the + // incremental one, which holds a single point per node. + // + // A static tree given points to add afterwards is rebuilt as an incremental + // one, so that building without the intention of adding points doesn't + // prevent it (see addPoints()). + void buildIndex( + const cv::Mat & features, + bool useDistanceL1, + bool incremental, + float removedRatio = 0.5f, + int leafMaxSize = 10); + + // Return an empty vector if the index cannot be serialized: when it is not + // built, or when points have been removed from it (the tree then indexes + // holes that the matrix given back to loadIndex() cannot reproduce). + std::vector serializeIndex() const; + + // features must hold the very same points, in the same order, than those + // that were indexed when the index was serialized. + bool loadIndex( + const cv::Mat & features, + bool useDistanceL1, + bool incremental, + const unsigned char * indexData, + size_t indexDataSize, + float removedRatio = 0.5f, + int leafMaxSize = 10, + std::string * errorMsg = 0); + + // The indexed points as an indexedFeatures()x"dim" CV_32FC1 matrix, copied + // out of the features they are referenced from. Empty if the index is not + // built. Note that points removed from the tree are still part of it. + cv::Mat indexedPoints() const; + + bool isBuilt() const {return index_ != 0;} + + // removed points excluded + size_t indexedFeatures() const; + + // return Bytes + size_t memoryUsed() const; + + // return the index assigned to each added point + std::vector addPoints(const cv::Mat & features); + + void removePoint(unsigned int index); + + // return squared distances, indices and distances are set to -1 for the + // neighbors that couldn't be found. + void knnSearch( + const cv::Mat & query, + cv::Mat & indices, + cv::Mat & dists, + int knn) const; + + // return squared distances + void radiusSearch( + const cv::Mat & query, + std::vector > & indices, + std::vector > & dists, + float radius, + int maxNeighbors, + float eps, + bool sorted) const; + +private: + // The metric (L2 or L1) and the compile-time dimension of the tree are only + // known when the index is built, so the tree type is erased behind this + // implementation, which also owns the points it indexes. Keeping nanoflann + // out of this header is a side effect, not the reason. + size_t appendPoints(const cv::Mat & features); + void makeIncremental(); + + NanoFlannIndexImpl * index_; + int featuresDim_; + // kept to rebuild the tree as an incremental one, see makeIncremental() + bool useDistanceL1_; + float removedRatio_; +}; + +} /* namespace rtabmap */ + +#endif /* CORELIB_SRC_NANOFLANN_NANOFLANNINDEX_H_ */ diff --git a/corelib/src/nanoflann/nanoflann.h b/corelib/src/nanoflann/nanoflann.h new file mode 100644 index 00000000..426ad309 --- /dev/null +++ b/corelib/src/nanoflann/nanoflann.h @@ -0,0 +1,4469 @@ +/*********************************************************************** + * Software License Agreement (BSD License) + * + * Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved. + * Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved. + * Copyright 2011-2026 Jose Luis Blanco (joseluisblancoc@gmail.com). + * All rights reserved. + * + * THE BSD LICENSE + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES + * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. + * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, + * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT + * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF + * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + *************************************************************************/ + +/** \mainpage nanoflann C++ API documentation + * nanoflann is a C++ header-only library for building KD-Trees, mostly + * optimized for 2D or 3D point clouds. + * + * nanoflann does not require compiling or installing, just an + * #include in your code. + * + * Macros that can be defined by the user to configure nanoflann: + * - NANOFLANN_NO_THREADS: If defined, multithreading is disabled, so the + * library can be used without linking against a threads library. Requesting a + * multi-threaded index build (`n_thread_build != 1`) then throws. + * - NANOFLANN_FIRST_MATCH: If defined, in case of a tie in distances the item + * with the smallest index will be returned. + * - NANOFLANN_NODE_ALIGNMENT: The memory alignment, in bytes, for kd-tree + * nodes. Default: 16. + * + * Macros defined internally by nanoflann (not meant to be set by the user): + * - NANOFLANN_RESTRICT: Expands to the compiler-specific `restrict` pointer + * qualifier (`__restrict__`, `__restrict`) when available, empty otherwise. + * - NANOFLANN_NODISCARD: Expands to `[[nodiscard]]` when the compiler supports + * it, empty otherwise. + * - NANOFLANN_VERSION: Library version as 0xMmP (M=Major, m=minor, P=patch). + * + * See the [README](https://github.com/jlblancoc/nanoflann#readme) for usage + * details and examples. + * + * See: + * - [Online README](https://github.com/jlblancoc/nanoflann) + * - [C++ API documentation](https://jlblancoc.github.io/nanoflann/) + */ + +#pragma once + +#include +#include +#include +#include +#include // std::chrono (async incremental index polling) +#include // for abs() +#include // rebuild worker of the async incremental index +#include +#include // snprintf +#include // for abs() +#include // std::exception_ptr (async incremental index) +#include // std::reference_wrapper +#include +#include +#include // std::numeric_limits +#include // std::unique_ptr (async incremental index) +#include // rebuild worker of the async incremental index +#include // placement new (incremental index node pool) +#include +#include +#include +#include +#include // std::is_trivially_destructible +#include +#include + +/** Library version as a decimal string "MAJOR.MINOR.PATCH" */ +#define NANOFLANN_VERSION_STRING "1.12.1" +/** Library version: 0xMMmmPP (MM=Major, mm=minor, PP=patch) */ +#define NANOFLANN_VERSION 0x010C01 + +// Avoid conflicting declaration of min/max macros in Windows headers +#if !defined(NOMINMAX) && (defined(_WIN32) || defined(_WIN32_) || defined(WIN32) || defined(_WIN64)) +#define NOMINMAX +#ifdef max +#undef max +#undef min +#endif +#endif +// Avoid conflicts with X11 headers +#ifdef None +#undef None +#endif + +// Handle restricted pointers +#if defined(__GNUC__) || defined(__clang__) +#define NANOFLANN_RESTRICT __restrict__ +#elif defined(_MSC_VER) +#define NANOFLANN_RESTRICT __restrict +#else +#define NANOFLANN_RESTRICT +#endif + +// [[nodiscard]] support +#if defined(__has_cpp_attribute) && __has_cpp_attribute(nodiscard) +#define NANOFLANN_NODISCARD [[nodiscard]] +#else +#define NANOFLANN_NODISCARD +#endif + +// [[fallthrough]] support for intentional switch fall-throughs +#if defined(__has_cpp_attribute) && __has_cpp_attribute(fallthrough) +#define NANOFLANN_FALLTHROUGH [[fallthrough]] +#else +#define NANOFLANN_FALLTHROUGH +#endif + +// Memory alignment of KD-tree nodes: +#ifndef NANOFLANN_NODE_ALIGNMENT +#define NANOFLANN_NODE_ALIGNMENT 16 +#endif + +namespace nanoflann +{ +/** @addtogroup nanoflann_grp nanoflann C++ library for KD-trees + * @{ */ + +/** the PI constant (required to avoid MSVC missing symbols) */ +template +constexpr T pi_const() +{ + return static_cast(3.14159265358979323846); +} + +/** + * Traits if object is resizable and assignable (typically has a resize | assign + * method) + */ +template +struct has_resize : std::false_type +{ +}; + +template +struct has_resize().resize(1), 0)> : std::true_type +{ +}; + +template +struct has_assign : std::false_type +{ +}; + +template +struct has_assign().assign(1, 0), 0)> : std::true_type +{ +}; + +/** + * Free function to resize a resizable object + */ +template +inline typename std::enable_if::value, void>::type resize( + Container& c, const size_t nElements) +{ + c.resize(nElements); +} + +/** + * Free function that has no effects on non resizable containers (e.g. + * std::array) It raises an exception if the expected size does not match + */ +template +inline typename std::enable_if::value, void>::type resize( + Container& c, const size_t nElements) +{ + if (nElements != c.size()) throw std::logic_error("Attempt to resize a fixed size container."); +} + +/** + * Free function to assign to a container + */ +template +inline typename std::enable_if::value, void>::type assign( + Container& c, const size_t nElements, const T& value) +{ + c.assign(nElements, value); +} + +/** + * Free function to assign to a std::array + */ +template +inline typename std::enable_if::value, void>::type assign( + Container& c, const size_t nElements, const T& value) +{ + for (size_t i = 0; i < nElements; i++) c[i] = value; +} + +/** operator "<" for std::sort() */ +struct IndexDist_Sorter +{ + /** PairType will be typically: ResultItem */ + template + bool operator()(const PairType& p1, const PairType& p2) const + { + return p1.second < p2.second; + } +}; + +/** + * Each result element in RadiusResultSet. Note that distances and indices + * are named `first` and `second` to keep backward-compatibility with the + * `std::pair<>` type used in the past. In contrast, this structure is ensured + * to be `std::is_standard_layout` so it can be used in wrappers to other + * languages. + * See: https://github.com/jlblancoc/nanoflann/issues/166 + */ +template +struct ResultItem +{ + ResultItem() = default; + ResultItem(const IndexType index, const DistanceType distance) : first(index), second(distance) + { + } + + IndexType first; //!< Index of the sample in the dataset + DistanceType second; //!< Distance from sample to query point +}; + +namespace detail +{ +/** Insert (dist, index) into a sorted result buffer (dists, indices) of the + * given capacity, keeping ascending distance order. Shared by KNNResultSet + * and RKNNResultSet, which are otherwise byte-for-byte identical. + * Always returns true (caller should continue searching). */ +template +bool addPointToSortedResultSet( + DistanceType* dists, IndexType* indices, CountType& count, CountType capacity, + DistanceType dist, IndexType index) +{ + CountType i; + for (i = count; i > 0; --i) + { +#ifdef NANOFLANN_FIRST_MATCH + if ((dists[i - 1] > dist) || ((dist == dists[i - 1]) && (indices[i - 1] > index))) + { +#else + if (dists[i - 1] > dist) + { +#endif + if (i < capacity) + { + dists[i] = dists[i - 1]; + indices[i] = indices[i - 1]; + } + } + else + break; + } + if (i < capacity) + { + dists[i] = dist; + indices[i] = index; + } + if (count < capacity) count++; + return true; +} +} // namespace detail + +/** @addtogroup result_sets_grp Result set classes + * @{ */ + +/** Result set for KNN searches (N-closest neighbors) */ +template +class KNNResultSet +{ + public: + using DistanceType = _DistanceType; + using IndexType = _IndexType; + using CountType = _CountType; + + private: + IndexType* indices; + DistanceType* dists; + CountType capacity; + CountType count; + + public: + explicit KNNResultSet(CountType capacity_) + : indices(nullptr), dists(nullptr), capacity(capacity_), count(0) + { + } + + void init(IndexType* indices_, DistanceType* dists_) + { + indices = indices_; + dists = dists_; + count = 0; + } + + NANOFLANN_NODISCARD CountType size() const noexcept { return count; } + NANOFLANN_NODISCARD bool empty() const noexcept { return count == 0; } + NANOFLANN_NODISCARD bool full() const noexcept { return count == capacity; } + + /** + * Called during search to add an element matching the criteria. + * @return true if the search should be continued, false if the results are + * sufficient + */ + bool addPoint(DistanceType dist, IndexType index) + { + return detail::addPointToSortedResultSet(dists, indices, count, capacity, dist, index); + } + + //! Returns the worst distance among found solutions if the search result is + //! full, or the maximum possible distance, if not full yet. + NANOFLANN_NODISCARD DistanceType worstDist() const noexcept + { + return (count < capacity || !count) ? std::numeric_limits::max() + : dists[count - 1]; + } + + void sort() + { + // already sorted + } +}; + +/** Result set for RKNN searches (N-closest neighbors with a maximum radius) */ +template +class RKNNResultSet +{ + public: + using DistanceType = _DistanceType; + using IndexType = _IndexType; + using CountType = _CountType; + + private: + IndexType* indices; + DistanceType* dists; + CountType capacity; + CountType count; + DistanceType maximumSearchDistanceSquared; + + public: + explicit RKNNResultSet(CountType capacity_, DistanceType maximumSearchDistanceSquared_) + : indices(nullptr), + dists(nullptr), + capacity(capacity_), + count(0), + maximumSearchDistanceSquared(maximumSearchDistanceSquared_) + { + } + + void init(IndexType* indices_, DistanceType* dists_) + { + indices = indices_; + dists = dists_; + count = 0; + if (capacity) dists[capacity - 1] = maximumSearchDistanceSquared; + } + + NANOFLANN_NODISCARD CountType size() const noexcept { return count; } + NANOFLANN_NODISCARD bool empty() const noexcept { return count == 0; } + NANOFLANN_NODISCARD bool full() const noexcept { return count == capacity; } + + /** + * Called during search to add an element matching the criteria. + * @return true if the search should be continued, false if the results are + * sufficient + */ + bool addPoint(DistanceType dist, IndexType index) + { + return detail::addPointToSortedResultSet(dists, indices, count, capacity, dist, index); + } + + //! Returns the worst distance among found solutions if the search result is + //! full, or the maximum possible distance, if not full yet. + NANOFLANN_NODISCARD DistanceType worstDist() const noexcept + { + return (count < capacity || !count) ? maximumSearchDistanceSquared : dists[count - 1]; + } + + void sort() + { + // already sorted + } +}; + +/** + * A result-set class used when performing a radius based search. + */ +template +class RadiusResultSet +{ + public: + using DistanceType = _DistanceType; + using IndexType = _IndexType; + + public: + const DistanceType radius; + + std::vector>& m_indices_dists; + + explicit RadiusResultSet( + DistanceType radius_, std::vector>& indices_dists) + : radius(radius_), m_indices_dists(indices_dists) + { + init(); + } + + void init() { clear(); } + void clear() { m_indices_dists.clear(); } + + NANOFLANN_NODISCARD size_t size() const noexcept { return m_indices_dists.size(); } + NANOFLANN_NODISCARD bool empty() const noexcept { return m_indices_dists.empty(); } + NANOFLANN_NODISCARD bool full() const noexcept { return true; } + + /** + * Called during search to add an element matching the criteria. + * @return true if the search should be continued, false if the results are + * sufficient + */ + bool addPoint(DistanceType dist, IndexType index) + { + if (dist < radius) m_indices_dists.emplace_back(index, dist); + return true; + } + + NANOFLANN_NODISCARD DistanceType worstDist() const noexcept { return radius; } + + /** + * Find the worst result (farthest neighbor) without copying or sorting + * Pre-conditions: size() > 0 + */ + ResultItem worst_item() const + { + if (m_indices_dists.empty()) + throw std::runtime_error( + "Cannot invoke RadiusResultSet::worst_item() on " + "an empty list of results."); + auto it = + std::max_element(m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter()); + return *it; + } + + void sort() { std::sort(m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter()); } +}; + +/** + * A result-set class used when collecting all points contained within an + * axis-aligned bounding box (see findWithinBox()). Distances are not used; + * matching point indices are appended to the user-provided vector. + */ +template +class BoxResultSet +{ + public: + using IndexType = _IndexType; + + std::vector& m_indices; + + explicit BoxResultSet(std::vector& indices) : m_indices(indices) + { + m_indices.clear(); + } + + NANOFLANN_NODISCARD size_t size() const noexcept { return m_indices.size(); } + NANOFLANN_NODISCARD bool empty() const noexcept { return m_indices.empty(); } + NANOFLANN_NODISCARD bool full() const noexcept { return true; } + + /** Called for each point found inside the query box. The distance argument + * is unused (always 0 for a box query). @return always true (keep going). */ + template + bool addPoint(DistanceType /*dist*/, IndexType index) + { + m_indices.push_back(index); + return true; + } + + void sort() { std::sort(m_indices.begin(), m_indices.end()); } +}; + +/** @} */ + +/** @addtogroup loadsave_grp Load/save auxiliary functions + * @{ */ +template +void save_value(std::ostream& stream, const T& value) +{ + stream.write(reinterpret_cast(&value), sizeof(T)); +} + +template +void save_value(std::ostream& stream, const std::vector& value) +{ + size_t size = value.size(); + stream.write(reinterpret_cast(&size), sizeof(size_t)); + stream.write(reinterpret_cast(value.data()), sizeof(T) * size); +} + +template +void load_value(std::istream& stream, T& value) +{ + stream.read(reinterpret_cast(&value), sizeof(T)); +} + +template +void load_value(std::istream& stream, std::vector& value) +{ + size_t size; + stream.read(reinterpret_cast(&size), sizeof(size_t)); + value.resize(size); + stream.read(reinterpret_cast(value.data()), sizeof(T) * size); +} +/** @} */ + +/** @addtogroup metric_grp Metric (distance) classes + * @{ */ + +struct Metric +{ +}; + +/** Manhattan distance functor (generic version, optimized for + * high-dimensionality data sets). Corresponding distance traits: + * nanoflann::metric_L1 + * + * \tparam T Type of the elements (e.g. double, float, uint8_t) + * \tparam DataSource Source of the data, i.e. where the vectors are stored + * \tparam _DistanceType Type of distance variables (must be signed) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +struct L1_Adaptor +{ + using ElementType = T; + using DistanceType = _DistanceType; + + const DataSource& data_source; + + L1_Adaptor(const DataSource& _data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric( + const T* NANOFLANN_RESTRICT a, const IndexType b_idx, size_t size) const + { + DistanceType result = DistanceType(); + const size_t multof4 = (size >> 2) << 2; // largest multiple of 4 + size_t d; + + for (d = 0; d < multof4; d += 4) + { + const DistanceType diff0 = std::abs(a[d + 0] - data_source.kdtree_get_pt(b_idx, d + 0)); + const DistanceType diff1 = std::abs(a[d + 1] - data_source.kdtree_get_pt(b_idx, d + 1)); + const DistanceType diff2 = std::abs(a[d + 2] - data_source.kdtree_get_pt(b_idx, d + 2)); + const DistanceType diff3 = std::abs(a[d + 3] - data_source.kdtree_get_pt(b_idx, d + 3)); + /* Parentheses break dependency chain: */ + result += (diff0 + diff1) + (diff2 + diff3); + } + /* Process last 0-3 components. Unrolled loop with fall-through switch. + */ + switch (size - multof4) + { + case 3: + result += std::abs(a[d + 2] - data_source.kdtree_get_pt(b_idx, d + 2)); + NANOFLANN_FALLTHROUGH; + case 2: + result += std::abs(a[d + 1] - data_source.kdtree_get_pt(b_idx, d + 1)); + NANOFLANN_FALLTHROUGH; + case 1: + result += std::abs(a[d + 0] - data_source.kdtree_get_pt(b_idx, d + 0)); + NANOFLANN_FALLTHROUGH; + case 0: + break; + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const + { + return std::abs(a - b); + } +}; + +/** **Squared** Euclidean distance functor (generic version, optimized for + * high-dimensionality data sets). Corresponding distance traits: + * nanoflann::metric_L2 + * + * \tparam T Type of the elements (e.g. double, float, uint8_t) + * \tparam DataSource Source of the data, i.e. where the vectors are stored + * \tparam _DistanceType Type of distance variables (must be signed) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +struct L2_Adaptor +{ + using ElementType = T; + using DistanceType = _DistanceType; + + const DataSource& data_source; + + L2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric( + const T* NANOFLANN_RESTRICT a, const IndexType b_idx, size_t size) const + { + DistanceType result = DistanceType(); + const size_t multof4 = (size >> 2) << 2; // largest multiple of 4 + size_t d; + + for (d = 0; d < multof4; d += 4) + { + const DistanceType diff0 = a[d + 0] - data_source.kdtree_get_pt(b_idx, d + 0); + const DistanceType diff1 = a[d + 1] - data_source.kdtree_get_pt(b_idx, d + 1); + const DistanceType diff2 = a[d + 2] - data_source.kdtree_get_pt(b_idx, d + 2); + const DistanceType diff3 = a[d + 3] - data_source.kdtree_get_pt(b_idx, d + 3); + /* Parentheses break dependency chain: */ + result += (diff0 * diff0 + diff1 * diff1) + (diff2 * diff2 + diff3 * diff3); + } + /* Process last 0-3 components. Unrolled loop with fall-through switch. + */ + DistanceType diff; + switch (size - multof4) + { + case 3: + diff = a[d + 2] - data_source.kdtree_get_pt(b_idx, d + 2); + result += diff * diff; + NANOFLANN_FALLTHROUGH; + case 2: + diff = a[d + 1] - data_source.kdtree_get_pt(b_idx, d + 1); + result += diff * diff; + NANOFLANN_FALLTHROUGH; + case 1: + diff = a[d + 0] - data_source.kdtree_get_pt(b_idx, d + 0); + result += diff * diff; + NANOFLANN_FALLTHROUGH; + case 0: + break; + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const + { + auto diff = a - b; + return diff * diff; + } +}; + +/** **Squared** Euclidean (L2) distance functor (suitable for low-dimensionality + * datasets, like 2D or 3D point clouds) Corresponding distance traits: + * nanoflann::metric_L2_Simple + * + * \tparam T Type of the elements (e.g. double, float, uint8_t) + * \tparam DataSource Source of the data, i.e. where the vectors are stored + * \tparam _DistanceType Type of distance variables (must be signed) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +struct L2_Simple_Adaptor +{ + using ElementType = T; + using DistanceType = _DistanceType; + + const DataSource& data_source; + + L2_Simple_Adaptor(const DataSource& _data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric(const T* a, const IndexType b_idx, size_t size) const + { + DistanceType result = DistanceType(); + for (size_t i = 0; i < size; ++i) + { + const DistanceType diff = a[i] - data_source.kdtree_get_pt(b_idx, i); + result += diff * diff; + } + return result; + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const + { + auto diff = a - b; + return diff * diff; + } +}; + +/** SO2 distance functor + * Corresponding distance traits: nanoflann::metric_SO2 + * + * \tparam T Type of the elements (e.g. double, float, uint8_t) + * \tparam DataSource Source of the data, i.e. where the vectors are stored + * \tparam _DistanceType Type of distance variables (must be signed) (e.g. + * float, double) orientation is constrained to be in [-pi, pi] + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +struct SO2_Adaptor +{ + using ElementType = T; + using DistanceType = _DistanceType; + + const DataSource& data_source; + + SO2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {} + + inline DistanceType evalMetric(const T* a, const IndexType b_idx, size_t size) const + { + return accum_dist(a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), size - 1); + } + + /** Returns the absolute shortest angular distance between a and b, + * assuming both are in [-pi, pi]. The result is in [0, pi], which + * satisfies the non-negativity requirement of a kd-tree metric and + * gives correct nearest-neighbour pruning. + */ + template + inline DistanceType accum_dist(const U a, const V b, const size_t) const + { + DistanceType diff = static_cast(b) - static_cast(a); + const DistanceType PI = pi_const(); + if (diff > PI) + diff -= 2 * PI; + else if (diff < -PI) + diff += 2 * PI; + return diff < DistanceType(0) ? -diff : diff; // abs without dependency + } +}; + +/** SO3 distance functor (Uses L2_Simple) + * Corresponding distance traits: nanoflann::metric_SO3 + * + * \tparam T Type of the elements (e.g. double, float, uint8_t) + * \tparam DataSource Source of the data, i.e. where the vectors are stored + * \tparam _DistanceType Type of distance variables (must be signed) (e.g. + * float, double) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +struct SO3_Adaptor +{ + using ElementType = T; + using DistanceType = _DistanceType; + + L2_Simple_Adaptor distance_L2_Simple; + + SO3_Adaptor(const DataSource& _data_source) : distance_L2_Simple(_data_source) {} + + inline DistanceType evalMetric(const T* a, const IndexType b_idx, size_t size) const + { + return distance_L2_Simple.evalMetric(a, b_idx, size); + } + + template + inline DistanceType accum_dist(const U a, const V b, const size_t idx) const + { + return distance_L2_Simple.accum_dist(a, b, idx); + } +}; + +/** Metaprogramming helper traits class for the L1 (Manhattan) metric */ +struct metric_L1 : public Metric +{ + template + struct traits + { + using distance_t = L1_Adaptor; + }; +}; +/** Metaprogramming helper traits class for the L2 (Euclidean) **squared** + * distance metric */ +struct metric_L2 : public Metric +{ + template + struct traits + { + using distance_t = L2_Adaptor; + }; +}; +/** Metaprogramming helper traits class for the L2_simple (Euclidean) + * **squared** distance metric */ +struct metric_L2_Simple : public Metric +{ + template + struct traits + { + using distance_t = L2_Simple_Adaptor; + }; +}; +/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */ +struct metric_SO2 : public Metric +{ + template + struct traits + { + using distance_t = SO2_Adaptor; + }; +}; +/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */ +struct metric_SO3 : public Metric +{ + template + struct traits + { + using distance_t = SO3_Adaptor; + }; +}; + +/** @} */ + +/** @addtogroup param_grp Parameter structs + * @{ */ + +enum class KDTreeSingleIndexAdaptorFlags +{ + None = 0, + SkipInitialBuildIndex = 1 +}; + +inline std::underlying_type::type operator&( + KDTreeSingleIndexAdaptorFlags lhs, KDTreeSingleIndexAdaptorFlags rhs) +{ + using underlying = typename std::underlying_type::type; + return static_cast(lhs) & static_cast(rhs); +} + +/** Returns true if \a f has the given \a flag bit set. + * Prefer this over the raw operator& in boolean contexts. */ +inline bool has_flag(KDTreeSingleIndexAdaptorFlags f, KDTreeSingleIndexAdaptorFlags flag) +{ + return (f & flag) != 0; +} + +/** Parameters (see README.md) */ +struct KDTreeSingleIndexAdaptorParams +{ + KDTreeSingleIndexAdaptorParams( + size_t _leaf_max_size = 10, + KDTreeSingleIndexAdaptorFlags _flags = KDTreeSingleIndexAdaptorFlags::None, + unsigned int _n_thread_build = 1) + : leaf_max_size(_leaf_max_size), flags(_flags), n_thread_build(_n_thread_build) + { + } + + size_t leaf_max_size; + KDTreeSingleIndexAdaptorFlags flags; + unsigned int n_thread_build; +}; + +/** Search options for KDTreeSingleIndexAdaptor::findNeighbors() */ +struct SearchParameters +{ + SearchParameters(float eps_ = 0, bool sorted_ = true) : eps(eps_), sorted(sorted_) {} + + float eps; //!< search for eps-approximate neighbors (default: 0) + bool sorted; //!< only for radius search, require neighbors sorted by + //!< distance (default: true) +}; +/** @} */ + +/** @addtogroup memalloc_grp Memory allocation + * @{ */ + +/** + * Pooled storage allocator + * + * The following routines allow for the efficient allocation of storage in + * small chunks from a specified pool. Rather than allowing each structure + * to be freed individually, an entire pool of storage is freed at once. + * This method has two advantages over just using malloc() and free(). First, + * it is far more efficient for allocating small objects, as there is + * no overhead for remembering all the information needed to free each + * object or consolidating fragmented memory. Second, the decision about + * how long to keep an object is made at the time of allocation, and there + * is no need to track down all the objects to free them. + * + */ +class PooledAllocator +{ + static constexpr size_t WORDSIZE = 16; // WORDSIZE must >= 8 + static constexpr size_t BLOCKSIZE = 8192; + + /* We maintain memory alignment to word boundaries by requiring that all + allocations be in multiples of the machine wordsize. */ + /* Size of machine word in bytes. Must be power of 2. */ + /* Minimum number of bytes requested at a time from the system. Must be + * multiple of WORDSIZE. */ + + using Size = size_t; + + Size remaining_ = 0; //!< Number of bytes left in current block of storage + void* base_ = nullptr; //!< Pointer to base of current block of storage + void* loc_ = nullptr; //!< Current location in block to next allocate + + void internal_init() + { + remaining_ = 0; + base_ = nullptr; + usedMemory = 0; + wastedMemory = 0; + } + + public: + Size usedMemory = 0; + Size wastedMemory = 0; + + /** + Default constructor. Initializes a new pool. + */ + PooledAllocator() { internal_init(); } + + /** + * Destructor. Frees all the memory allocated in this pool. + */ + ~PooledAllocator() { free_all(); } + + /** Frees all allocated memory chunks */ + void free_all() + { + while (base_ != nullptr) + { + // Get pointer to prev block + void* prev = *(static_cast(base_)); + ::free(base_); + base_ = prev; + } + internal_init(); + } + + /** + * Returns a pointer to a piece of new memory of the given size in bytes + * allocated from the pool. + */ + void* allocateBytes(const size_t req_size) + { + /* Round size up to a multiple of wordsize. The following expression + only works for WORDSIZE that is a power of 2, by masking last bits + of incremented size to zero. + */ + const Size size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1); + + /* Check whether a new block must be allocated. Note that the first + word of a block is reserved for a pointer to the previous block. + */ + if (size > remaining_) + { + wastedMemory += remaining_; + + /* Allocate new storage. */ + const Size blocksize = size > BLOCKSIZE ? size + WORDSIZE : BLOCKSIZE + WORDSIZE; + + // use the standard C malloc to allocate memory + void* m = ::malloc(blocksize); + if (!m) + { + throw std::bad_alloc(); + } + + /* Fill first word of new block with pointer to previous block. */ + static_cast(m)[0] = base_; + base_ = m; + + remaining_ = blocksize - WORDSIZE; + loc_ = static_cast(m) + WORDSIZE; + } + void* rloc = loc_; + loc_ = static_cast(loc_) + size; + remaining_ -= size; + + usedMemory += size; + + return rloc; + } + + /** + * Allocates (using this pool) a generic type T. + * + * Params: + * count = number of instances to allocate. + * Returns: pointer (of type T*) to memory buffer + */ + template + T* allocate(const size_t count = 1) + { + T* mem = static_cast(this->allocateBytes(sizeof(T) * count)); + return mem; + } +}; +/** @} */ + +/** @addtogroup nanoflann_metaprog_grp Auxiliary metaprogramming stuff + * @{ */ + +/** Used to declare fixed-size arrays when DIM>0, dynamically-allocated vectors + * when DIM=-1. Fixed size version for a generic DIM: + */ +template +struct array_or_vector +{ + using type = std::array; +}; +/** Dynamic size version */ +template +struct array_or_vector<-1, T> +{ + using type = std::vector; +}; + +/** @} */ + +/** kd-tree base-class + * + * Contains the member functions common to the classes KDTreeSingleIndexAdaptor + * and KDTreeSingleIndexDynamicAdaptor_. + * + * \tparam Derived The name of the class which inherits this class. + * \tparam DatasetAdaptor The user-provided adaptor, which must be ensured to + * have a lifetime equal or longer than the instance of this class. + * \tparam Distance The distance metric to use, these are all classes derived + * from nanoflann::Metric + * \tparam DIM Dimensionality of data points (e.g. 3 for 3D points) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template < + class Derived, typename Distance, class DatasetAdaptor, int32_t DIM = -1, + typename index_t = uint32_t> +class KDTreeBaseClass +{ + public: + /** Frees the previously-built index. Automatically called within + * buildIndex(). */ + void freeIndex(Derived& obj) + { + obj.pool_.free_all(); + obj.root_node_ = nullptr; + obj.size_at_index_build_ = 0; + } + + using ElementType = typename Distance::ElementType; + using DistanceType = typename Distance::DistanceType; + using IndexType = index_t; + + /** + * Array of indices to vectors in the dataset_. + */ + std::vector vAcc_; + + using Offset = typename decltype(vAcc_)::size_type; + using Size = typename decltype(vAcc_)::size_type; + using Dimension = int32_t; + + /*------------------------------------------------------------------- + * Internal Data Structures + * + * "Node" below can be declared with alignas(N) to improve + * cache friendliness and SIMD load/store performance. + * + * The optimal N depends on the underlying hardware: + * + Intel x86-64: 16 for SSE, 32 for AVX/AVX2 and 64 for AVX-512 + * + NVIDIA Jetson: 16 for ARM + NEON and CUDA float4/ + * To avoid unnecessary padding, the smallest alignment + * compatible with a platform's vector width should be chosen. + * ------------------------------------------------------------------*/ + struct alignas(NANOFLANN_NODE_ALIGNMENT) Node + { + /** Union used because a node can be either a LEAF node or a non-leaf + * node, so both data fields are never used simultaneously */ + union + { + struct leaf + { + Offset left, right; //!< Indices of points in leaf node + } lr; + struct nonleaf + { + Dimension divfeat; //!< Dimension used for subdivision. + /// The values used for subdivision. + DistanceType divlow, divhigh; + } sub; + } node_type; + + /** Child nodes (both=nullptr mean its a leaf node) */ + Node *child1 = nullptr, *child2 = nullptr; + }; + + using NodePtr = Node*; + using NodeConstPtr = const Node*; + + struct Interval + { + ElementType low, high; + }; + + NodePtr root_node_ = nullptr; + + Size leaf_max_size_ = 0; + + /// Number of thread for concurrent tree build + Size n_thread_build_ = 1; + /// Number of current points in the dataset + Size size_ = 0; + /// Number of points in the dataset when the index was built + Size size_at_index_build_ = 0; + Dimension dim_ = 0; //!< Dimensionality of each data point + + /** Define "BoundingBox" as a fixed-size or variable-size container + * depending on "DIM" */ + using BoundingBox = typename array_or_vector::type; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + using distance_vector_t = typename array_or_vector::type; + + /** The KD-tree used to find neighbors */ + BoundingBox root_bbox_; + + /** + * Pooled memory allocator. + * + * Using a pooled memory allocator is more efficient + * than allocating memory directly when there is a large + * number small of memory allocations. + */ + PooledAllocator pool_; + + /** Returns number of points in dataset */ + NANOFLANN_NODISCARD Size size(const Derived& obj) const noexcept { return obj.size_; } + + /** Returns the length of each point in the dataset. + * For a fixed-size tree (DIM > 0) this is a compile-time constant; under + * C++17 the `if constexpr` lets the compiler drop the runtime read of + * `dim_` entirely. The C++11 path keeps the equivalent ternary. */ + NANOFLANN_NODISCARD Size veclen(const Derived& obj) const noexcept + { +#if defined(__cpp_if_constexpr) && __cpp_if_constexpr >= 201606L + if constexpr (DIM > 0) + { + return DIM; + } + else + { + return obj.dim_; + } +#else + return DIM > 0 ? DIM : obj.dim_; +#endif + } + + /// Helper accessor to the dataset points: + ElementType dataset_get(const Derived& obj, IndexType element, Dimension component) const + { + return obj.dataset_.kdtree_get_pt(element, component); + } + + /** + * Computes the index memory usage + * Returns: memory used by the index + */ + NANOFLANN_NODISCARD Size usedMemory(const Derived& obj) const + { + return obj.pool_.usedMemory + obj.pool_.wastedMemory + + obj.dataset_.kdtree_get_point_count() * + sizeof(IndexType); // pool memory and vind array memory + } + + /** + * Compute the minimum and maximum element values in the specified dimension + */ + void computeMinMax( + const Derived& obj, Offset ind, Size count, Dimension element, ElementType& min_elem, + ElementType& max_elem) const + { + min_elem = dataset_get(obj, vAcc_[ind], element); + max_elem = min_elem; + for (Offset i = 1; i < count; ++i) + { + ElementType val = dataset_get(obj, vAcc_[ind + i], element); + if (val < min_elem) min_elem = val; + if (val > max_elem) max_elem = val; + } + } + + /** Returns true if the point at index idx should be visited during search. + * The static adaptor always returns true; the dynamic adaptor overrides + * this to skip tombstoned (removed) points. */ + NANOFLANN_NODISCARD bool isActive(IndexType /*idx*/) const { return true; } + + /** Computes the bounding box of the points currently in the index. + * Uses size_ (set by buildIndex before this is called) so the result is + * correct for both the static and dynamic adaptors. */ + void computeBoundingBox(BoundingBox& bbox) + { + Derived& obj = static_cast(*this); + const Dimension dims = static_cast(veclen(obj)); + resize(bbox, dims); + if (obj.dataset_.kdtree_get_bbox(bbox)) return; + if (!size_) + throw std::runtime_error( + "[nanoflann] computeBoundingBox() called but " + "no data points found."); + for (Dimension i = 0; i < dims; ++i) + bbox[i].low = bbox[i].high = dataset_get(obj, vAcc_[0], i); + for (Offset k = 1; k < size_; ++k) + for (Dimension i = 0; i < dims; ++i) + { + const auto val = dataset_get(obj, vAcc_[k], i); + if (val < bbox[i].low) bbox[i].low = val; + if (val > bbox[i].high) bbox[i].high = val; + } + } + + /** + * Performs an exact search in the tree starting from a node. + * Uses the CRTP-dispatched isActive() hook to skip removed points (no-op + * in the static adaptor, checks treeIndex_ in the dynamic adaptor). + * \tparam RESULTSET Should be any ResultSet + * \return true if the search should be continued, false if the results are + * sufficient + */ + template + bool searchLevel( + RESULTSET& result_set, const ElementType* vec, const NodePtr node, DistanceType mindist, + distance_vector_t& dists, const DistanceType epsError) const + { + const Derived& obj = static_cast(*this); + // If this is a leaf node, then do check and return. + if (!node->child1) // (if one node is nullptr, both are) + { + // Hoist the point length out of the per-point loop. For a + // fixed-size tree (DIM > 0) this is a compile-time constant; for a + // runtime dimension it avoids re-reading obj.dim_ on every point. + const Size dim = veclen(obj); + for (Offset i = node->node_type.lr.left; i < node->node_type.lr.right; ++i) + { + const IndexType accessor = vAcc_[i]; + if (!obj.isActive(accessor)) continue; + DistanceType dist = obj.distance_.evalMetric(vec, accessor, dim); + if (dist < result_set.worstDist()) + { + if (!result_set.addPoint( + static_cast(dist), + static_cast(accessor))) + return false; + } + } + return true; + } + + /* Which child branch should be taken first? */ + Dimension idx = node->node_type.sub.divfeat; + ElementType val = vec[idx]; + DistanceType diff1 = val - node->node_type.sub.divlow; + DistanceType diff2 = val - node->node_type.sub.divhigh; + + NodePtr bestChild; + NodePtr otherChild; + DistanceType cut_dist; + if ((diff1 + diff2) < 0) + { + bestChild = node->child1; + otherChild = node->child2; + cut_dist = obj.distance_.accum_dist(val, node->node_type.sub.divhigh, idx); + } + else + { + bestChild = node->child2; + otherChild = node->child1; + cut_dist = obj.distance_.accum_dist(val, node->node_type.sub.divlow, idx); + } + + /* Call recursively to search next level down. */ + if (!searchLevel(result_set, vec, bestChild, mindist, dists, epsError)) return false; + + DistanceType dst = dists[idx]; + mindist = mindist + cut_dist - dst; + dists[idx] = cut_dist; + if (mindist * epsError <= result_set.worstDist()) + { + if (!searchLevel(result_set, vec, otherChild, mindist, dists, epsError)) return false; + } + dists[idx] = dst; + return true; + } + + /** + * Create a tree node that subdivides the list of vecs from vind[first] + * to vind[last]. The routine is called recursively on each sublist. + * + * @param left index of the first vector + * @param right index of the last vector + * @param bbox bounding box used as input for splitting and output for + * parent node + */ + /** + * Initialize a freshly-allocated node while building the tree: either turn + * it into a leaf node (computing the leaf bounding-box) or compute the + * split plane for an interior node. Shared by the sequential and concurrent + * builders, which differ only in how they recurse. + * + * @return true if the node became a leaf (no further recursion needed), + * false if it is an interior node and \a idx / \a cutfeat / \a cutval + * describe the split plane. + */ + bool makeNode( + Derived& obj, NodePtr node, const Offset left, const Offset right, BoundingBox& bbox, + Offset& idx, Dimension& cutfeat, DistanceType& cutval) + { + const Dimension dims = static_cast(veclen(obj)); + + /* If too few exemplars remain, then make this a leaf node. */ + if ((right - left) <= static_cast(obj.leaf_max_size_)) + { + node->child1 = node->child2 = nullptr; /* Mark as leaf node. */ + node->node_type.lr.left = left; + node->node_type.lr.right = right; + + // compute bounding-box of leaf points + for (Dimension i = 0; i < dims; ++i) + { + bbox[i].low = dataset_get(obj, obj.vAcc_[left], i); + bbox[i].high = dataset_get(obj, obj.vAcc_[left], i); + } + for (Offset k = left + 1; k < right; ++k) + { + for (Dimension i = 0; i < dims; ++i) + { + const auto val = dataset_get(obj, obj.vAcc_[k], i); + if (bbox[i].low > val) bbox[i].low = val; + if (bbox[i].high < val) bbox[i].high = val; + } + } + return true; + } + + /* Determine the index, dimension and value for split plane */ + middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox); + node->node_type.sub.divfeat = cutfeat; + return false; + } + + /** + * After both children of an interior node have been built, record the split + * planes and expand \a bbox to the union of the children bounding-boxes. + * Shared by the sequential and concurrent builders. + */ + void finalizeSplitNode( + Derived& obj, NodePtr node, const Dimension cutfeat, const BoundingBox& left_bbox, + const BoundingBox& right_bbox, BoundingBox& bbox) + { + node->node_type.sub.divlow = left_bbox[cutfeat].high; + node->node_type.sub.divhigh = right_bbox[cutfeat].low; + + const Dimension dims = static_cast(veclen(obj)); + for (Dimension i = 0; i < dims; ++i) + { + bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low); + bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high); + } + } + + NodePtr divideTree(Derived& obj, const Offset left, const Offset right, BoundingBox& bbox) + { + assert(static_cast(obj.vAcc_.at(left)) < obj.dataset_.kdtree_get_point_count()); + + NodePtr node = obj.pool_.template allocate(); // allocate memory + Offset idx; + Dimension cutfeat; + DistanceType cutval; + if (makeNode(obj, node, left, right, bbox, idx, cutfeat, cutval)) return node; + + /* Recurse on left */ + BoundingBox left_bbox(bbox); + left_bbox[cutfeat].high = cutval; + node->child1 = this->divideTree(obj, left, left + idx, left_bbox); + + /* Recurse on right */ + BoundingBox right_bbox(bbox); + right_bbox[cutfeat].low = cutval; + node->child2 = this->divideTree(obj, left + idx, right, right_bbox); + + finalizeSplitNode(obj, node, cutfeat, left_bbox, right_bbox, bbox); + + return node; + } + + /** + * Create a tree node that subdivides the list of vecs from vind[first] to + * vind[last] concurrently. The routine is called recursively on each + * sublist. + * + * @param left index of the first vector + * @param right index of the last vector + * @param bbox bounding box used as input for splitting and output for + * parent node + * @param thread_count count of std::async threads + * @param mutex mutex for mempool allocation + */ + NodePtr divideTreeConcurrent( + Derived& obj, const Offset left, const Offset right, BoundingBox& bbox, + std::atomic& thread_count, std::mutex& mutex) + { + std::unique_lock lock(mutex); + NodePtr node = obj.pool_.template allocate(); // allocate memory + lock.unlock(); + + Offset idx; + Dimension cutfeat; + DistanceType cutval; + if (makeNode(obj, node, left, right, bbox, idx, cutfeat, cutval)) return node; + + std::future right_future; + + /* Recurse on right concurrently, if possible */ + + BoundingBox right_bbox(bbox); + right_bbox[cutfeat].low = cutval; + if (++thread_count < n_thread_build_) + { + /* Concurrent thread for right recursion */ + + right_future = std::async( + std::launch::async, &KDTreeBaseClass::divideTreeConcurrent, this, std::ref(obj), + left + idx, right, std::ref(right_bbox), std::ref(thread_count), std::ref(mutex)); + } + else + { + --thread_count; + } + + /* Recurse on left in this thread */ + + BoundingBox left_bbox(bbox); + left_bbox[cutfeat].high = cutval; + node->child1 = + this->divideTreeConcurrent(obj, left, left + idx, left_bbox, thread_count, mutex); + + if (right_future.valid()) + { + /* Block and wait for concurrent right from above */ + + node->child2 = right_future.get(); + --thread_count; + } + else + { + /* Otherwise, recurse on right in this thread */ + + node->child2 = + this->divideTreeConcurrent(obj, left + idx, right, right_bbox, thread_count, mutex); + } + + finalizeSplitNode(obj, node, cutfeat, left_bbox, right_bbox, bbox); + + return node; + } + + void middleSplit_( + const Derived& obj, const Offset ind, const Size count, Offset& index, Dimension& cutfeat, + DistanceType& cutval, const BoundingBox& bbox) + { + const Dimension dims = static_cast(veclen(obj)); + const auto EPS = static_cast(0.00001); + + // Pre-compute max_span once + ElementType max_span = bbox[0].high - bbox[0].low; + for (Dimension i = 1; i < dims; ++i) + { + ElementType span = bbox[i].high - bbox[i].low; + if (span > max_span) max_span = span; + } + + // Two-pass: first find max_span (done above), then scan candidate dims + // inline — no heap allocation for a candidates vector. + cutfeat = 0; + ElementType max_spread = -1; + ElementType min_elem = 0, max_elem = 0; + const ElementType threshold = (1 - EPS) * max_span; + + for (Dimension dim = 0; dim < dims; ++dim) + { + if (bbox[dim].high - bbox[dim].low < threshold) continue; + + ElementType local_min = dataset_get(obj, vAcc_[ind], dim); + ElementType local_max = local_min; + + // Unrolled loop for better performance + constexpr size_t UNROLL = 4; + Offset k = 1; + for (; k + UNROLL <= count; k += UNROLL) + { + ElementType v0 = dataset_get(obj, vAcc_[ind + k], dim); + ElementType v1 = dataset_get(obj, vAcc_[ind + k + 1], dim); + ElementType v2 = dataset_get(obj, vAcc_[ind + k + 2], dim); + ElementType v3 = dataset_get(obj, vAcc_[ind + k + 3], dim); + + local_min = std::min({local_min, v0, v1, v2, v3}); + local_max = std::max({local_max, v0, v1, v2, v3}); + } + + // Handle remainder + for (; k < count; ++k) + { + ElementType val = dataset_get(obj, vAcc_[ind + k], dim); + local_min = std::min(local_min, val); + local_max = std::max(local_max, val); + } + + ElementType spread = local_max - local_min; + if (spread > max_spread) + { + cutfeat = dim; + max_spread = spread; + min_elem = local_min; + max_elem = local_max; + } + } + + // Median-of-three for better balance + DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2; + if (split_val < min_elem) split_val = min_elem; + if (split_val > max_elem) split_val = max_elem; + + cutval = split_val; + + // Optimized partitioning + Offset lim1, lim2; + planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2); + + index = (lim1 > count / 2) ? lim1 : (lim2 < count / 2) ? lim2 : count / 2; + } + + /** + * Subdivide the list of points by a plane perpendicular on the axis + * corresponding to the 'cutfeat' dimension at 'cutval' position. + * + * On return: + * dataset[ind[0..lim1-1]][cutfeat] < cutval + * dataset[ind[lim1..lim2-1]][cutfeat] == cutval + * dataset[ind[lim2..count]][cutfeat] > cutval + */ + void planeSplit( + const Derived& obj, const Offset ind, const Size count, const Dimension cutfeat, + const DistanceType& cutval, Offset& lim1, Offset& lim2) + { + // Dutch National Flag algorithm for three-way partitioning + Offset left = 0; + Offset mid = 0; + Offset right = count - 1; + + while (mid <= right) + { + ElementType val = dataset_get(obj, vAcc_[ind + mid], cutfeat); + + if (val < cutval) + { + std::swap(vAcc_[ind + left], vAcc_[ind + mid]); + left++; + mid++; + } + else if (val > cutval) + { + std::swap(vAcc_[ind + mid], vAcc_[ind + right]); + right--; + } + else + { + mid++; + } + } + + lim1 = left; + lim2 = mid; + } + + DistanceType computeInitialDistances( + const Derived& obj, const ElementType* vec, distance_vector_t& dists) const + { + assert(vec); + DistanceType dist = DistanceType(); + + const Dimension dims = static_cast(veclen(obj)); + for (Dimension i = 0; i < dims; ++i) + { + if (vec[i] < obj.root_bbox_[i].low) + { + dists[i] = obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].low, i); + dist += dists[i]; + } + else if (vec[i] > obj.root_bbox_[i].high) + { + dists[i] = obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].high, i); + dist += dists[i]; + } + } + return dist; + } + + static void save_tree(const Derived& obj, std::ostream& stream, const NodeConstPtr tree) + { + save_value(stream, *tree); + if (tree->child1 != nullptr) + { + save_tree(obj, stream, tree->child1); + } + if (tree->child2 != nullptr) + { + save_tree(obj, stream, tree->child2); + } + } + + static void load_tree(Derived& obj, std::istream& stream, NodePtr& tree) + { + tree = obj.pool_.template allocate(); + load_value(stream, *tree); + if (tree->child1 != nullptr) + { + load_tree(obj, stream, tree->child1); + } + if (tree->child2 != nullptr) + { + load_tree(obj, stream, tree->child2); + } + } + + /** Magic number written at the start of every saveIndex() stream. + * Spells 'NFLN' in ASCII. */ + static constexpr uint32_t SAVE_MAGIC = 0x4E464C4E; + + /** Stores the index in a binary stream. + * + * The set of data points is NOT stored; when reloading, the index object + * must be constructed with the same dataset. See: examples/saveload_example.cpp + * + * \note **Portability limitations** (by design -- fixing them would require + * a breaking format change): + * - Files are NOT portable across different endianness (e.g. x86 little-endian + * vs. big-endian SPARC/PowerPC). No byte-swapping is performed. + * - Files are NOT portable across 32-bit vs. 64-bit platforms (sizeof(size_t) + * differs). + * - Files are NOT portable across different nanoflann versions; loadIndex() + * throws if the version in the file does not match the library. + * - Files are NOT portable across different template instantiations (e.g. + * float vs. double IndexType/ElementType); loadIndex() throws on mismatch. + * + * \sa loadIndex + */ + void saveIndex(const Derived& obj, std::ostream& stream) const + { + // 10-byte header: magic | version | sizeof_size_t | sizeof_IndexType + // | sizeof_ElementType | sizeof_DistanceType + // Use local copies: passing a static constexpr by const-ref ODR-uses it + // in C++11/14, which requires an out-of-class definition we cannot provide + // in a header-only library. + const uint32_t hdr_magic = SAVE_MAGIC; + const uint32_t hdr_version = static_cast(NANOFLANN_VERSION); + const uint8_t hdr_sz_st = static_cast(sizeof(size_t)); + const uint8_t hdr_sz_idx = static_cast(sizeof(IndexType)); + const uint8_t hdr_sz_elem = static_cast(sizeof(ElementType)); + const uint8_t hdr_sz_dist = static_cast(sizeof(DistanceType)); + save_value(stream, hdr_magic); + save_value(stream, hdr_version); + save_value(stream, hdr_sz_st); + save_value(stream, hdr_sz_idx); + save_value(stream, hdr_sz_elem); + save_value(stream, hdr_sz_dist); + + save_value(stream, obj.size_); + save_value(stream, obj.dim_); + save_value(stream, obj.root_bbox_); + save_value(stream, obj.leaf_max_size_); + save_value(stream, obj.vAcc_); + if (obj.root_node_) + { + save_tree(obj, stream, obj.root_node_); + } + } + + /** Loads an index previously saved with saveIndex() from a binary stream. + * + * The index object must be constructed associated to the same dataset that + * was used when building the saved index. See: examples/saveload_example.cpp + * + * \throws std::runtime_error if the stream does not start with the expected + * magic number (wrong file or corrupt data), if the nanoflann version in + * the file differs from the current library version, if the saved type + * sizes (size_t, IndexType, ElementType, DistanceType) do not match the + * current template instantiation, or if a read error occurs. + * + * \note See saveIndex() for portability limitations. + * + * \sa saveIndex + */ + void loadIndex(Derived& obj, std::istream& stream) + { + // Validate header + uint32_t magic = 0; + load_value(stream, magic); + if (stream.fail() || magic != SAVE_MAGIC) + { + throw std::runtime_error( + "nanoflann loadIndex: invalid file (wrong magic number). " + "The stream was not written by nanoflann saveIndex()."); + } + + uint32_t file_version = 0; + load_value(stream, file_version); + if (file_version != static_cast(NANOFLANN_VERSION)) + { + char msg[200]; + snprintf( + msg, sizeof(msg), + "nanoflann loadIndex: version mismatch " + "(file=0x%03X, library=0x%03X). Rebuild the index.", + file_version, static_cast(NANOFLANN_VERSION)); + throw std::runtime_error(msg); + } + + uint8_t sz_size_t = 0; + uint8_t sz_idx = 0; + uint8_t sz_elem = 0; + uint8_t sz_dist = 0; + load_value(stream, sz_size_t); + load_value(stream, sz_idx); + load_value(stream, sz_elem); + load_value(stream, sz_dist); + if (sz_size_t != static_cast(sizeof(size_t)) || + sz_idx != static_cast(sizeof(IndexType)) || + sz_elem != static_cast(sizeof(ElementType)) || + sz_dist != static_cast(sizeof(DistanceType))) + { + throw std::runtime_error( + "nanoflann loadIndex: type-size mismatch between saved index and " + "current template instantiation (sizeof size_t / IndexType / " + "ElementType / DistanceType differ). Rebuild the index."); + } + + load_value(stream, obj.size_); + load_value(stream, obj.dim_); + load_value(stream, obj.root_bbox_); + load_value(stream, obj.leaf_max_size_); + load_value(stream, obj.vAcc_); + + if (obj.size_ > 0) + { + load_tree(obj, stream, obj.root_node_); + } + + if (stream.fail()) + { + throw std::runtime_error( + "nanoflann loadIndex: unexpected end of stream or read error."); + } + } +}; + +/** @addtogroup kdtrees_grp KD-tree classes and adaptors + * @{ */ + +/** kd-tree static index + * + * Contains the k-d trees and other information for indexing a set of points + * for nearest-neighbor matching. + * + * The class "DatasetAdaptor" must provide the following interface (can be + * non-virtual, inlined methods): + * + * \code + * // Must return the number of data points + * size_t kdtree_get_point_count() const { ... } + * + * + * // Must return the dim'th component of the idx'th point in the class: + * T kdtree_get_pt(const size_t idx, const size_t dim) const { ... } + * + * // Optional bounding-box computation: return false to default to a standard + * bbox computation loop. + * // Return true if the BBOX was already computed by the class and returned + * in "bb" so it can be avoided to redo it again. + * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 + * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const + * { + * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits + * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits + * ... + * return true; + * } + * + * \endcode + * + * \tparam DatasetAdaptor The user-provided adaptor, which must be ensured to + * have a lifetime equal or longer than the instance of this class. + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam DIM + * Dimensionality of data points (e.g. 3 for 3D points) \tparam IndexType Will + * be typically size_t or int + * + * \note Threading guarantees: + * - Index build: passing `n_thread_build > 1` in the params parallelizes the + * build using `std::async` (unless NANOFLANN_NO_THREADS is defined, in which + * case requesting more than one thread throws). + * - Queries (`findNeighbors`, `knnSearch`, `radiusSearch`, `rknnSearch`) are + * `const` and thread-safe for concurrent readers: multiple threads may query + * the same index simultaneously, as long as no thread is concurrently + * (re)building or modifying it. + * - The internal `PooledAllocator` is NOT thread-safe; building an index from + * multiple threads, or mixing queries with a concurrent build, is not + * supported. + */ +template +class KDTreeSingleIndexAdaptor + : public KDTreeBaseClass< + KDTreeSingleIndexAdaptor, Distance, + DatasetAdaptor, DIM, index_t> +{ + public: + /** Deleted copy constructor*/ + explicit KDTreeSingleIndexAdaptor( + const KDTreeSingleIndexAdaptor&) = delete; + + /** The data source used by this index */ + const DatasetAdaptor& dataset_; + + const KDTreeSingleIndexAdaptorParams indexParams; + + Distance distance_; + + using Base = typename nanoflann::KDTreeBaseClass< + nanoflann::KDTreeSingleIndexAdaptor, Distance, + DatasetAdaptor, DIM, index_t>; + + using Offset = typename Base::Offset; + using Size = typename Base::Size; + using Dimension = typename Base::Dimension; + + using ElementType = typename Base::ElementType; + using DistanceType = typename Base::DistanceType; + using IndexType = typename Base::IndexType; + + using Node = typename Base::Node; + using NodePtr = Node*; + + using Interval = typename Base::Interval; + + /** Define "BoundingBox" as a fixed-size or variable-size container + * depending on "DIM" */ + using BoundingBox = typename Base::BoundingBox; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + using distance_vector_t = typename Base::distance_vector_t; + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the dataset, e.g. + * 3 for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features. Its lifetime must be + * equal or longer than that of the instance of this class. + * @param params Basically, the maximum leaf node size + * + * Note that there is a variable number of optional additional parameters + * which will be forwarded to the metric class constructor. Refer to example + * `examples/pointcloud_custom_metric.cpp` for a use case. + * + */ + template + explicit KDTreeSingleIndexAdaptor( + const Dimension dimensionality, const DatasetAdaptor& inputData, + const KDTreeSingleIndexAdaptorParams& params, Args&&... args) + : dataset_(inputData), + indexParams(params), + distance_(inputData, std::forward(args)...) + { + init(dimensionality, params); + } + + explicit KDTreeSingleIndexAdaptor( + const Dimension dimensionality, const DatasetAdaptor& inputData, + const KDTreeSingleIndexAdaptorParams& params = {}) + : dataset_(inputData), indexParams(params), distance_(inputData) + { + init(dimensionality, params); + } + + private: + void init(const Dimension dimensionality, const KDTreeSingleIndexAdaptorParams& params) + { + Base::size_ = dataset_.kdtree_get_point_count(); + Base::size_at_index_build_ = Base::size_; + Base::dim_ = dimensionality; + if (DIM > 0) Base::dim_ = DIM; + Base::leaf_max_size_ = params.leaf_max_size; + if (params.n_thread_build > 0) + { + Base::n_thread_build_ = params.n_thread_build; + } + else + { + Base::n_thread_build_ = std::max(std::thread::hardware_concurrency(), 1u); + } + + if (!has_flag(params.flags, KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex)) + { + // Build KD-tree: + buildIndex(); + } + } + + public: + /** + * Builds the index + */ + void buildIndex() + { + Base::size_ = dataset_.kdtree_get_point_count(); + Base::size_at_index_build_ = Base::size_; + init_vind(); + this->freeIndex(*this); + Base::size_at_index_build_ = Base::size_; + if (Base::size_ == 0) return; + this->computeBoundingBox(Base::root_bbox_); + // construct the tree + if (Base::n_thread_build_ == 1) + { + Base::root_node_ = this->divideTree(*this, 0, Base::size_, Base::root_bbox_); + } + else + { +#ifndef NANOFLANN_NO_THREADS + std::atomic thread_count(0u); + std::mutex mutex; + Base::root_node_ = this->divideTreeConcurrent( + *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex); +#else /* NANOFLANN_NO_THREADS */ + throw std::runtime_error("Multithreading is disabled"); +#endif /* NANOFLANN_NO_THREADS */ + } + } + + /** \name Query methods + * @{ */ + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * + * Params: + * result = the result object in which the indices of the + * nearest-neighbors are stored vec = the vector for which to search the + * nearest neighbors + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * \sa knnSearch, radiusSearch + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + */ + template + bool findNeighbors( + RESULTSET& result, const ElementType* vec, const SearchParameters& searchParams = {}) const + { + assert(vec); + if (this->size(*this) == 0) return false; + if (!Base::root_node_) + throw std::runtime_error( + "[nanoflann] findNeighbors() called before building the " + "index."); + DistanceType epsError = 1 + static_cast(searchParams.eps); + + // fixed or variable-sized container (depending on DIM) + distance_vector_t dists; + // Fill it with zeros. + auto zero = static_cast(0); + assign(dists, this->veclen(*this), zero); + DistanceType dist = this->computeInitialDistances(*this, vec, dists); + this->searchLevel(result, vec, Base::root_node_, dist, dists, epsError); + + if (searchParams.sorted) result.sort(); + + return result.full(); + } + + /** + * Find all points contained within the specified bounding box. Their + * indices are stored inside the result object. + * + * Params: + * result = the result object in which the indices of the points + * within the bounding box are stored + * bbox = the bounding box defining the search region + * + * \tparam RESULTSET Should be any ResultSet + * \return Number of points found within the bounding box. + * \sa findNeighbors, knnSearch, radiusSearch + * + * \note The search is inclusive - points on the boundary are included. + */ + template + NANOFLANN_NODISCARD Size findWithinBox(RESULTSET& result, const BoundingBox& bbox) const + { + if (this->size(*this) == 0) return 0; + if (!Base::root_node_) + throw std::runtime_error( + "[nanoflann] findWithinBox() called before building the " + "index."); + + std::stack stack; + stack.push(Base::root_node_); + + while (!stack.empty()) + { + const NodePtr node = stack.top(); + stack.pop(); + + // If this is a leaf node, then do check and return. + if (!node->child1) // (if one node is nullptr, both are) + { + for (Offset i = node->node_type.lr.left; i < node->node_type.lr.right; ++i) + { + if (contains(bbox, Base::vAcc_[i])) + { + if (!result.addPoint(0, Base::vAcc_[i])) + { + // the resultset doesn't want to receive any more + // points, we're done searching! + return result.size(); + } + } + } + } + else + { + const Dimension idx = node->node_type.sub.divfeat; + const auto low_bound = node->node_type.sub.divlow; + const auto high_bound = node->node_type.sub.divhigh; + + if (bbox[idx].low <= low_bound) stack.push(node->child1); + if (bbox[idx].high >= high_bound) stack.push(node->child2); + } + } + + return result.size(); + } + + /** + * Find the "num_closest" nearest neighbors to the \a query_point[0:dim-1]. + * Their indices and distances are stored in the provided pointers to + * array/vector. + * + * \sa radiusSearch, findNeighbors + * \return Number `N` of valid points in the result set. + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + * + * \note Only the first `N` entries in `out_indices` and `out_distances` + * will be valid. Return is less than `num_closest` only if the + * number of elements in the tree is less than `num_closest`. + */ + NANOFLANN_NODISCARD Size knnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances) const + { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances); + findNeighbors(resultSet, query_point); + return resultSet.size(); + } + + /** + * Find all the neighbors to \a query_point[0:dim-1] within a maximum + * radius. The output is given as a vector of pairs, of which the first + * element is a point index and the second the corresponding distance. + * Previous contents of \a IndicesDists are cleared. + * + * If searchParams.sorted==true, the output list is sorted by ascending + * distances. + * + * For a better performance, it is advisable to do a .reserve() on the + * vector if you have any wild guess about the number of expected matches. + * + * \sa knnSearch, findNeighbors, radiusSearchCustomCallback + * \return The number of points within the given radius (i.e. indices.size() + * or dists.size() ) + * + * \note If L2 norms are used, search radius and all returned distances + * are actually squared distances. + */ + NANOFLANN_NODISCARD Size radiusSearch( + const ElementType* query_point, const DistanceType& radius, + std::vector>& IndicesDists, + const SearchParameters& searchParams = {}) const + { + RadiusResultSet resultSet(radius, IndicesDists); + const Size nFound = radiusSearchCustomCallback(query_point, resultSet, searchParams); + return nFound; + } + + /** + * Just like radiusSearch() but with a custom callback class for each point + * found in the radius of the query. See the source of RadiusResultSet<> as + * a start point for your own classes. \sa radiusSearch + */ + template + NANOFLANN_NODISCARD Size radiusSearchCustomCallback( + const ElementType* query_point, SEARCH_CALLBACK& resultSet, + const SearchParameters& searchParams = {}) const + { + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** + * Find the N closest neighbors to \a query_point[0:dim-1] that are also + * within the given maximum radius. Results are stored in the provided + * output arrays; previous contents are overwritten. + * + * \sa radiusSearch, findNeighbors + * \return Number of valid points written (at most `num_closest`). May be + * less if fewer than `num_closest` points lie within the radius. + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + * + * \note Only the first `N` entries in `out_indices` and `out_distances` + * will be valid. + */ + NANOFLANN_NODISCARD Size rknnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const DistanceType& radius) const + { + nanoflann::RKNNResultSet resultSet(num_closest, radius); + resultSet.init(out_indices, out_distances); + findNeighbors(resultSet, query_point); + return resultSet.size(); + } + + /** @} */ + + public: + /** Make sure the auxiliary list \a vind has the same size as the + * current dataset, and re-generate if size has changed. */ + void init_vind() + { + // Create a permutable array of indices to the input vectors. + Base::size_ = dataset_.kdtree_get_point_count(); + if (Base::vAcc_.size() != Base::size_) Base::vAcc_.resize(Base::size_); + for (IndexType i = 0; i < static_cast(Base::size_); i++) Base::vAcc_[i] = i; + } + + bool contains(const BoundingBox& bbox, IndexType idx) const + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + { + const auto point = this->dataset_.kdtree_get_pt(idx, i); + if (point < bbox[i].low || point > bbox[i].high) return false; + } + return true; + } + + public: + /** Stores the index in a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so + * when loading the index object it must be constructed associated to the + * same source of data points used while building it. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void saveIndex(std::ostream& stream) const { Base::saveIndex(*this, stream); } + + /** Loads a previous index from a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so + * the index object must be constructed associated to the same source of + * data points used while building the index. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void loadIndex(std::istream& stream) { Base::loadIndex(*this, stream); } + +}; // class KDTree + +/** kd-tree dynamic index + * + * Contains the k-d trees and other information for indexing a set of points + * for nearest-neighbor matching. + * + * The class "DatasetAdaptor" must provide the following interface (can be + * non-virtual, inlined methods): + * + * \code + * // Must return the number of data points + * size_t kdtree_get_point_count() const { ... } + * + * // Must return the dim'th component of the idx'th point in the class: + * T kdtree_get_pt(const size_t idx, const size_t dim) const { ... } + * + * // Optional bounding-box computation: return false to default to a standard + * bbox computation loop. + * // Return true if the BBOX was already computed by the class and returned + * in "bb" so it can be avoided to redo it again. + * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 + * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const + * { + * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits + * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits + * ... + * return true; + * } + * + * \endcode + * + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. + * \tparam DIM Dimensionality of data points (e.g. 3 for 3D points) + * \tparam IndexType Type of the arguments with which the data can be + * accessed (e.g. float, double, int64_t, T*) + */ +template +class KDTreeSingleIndexDynamicAdaptor_ + : public KDTreeBaseClass< + KDTreeSingleIndexDynamicAdaptor_, Distance, + DatasetAdaptor, DIM, IndexType> +{ + public: + /** + * The dataset used by this index + */ + const DatasetAdaptor& dataset_; //!< The source of our data + + KDTreeSingleIndexAdaptorParams index_params_; + + std::vector& treeIndex_; + + Distance distance_; + + using Base = typename nanoflann::KDTreeBaseClass< + nanoflann::KDTreeSingleIndexDynamicAdaptor_, + Distance, DatasetAdaptor, DIM, IndexType>; + + using ElementType = typename Base::ElementType; + using DistanceType = typename Base::DistanceType; + + using Offset = typename Base::Offset; + using Size = typename Base::Size; + using Dimension = typename Base::Dimension; + + using Node = typename Base::Node; + using NodePtr = Node*; + + using Interval = typename Base::Interval; + /** Define "BoundingBox" as a fixed-size or variable-size container + * depending on "DIM" */ + using BoundingBox = typename Base::BoundingBox; + + /** Define "distance_vector_t" as a fixed-size or variable-size container + * depending on "DIM" */ + using distance_vector_t = typename Base::distance_vector_t; + + /** Returns false for points that have been removed (lazy deletion). */ + NANOFLANN_NODISCARD bool isActive(IndexType idx) const { return treeIndex_[idx] != -1; } + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the dataset, e.g. + * 3 for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features. Its lifetime must be + * equal or longer than that of the instance of this class. + * @param params Basically, the maximum leaf node size + */ + KDTreeSingleIndexDynamicAdaptor_( + const Dimension dimensionality, const DatasetAdaptor& inputData, + std::vector& treeIndex, + const KDTreeSingleIndexAdaptorParams& params = KDTreeSingleIndexAdaptorParams()) + : dataset_(inputData), index_params_(params), treeIndex_(treeIndex), distance_(inputData) + { + Base::size_ = 0; + Base::size_at_index_build_ = 0; + for (auto& v : Base::root_bbox_) v = {}; + Base::dim_ = dimensionality; + if (DIM > 0) Base::dim_ = DIM; + Base::leaf_max_size_ = params.leaf_max_size; + if (params.n_thread_build > 0) + { + Base::n_thread_build_ = params.n_thread_build; + } + else + { + Base::n_thread_build_ = std::max(std::thread::hardware_concurrency(), 1u); + } + } + + /** Explicitly default the copy constructor */ + KDTreeSingleIndexDynamicAdaptor_(const KDTreeSingleIndexDynamicAdaptor_& rhs) = default; + + /** Assignment operator definition */ + KDTreeSingleIndexDynamicAdaptor_& operator=(const KDTreeSingleIndexDynamicAdaptor_& rhs) + { + if (this == &rhs) return *this; + KDTreeSingleIndexDynamicAdaptor_ tmp(rhs); + std::swap(Base::vAcc_, tmp.Base::vAcc_); + std::swap(Base::leaf_max_size_, tmp.Base::leaf_max_size_); + std::swap(index_params_, tmp.index_params_); + // treeIndex_ is a reference member and cannot be rebound; do not swap. + std::swap(Base::size_, tmp.Base::size_); + std::swap(Base::size_at_index_build_, tmp.Base::size_at_index_build_); + std::swap(Base::root_node_, tmp.Base::root_node_); + std::swap(Base::root_bbox_, tmp.Base::root_bbox_); + std::swap(Base::pool_, tmp.Base::pool_); + return *this; + } + + /** + * Builds the index + */ + void buildIndex() + { + Base::size_ = Base::vAcc_.size(); + this->freeIndex(*this); + Base::size_at_index_build_ = Base::size_; + if (Base::size_ == 0) return; + this->computeBoundingBox(Base::root_bbox_); + // construct the tree + if (Base::n_thread_build_ == 1) + { + Base::root_node_ = this->divideTree(*this, 0, Base::size_, Base::root_bbox_); + } + else + { +#ifndef NANOFLANN_NO_THREADS + std::atomic thread_count(0u); + std::mutex mutex; + Base::root_node_ = this->divideTreeConcurrent( + *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex); +#else /* NANOFLANN_NO_THREADS */ + throw std::runtime_error("Multithreading is disabled"); +#endif /* NANOFLANN_NO_THREADS */ + } + } + + /** \name Query methods + * @{ */ + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * This is the core search function, all others are wrappers around this + * one. + * + * \param result The result object in which the indices of the + * nearest-neighbors are stored. + * \param vec The vector of the query point for which to search the + * nearest neighbors. + * \param searchParams Optional parameters for the search. + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * + * \sa knnSearch(), radiusSearch(), radiusSearchCustomCallback() + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + */ + template + bool findNeighbors( + RESULTSET& result, const ElementType* vec, const SearchParameters& searchParams = {}) const + { + assert(vec); + if (this->size(*this) == 0) return false; + if (!Base::root_node_) return false; + DistanceType epsError = 1 + static_cast(searchParams.eps); + + // fixed or variable-sized container (depending on DIM) + distance_vector_t dists; + // Fill it with zeros. + assign(dists, this->veclen(*this), static_cast(0)); + DistanceType dist = this->computeInitialDistances(*this, vec, dists); + this->searchLevel(result, vec, Base::root_node_, dist, dists, epsError); + + if (searchParams.sorted) result.sort(); + + return result.full(); + } + + /** + * Find the "num_closest" nearest neighbors to the \a query_point[0:dim-1]. + * Their indices are stored inside the result object. \sa radiusSearch, + * findNeighbors + * \return Number `N` of valid points in + * the result set. + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + * + * \note Only the first `N` entries in `out_indices` and `out_distances` + * will be valid. Return may be less than `num_closest` only if the + * number of elements in the tree is less than `num_closest`. + */ + NANOFLANN_NODISCARD Size knnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const SearchParameters& searchParams = {}) const + { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances); + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** + * Find all the neighbors to \a query_point[0:dim-1] within a maximum + * radius. The output is given as a vector of pairs, of which the first + * element is a point index and the second the corresponding distance. + * Previous contents of \a IndicesDists are cleared. + * + * If searchParams.sorted==true, the output list is sorted by ascending + * distances. + * + * For a better performance, it is advisable to do a .reserve() on the + * vector if you have any wild guess about the number of expected matches. + * + * \sa knnSearch, findNeighbors, radiusSearchCustomCallback + * \return The number of points within the given radius (i.e. indices.size() + * or dists.size() ) + * + * \note If L2 norms are used, search radius and all returned distances + * are actually squared distances. + */ + NANOFLANN_NODISCARD Size radiusSearch( + const ElementType* query_point, const DistanceType& radius, + std::vector>& IndicesDists, + const SearchParameters& searchParams = {}) const + { + RadiusResultSet resultSet(radius, IndicesDists); + const Size nFound = radiusSearchCustomCallback(query_point, resultSet, searchParams); + return nFound; + } + + /** + * Just like radiusSearch() but with a custom callback class for each point + * found in the radius of the query. See the source of RadiusResultSet<> as + * a start point for your own classes. \sa radiusSearch + */ + template + NANOFLANN_NODISCARD Size radiusSearchCustomCallback( + const ElementType* query_point, SEARCH_CALLBACK& resultSet, + const SearchParameters& searchParams = {}) const + { + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** @} */ + + public: + public: + /** Stores the index in a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so + * when loading the index object it must be constructed associated to the + * same source of data points used while building it. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void saveIndex(std::ostream& stream) { Base::saveIndex(*this, stream); } + + /** Loads a previous index from a binary file. + * IMPORTANT NOTE: The set of data points is NOT stored in the file, so + * the index object must be constructed associated to the same source of + * data points used while building the index. See the example: + * examples/saveload_example.cpp \sa loadIndex */ + void loadIndex(std::istream& stream) { Base::loadIndex(*this, stream); } +}; + +/** kd-tree dynamic index + * + * class to create multiple static index and merge their results to behave as + * single dynamic index as proposed in Logarithmic Approach. + * + * Example of usage: + * examples/dynamic_pointcloud_example.cpp + * + * \tparam DatasetAdaptor The user-provided adaptor (see comments above). + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \tparam DIM + * Dimensionality of data points (e.g. 3 for 3D points) \tparam IndexType + * Will be typically size_t or int + */ +template +class KDTreeSingleIndexDynamicAdaptor +{ + public: + using ElementType = typename Distance::ElementType; + using DistanceType = typename Distance::DistanceType; + + using Offset = typename KDTreeSingleIndexDynamicAdaptor_::Offset; + using Size = typename KDTreeSingleIndexDynamicAdaptor_::Size; + using Dimension = + typename KDTreeSingleIndexDynamicAdaptor_::Dimension; + + protected: + Size leaf_max_size_; + Size treeCount_; + Size pointCount_; + + /** + * The dataset used by this index + */ + const DatasetAdaptor& dataset_; //!< The source of our data + + /** treeIndex[idx] is the index of tree in which point at idx is stored. + * treeIndex[idx]=-1 means that point has been removed. */ + std::vector treeIndex_; + /** Maps each currently-removed point index to the sub-tree that still + * physically holds it (removal is lazy). Used to reactivate the point in + * place if it is later re-added, instead of inserting a duplicate. */ + std::unordered_map removedPoints_; + + KDTreeSingleIndexAdaptorParams index_params_; + + Dimension dim_; //!< Dimensionality of each data point + + using index_container_t = + KDTreeSingleIndexDynamicAdaptor_; + std::vector index_; + + public: + /** Get a const ref to the internal list of indices; the number of indices + * is adapted dynamically as the dataset grows in size. */ + const std::vector& getAllIndices() const { return index_; } + + private: + /** finds position of least significant unset bit */ + int First0Bit(Size num) + { + int pos = 0; + while (num & 1) + { + num = num >> 1; + pos++; + } + return pos; + } + + /** Creates multiple empty trees to handle dynamic support */ + void init() + { + using my_kd_tree_t = + KDTreeSingleIndexDynamicAdaptor_; + std::vector index( + treeCount_, my_kd_tree_t(dim_ /*dim*/, dataset_, treeIndex_, index_params_)); + index_ = index; + } + + public: + Distance distance_; + + /** + * KDTree constructor + * + * Refer to docs in README.md or online in + * https://github.com/jlblancoc/nanoflann + * + * The KD-Tree point dimension (the length of each point in the dataset, e.g. + * 3 for 3D points) is determined by means of: + * - The \a DIM template parameter if >0 (highest priority) + * - Otherwise, the \a dimensionality parameter of this constructor. + * + * @param inputData Dataset with the input features. Its lifetime must be + * equal or longer than that of the instance of this class. + * @param params Basically, the maximum leaf node size + */ + explicit KDTreeSingleIndexDynamicAdaptor( + const int dimensionality, const DatasetAdaptor& inputData, + const KDTreeSingleIndexAdaptorParams& params = KDTreeSingleIndexAdaptorParams(), + const size_t maximumPointCount = 1000000000U) + : dataset_(inputData), index_params_(params), distance_(inputData) + { + treeCount_ = static_cast(std::log2(maximumPointCount)) + 1; + pointCount_ = 0U; + dim_ = dimensionality; + treeIndex_.clear(); + if (DIM > 0) dim_ = DIM; + leaf_max_size_ = params.leaf_max_size; + init(); + const size_t num_initial_points = dataset_.kdtree_get_point_count(); + if (num_initial_points > 0) + { + addPoints(0, static_cast(num_initial_points - 1)); + } + } + + /** Deleted copy constructor*/ + explicit KDTreeSingleIndexDynamicAdaptor( + const KDTreeSingleIndexDynamicAdaptor&) = delete; + + /** Add points to the set, Inserts all points from [start, end] */ + void addPoints(IndexType start, IndexType end) + { + int maxIndex = 0; + for (IndexType idx = start; idx <= end; idx++) + { + // If this index was previously removed, its point is still + // physically present in its sub-tree (removal is lazy and never + // deletes from vAcc_). Just clear the "removed" mark and restore its + // tree index. Re-inserting it would create a duplicate entry that + // grows the trees without bound and yields duplicate search results. + const auto it = removedPoints_.find(idx); + if (it != removedPoints_.end()) + { + treeIndex_[idx] = it->second; + removedPoints_.erase(it); + continue; + } + + const int pos = First0Bit(pointCount_); + maxIndex = std::max(pos, maxIndex); + if (treeIndex_.size() <= static_cast(pointCount_)) + treeIndex_.resize(static_cast(pointCount_) + 1); + treeIndex_[pointCount_] = pos; + + for (int i = 0; i < pos; i++) + { + for (size_t j = 0; j < index_[i].vAcc_.size(); j++) + { + const IndexType e = index_[i].vAcc_[j]; + index_[pos].vAcc_.push_back(e); + if (treeIndex_[e] != -1) + treeIndex_[e] = pos; + else + removedPoints_[e] = pos; // keep tombstone's tree index current + } + index_[i].vAcc_.clear(); + } + index_[pos].vAcc_.push_back(idx); + pointCount_++; + } + + for (int i = 0; i <= maxIndex; ++i) + { + index_[i].freeIndex(index_[i]); + if (!index_[i].vAcc_.empty()) index_[i].buildIndex(); + } + } + + /** Remove a point from the set (Lazy Deletion) */ + void removePoint(size_t idx) + { + if (idx >= pointCount_) return; + if (treeIndex_[idx] == -1) return; // already removed + // Remember which sub-tree still physically holds this point, so it can + // be reactivated in place if re-added later (see addPoints). + removedPoints_[static_cast(idx)] = treeIndex_[idx]; + treeIndex_[idx] = -1; + } + + /** + * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored + * inside the result object. + * + * Params: + * result = the result object in which the indices of the + * nearest-neighbors are stored vec = the vector for which to search the + * nearest neighbors + * + * \tparam RESULTSET Should be any ResultSet + * \return True if the requested neighbors could be found. + * \sa knnSearch, radiusSearch + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + */ + template + bool findNeighbors( + RESULTSET& result, const ElementType* vec, const SearchParameters& searchParams = {}) const + { + for (size_t i = 0; i < treeCount_; i++) + { + index_[i].findNeighbors(result, &vec[0], searchParams); + } + return result.full(); + } +}; + +/** Parameters for KDTreeSingleIndexIncrementalAdaptor. + * + * The two \a alpha_* thresholds drive the weight-balanced (scapegoat-style) + * partial rebuilds: + * - \a alpha_balance : a subtree is rebuilt when the larger of its two child + * subtrees holds more than this fraction of the subtree's points. Lower + * values keep the tree better balanced (faster queries) at the price of more + * frequent rebuilds. Typical range [0.55, 0.85]. + * - \a alpha_deleted : a subtree is rebuilt (physically dropping tombstoned + * points) when the fraction of removed points in it exceeds this value. Lower + * values reclaim memory more aggressively. Typical range [0.3, 0.7]. + */ +struct KDTreeIncrementalIndexParams +{ + KDTreeIncrementalIndexParams(float alpha_balance_ = 0.75f, float alpha_deleted_ = 0.5f) + : alpha_balance(alpha_balance_), alpha_deleted(alpha_deleted_) + { + } + + float alpha_balance; + float alpha_deleted; +}; + +/** kd-tree incremental dynamic index — a single, self-balancing k-d tree. + * + * This is an additive alternative to KDTreeSingleIndexDynamicAdaptor (the + * "logarithmic forest"). Instead of maintaining O(log N) static sub-trees, it + * keeps a *single* weight-balanced k-d tree that supports cheap incremental + * point insertion, lazy point removal, and pruned axis-aligned box deletions + * (removeBox / removeOutsideBox), the latter being the primary map-maintenance + * primitive used in LiDAR odometry (keep a local cube around the sensor and + * discard everything outside it as the platform moves). + * + * Compared with the forest, the single tree avoids the multiplicative + * O(log N) query penalty of querying every sub-tree, and keeps deletion garbage + * bounded (a subtree is rebuilt once its dead fraction crosses + * `alpha_deleted`), at the price of synchronous O(N) rebuild spikes near the + * root. + * + * Like the other adaptors this is zero-copy: the data points live in the + * user-provided dataset; the tree only stores point indices. Each tree node + * holds a single point plus augmentation (subtree size, tombstone count and an + * axis-aligned bounding box) used to prune the box-region operations. + * + * \note Threading: like the rest of nanoflann, const queries are safe for + * concurrent readers, but mutating operations (addPoints / removePoint / + * removeBox / removeOutsideBox) must not run concurrently with any query. + * + * \note A compile-time fixed \a DIM (e.g. 3 for 3D LiDAR) is recommended: the + * per-node bounding box is then a stack std::array and no per-node heap + * allocation occurs. With DIM=-1 each node's box is a std::vector. + * + * \tparam Distance The distance metric (nanoflann::metric_L2_Simple, etc). + * \tparam DatasetAdaptor The user-provided dataset adaptor. + * \tparam DIM Dimensionality of the data points (e.g. 3), or -1 for runtime. + * \tparam IndexType Type used to index data points (e.g. uint32_t). + */ +template +class KDTreeSingleIndexIncrementalAdaptor + : public KDTreeBaseClass< + KDTreeSingleIndexIncrementalAdaptor, Distance, + DatasetAdaptor, DIM, IndexType> +{ + public: + using Base = typename nanoflann::KDTreeBaseClass< + KDTreeSingleIndexIncrementalAdaptor, Distance, + DatasetAdaptor, DIM, IndexType>; + + using ElementType = typename Base::ElementType; + using DistanceType = typename Base::DistanceType; + + using Offset = typename Base::Offset; + using Size = typename Base::Size; + using Dimension = typename Base::Dimension; + + using Interval = typename Base::Interval; + using BoundingBox = typename Base::BoundingBox; + using distance_vector_t = typename Base::distance_vector_t; + + /** The data source used by this index. */ + const DatasetAdaptor& dataset_; + + Distance distance_; + + /** Augmented tree node: stores a single data point plus the maintenance + * metadata. Children pointers are nullptr at the leaves. */ + struct INode + { + IndexType ptIdx = 0; //!< index of the stored data point + Dimension divfeat = 0; //!< splitting axis at this node + bool deleted = false; //!< this node's point is tombstoned + bool treeDeleted = false; //!< whole subtree lazily tombstoned + INode* child1 = nullptr; //!< "< split" child (also free-list link) + INode* child2 = nullptr; //!< ">= split" child + INode* parent = nullptr; //!< parent (nullptr at the root) + Size subtree_size = 0; //!< number of nodes in this subtree + Size invalid_count = 0; //!< number of tombstoned nodes in subtree + BoundingBox box; //!< AABB of all points (live+dead) in this subtree + //! Cache of this node's own point coordinates, kept in-node to avoid the + //! dataset_get() indirection on the hot query / insert / box paths. Only + //! populated for a compile-time fixed DIM (`kCacheCoords`); for DIM=-1 it + //! stays an empty vector and the code falls back to the dataset. + typename array_or_vector::type pcoord; + }; + + /** Whether per-node coordinate caching is active: only for a compile-time + * fixed DIM, so the cache is a stack array and adds no per-node heap. + * Define NANOFLANN_INCREMENTAL_NO_COORD_CACHE to opt out (e.g. a very large + * ElementType) and always read coordinates from the dataset. */ +#if defined(NANOFLANN_INCREMENTAL_NO_COORD_CACHE) + static constexpr bool kCacheCoords = false; +#else + static constexpr bool kCacheCoords = (DIM > 0); +#endif + + private: + INode* iroot_ = nullptr; //!< root of the incremental tree + INode* freeList_ = nullptr; //!< recycled nodes (linked via child1) + + Size liveCount_ = 0; //!< number of live (non-tombstoned) points + Size totalCount_ = 0; //!< number of nodes physically in the tree + + float alphaBal_ = 0.75f; + float alphaDel_ = 0.5f; + /// Subtrees smaller than this are never rebuilt for *balance* reasons. + static constexpr Size kMinBalanceRebuild = 4; + /// addPoints() bulk-builds instead of inserting point-by-point when the + /// batch is at least this fraction of the current live count (see addPoints). + static constexpr double kBulkInsertFraction = 0.5; + + /// Highest unbalanced node found during the current insertion (rebuilt once). + INode* pendingRebuild_ = nullptr; + + /// When false, the index never performs inline (synchronous) balance or + /// deletion rebuilds: it only appends / lazily tombstones, and balance is + /// restored externally by full bulk rebuilds. Used by the multi-threaded + /// wrapper, which offloads the expensive rebuilds to a background thread. + bool inlineRebuild_ = true; + + /// idx -> node holding that point (nullptr if the point is not present). + std::vector nodeOfPoint_; + + /// Scratch buffer reused across rebuilds (live indices being re-balanced). + std::vector buildBuf_; + + /// Optional sink for physically-evicted point indices (acquireRemovedPoints). + bool collectRemoved_ = false; + std::vector removedSink_; + + public: + /** Constructor. + * @param dimensionality Runtime dimensionality (ignored if DIM>0). + * @param inputData Dataset adaptor; its lifetime must outlive this index. + * @param params Balancing thresholds (see KDTreeIncrementalIndexParams). + * + * The tree starts empty regardless of the dataset size; call addPoints() + * to insert points (their indices must be valid in \a inputData). + */ + explicit KDTreeSingleIndexIncrementalAdaptor( + const Dimension dimensionality, const DatasetAdaptor& inputData, + const KDTreeIncrementalIndexParams& params = {}) + : dataset_(inputData), distance_(inputData) + { + Base::dim_ = dimensionality; + if (DIM > 0) Base::dim_ = DIM; + alphaBal_ = params.alpha_balance; + alphaDel_ = params.alpha_deleted; + resize(Base::root_bbox_, static_cast(this->veclen(*this))); + } + + /** Deleted copy constructor (owns raw node memory). */ + KDTreeSingleIndexIncrementalAdaptor(const KDTreeSingleIndexIncrementalAdaptor&) = delete; + KDTreeSingleIndexIncrementalAdaptor& operator=(const KDTreeSingleIndexIncrementalAdaptor&) = + delete; + + ~KDTreeSingleIndexIncrementalAdaptor() { destroyNodeObjects(); } + + /** \name Modifiers + * @{ */ + + /** Insert a single point (by its index in the dataset). */ + void addPoint(IndexType idx) + { + ensureNodeMap(idx); + insertOne(idx); + syncRootBox(); + } + + /** Insert all points with indices in the inclusive range [start, end]. + * + * When the batch is large relative to the current tree (an empty tree, or + * a batch comparable to the live count — e.g. a full LiDAR scan rebuilding + * a heavily-trimmed map) it is cheaper to flatten the live points and + * bulk-build one balanced tree than to descend-insert each point and + * trigger near-root scapegoat rebuilds. Small batches relative to a large + * map take the incremental per-point path. */ + void addPoints(IndexType start, IndexType end) + { + if (end < start) return; + ensureNodeMap(end); + const Size batch = static_cast(end - start) + 1; + if (!iroot_ || + static_cast(batch) >= kBulkInsertFraction * static_cast(liveCount_)) + { + buildBuf_.clear(); + if (iroot_) collectLiveAndFree(iroot_, buildBuf_); // keep existing live points + for (IndexType idx = start; idx <= end; ++idx) buildBuf_.push_back(idx); + iroot_ = buildBalanced(buildBuf_, 0, buildBuf_.size(), 0, nullptr); + liveCount_ = buildBuf_.size(); + totalCount_ = liveCount_; + } + else + { + for (IndexType idx = start; idx <= end; ++idx) insertOne(idx); + } + syncRootBox(); + } + + /** Lazily remove the point with the given index (no-op if absent/removed). */ + void removePoint(IndexType idx) + { + if (idx >= nodeOfPoint_.size()) return; + INode* n = nodeOfPoint_[idx]; + if (!n || n->deleted) return; + // A point can also be logically dead via a treeDeleted ancestor (a + // lazily-killed box region). Detect that and treat it as already gone. + for (INode* p = n->parent; p; p = p->parent) + if (p->treeDeleted) return; + + n->deleted = true; + for (INode* p = n; p; p = p->parent) ++p->invalid_count; + --liveCount_; + maybeRebuildForDeletion(); + syncRootBox(); + } + + /** Remove every live point lying inside the axis-aligned box \a box. */ + void removeBox(const BoundingBox& box) + { + if (iroot_) removeBoxRec(iroot_, box); + maybeRebuildForDeletion(); + syncRootBox(); + } + + /** Remove every live point lying *outside* the axis-aligned box \a keep. + * This is the LiDAR sliding-window map-trimming primitive. */ + void removeOutsideBox(const BoundingBox& keep) + { + if (iroot_) removeOutsideBoxRec(iroot_, keep); + maybeRebuildForDeletion(); + syncRootBox(); + } + + /** Enable/disable recording of physically-evicted point indices, returned + * by acquireRemovedPoints(). Off by default (cost-free when unused). */ + void setCollectRemovedPoints(bool enable) + { + collectRemoved_ = enable; + if (!enable) std::vector().swap(removedSink_); + } + + /** Move out the list of point indices physically dropped (during rebuilds) + * since the last call. Requires setCollectRemovedPoints(true). */ + std::vector acquireRemovedPoints() + { + std::vector out; + out.swap(removedSink_); + return out; + } + + /** Enable/disable inline (synchronous) rebalancing. When disabled the index + * only appends and lazily tombstones; balance must be restored externally + * via buildFromIndices(). Used by the multi-threaded wrapper. */ + void setInlineRebuild(bool enable) { inlineRebuild_ = enable; } + + /** Append the live point indices (DFS, skipping tombstones) into \a out. + * Non-destructive; used to snapshot the tree for a background rebuild. */ + void snapshotLiveIndices(std::vector& out) const { snapshotRec(iroot_, out); } + + /** Append EVERY physically-stored point index (live and tombstoned) into + * \a out. Used by the multi-threaded wrapper to detect which dataset slots + * become free after a background rebuild swap. */ + void collectPhysicalIndices(std::vector& out) const { collectAllRec(iroot_, out); } + + /** True if some tree node currently references the given point index (i.e. + * the dataset slot is in use and must not be recycled). */ + NANOFLANN_NODISCARD bool referencesIndex(IndexType idx) const + { + return idx < nodeOfPoint_.size() && nodeOfPoint_[idx] != nullptr; + } + + /** Discard the current tree and bulk-build a fresh, balanced tree over the + * given point indices. O(M log M). Reuses recycled nodes via the pool. */ + void buildFromIndices(const std::vector& idxs) + { + // Validate (and grow the point->node map) before touching the current + // tree, so a rejected index list leaves the index untouched. + IndexType maxIdx = 0; + for (IndexType v : idxs) maxIdx = std::max(maxIdx, v); + if (!idxs.empty()) ensureNodeMap(maxIdx); + + if (iroot_) + { + buildBuf_.clear(); + collectLiveAndFree(iroot_, buildBuf_); // recycle existing nodes + iroot_ = nullptr; + } + buildBuf_.assign(idxs.begin(), idxs.end()); + iroot_ = buildBalanced(buildBuf_, 0, buildBuf_.size(), 0, nullptr); + liveCount_ = buildBuf_.size(); + totalCount_ = liveCount_; + syncRootBox(); + } + + /** @} */ + + /** \name Capacity / observers + * @{ */ + + /** Number of live (non-removed) points currently in the index. */ + NANOFLANN_NODISCARD Size size() const noexcept { return liveCount_; } + NANOFLANN_NODISCARD bool empty() const noexcept { return liveCount_ == 0; } + + /** Number of nodes physically stored (live + not-yet-reclaimed tombstones). */ + NANOFLANN_NODISCARD Size physicalSize() const noexcept { return totalCount_; } + + /** Approximate bytes used by the node pool and the index->node map. */ + NANOFLANN_NODISCARD Size usedMemory() const + { + return Base::pool_.usedMemory + Base::pool_.wastedMemory + + nodeOfPoint_.capacity() * sizeof(INode*); + } + + /** Axis-aligned bounding box of all points currently in the index (live and + * not-yet-reclaimed tombstones — a conservative superset of the live set). + * O(1): returns the cached root box. Meaningless if empty() (all zeros). */ + NANOFLANN_NODISCARD BoundingBox boundingBox() const { return Base::root_bbox_; } + + /** Pre-size the internal index->node map (and the rebuild scratch buffer) to + * avoid reallocations while the point count grows toward \a n. */ + void reserve(Size n) + { + nodeOfPoint_.reserve(n); + buildBuf_.reserve(n); + } + + /** @} */ + + /** \name Query methods + * @{ */ + + /** Core search: find neighbors of \a vec, storing them in \a result. */ + template + bool findNeighbors( + RESULTSET& result, const ElementType* vec, const SearchParameters& searchParams = {}) const + { + assert(vec); + if (!iroot_ || liveCount_ == 0) return false; + const DistanceType epsError = 1 + static_cast(searchParams.eps); + + distance_vector_t dists; + assign(dists, this->veclen(*this), static_cast(0)); + const DistanceType dist = this->computeInitialDistances(*this, vec, dists); + searchLevelInc(result, vec, iroot_, dist, dists, epsError, this->veclen(*this)); + if (searchParams.sorted) result.sort(); + return result.full(); + } + + /** Find the \a num_closest nearest neighbors to \a query_point. */ + NANOFLANN_NODISCARD Size knnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const SearchParameters& searchParams = {}) const + { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances); + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** Radius search around \a query_point. */ + NANOFLANN_NODISCARD Size radiusSearch( + const ElementType* query_point, const DistanceType& radius, + std::vector>& IndicesDists, + const SearchParameters& searchParams = {}) const + { + RadiusResultSet resultSet(radius, IndicesDists); + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** Custom-callback radius search. */ + template + NANOFLANN_NODISCARD Size radiusSearchCustomCallback( + const ElementType* query_point, SEARCH_CALLBACK& resultSet, + const SearchParameters& searchParams = {}) const + { + findNeighbors(resultSet, query_point, searchParams); + return resultSet.size(); + } + + /** Radius-limited KNN around \a query_point. */ + NANOFLANN_NODISCARD Size rknnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const DistanceType& radius) const + { + nanoflann::RKNNResultSet resultSet(num_closest, radius); + resultSet.init(out_indices, out_distances); + findNeighbors(resultSet, query_point); + return resultSet.size(); + } + + /** Find all live points contained within the box \a bbox. */ + template + NANOFLANN_NODISCARD Size findWithinBox(RESULTSET& result, const BoundingBox& bbox) const + { + if (iroot_) findWithinBoxRec(result, iroot_, bbox); + return result.size(); + } + + /** @} */ + + /** \name Persistence (index topology only; the dataset is NOT stored) + * @{ */ + + /** Magic number written at the start of every saveIndex() stream of this + * incremental index. Distinct from the static KDTreeSingleIndexAdaptor's + * SAVE_MAGIC so that loading the wrong kind of file fails fast with a + * clear error instead of silently misinterpreting the bytes. + * Spells 'NFLI' (nanoflann incremental) in ASCII. */ + static constexpr uint32_t INCREMENTAL_SAVE_MAGIC = 0x4E464C49; + + /** Serializes the tree *topology* (split axis, tombstone flags, and the + * live/dead tree shape) to a binary stream. Point coordinates are NOT + * stored: as with the static index's saveIndex(), the object must be + * reattached to a dataset with the very same content before loadIndex() + * is called on it. + * + * Every other per-node field (bounding box, subtree size, tombstone + * counts, coordinate cache, parent pointer) is recomputed on load instead + * of stored, since they are all structural invariants derivable from the + * fields above; this keeps the format independent of DIM being + * compile-time fixed or runtime, unlike a raw struct byte-dump. + * + * \note Portability limitations as in the static index's saveIndex(): not + * portable across endianness, 32- vs 64-bit size_t, nanoflann versions, + * or IndexType/ElementType/DistanceType instantiations (checked, throws + * on mismatch). + * + * \sa loadIndex + */ + void saveIndex(std::ostream& stream) const + { + const uint32_t hdr_magic = INCREMENTAL_SAVE_MAGIC; + const uint32_t hdr_version = static_cast(NANOFLANN_VERSION); + const uint8_t hdr_sz_st = static_cast(sizeof(size_t)); + const uint8_t hdr_sz_idx = static_cast(sizeof(IndexType)); + const uint8_t hdr_sz_elem = static_cast(sizeof(ElementType)); + const uint8_t hdr_sz_dist = static_cast(sizeof(DistanceType)); + save_value(stream, hdr_magic); + save_value(stream, hdr_version); + save_value(stream, hdr_sz_st); + save_value(stream, hdr_sz_idx); + save_value(stream, hdr_sz_elem); + save_value(stream, hdr_sz_dist); + + const Dimension dims = static_cast(this->veclen(*this)); + save_value(stream, dims); + + const uint8_t hasRoot = iroot_ ? 1 : 0; + save_value(stream, hasRoot); + if (iroot_) saveNode(stream, iroot_); + } + + /** Loads a tree topology previously saved with saveIndex(). + * + * \note Must be called on a freshly constructed index (mirrors the static + * index's loadIndex() precondition): loading into an already-populated + * tree is not supported and would leak its nodes. The object must + * already be attached (via the constructor) to a dataset holding the + * very same points that were indexed when saveIndex() was called. + * + * \throws std::runtime_error on a magic-number/version/type-size/ + * dimensionality mismatch, or a stream read error. + * + * \sa saveIndex + */ + void loadIndex(std::istream& stream) + { + uint32_t magic = 0; + load_value(stream, magic); + if (stream.fail() || magic != INCREMENTAL_SAVE_MAGIC) + { + throw std::runtime_error( + "KDTreeSingleIndexIncrementalAdaptor::loadIndex: invalid file (wrong magic " + "number). The stream was not written by this class' saveIndex(), or was written " + "by the static KDTreeSingleIndexAdaptor instead."); + } + + uint32_t file_version = 0; + load_value(stream, file_version); + if (file_version != static_cast(NANOFLANN_VERSION)) + { + char msg[200]; + snprintf( + msg, sizeof(msg), + "KDTreeSingleIndexIncrementalAdaptor::loadIndex: version mismatch " + "(file=0x%03X, library=0x%03X). Rebuild the index.", + file_version, static_cast(NANOFLANN_VERSION)); + throw std::runtime_error(msg); + } + + uint8_t sz_size_t = 0; + uint8_t sz_idx = 0; + uint8_t sz_elem = 0; + uint8_t sz_dist = 0; + load_value(stream, sz_size_t); + load_value(stream, sz_idx); + load_value(stream, sz_elem); + load_value(stream, sz_dist); + if (sz_size_t != static_cast(sizeof(size_t)) || + sz_idx != static_cast(sizeof(IndexType)) || + sz_elem != static_cast(sizeof(ElementType)) || + sz_dist != static_cast(sizeof(DistanceType))) + { + throw std::runtime_error( + "KDTreeSingleIndexIncrementalAdaptor::loadIndex: type-size mismatch between " + "saved index and current template instantiation (sizeof size_t / IndexType / " + "ElementType / DistanceType differ). Rebuild the index."); + } + + Dimension dims = 0; + load_value(stream, dims); + if (dims != static_cast(this->veclen(*this))) + { + throw std::runtime_error( + "KDTreeSingleIndexIncrementalAdaptor::loadIndex: dimensionality mismatch " + "between the saved index and this object's dataset."); + } + + uint8_t hasRoot = 0; + load_value(stream, hasRoot); + iroot_ = hasRoot ? loadNode(stream, nullptr) : nullptr; + + if (stream.fail()) + { + throw std::runtime_error( + "KDTreeSingleIndexIncrementalAdaptor::loadIndex: unexpected end of stream or " + "read error."); + } + + totalCount_ = iroot_ ? iroot_->subtree_size : 0; + liveCount_ = iroot_ ? iroot_->subtree_size - iroot_->invalid_count : 0; + syncRootBox(); + } + + /** @} */ + + private: + // -------------------------------------------------------------------- + // Node allocation (bump-allocate from the pool, recycle via free-list) + // -------------------------------------------------------------------- + INode* allocNode() + { + if (freeList_) + { + INode* n = freeList_; + freeList_ = n->child1; + return n; // already constructed; box storage reused + } + INode* n = Base::pool_.template allocate(); + // Placement-new so that, for DIM=-1, the std::vector box is constructed. + ::new (static_cast(n)) INode(); + resize(n->box, static_cast(this->veclen(*this))); + if (kCacheCoords) resize(n->pcoord, static_cast(this->veclen(*this))); + return n; + } + + /** Fill the node's cached coordinates from the dataset (fixed DIM only). */ + void cacheCoords(INode* n) + { + if (!kCacheCoords) return; + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) n->pcoord[i] = pt(n->ptIdx, i); + } + + /** Coordinate of a node's own point along axis \a d (cached for fixed DIM). */ + ElementType nodeCoord(const INode* n, Dimension d) const + { + return kCacheCoords ? n->pcoord[d] : pt(n->ptIdx, d); + } + + /** Whether a node's own point lies inside box \a b (uses the coord cache). */ + bool nodeInBox(const INode* n, const BoundingBox& b) const + { + if (!kCacheCoords) return pointInBox(n->ptIdx, b); + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + if (n->pcoord[i] < b[i].low || n->pcoord[i] > b[i].high) return false; + return true; + } + + void recycleNode(INode* n) + { + n->child1 = freeList_; + freeList_ = n; + } + + /** Destroy all constructed INode objects (only needed when the box type is + * not trivially destructible, i.e. DIM=-1 where it owns a std::vector). */ + void destroyNodeObjects() + { + if (std::is_trivially_destructible::value) return; + destroySubtree(iroot_); + iroot_ = nullptr; + while (freeList_) + { + INode* n = freeList_; + freeList_ = n->child1; + n->~INode(); + } + } + + void destroySubtree(INode* n) + { + if (!n) return; + destroySubtree(n->child1); + destroySubtree(n->child2); + n->~INode(); + } + + // -------------------------------------------------------------------- + // Helpers + // -------------------------------------------------------------------- + ElementType pt(IndexType idx, Dimension d) const { return dataset_.kdtree_get_pt(idx, d); } + + void ensureNodeMap(IndexType idx) + { + // The point->node map is grown to hold `idx`, so a bogus index is not a + // wrong result but an out-of-memory: the maximum IndexType value alone + // asks for a 2^32-entry (32 GB) map on the default uint32_t. That value + // is what `static_cast(n - 1)` produces when a caller forgets + // to special-case an empty (n == 0) dataset, and it would additionally + // make the [start,end] loops of addPoints() wrap around forever, so + // reject it here, where every index-taking entry point passes through. + if (idx == (std::numeric_limits::max)()) + { + throw std::invalid_argument( + "[nanoflann] KDTreeSingleIndexIncrementalAdaptor: point index equal to the " + "maximum IndexType value; this is almost certainly an underflowed 'size - 1' " + "on an empty dataset."); + } + if (idx >= nodeOfPoint_.size()) nodeOfPoint_.resize(static_cast(idx) + 1, nullptr); + } + + void syncRootBox() + { + const Dimension dims = static_cast(this->veclen(*this)); + if (iroot_) + for (Dimension i = 0; i < dims; ++i) Base::root_bbox_[i] = iroot_->box[i]; + else + for (Dimension i = 0; i < dims; ++i) Base::root_bbox_[i] = Interval{0, 0}; + } + + void initBoxToPoint(INode* n) + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + { + const ElementType v = pt(n->ptIdx, i); + n->box[i].low = n->box[i].high = v; + } + } + + void expandBoxToPoint(INode* n, IndexType idx) + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + { + const ElementType v = pt(idx, i); + if (v < n->box[i].low) n->box[i].low = v; + if (v > n->box[i].high) n->box[i].high = v; + } + } + + void unionBox(INode* n, const INode* c) + { + if (!c) return; + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + { + if (c->box[i].low < n->box[i].low) n->box[i].low = c->box[i].low; + if (c->box[i].high > n->box[i].high) n->box[i].high = c->box[i].high; + } + } + + bool pointInBox(IndexType idx, const BoundingBox& b) const + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + { + const ElementType v = pt(idx, i); + if (v < b[i].low || v > b[i].high) return false; + } + return true; + } + + bool boxFullyInside(const BoundingBox& inner, const BoundingBox& outer) const + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + if (inner[i].low < outer[i].low || inner[i].high > outer[i].high) return false; + return true; + } + + bool boxDisjoint(const BoundingBox& a, const BoundingBox& b) const + { + const Dimension dims = static_cast(this->veclen(*this)); + for (Dimension i = 0; i < dims; ++i) + if (a[i].high < b[i].low || a[i].low > b[i].high) return true; + return false; + } + + // -------------------------------------------------------------------- + // Insertion + // -------------------------------------------------------------------- + INode* makeLeaf(IndexType idx, Dimension depth, INode* parent) + { + INode* n = allocNode(); + const Dimension dims = static_cast(this->veclen(*this)); + n->ptIdx = idx; + n->divfeat = static_cast(depth % dims); + n->deleted = false; + n->treeDeleted = false; + n->child1 = n->child2 = nullptr; + n->parent = parent; + n->subtree_size = 1; + n->invalid_count = 0; + initBoxToPoint(n); + cacheCoords(n); + nodeOfPoint_[idx] = n; + return n; + } + + /** Insert one point and, in the same pass, rebuild the highest unbalanced + * node on the insertion path (BB[alpha] scapegoat rebalancing). */ + void insertOne(IndexType idx) + { + pendingRebuild_ = nullptr; + iroot_ = insertRec(iroot_, idx, 0, nullptr); + ++liveCount_; + ++totalCount_; + if (pendingRebuild_ && inlineRebuild_) rebuildAt(pendingRebuild_); + } + + INode* insertRec(INode* node, IndexType idx, Dimension depth, INode* parent) + { + if (!node) return makeLeaf(idx, depth, parent); + if (node->treeDeleted) pushDownDelete(node); + + ++node->subtree_size; + expandBoxToPoint(node, idx); + + const Dimension axis = node->divfeat; + if (pt(idx, axis) < nodeCoord(node, axis)) + node->child1 = insertRec(node->child1, idx, static_cast(depth + 1), node); + else + node->child2 = insertRec(node->child2, idx, static_cast(depth + 1), node); + + // On the unwind, remember the *highest* unbalanced node seen on the + // path (ancestors are visited after descendants, so the last write + // wins). insertOne() rebuilds it once, avoiding a second descent. + if (isBalanceScapegoat(node)) pendingRebuild_ = node; + return node; + } + + /** Make the lazy whole-subtree tombstone one level explicit so the subtree + * is consistent before we descend into it for an insertion. */ + void pushDownDelete(INode* node) + { + node->deleted = true; + if (node->child1) + { + node->child1->treeDeleted = true; + node->child1->invalid_count = node->child1->subtree_size; + } + if (node->child2) + { + node->child2->treeDeleted = true; + node->child2->invalid_count = node->child2->subtree_size; + } + node->treeDeleted = false; // invalid_count already == subtree_size + } + + Size maxChildSize(const INode* node) const + { + const Size l = node->child1 ? node->child1->subtree_size : 0; + const Size r = node->child2 ? node->child2->subtree_size : 0; + return l > r ? l : r; + } + + bool isBalanceScapegoat(const INode* node) const + { + if (node->subtree_size < kMinBalanceRebuild) return false; + return static_cast(maxChildSize(node)) > + alphaBal_ * static_cast(node->subtree_size); + } + + // -------------------------------------------------------------------- + // Deletion (lazy) + box-region deletion + // -------------------------------------------------------------------- + /** Kill an entire subtree in O(1): mark it as wholly tombstoned. */ + void killSubtree(INode* node) + { + node->treeDeleted = true; + node->invalid_count = node->subtree_size; + } + + /** Remove points outside \a keep. Returns the number newly tombstoned. */ + Size removeOutsideBoxRec(INode* node, const BoundingBox& keep) + { + if (!node) return 0; + if (node->invalid_count == node->subtree_size) return 0; // already all dead + if (boxFullyInside(node->box, keep)) return 0; // keep entire subtree + if (boxDisjoint(node->box, keep)) + { + const Size newly = node->subtree_size - node->invalid_count; + killSubtree(node); + liveCount_ -= newly; + return newly; + } + Size newly = 0; + if (!node->deleted && !nodeInBox(node, keep)) + { + node->deleted = true; + ++newly; + --liveCount_; + } + newly += removeOutsideBoxRec(node->child1, keep); + newly += removeOutsideBoxRec(node->child2, keep); + node->invalid_count += newly; + return newly; + } + + /** Remove points inside \a box. Returns the number newly tombstoned. */ + Size removeBoxRec(INode* node, const BoundingBox& box) + { + if (!node) return 0; + if (node->invalid_count == node->subtree_size) return 0; + if (boxDisjoint(node->box, box)) return 0; // nothing inside + if (boxFullyInside(node->box, box)) + { + const Size newly = node->subtree_size - node->invalid_count; + killSubtree(node); + liveCount_ -= newly; + return newly; + } + Size newly = 0; + if (!node->deleted && nodeInBox(node, box)) + { + node->deleted = true; + ++newly; + --liveCount_; + } + newly += removeBoxRec(node->child1, box); + newly += removeBoxRec(node->child2, box); + node->invalid_count += newly; + return newly; + } + + bool isDeletionScapegoat(const INode* node) const + { + if (node->subtree_size == 0) return false; + return static_cast(node->invalid_count) > + alphaDel_ * static_cast(node->subtree_size); + } + + INode* findDeletionScapegoat(INode* node) const + { + if (!node) return nullptr; + if (isDeletionScapegoat(node)) return node; // highest wins + if (INode* l = findDeletionScapegoat(node->child1)) return l; + return findDeletionScapegoat(node->child2); + } + + void maybeRebuildForDeletion() + { + if (!iroot_ || !inlineRebuild_) return; + if (INode* sg = findDeletionScapegoat(iroot_)) rebuildAt(sg); + } + + // -------------------------------------------------------------------- + // Partial rebuild (scapegoat): flatten live points, rebuild balanced + // -------------------------------------------------------------------- + void rebuildAt(INode* node) + { + INode* par = node->parent; + INode** link = par ? (par->child1 == node ? &par->child1 : &par->child2) : &iroot_; + + const Size oldSize = node->subtree_size; + const Size oldInvalid = node->invalid_count; + + buildBuf_.clear(); + collectLiveAndFree(node, buildBuf_); + + INode* nb = buildBalanced(buildBuf_, 0, buildBuf_.size(), 0, par); + *link = nb; + + const Size newSize = nb ? nb->subtree_size : 0; // == number of live pts + // Propagate the change in (size, invalid) up to the ancestors. + for (INode* p = par; p; p = p->parent) + { + p->subtree_size = p->subtree_size - oldSize + newSize; + p->invalid_count = p->invalid_count - oldInvalid; + } + totalCount_ = totalCount_ - oldSize + newSize; + } + + /** Collect the live point indices under \a node (DFS) and recycle every + * node to the free-list. Tombstoned points are dropped (and optionally + * recorded for acquireRemovedPoints()). */ + void collectLiveAndFree(INode* node, std::vector& out) + { + if (!node) return; + if (node->treeDeleted) + { + freeDeadSubtree(node); + return; + } + if (!node->deleted) + out.push_back(node->ptIdx); + else + dropDeadPoint(node->ptIdx); + collectLiveAndFree(node->child1, out); + collectLiveAndFree(node->child2, out); + recycleNode(node); + } + + void freeDeadSubtree(INode* node) + { + if (!node) return; + dropDeadPoint(node->ptIdx); + freeDeadSubtree(node->child1); + freeDeadSubtree(node->child2); + recycleNode(node); + } + + void dropDeadPoint(IndexType idx) + { + if (idx < nodeOfPoint_.size()) nodeOfPoint_[idx] = nullptr; + if (collectRemoved_) removedSink_.push_back(idx); + } + + /** DFS collecting live point indices (skips tombstoned points/subtrees). */ + void snapshotRec(const INode* node, std::vector& out) const + { + if (!node) return; + if (node->invalid_count == node->subtree_size) return; // whole subtree dead + if (!node->deleted) out.push_back(node->ptIdx); + snapshotRec(node->child1, out); + snapshotRec(node->child2, out); + } + + /** DFS collecting every physically-stored index (live and tombstoned). */ + void collectAllRec(const INode* node, std::vector& out) const + { + if (!node) return; + out.push_back(node->ptIdx); + collectAllRec(node->child1, out); + collectAllRec(node->child2, out); + } + + // -------------------------------------------------------------------- + // Persistence (see saveIndex()/loadIndex()) + // -------------------------------------------------------------------- + /** Writes one node and its subtree. Only the fields that cannot be + * derived from the rest are stored; see loadNode(). */ + void saveNode(std::ostream& stream, const INode* n) const + { + save_value(stream, n->ptIdx); + save_value(stream, n->divfeat); + save_value(stream, n->deleted); + save_value(stream, n->treeDeleted); + + const uint8_t hasChild1 = n->child1 ? 1 : 0; + save_value(stream, hasChild1); + if (n->child1) saveNode(stream, n->child1); + + const uint8_t hasChild2 = n->child2 ? 1 : 0; + save_value(stream, hasChild2); + if (n->child2) saveNode(stream, n->child2); + } + + /** Loads one node and its subtree. Reconstructs every field not written + * by saveNode() (box, subtree_size, invalid_count, coordinate cache, + * parent link, and the idx->node map entry) from invariants that hold + * regardless of how the tree was originally built: + * - subtree_size is always 1 + the children's (0 if absent). + * - invalid_count equals subtree_size for a treeDeleted node (see + * killSubtree()), otherwise it is this node's own deleted flag plus + * the children's invalid_count. + * - box is this node's point, unioned with the children's boxes. + */ + INode* loadNode(std::istream& stream, INode* parent) + { + INode* n = allocNode(); + load_value(stream, n->ptIdx); + load_value(stream, n->divfeat); + load_value(stream, n->deleted); + load_value(stream, n->treeDeleted); + n->parent = parent; + + uint8_t hasChild1 = 0; + load_value(stream, hasChild1); + n->child1 = hasChild1 ? loadNode(stream, n) : nullptr; + + uint8_t hasChild2 = 0; + load_value(stream, hasChild2); + n->child2 = hasChild2 ? loadNode(stream, n) : nullptr; + + n->subtree_size = 1 + (n->child1 ? n->child1->subtree_size : 0) + + (n->child2 ? n->child2->subtree_size : 0); + n->invalid_count = n->treeDeleted ? n->subtree_size + : static_cast(n->deleted ? 1 : 0) + + (n->child1 ? n->child1->invalid_count : 0) + + (n->child2 ? n->child2->invalid_count : 0); + + initBoxToPoint(n); + unionBox(n, n->child1); + unionBox(n, n->child2); + cacheCoords(n); + + ensureNodeMap(n->ptIdx); + nodeOfPoint_[n->ptIdx] = n; + + return n; + } + + /** Build a balanced subtree over buf[lo,hi) (median split on widest axis). */ + INode* buildBalanced( + std::vector& buf, size_t lo, size_t hi, Dimension depth, INode* parent) + { + if (lo >= hi) return nullptr; + const Dimension dims = static_cast(this->veclen(*this)); + + // Widest-spread axis over buf[lo,hi). + Dimension axis = static_cast(depth % dims); + ElementType bestSpan = -1; + for (Dimension d = 0; d < dims; ++d) + { + ElementType mn = pt(buf[lo], d), mx = mn; + for (size_t k = lo + 1; k < hi; ++k) + { + const ElementType v = pt(buf[k], d); + if (v < mn) mn = v; + if (v > mx) mx = v; + } + const ElementType span = mx - mn; + if (span > bestSpan) + { + bestSpan = span; + axis = d; + } + } + + const size_t mid = lo + (hi - lo) / 2; + std::nth_element( + buf.begin() + lo, buf.begin() + mid, buf.begin() + hi, + [this, axis](IndexType a, IndexType b) { return pt(a, axis) < pt(b, axis); }); + + INode* node = allocNode(); + node->ptIdx = buf[mid]; + node->divfeat = axis; + node->deleted = false; + node->treeDeleted = false; + node->parent = parent; + cacheCoords(node); + nodeOfPoint_[buf[mid]] = node; + + node->child1 = buildBalanced(buf, lo, mid, static_cast(depth + 1), node); + node->child2 = buildBalanced(buf, mid + 1, hi, static_cast(depth + 1), node); + + node->subtree_size = hi - lo; + node->invalid_count = 0; + initBoxToPoint(node); + unionBox(node, node->child1); + unionBox(node, node->child2); + return node; + } + + // -------------------------------------------------------------------- + // Search + // -------------------------------------------------------------------- + template + void searchLevelInc( + RESULTSET& rs, const ElementType* vec, const INode* node, DistanceType mindist, + distance_vector_t& dists, const DistanceType epsError, const Size dim) const + { + if (!node) return; + if (node->invalid_count == node->subtree_size) return; // whole subtree dead + + if (!node->deleted) + { +#if defined(NANOFLANN_INCREMENTAL_INNODE_DISTANCE) + // Opt-in: compute the node distance from the in-node coordinate cache + // as a sum of per-axis accum_dist contributions. This avoids the + // dataset_get() indirection and is ~12% faster on KNN, but is only + // valid for *additive* (axis-decomposable) metrics — L1, L2, + // L2_Simple. Do NOT enable it for SO2/SO3. + DistanceType d = DistanceType(); + if (kCacheCoords) + for (Size i = 0; i < dim; ++i) + d += distance_.accum_dist( + vec[i], node->pcoord[static_cast(i)], static_cast(i)); + else + d = distance_.evalMetric(vec, node->ptIdx, dim); +#else + const DistanceType d = distance_.evalMetric(vec, node->ptIdx, dim); +#endif + if (d < rs.worstDist()) + rs.addPoint( + static_cast(d), + static_cast(node->ptIdx)); + } + + const Dimension axis = node->divfeat; + const ElementType splitval = nodeCoord(node, axis); + const ElementType val = vec[axis]; + const DistanceType cut = distance_.accum_dist(val, splitval, axis); + + const INode* nearChild; + const INode* farChild; + if (val < splitval) + { + nearChild = node->child1; + farChild = node->child2; + } + else + { + nearChild = node->child2; + farChild = node->child1; + } + + searchLevelInc(rs, vec, nearChild, mindist, dists, epsError, dim); + + const DistanceType dst = dists[axis]; + const DistanceType newmin = mindist + cut - dst; + dists[axis] = cut; + if (newmin * epsError <= rs.worstDist()) + searchLevelInc(rs, vec, farChild, newmin, dists, epsError, dim); + dists[axis] = dst; + } + + template + void findWithinBoxRec(RESULTSET& result, const INode* node, const BoundingBox& bbox) const + { + if (!node) return; + if (node->invalid_count == node->subtree_size) return; + if (boxDisjoint(node->box, bbox)) return; + if (!node->deleted && nodeInBox(node, bbox)) result.addPoint(0, node->ptIdx); + findWithinBoxRec(result, node->child1, bbox); + findWithinBoxRec(result, node->child2, bbox); + } +}; + +#ifndef NANOFLANN_NO_THREADS +/** Multi-threaded variant of KDTreeSingleIndexIncrementalAdaptor that hides the + * O(N) near-root rebuild spike: the large balancing rebuild is performed on a + * background thread while the foreground tree keeps serving inserts, deletions + * and queries. + * + * Model (single foreground thread + one background rebuild thread): + * - The live ("active") tree never rebalances inline; it only appends and + * lazily tombstones, so every foreground call returns quickly. + * - When the active tree has grown / accumulated tombstones past a threshold, + * a snapshot of its live point indices is taken and a background thread + * bulk-builds a fresh, balanced tree from it. Foreground operations meanwhile + * keep mutating the active tree and are appended to a small op-log. + * That background thread is **one persistent worker** started on the first + * rebuild and reused for every later one, not one thread per rebuild: a + * long-running incremental map triggers rebuilds continuously, and a thread + * per rebuild would make the host process churn through thousands of OS + * threads, multiplying every per-thread cost it carries (malloc arenas, + * sanitizer bookkeeping) for no benefit. + * - At the next foreground call after the build finishes, the op-log is + * replayed onto the fresh tree and it atomically replaces the active tree. + * + * This keeps the foreground tail latency bounded (snapshot + replay) instead of + * paying the full O(N) rebuild inline. It matches ikd-Tree's async-rebuild idea + * but with a thread-isolated build (no per-node locks), a bounded `std::deque`- + * style op-log (no fixed 10⁶ queue), and no PCL/pthread dependency. + * + * \warning Same threading contract as the synchronous index for the *foreground* + * thread (const queries are safe for concurrent readers; no concurrent writer). + * Additionally, because the background thread reads point coordinates from the + * dataset adaptor, the **dataset must keep stable element storage while a + * rebuild is in flight** (e.g. `reserve()` the backing vector, or use a + * std::deque) so that appends on the foreground thread do not reallocate it. + * Disabled entirely under NANOFLANN_NO_THREADS. + */ +template +class KDTreeSingleIndexIncrementalAdaptorMT +{ + public: + using Inner = KDTreeSingleIndexIncrementalAdaptor; + using ElementType = typename Inner::ElementType; + using DistanceType = typename Inner::DistanceType; + using Size = typename Inner::Size; + using Dimension = typename Inner::Dimension; + using BoundingBox = typename Inner::BoundingBox; + + /** Constructor. + * @param rebuild_growth Trigger a background rebuild once the physical node + * count exceeds this factor of the live count at the last rebuild + * (captures both appends and tombstone accumulation). + * @param min_rebuild_size Never trigger below this many physical nodes. */ + explicit KDTreeSingleIndexIncrementalAdaptorMT( + const Dimension dimensionality, const DatasetAdaptor& inputData, + const KDTreeIncrementalIndexParams& params = {}, double rebuild_growth = 1.3, + Size min_rebuild_size = 10000) + : dataset_(inputData), + dim_(dimensionality), + params_(params), + rebuildGrowth_(rebuild_growth), + minRebuildSize_(min_rebuild_size) + { + active_.reset(new Inner(dimensionality, inputData, params)); + active_->setInlineRebuild(false); + } + + KDTreeSingleIndexIncrementalAdaptorMT(const KDTreeSingleIndexIncrementalAdaptorMT&) = delete; + KDTreeSingleIndexIncrementalAdaptorMT& operator=(const KDTreeSingleIndexIncrementalAdaptorMT&) = + delete; + + ~KDTreeSingleIndexIncrementalAdaptorMT() + { + // Lets any in-flight build finish before teardown: the worker reads the + // caller's dataset, which is typically freed right after this returns. + stopWorker(); + } + + /** \name Modifiers @{ */ + void addPoints(IndexType start, IndexType end) + { + integrateIfReady(); + active_->addPoints(start, end); + if (building_) log_.push_back({OpKind::Add, start, end, {}}); + maybeTriggerRebuild(); + } + void addPoint(IndexType idx) { addPoints(idx, idx); } + + void removePoint(IndexType idx) + { + integrateIfReady(); + active_->removePoint(idx); + if (building_) log_.push_back({OpKind::Remove, idx, idx, {}}); + maybeTriggerRebuild(); + } + void removeBox(const BoundingBox& box) + { + integrateIfReady(); + active_->removeBox(box); + if (building_) log_.push_back({OpKind::RemoveBox, 0, 0, box}); + maybeTriggerRebuild(); + } + void removeOutsideBox(const BoundingBox& keep) + { + integrateIfReady(); + active_->removeOutsideBox(keep); + if (building_) log_.push_back({OpKind::RemoveOutsideBox, 0, 0, keep}); + maybeTriggerRebuild(); + } + /** @} */ + + /** \name Query methods (forwarded to the active tree) @{ */ + template + bool findNeighbors( + RESULTSET& result, const ElementType* vec, const SearchParameters& sp = {}) const + { + return active_->findNeighbors(result, vec, sp); + } + Size knnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const SearchParameters& sp = {}) const + { + return active_->knnSearch(query_point, num_closest, out_indices, out_distances, sp); + } + Size radiusSearch( + const ElementType* query_point, const DistanceType& radius, + std::vector>& IndicesDists, + const SearchParameters& sp = {}) const + { + return active_->radiusSearch(query_point, radius, IndicesDists, sp); + } + Size rknnSearch( + const ElementType* query_point, const Size num_closest, IndexType* out_indices, + DistanceType* out_distances, const DistanceType& radius) const + { + return active_->rknnSearch(query_point, num_closest, out_indices, out_distances, radius); + } + template + Size findWithinBox(RESULTSET& result, const BoundingBox& bbox) const + { + return active_->findWithinBox(result, bbox); + } + /** @} */ + + /** \name Observers @{ */ + Size size() const noexcept { return active_->size(); } + bool empty() const noexcept { return active_->empty(); } + Size physicalSize() const noexcept { return active_->physicalSize(); } + bool isRebuilding() const noexcept { return building_; } + + /** Live AABB of the active tree (see the synchronous index). */ + NANOFLANN_NODISCARD BoundingBox boundingBox() const { return active_->boundingBox(); } + + /** Append the live point indices of the active tree into \a out. */ + void snapshotLiveIndices(std::vector& out) const + { + active_->snapshotLiveIndices(out); + } + + /** Pre-size the active tree's internal map (see the synchronous index). */ + void reserve(Size n) { active_->reserve(n); } + + /** Block until any in-flight background rebuild has been integrated. */ + void sync() + { + if (building_) + { + std::unique_lock lk(workerMtx_); + workerCvDone_.wait(lk, [this] { return resultReady_; }); + } + integrateIfReady(); + } + + /** Access the underlying active tree (e.g. for further query types). */ + const Inner& activeIndex() const { return *active_; } + /** @} */ + + /** \name Persistence (see the synchronous index's saveIndex()/loadIndex()) + * @{ */ + + /** Blocks until any in-flight background rebuild is integrated (see + * sync()), then serializes the active tree's topology. */ + void saveIndex(std::ostream& stream) + { + sync(); + active_->saveIndex(stream); + } + + /** Blocks until any in-flight background rebuild is integrated, then loads + * a previously saved topology into the active tree (see the synchronous + * index's loadIndex() precondition: the active tree must be freshly + * constructed / empty). */ + void loadIndex(std::istream& stream) + { + sync(); + active_->loadIndex(stream); + lastBuildLive_ = active_->size(); + } + /** @} */ + + /** Set a callback invoked **on the background worker thread** on each freshly + * built tree, right after it is balanced and before it is handed back for + * integration. Lets the caller recompute per-point auxiliary data (e.g. + * covariances) off the foreground thread. The callback runs concurrently + * with foreground queries on the *old* tree, so it must only touch the + * passed-in fresh index and its own/snapshot data — never foreground-shared + * state without external synchronization. Pass {} to clear. */ + void setRebuildCallback(std::function cb) { rebuildCallback_ = std::move(cb); } + + /** \name Dataset-storage reclamation @{ */ + + /** Enable recording of dataset slots that become free when a background + * rebuild drops tombstoned points (so the caller can recycle them and keep + * the dataset bounded). Off by default (cost-free when unused). */ + void setCollectRemovedPoints(bool enable) + { + collectRemoved_ = enable; + if (!enable) std::vector().swap(removedSink_); + } + + /** Move out the point indices whose dataset slots became free since the last + * call (no live OR tombstoned tree node references them — safe to recycle). + * Requires setCollectRemovedPoints(true). */ + std::vector acquireRemovedPoints() + { + std::vector out; + out.swap(removedSink_); + return out; + } + /** @} */ + + private: + enum class OpKind + { + Add, + Remove, + RemoveBox, + RemoveOutsideBox + }; + struct LoggedOp + { + OpKind kind; + IndexType a, b; + BoundingBox box; + }; + + void maybeTriggerRebuild() + { + if (building_) return; + const Size phys = active_->physicalSize(); + if (phys < minRebuildSize_) return; + const Size base = lastBuildLive_ ? lastBuildLive_ : Size(1); + if (static_cast(phys) < rebuildGrowth_ * static_cast(base)) return; + + // Snapshot the live indices on the foreground thread, then hand them to + // the background worker, which bulk-builds a fresh balanced tree. + auto snapshot = std::make_shared>(); + active_->snapshotLiveIndices(*snapshot); + + // Started on the first rebuild only, so an index that never rebuilds + // costs no thread at all: + if (!workerThread_.joinable()) + { + workerThread_ = std::thread(&KDTreeSingleIndexIncrementalAdaptorMT::workerLoop, this); + } + + { + std::lock_guard lk(workerMtx_); + pendingJob_ = std::move(snapshot); + // Copied per job, so the worker never reads rebuildCallback_ while + // setRebuildCallback() writes it: + pendingCallback_ = rebuildCallback_; + builtTree_.reset(); + buildError_ = nullptr; + resultReady_ = false; + } + workerCvJob_.notify_one(); + + building_ = true; + log_.clear(); + } + + /** Body of the single persistent rebuild worker: wait for a job, bulk-build + * a balanced tree from it, publish the result, repeat. */ + void workerLoop() + { + for (;;) + { + std::shared_ptr> job; + std::function cb; + { + std::unique_lock lk(workerMtx_); + workerCvJob_.wait(lk, [this] { return workerStop_ || pendingJob_ != nullptr; }); + // A job already handed over is always finished before stopping: + // the destructor relies on that to keep the caller's dataset + // alive for as long as the build reads it. + if (workerStop_ && !pendingJob_) return; + job = std::move(pendingJob_); + cb = std::move(pendingCallback_); + } + + // dim_, dataset_ and params_ are set at construction and never + // mutated, so the worker reads them without synchronization. + std::unique_ptr t; + std::exception_ptr err; + try + { + t.reset(new Inner(dim_, dataset_, params_)); + t->setInlineRebuild(false); + t->buildFromIndices(*job); + if (cb) cb(*t); // background post-rebuild hook (e.g. recompute covariances) + } + catch (...) + { + // Handed to the foreground thread instead of terminating the + // process; see integrateIfReady(). + err = std::current_exception(); + t.reset(); + } + + { + std::lock_guard lk(workerMtx_); + builtTree_ = std::move(t); + buildError_ = err; + resultReady_ = true; + } + workerCvDone_.notify_all(); + } + } + + /** Stops the worker once its current build (if any) is done, and joins it. */ + void stopWorker() + { + if (!workerThread_.joinable()) return; + { + std::lock_guard lk(workerMtx_); + workerStop_ = true; + } + workerCvJob_.notify_all(); + workerThread_.join(); + } + + void integrateIfReady() + { + if (!building_) return; + + std::unique_ptr fresh; + std::exception_ptr err; + { + std::lock_guard lk(workerMtx_); + if (!resultReady_) return; // the rebuild is still running + resultReady_ = false; + fresh = std::move(builtTree_); + err = buildError_; + buildError_ = nullptr; + } + + // However this attempt ends (integrated, failed build, or a failed + // replay below), it is over: return to the "not rebuilding" state so a + // later rebuild can be triggered again. Latching `building_` would + // silently disable rebuilding for good, and with it the reclaiming of + // tombstoned nodes and of the dataset slots reported by + // acquireRemovedPoints() (i.e. unbounded growth), and it would make a + // later sync() wait forever for a result nobody is producing. + struct EndOfRebuild + { + KDTreeSingleIndexIncrementalAdaptorMT& self; + ~EndOfRebuild() + { + self.log_.clear(); + self.building_ = false; + } + } endOfRebuild{*this}; + + // The build failed (e.g. std::bad_alloc). The active tree was never + // touched by it, so it is still correct: just drop the attempt. + if (err) std::rethrow_exception(err); + + fresh->setInlineRebuild(false); + // Replay the operations buffered while the background build was running. + for (const auto& op : log_) + { + switch (op.kind) + { + case OpKind::Add: + fresh->addPoints(op.a, op.b); + break; + case OpKind::Remove: + fresh->removePoint(op.a); + break; + case OpKind::RemoveBox: + fresh->removeBox(op.box); + break; + case OpKind::RemoveOutsideBox: + fresh->removeOutsideBox(op.box); + break; + } + } + // Dataset slots referenced by the OLD tree but not by the fresh one are + // now free for the caller to recycle (no node references them anymore). + if (collectRemoved_) + { + std::vector oldPhysical; + active_->collectPhysicalIndices(oldPhysical); + for (IndexType idx : oldPhysical) + if (!fresh->referencesIndex(idx)) removedSink_.push_back(idx); + } + active_ = std::move(fresh); + lastBuildLive_ = active_->size(); + } + + const DatasetAdaptor& dataset_; + Dimension dim_; + KDTreeIncrementalIndexParams params_; + double rebuildGrowth_; + Size minRebuildSize_; + + std::unique_ptr active_; + /// Foreground-only: a rebuild has been handed to the worker and not + /// integrated back yet. + bool building_ = false; + Size lastBuildLive_ = 0; + std::vector log_; + + /** @name Background rebuild worker + * One persistent thread, created on the first rebuild (see + * maybeTriggerRebuild()) and joined by the destructor. Everything below is + * guarded by workerMtx_, except workerThread_ itself, which is only + * touched by the foreground thread. + * @{ */ + std::thread workerThread_; + std::mutex workerMtx_; + std::condition_variable workerCvJob_; // foreground -> worker + std::condition_variable workerCvDone_; // worker -> foreground + std::shared_ptr> pendingJob_; // live indices to build from + std::function pendingCallback_; + std::unique_ptr builtTree_; // result of the last build + std::exception_ptr buildError_; // ...or how it failed + bool resultReady_ = false; + bool workerStop_ = false; + /** @} */ + + bool collectRemoved_ = false; + std::vector removedSink_; + std::function rebuildCallback_; +}; +#endif // NANOFLANN_NO_THREADS + +/** An L2-metric KD-tree adaptor for working with data directly stored in an + * Eigen Matrix, without duplicating the data storage. You can select whether a + * row or column in the matrix represents a point in the state space. + * + * Example of usage: + * \code + * Eigen::Matrix mat; + * + * // Fill out "mat"... + * using my_kd_tree_t = nanoflann::KDTreeEigenMatrixAdaptor< + * Eigen::Matrix>; + * + * const int max_leaf = 10; + * my_kd_tree_t mat_index(mat, max_leaf); + * mat_index.index->... + * \endcode + * + * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality + * for the points in the data set, allowing more compiler optimizations. + * \tparam Distance The distance metric to use: nanoflann::metric_L1, + * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. + * \tparam row_major If set to true the rows of the matrix are used as the + * points, if set to false the columns of the matrix are used as the + * points. + */ +template < + class MatrixType, int32_t DIM = -1, class Distance = nanoflann::metric_L2, + bool row_major = true> +struct KDTreeEigenMatrixAdaptor +{ + using self_t = KDTreeEigenMatrixAdaptor; + using num_t = typename MatrixType::Scalar; + using IndexType = typename MatrixType::Index; + using metric_t = typename Distance::template traits::distance_t; + + using index_t = KDTreeSingleIndexAdaptor< + metric_t, self_t, row_major ? MatrixType::ColsAtCompileTime : MatrixType::RowsAtCompileTime, + IndexType>; + + index_t* index_; //! The kd-tree index for the user to call its methods as + //! usual with any other FLANN index. + + using Offset = typename index_t::Offset; + using Size = typename index_t::Size; + using Dimension = typename index_t::Dimension; + + /// Constructor: takes a const ref to the matrix object with the data points + explicit KDTreeEigenMatrixAdaptor( + const Dimension dimensionality, const std::reference_wrapper& mat, + const int leaf_max_size = 10, const unsigned int n_thread_build = 1) + : m_data_matrix(mat) + { + const auto dims = row_major ? mat.get().cols() : mat.get().rows(); + if (static_cast(dims) != dimensionality) + throw std::runtime_error( + "Error: 'dimensionality' must match column count in data " + "matrix"); + if (DIM > 0 && static_cast(dims) != DIM) + throw std::runtime_error( + "Data set dimensionality does not match the 'DIM' template " + "argument"); + index_ = new index_t( + static_cast(dims), *this /* adaptor */, + nanoflann::KDTreeSingleIndexAdaptorParams( + leaf_max_size, nanoflann::KDTreeSingleIndexAdaptorFlags::None, n_thread_build)); + } + + public: + /** Deleted copy constructor */ + KDTreeEigenMatrixAdaptor(const self_t&) = delete; + self_t& operator=(const self_t&) = delete; + + /** Move operations are deleted: the owned index_ stores a reference back to + * this adaptor object (passed as the dataset adaptor at construction), so + * moving would leave that reference dangling. Deleting them also prevents + * a double-free of the raw index_ pointer. */ + KDTreeEigenMatrixAdaptor(self_t&&) = delete; + self_t& operator=(self_t&&) = delete; + + ~KDTreeEigenMatrixAdaptor() { delete index_; } + + const std::reference_wrapper m_data_matrix; + + /** Query for the \a num_closest closest points to a given point (entered as + * query_point[0:dim-1]). Note that this is a short-cut method for + * index->findNeighbors(). The user can also call index->... methods as + * desired. + * + * \note If L2 norms are used, all returned distances are actually squared + * distances. + */ + void query( + const num_t* query_point, const Size num_closest, IndexType* out_indices, + num_t* out_distances) const + { + nanoflann::KNNResultSet resultSet(num_closest); + resultSet.init(out_indices, out_distances); + index_->findNeighbors(resultSet, query_point); + } + + /** @name Interface expected by KDTreeSingleIndexAdaptor + * @{ */ + + inline const self_t& derived() const noexcept { return *this; } + inline self_t& derived() noexcept { return *this; } + + // Must return the number of data points + inline Size kdtree_get_point_count() const + { + if (row_major) + return m_data_matrix.get().rows(); + else + return m_data_matrix.get().cols(); + } + + // Returns the dim'th component of the idx'th point in the class: + inline num_t kdtree_get_pt(const IndexType idx, size_t dim) const + { + if (row_major) + return m_data_matrix.get().coeff(idx, IndexType(dim)); + else + return m_data_matrix.get().coeff(IndexType(dim), idx); + } + + // Optional bounding-box computation: return false to default to a standard + // bbox computation loop. + // Return true if the BBOX was already computed by the class and returned + // in "bb" so it can be avoided to redo it again. Look at bb.size() to + // find out the expected dimensionality (e.g. 2 or 3 for point clouds) + template + inline bool kdtree_get_bbox(BBOX& /*bb*/) const + { + return false; + } + + /** @} */ + +}; // end of KDTreeEigenMatrixAdaptor +/** @} */ + +/** @} */ // end of grouping +} // namespace nanoflann + +#undef NANOFLANN_RESTRICT diff --git a/corelib/src/nanoflann/readme.txt b/corelib/src/nanoflann/readme.txt new file mode 100644 index 00000000..af979305 --- /dev/null +++ b/corelib/src/nanoflann/readme.txt @@ -0,0 +1,7 @@ + +nanoflann is included in rtabmap for convenience + +Source: https://github.com/jlblancoc/nanoflann +Version: 1.12.1 +Commit: 7812aa08260b6971af2230b1ac446d54f7939822 +License: BSD diff --git a/corelib/src/rtflann/config.h b/corelib/src/rtflann/config.h.in similarity index 82% rename from corelib/src/rtflann/config.h rename to corelib/src/rtflann/config.h.in index 29c070b7..c90c8fb4 100644 --- a/corelib/src/rtflann/config.h +++ b/corelib/src/rtflann/config.h.in @@ -35,4 +35,11 @@ #endif #define FLANN_VERSION_ "1.8.4" +// Generated by CMake from config.h.in, as upstream flann does. Carrying the +// option here rather than in a compile definition keeps a toggle of it from +// rebuilding the whole library: a definition given to the compiler lands in the +// target's flags, which every object file depends on, while this header is only +// included where rtflann is. +#cmakedefine FLANN_KDTREE_MEM_OPT + #endif /* RTABMAP_FLANN_CONFIG_H_ */ diff --git a/corelib/src/rtflann/defines.h b/corelib/src/rtflann/defines.h index c3e885d2..8cfc7356 100644 --- a/corelib/src/rtflann/defines.h +++ b/corelib/src/rtflann/defines.h @@ -29,7 +29,7 @@ #ifndef RTABMAP_FLANN_DEFINES_H_ #define RTABMAP_FLANN_DEFINES_H_ -#include "config.h" +#include "rtflann/config.h" // generated by CMake #ifdef FLANN_EXPORT #undef FLANN_EXPORT diff --git a/corelib/test/CMakeLists.txt b/corelib/test/CMakeLists.txt index 005ddcf9..fa754022 100644 --- a/corelib/test/CMakeLists.txt +++ b/corelib/test/CMakeLists.txt @@ -28,6 +28,7 @@ set(corelib_test_sources test_util3d_surface.cpp #util3d_surface.h test_visualword.cpp #VisualWord.h test_vwdictionary.cpp #VWDictionary.h + test_flann_index.cpp #FlannIndex.h test_transform.cpp #Transform.h test_stereo_dense.cpp #StereoDense.h (BM and SGBM strategies) test_stereo.cpp #Stereo.h (BlockMatching and OpticalFlow) @@ -102,6 +103,24 @@ foreach(shard RANGE 0 ${_corelib_last_shard}) ENVIRONMENT "GTEST_TOTAL_SHARDS=${CORELIB_TEST_SHARDS};GTEST_SHARD_INDEX=${shard}") endforeach() +# Comparison of the FlannIndex backends (times, memory, recall). Its own +# executable so that it can be run on demand: +# bin/test_flann_index_perf +# bin/test_flann_index_perf --gtest_filter=*Descriptors* +# and so that its seconds of benchmarking stay out of the unit test shards. +IF(BUILD_PERF_TESTS) + add_executable(test_flann_index_perf perf_flann_index.cpp) + target_link_libraries(test_flann_index_perf gtest_main rtabmap_core) + + # Labelled "performance" so CI jobs can opt out via `ctest -LE performance`, + # and run them alone with `ctest -L performance`. + add_test(NAME test_flann_index_perf COMMAND test_flann_index_perf) + math(EXPR _perf_timeout "600 * ${_test_timeout_scale}") + set_tests_properties(test_flann_index_perf PROPERTIES + TIMEOUT ${_perf_timeout} + LABELS "performance") +ENDIF(BUILD_PERF_TESTS) + # Rtabmap end-to-end replay of sample DBs (test data fetched by # scripts/fetch_test_data.sh into data/tests/*.db). Skips at runtime if the # assets are absent, so the build stays green for contributors without them. diff --git a/corelib/test/FlannIndexBackends.h b/corelib/test/FlannIndexBackends.h new file mode 100644 index 00000000..1e403523 --- /dev/null +++ b/corelib/test/FlannIndexBackends.h @@ -0,0 +1,313 @@ +#ifndef RTABMAP_CORELIB_TEST_FLANNINDEXBACKENDS_H_ +#define RTABMAP_CORELIB_TEST_FLANNINDEXBACKENDS_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace rtabmap; + +// Everything here is inline rather than static: the two test files including +// this header each use a subset of it, and unused functions with internal +// linkage warn. +namespace { + +struct Backend +{ + const char * name; + FlannIndex::flann_algorithm_t algorithm; + // The nanoflann structures differ by their rebalancing factor: 1 for the one + // built once, more for the one accepting points afterwards. + float rebalancingFactor = 2.0f; + // Not a FlannIndex at all: cv::BFMatcher, what the brute force strategies of + // VWDictionary and RegistrationVis use. Kept in the comparisons as the + // baseline every index has to beat. OpenCV threads its search where the + // indexes here search on one core, so it comes in two flavours: as the + // application gets it, and held to one core to compare the work done rather + // than the time it takes on an idle machine. + bool bruteForce = false; + bool singleCore = false; +}; + +// Every algorithm that indexes float features. The exhaustive search comes +// first: it is the reference the others are compared to, both for the neighbors +// found and for the time taken. +const Backend FLOAT_BACKENDS[] = { + {"linear exhaustive ", FlannIndex::FLANN_INDEX_LINEAR}, + // No single core row for the float features: OpenCV doesn't thread that + // match at these sizes, it measures the same thing as the one above. + {"cv BFMatcher ", FlannIndex::FLANN_INDEX_LINEAR, 1.0f, true}, + {"rtflann kd-tree (4 randomized) ", FlannIndex::FLANN_INDEX_KDTREE}, + {"rtflann kd-tree single ", FlannIndex::FLANN_INDEX_KDTREE_SINGLE}, + {"nanoflann kd-tree single ", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 1.0f}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, +}; + +// Those of them that search exactly, so are expected to return the very same +// neighbors. The randomized kd-tree is left out: it trades recall for speed. +const Backend EXACT_BACKENDS[] = { + {"linear exhaustive ", FlannIndex::FLANN_INDEX_LINEAR}, + {"rtflann kd-tree single ", FlannIndex::FLANN_INDEX_KDTREE_SINGLE}, + {"nanoflann kd-tree single ", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 1.0f}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, +}; + +// Binary features only have the Hamming based algorithms: nanoflann has no +// Hamming metric, and a kd-tree splitting on bits doesn't work, which is what +// LSH is for. +const Backend BINARY_BACKENDS[] = { + {"linear exhaustive (hamming) ", FlannIndex::FLANN_INDEX_LINEAR}, + {"cv BFMatcher (hamming) ", FlannIndex::FLANN_INDEX_LINEAR, 1.0f, true}, + {"cv BFMatcher (hamming,1 core)", FlannIndex::FLANN_INDEX_LINEAR, 1.0f, true, true}, + {"rtflann LSH ", FlannIndex::FLANN_INDEX_LSH}, +}; + +// Those compared when points are added after the index is built. +const Backend INCREMENTAL_BACKENDS[] = { + {"linear exhaustive ", FlannIndex::FLANN_INDEX_LINEAR}, + {"rtflann kd-tree (4 randomized) ", FlannIndex::FLANN_INDEX_KDTREE}, + {"rtflann kd-tree single ", FlannIndex::FLANN_INDEX_KDTREE_SINGLE}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, +}; + +// Note when reading the times of the first configuration of a comparison: it +// pays the warm up of the allocator, which the runs after it reuse. It shows on +// the nanoflann trees, which allocate a node per point, and the smaller the +// index the more it weighs. +// +// The rebalancing factor doesn't change how a search is done, only how the +// index is built and grown: with a factor over 1 the rtflann backend keeps one +// cv::Mat header per indexed point so that it can rebuild itself later, and the +// nanoflann incremental tree turns it into the fraction of removed points it +// tolerates. Both are compared where that matters, building and inserting. +const float REBALANCING_FACTORS[] = {1.0f, 2.0f}; + +// Synthetic point cloud, uniformly distributed in a 100 m box. The seeds are +// fixed so that a comparison that passes keeps passing. +inline cv::Mat makeCloud(int points, int dim, uint64_t seed) +{ + cv::RNG rng(seed); + cv::Mat cloud(points, dim, CV_32FC1); + rng.fill(cloud, cv::RNG::UNIFORM, 0.0f, 100.0f); + return cloud; +} + +// Float descriptors drawn around a limited number of centers. Uniformly +// distributed descriptors would be the worst case there is for any of the +// approximate algorithms, and nothing like the clustered distribution real +// descriptors have. +inline cv::Mat makeDescriptors(int count, int dim, int clusters, uint64_t seed) +{ + cv::RNG rng(seed); + cv::Mat centers(clusters, dim, CV_32FC1); + rng.fill(centers, cv::RNG::UNIFORM, 0.0f, 255.0f); + + // Filled in one go: a fill() per row costs seconds on a million descriptors. + cv::Mat descriptors(count, dim, CV_32FC1); + rng.fill(descriptors, cv::RNG::NORMAL, 0.0f, 12.0f); + for(int i=0; i(i, rng.uniform(0, bytes)) ^= (1 << rng.uniform(0, 8)); + } + } + return descriptors; +} + +// Queries taken from the indexed descriptors and perturbed, the way the same +// feature observed twice would be. Their nearest neighbor is unambiguous, which +// is what makes the recall of the approximate algorithms meaningful. +inline cv::Mat perturbedQueries(const cv::Mat & descriptors, int count, uint64_t seed) +{ + cv::RNG rng(seed); + cv::Mat queries(count, descriptors.cols, descriptors.type()); + for(int i=0; i(i, rng.uniform(0, descriptors.cols)) ^= (1 << rng.uniform(0, 8)); + } + } + } + return queries; +} + +struct Result +{ + double buildTime; + double knnTime; + double radiusTime; // negative when no radius search was done + size_t memory; + cv::Mat indices; +}; + +inline Result run( + const Backend & backend, + const cv::Mat & data, + const cv::Mat & queries, + int knn, + float radius, + float rebalancingFactor) +{ + Result result; + result.radiusTime = -1.0; + cv::Mat dists; + + if(backend.bruteForce) + { + // cv::setNumThreads() is global, put it back before leaving. + const int threads = cv::getNumThreads(); + if(backend.singleCore) + { + cv::setNumThreads(1); + } + + UTimer timer; + cv::BFMatcher matcher(data.type()==CV_8U?cv::NORM_HAMMING:cv::NORM_L2SQR); + matcher.add(std::vector(1, data)); + matcher.train(); // nothing to build, kept for the symmetry of the times + result.buildTime = timer.ticks(); + + std::vector > matches; + matcher.knnMatch(queries, matches, knn); + result.knnTime = timer.ticks(); + + result.indices = cv::Mat(queries.rows, knn, CV_32SC1, cv::Scalar(-1)); + for(size_t i=0; i((int)i, (int)j) = matches[i][j].trainIdx; + } + } + + if(radius > 0.0f) + { + std::vector > radiusMatches; + matcher.radiusMatch(queries, radiusMatches, radius); + result.radiusTime = timer.ticks(); + } + result.memory = 0; // it indexes nothing + if(backend.singleCore) + { + cv::setNumThreads(threads); + } + return result; + } + + FlannIndex index; + UTimer timer; + index.buildIndex(backend.algorithm, data, false, rebalancingFactor); + result.buildTime = timer.ticks(); + + index.knnSearch(queries, result.indices, dists, knn); + result.knnTime = timer.ticks(); + + if(radius > 0.0f) + { + std::vector > radiusIndices; + std::vector > radiusDists; + index.radiusSearch(queries, radiusIndices, radiusDists, radius); + result.radiusTime = timer.ticks(); + } + + result.memory = index.memoryUsed(); + return result; +} + +// Fraction of the neighbors that are the ones the exhaustive search finds. +inline float recall(const cv::Mat & indices, const cv::Mat & reference) +{ + UASSERT(indices.size() == reference.size()); + int found = 0; + for(int i=0; i(i, j) == reference.at(i, j)) + { + ++found; + } + } + } + return float(found)/float(indices.total()); +} + +// Note on the memory column: the rtflann indexes point into the features they +// were built with, while the nanoflann ones copy them into their own storage, +// which their memoryUsed() includes. +inline void report(const char * name, const Result & result, float recallRatio) +{ + std::cout << "[ ] " << name + << " build=" << uFormat("%7.1f", result.buildTime*1000.0) << " ms" + << " knn=" << uFormat("%7.1f", result.knnTime*1000.0) << " ms"; + if(result.radiusTime >= 0.0) + { + std::cout << " radius=" << uFormat("%7.1f", result.radiusTime*1000.0) << " ms"; + } + std::cout << " memory=" << uFormat("%6d", (int)(result.memory/1024)) << " KB" + << " recall=" << uFormat("%5.1f", recallRatio*100.0f) << " %" + << std::endl; +} + +// Run every backend of a list on the same data, the first one being the +// reference the others are compared to. +inline void compare( + const Backend * backends, + size_t count, + const cv::Mat & data, + const cv::Mat & queries, + int knn, + float radius = 0.0f, + float rebalancingFactor = 1.0f) +{ + cv::Mat reference; + for(size_t i=0; i & growths) +{ + std::cout << "[ ] " << DIM << "D float descriptors, " << finalCount + << " indexed, " << queryCount << " queries, knn=" << KNN + << ", averaged over " << seeds << " seed" << (seeds>1?"s":"") << std::endl; + + std::map grown; // growth factor -> average + Average fresh; + for(int seed=0; seed & removedPercents) +{ + std::cout << "[ ] " << DIM << "D float descriptors, " << count + << " indexed then partly removed, " << queryCount << " queries, knn=" + << KNN << ", averaged over " << seeds << " seed" << (seeds>1?"s":"") << std::endl; + + std::map kept; // removed % -> index that kept the removed points + std::map rebuilt; // removed % -> index built on the live ones only + for(int seed=0; seed live; + std::vector indexOfLive(count, -1); + for(int i=0; i= removed) + { + indexOfLive[i] = (int)live.size(); + live.push_back(i); + } + } + cv::Mat liveDescriptors((int)live.size(), DIM, CV_32FC1); + for(size_t i=0; i(i, j); + translated.at(i, j) = found>=0?indexOfLive[found]:-1; + } + } + kept[removed].add(recall(translated, reference), buildTime, knnTime, index.memoryUsed()); + } + + // Only the live points indexed: what rebuilding gives. + { + FlannIndex index; + UTimer timer; + index.buildIndex(FlannIndex::FLANN_INDEX_KDTREE, liveDescriptors, false, 1.0f); + const double buildTime = timer.ticks(); + index.knnSearch(queries, indices, dists, KNN); + rebuilt[removed].add(recall(indices, reference), buildTime, timer.ticks(), index.memoryUsed()); + } + } + } + + for(const auto & iter: kept) + { + reportAverage(uFormat("%3d%% removed, kept in the index ", iter.first), iter.second); + reportAverage(uFormat("%3d%% removed, rebuilt without them ", iter.first), rebuilt[iter.first]); + } +} + +} // namespace + +TEST(FlannIndexPerfTest, GrownIndexAgainstFreshlyBuiltOne) +{ + compareGrownAndFreshlyBuilt(20000, 3, 200, {2, 10, 100, 1000}); +} + +TEST(FlannIndexPerfTest, RecallAgainstTheFractionOfRemovedPoints) +{ + compareRemovedFractions(10000, 3, 200, {0, 25, 50, 75, 90}); +} + +// The same two comparisons on a dictionary of a million words, where a single +// descriptor matrix is already 512 MB and each of them takes minutes. Disabled +// so that a plain run of this executable stays in the seconds, run them with: +// bin/test_flann_index_perf --gtest_also_run_disabled_tests +// They stay compiled, so they cannot rot as FlannIndex changes. +TEST(FlannIndexPerfTest, DISABLED_GrownIndexAgainstFreshlyBuiltOneOnAMillionWords) +{ + compareGrownAndFreshlyBuilt(1000000, 3, 100, {2, 100, 1000}); +} + +TEST(FlannIndexPerfTest, DISABLED_RecallAgainstTheFractionOfRemovedPointsOnAMillionWords) +{ + compareRemovedFractions(1000000, 3, 100, {50, 90}); +} + +// The search RegistrationVis does per frame when a guess transform is given +// (Vis/CorGuessWinSize): the keypoints of the frame are indexed, and the points +// projected from the previous frame are looked up around their projection. The +// index is built and thrown away every frame, so its build time weighs as much +// as its search time. Before the nanoflann backend, this was a rtflann +// randomized kd-tree forest. +TEST(FlannIndexPerfTest, RegistrationGuessMatching) +{ + const int keypoints = 1000; // Vis/MaxFeatures + const float radius = 40.0f; // Vis/CorGuessWinSize + const int frames = 1000; // ~ a 50 s sequence at 20 Hz + + // Image points rather than a cube of them. + cv::RNG rng(140); + cv::Mat points(keypoints, 2, CV_32FC1); + cv::Mat projected(keypoints, 2, CV_32FC1); + for(int i=0; i(i, 0) = rng.uniform(0.0f, 640.0f); + points.at(i, 1) = rng.uniform(0.0f, 480.0f); + projected.at(i, 0) = rng.uniform(0.0f, 640.0f); + projected.at(i, 1) = rng.uniform(0.0f, 480.0f); + } + + // A factor of 1 for the rtflann rows keeps their per-point bookkeeping out + // of the measurement, and picks the nanoflann tree that is built once. + const Backend backends[] = { + {"cv BFMatcher ", FlannIndex::FLANN_INDEX_LINEAR, 1.0f, true}, + {"rtflann kd-tree (4 randomized) ", FlannIndex::FLANN_INDEX_KDTREE, 1.0f}, + {"rtflann kd-tree single ", FlannIndex::FLANN_INDEX_KDTREE_SINGLE, 1.0f}, + {"nanoflann kd-tree single ", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 1.0f}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, + }; + + std::cout << "[ ] " << keypoints << " keypoints indexed and as many looked up in a " + << radius << " px radius, per frame" << std::endl; + + for(const Backend & backend: backends) + { + std::vector > indices; + std::vector > dists; + + UTimer timer; + for(int frame=0; frame(i, j), 0) << "dim=" << dim << " knn=" << knn; + } + } + continue; + } + + for(int i=0; i(i, j), referenceIndices.at(i, j)) + << backend.name << " dim=" << dim << " knn=" << knn << " query=" << i << " n=" << j; + EXPECT_NEAR(dists.at(i, j), referenceDists.at(i, j), 1e-3f) + << backend.name << " dim=" << dim << " knn=" << knn << " query=" << i << " n=" << j; + } + } + } + } + } +} + +TEST(FlannIndexTest, ExactBackendsFindTheSamePointsInRadius) +{ + for(int dim = 2; dim <= 3; ++dim) + { + const cv::Mat cloud = makeCloud(5000, dim, 3); + const cv::Mat queries = makeCloud(200, dim, 4); + const float radius = 5.0f; + + std::vector > referenceIndices; + for(const Backend & backend: EXACT_BACKENDS) + { + FlannIndex index; + index.buildIndex(backend.algorithm, cloud, false, backend.rebalancingFactor); + + std::vector > indices; + std::vector > dists; + index.radiusSearch(queries, indices, dists, radius); + ASSERT_EQ(indices.size(), (size_t)queries.rows) << backend.name; + + // The backends don't return the points in the same order when they + // are not sorted by distance, compare them as sets. + for(size_t i=0; i(i, j), indices.at(i, j)) + << backend.name << " checksum=" << checksum << " query=" << i; + } + } + + // The raw pointer overload takes the same data. + FlannIndex loadedRaw; + EXPECT_TRUE(loadedRaw.loadIndex(data.data(), data.size(), backend.algorithm, cloud, false, backend.rebalancingFactor, &error)) + << backend.name << ": " << error; + } + } +} + +TEST(FlannIndexTest, LoadIndexRefusesDataItCannotUse) +{ + const cv::Mat cloud = makeCloud(500, 3, 22); + FlannIndex index; + index.buildIndex(FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, cloud); + const std::vector data = index.serializeIndex(true); + ASSERT_FALSE(data.empty()); + + const FlannIndex::flann_algorithm_t algorithm = FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE; + FlannIndex loaded; + std::string error; + + // Another algorithm than the one it was built with + EXPECT_FALSE(loaded.loadIndex(data, FlannIndex::FLANN_INDEX_KDTREE_SINGLE, cloud, false, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Another number of features + error.clear(); + EXPECT_FALSE(loaded.loadIndex(data, algorithm, cloud.rowRange(0, cloud.rows-1), false, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Another dimension + error.clear(); + EXPECT_FALSE(loaded.loadIndex(data, algorithm, makeCloud(cloud.rows, 2, 22), false, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Another distance + error.clear(); + EXPECT_FALSE(loaded.loadIndex(data, algorithm, cloud, true, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Same shape, other content: caught by the checksum + error.clear(); + EXPECT_FALSE(loaded.loadIndex(data, algorithm, makeCloud(cloud.rows, cloud.cols, 23), false, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Truncated + error.clear(); + EXPECT_FALSE(loaded.loadIndex(data.data(), data.size()/2, algorithm, cloud, false, 2.0f, &error)); + EXPECT_FALSE(error.empty()); + + // Nothing at all + error.clear(); + std::vector empty; + EXPECT_FALSE(loaded.loadIndex(empty, algorithm, cloud, false, 2.0f, &error)); + + // None of it left a half loaded index behind + EXPECT_FALSE(loaded.isBuilt()); +} + +// A descriptor header can cover a whole batch of points when the index is never +// rebuilt (rebalancing factor of 1), which used to make serializeIndex() look +// them up one by one and throw. +TEST(FlannIndexTest, SerializesAnIndexWithBatchedHeadersAndRemovedPoints) +{ + const cv::Mat cloud = makeCloud(500, 3, 24); + const cv::Mat added = makeCloud(100, 3, 25); + + for(float factor: {1.0f, 2.0f}) + { + FlannIndex index; + index.buildIndex(FlannIndex::FLANN_INDEX_KDTREE, cloud, false, factor); + const std::vector indexes = index.addPoints(added); + ASSERT_EQ(indexes.size(), (size_t)added.rows); + + for(size_t i=0; i<10; ++i) + { + index.removePoint(indexes[i]); + } + EXPECT_EQ(index.indexedFeatures(), (size_t)(cloud.rows + added.rows - 10)) << "factor=" << factor; + + EXPECT_NO_THROW(index.serializeIndex(true)) << "factor=" << factor; + EXPECT_NO_THROW(index.serializeIndex(false)) << "factor=" << factor; + } +} + +TEST(FlannIndexTest, AddedPointsAreFoundAndRemovedOnesAreNot) +{ + const cv::Mat cloud = makeCloud(500, 3, 26); + const cv::Mat added = makeCloud(50, 3, 27); + + const Backend backends[] = { + {"rtflann kd-tree (4 randomized) ", FlannIndex::FLANN_INDEX_KDTREE}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, + }; + + for(const Backend & backend: backends) + { + FlannIndex index; + index.buildIndex(backend.algorithm, cloud, false, backend.rebalancingFactor); + ASSERT_EQ(index.indexedFeatures(), (size_t)cloud.rows) << backend.name; + + const std::vector indexes = index.addPoints(added); + ASSERT_EQ(indexes.size(), (size_t)added.rows) << backend.name; + for(size_t i=0; i(0, 0), (int)indexes[0]) << backend.name; + EXPECT_NEAR(dists.at(0, 0), 0.0f, 1e-3f) << backend.name; + + // Once removed, it is not returned anymore, by either search. + index.removePoint(indexes[0]); + EXPECT_EQ(index.indexedFeatures(), (size_t)(cloud.rows + added.rows - 1)) << backend.name; + + index.knnSearch(added.row(0), indices, dists, 1); + EXPECT_NE(indices.at(0, 0), (int)indexes[0]) << backend.name; + + std::vector > radiusIndices; + std::vector > radiusDists; + index.radiusSearch(added.row(0), radiusIndices, radiusDists, 1.0f); + ASSERT_EQ(radiusIndices.size(), 1u) << backend.name; + for(size_t neighbor: radiusIndices[0]) + { + EXPECT_NE(neighbor, (size_t)indexes[0]) << backend.name; + } + } +} + +// The rebuild triggered by the removals renumbers nothing: the indexes handed +// out before it still designate the same points, which VWDictionary relies on. +TEST(FlannIndexTest, IndexesSurviveARebuild) +{ + const cv::Mat cloud = makeCloud(400, 3, 28); + const cv::Mat added = makeCloud(400, 3, 29); + + const Backend backends[] = { + {"rtflann kd-tree (4 randomized) ", FlannIndex::FLANN_INDEX_KDTREE}, + {"nanoflann kd-tree single incremental", FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, 2.0f}, + }; + + for(const Backend & backend: backends) + { + FlannIndex index; + index.buildIndex(backend.algorithm, cloud, false, 2.0f); + + std::vector indexes; + for(int i=0; i added1 = index.addPoints(added.row(i)); + ASSERT_EQ(added1.size(), 1u) << backend.name; + indexes.push_back(added1[0]); + } + + // Enough removals to get over the ratio a factor of 2 tolerates. + for(int i=0; i(i, 0), (int)indexes[i]) << backend.name << " point=" << i; + } + } +} + +TEST(FlannIndexTest, UnsupportedOperationsAreRefused) +{ + const cv::Mat cloud = makeCloud(200, 3, 30); + const cv::Mat added = makeCloud(10, 3, 31); + + // A nanoflann index built to never be rebuilt (factor 1) still takes points, + // rebuilding itself as the tree that accepts them. + FlannIndex staticIndex; + staticIndex.buildIndex(FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, cloud, false, 1.0f); + const std::vector addedIndexes = staticIndex.addPoints(added); + ASSERT_EQ(addedIndexes.size(), (size_t)added.rows); + EXPECT_EQ(addedIndexes[0], (unsigned int)cloud.rows); + EXPECT_EQ(staticIndex.indexedFeatures(), (size_t)(cloud.rows + added.rows)); + staticIndex.removePoint(0); + EXPECT_EQ(staticIndex.indexedFeatures(), (size_t)(cloud.rows + added.rows - 1)); + + // The points it held are still there, under the same indexes. + cv::Mat indices; + cv::Mat dists; + staticIndex.knnSearch(cloud.row(1), indices, dists, 1); + EXPECT_EQ(indices.at(0, 0), 1); + staticIndex.knnSearch(added.row(0), indices, dists, 1); + EXPECT_EQ(indices.at(0, 0), (int)addedIndexes[0]); + + // An index with removed points refers to holes in the features it was built + // with, which the ones given back to loadIndex() cannot reproduce. + FlannIndex incremental; + incremental.buildIndex(FlannIndex::NANOFLANN_INDEX_KDTREE_SINGLE, cloud, false, 2.0f); + EXPECT_FALSE(incremental.serializeIndex(true).empty()); + incremental.removePoint(0); + EXPECT_TRUE(incremental.serializeIndex(true).empty()); +} + +TEST(FlannIndexTest, DistanceL1BackendsAgree) +{ + const cv::Mat cloud = makeCloud(1000, 8, 32); + const cv::Mat queries = makeCloud(51, 8, 33); + const int knn = 2; + + cv::Mat reference; + cv::Mat referenceDists; + for(const Backend & backend: EXACT_BACKENDS) + { + FlannIndex index; + index.buildIndex(backend.algorithm, cloud, true /* useDistanceL1 */, backend.rebalancingFactor); + + cv::Mat indices; + cv::Mat dists; + index.knnSearch(queries, indices, dists, knn); + ASSERT_EQ(indices.rows, queries.rows) << backend.name; + + if(reference.empty()) + { + reference = indices; + referenceDists = dists; + // The exhaustive reference is L1 indeed, not L2. + const float * query = queries.ptr(0); + const float * neighbor = cloud.ptr(indices.at(0, 0)); + float l1 = 0.0f; + for(int i=0; i(0, 0), l1, 1e-2f); + continue; + } + + for(int i=0; i(i, j), reference.at(i, j)) << backend.name << " query=" << i; + EXPECT_NEAR(dists.at(i, j), referenceDists.at(i, j), 1e-2f) << backend.name; + } + } + } +} + +TEST(FlannIndexTest, BinaryDescriptorsUseHammingDistances) +{ + const cv::Mat descriptors = makeBinaryDescriptors(1000, 32, 100, 34); + const cv::Mat queries = descriptors.rowRange(0, 20); // the indexed ones + + const Backend backends[] = { + {"linear exhaustive (hamming) ", FlannIndex::FLANN_INDEX_LINEAR}, + {"rtflann LSH ", FlannIndex::FLANN_INDEX_LSH}, + }; + + for(const Backend & backend: backends) + { + FlannIndex index; + index.buildIndex(backend.algorithm, descriptors, false, backend.rebalancingFactor); + EXPECT_EQ(index.featuresType(), CV_8UC1) << backend.name; + EXPECT_EQ(index.featuresDim(), descriptors.cols) << backend.name; + + cv::Mat indices; + cv::Mat dists; + index.knnSearch(queries, indices, dists, 1); + // Hamming distances are integers + ASSERT_EQ(dists.type(), CV_32S) << backend.name; + for(int i=0; i(i, 0), i) << backend.name << " query=" << i; + EXPECT_EQ(dists.at(i, 0), 0) << backend.name << " query=" << i; + } + } +} + +TEST(FlannIndexTest, RadiusSearchKeepsTheNearestMaxNeighbors) +{ + const cv::Mat cloud = makeCloud(2000, 2, 35); + const cv::Mat queries = makeCloud(50, 2, 36); + const float radius = 10.0f; + const int maxNeighbors = 3; + + std::vector > reference; + for(const Backend & backend: EXACT_BACKENDS) + { + FlannIndex index; + index.buildIndex(backend.algorithm, cloud, false, backend.rebalancingFactor); + + std::vector > indices; + std::vector > dists; + index.radiusSearch(queries, indices, dists, radius, maxNeighbors, 32, 0.0f, true); + ASSERT_EQ(indices.size(), (size_t)queries.rows) << backend.name; + + for(size_t i=0; i data = dict->serializeIndex(); #ifdef _WIN32 - // FlannIndex::serializeIndex() is not implemented on Windows - // (see corelib/src/FlannIndex.cpp), so it always returns empty - // data regardless of the strategy. Skip the rest of the - // round-trip assertions on Windows. - EXPECT_EQ(data.size(), 0u) << "Strategy: " << VWDictionary::nnStrategyName(strategy); - continue; -#else - if(strategy < VWDictionary::kNNBruteForce) + // The rtflann serialization needs fmemopen, which Windows doesn't have + // (see FlannIndex::serializeIndex()): there, only the nanoflann index + // gives data back, the others are left out of the round trip below. + if(strategy != VWDictionary::kNNNanoFlannKdTree) + { + EXPECT_EQ(data.size(), 0u) << "Strategy: " << VWDictionary::nnStrategyName(strategy); + continue; + } +#endif + if(hasFlannIndex(strategy)) { - // flann strategies EXPECT_GT(data.size(), 0u) << "Strategy: " << VWDictionary::nnStrategyName(strategy); } else { + // brute force strategies have no index to serialize EXPECT_EQ(data.size(), 0u) << "Strategy: " << VWDictionary::nnStrategyName(strategy); } -#endif // Create new dictionary and deserialize VWDictionary dict2; @@ -722,9 +743,8 @@ TEST_F(VWDictionaryTest, SerializeDeserializeIndex) dict3.setNNStrategy(strategy); dict3.addNewWords(descriptors, 1); success = dict3.deserializeIndex(data); - if(strategy < VWDictionary::kNNBruteForce) + if(hasFlannIndex(strategy)) { - // flann strategies EXPECT_TRUE(success) << "Strategy: " << VWDictionary::nnStrategyName(strategy); // Index should be loaded @@ -1001,3 +1021,108 @@ TEST_F(VWDictionaryTest, FindNNWithVisualWords) } } + +// What Kp/ByteToFloat costs in matching quality. Both conversions feed the same +// exact index, so any difference comes from the distance they induce: expanding +// each bit to a float keeps the L1 distance equal to the Hamming distance, +// while converting each byte to a float doesn't, one flipped bit moving a byte +// by 1 or by 128 depending on which bit it is. +TEST_F(VWDictionaryTest, ByteToFloatMatchingQuality) +{ + const int words = 200; + const int bytes = 32; // ORB + const int queries = 100; + + cv::RNG rng(42); + cv::Mat descriptors(words, bytes, CV_8U); + rng.fill(descriptors, cv::RNG::UNIFORM, 0, 256); + + int totalBitExpansion = 0; + int totalByteToFloat = 0; + // From a query a few bits away from its word, which any distance finds, to + // one almost as far as the others are from each other (two random 256 bits + // descriptors differ by about 128). + for(int flippedBits: {8, 32, 64, 96}) + { + // Queries are indexed descriptors with bits flipped, the way the same + // feature looks when seen again. + cv::Mat queryDescriptors(queries, bytes, CV_8U); + for(int i=0; i(i, rng.uniform(0, bytes)) ^= (1 << rng.uniform(0, 8)); + } + } + + // Ground truth: the closest descriptor in Hamming distance. + std::vector hammingNN(queries); + for(int i=0; i addedIds = dictionary.addNewWords(descriptors, 1); + ASSERT_EQ(addedIds.size(), (size_t)words) << "byteToFloat=" << byteToFloat; + dictionary.update(); + ASSERT_EQ(dictionary.getVisualWords().size(), (size_t)words) << "byteToFloat=" << byteToFloat; + + const std::vector ids(addedIds.begin(), addedIds.end()); + const std::vector matched = dictionary.findNN(queryDescriptors); + ASSERT_EQ(matched.size(), (size_t)queries) << "byteToFloat=" << byteToFloat; + for(int i=0; i - Factor used when rebuilding the incremental FLANN index. Set 1 to disable. + Rebuild the incremental FLANN index once the ratio (factor-1)/factor of its features has been removed, e.g. half of them for a factor of 2. Rebuilding frees the memory of the removed features and speeds up the searches. Features are mostly removed when memory management is enabled. Set to 1 to never rebuild, which also uses less memory as the features don't have to be referenced one by one. true @@ -11930,7 +11930,7 @@ Lower the ratio -> higher the precision. - When using a FLANN-based nearest neighbor strategy, add/remove points to its index without always rebuilding the index (the index is built only when the dictionary increases of the factor below in size). + When using a FLANN-based nearest neighbor strategy, add/remove points to its index without always rebuilding the index (the index is only rebuilt when too many of its features have been removed, see the factor below). true @@ -11980,6 +11980,16 @@ Lower the ratio -> higher the precision. Brute Force GPU + + + NanoFLANN KdTree + + + + + FLANN KdTree Single + + @@ -23538,6 +23548,16 @@ With <0, the length is estimated once for each unique marker, then re-used fo GMS + + + NanoFLANN KdTree + + + + + FLANN KdTree Single + + diff --git a/tools/Matcher/main.cpp b/tools/Matcher/main.cpp index ab0a3b5a..f975187c 100644 --- a/tools/Matcher/main.cpp +++ b/tools/Matcher/main.cpp @@ -414,6 +414,24 @@ int main(int argc, char * argv[]) std::string pyMatcherPath; Parameters::parse(parameters, Parameters::kVisPnPReprojError(), reprojError); Parameters::parse(parameters, Parameters::kPyMatcherPath(), pyMatcherPath); + + // PyMatcher cannot match binary features, RegistrationVis falls back to + // brute force with cross check (see the warning above). + const bool pyMatcherOnBinaryFeatures = + reg.getNNType()==6 && + !dataFrom.getWordsDescriptors().empty() && + dataFrom.getWordsDescriptors().type()!=CV_32F; + QString nnTypeName = (pyMatcherOnBinaryFeatures? + RegistrationVis::getNNTypeName(5):reg.getNNTypeName()).c_str(); + // 5, 6 and 7 are the approaches RegistrationVis matches with itself, + // the others search for the k nearest neighbors and do the ratio test. + const bool nndrUsed = reg.getNNType()<5 || reg.getNNType()>7; + if(reg.getNNType()==6 && !pyMatcherOnBinaryFeatures) + { + // the script actually used is more telling than the generic name + nnTypeName = QString(uSplit(UFile::getName(pyMatcherPath), '.').front().c_str()).replace("rtabmap_", ""); + } + dialog.setWindowTitle(QString("Matches (%1/%2) %3 sec [%4=%5 (%6) %7=%8 (%9)%10 %11=%12 (%13) %14=%15]") .arg(info.inliers) .arg(info.matches) @@ -423,11 +441,8 @@ int main(int argc, char * argv[]) .arg(reg.getDetector()?Feature2D::typeName(reg.getDetector()->getType()).c_str():"?") .arg(Parameters::kVisCorNNType().c_str()) .arg(reg.getNNType()) - .arg(reg.getNNType()3D":reg.getEstimationType()==1?"3D->2D":reg.getEstimationType()==2?"2D->2D":"?")