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 <[email protected]>
2026-10-07 02:27:47 +08:00 · 2026-08-16 09:51:39 -07:00
parent df52523a0c
commit f647014f54
28 changed files with 7537 additions and 179 deletions
@@ -34,12 +34,14 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <rtabmap/core/Parameters.h>

 #include "rtflann/flann.hpp"
+#include "nanoflann/NanoFlannIndex.h"
 #include <boost/crc.hpp>

 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<class T>
+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<unsigned char> FlannIndex::serializeIndex(bool computeChecksum) const {
+	if(nanoIndex_)
+	{
+		std::vector<unsigned char> nanoIndexData = nanoIndex_->serializeIndex();
+		if(!nanoIndexData.empty())
+		{
+			const size_t headerSizeBytes = sizeof(int)*FLANN_INDEX_HEADER_SIZE;
+			const cv::Mat dataset = nanoIndex_->indexedPoints();
+			std::vector<unsigned char> 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<unsigned char>();
+	}
 	if(index_ && !addedDescriptors_.empty())
 	{
 #ifdef WIN32
@@ -147,6 +261,10 @@ std::vector<unsigned char> 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<unsigned char> 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<iter.second.rows; ++r)
+					{
+						if(removedDescriptors.find(iter.first+r) == removedDescriptors.end())
+						{
+							UASSERT(row < dataRows);
+							iter.second.row(r).copyTo(dataset.row(row++));
+						}
+					}
 				}
+				UASSERT(row == dataRows);
 			}

-			unsigned int crcValue = 0;
-			if(computeChecksum) {
-				boost::crc_32_type result;
-				result.process_bytes(dataset.data, dataset.total()*dataset.elemSize());
-				crcValue = result.checksum(); 
-			}
-			
 			indexData.resize(bytes_written+headerSizeBytes);
 			indexData.shrink_to_fit();
-			int rebalancingFactorAsInt; // Deprecated
-			memcpy(&rebalancingFactorAsInt, &rebalancingFactor_, sizeof(rebalancingFactor_)); // Deprecated
-			int crcValueAsInt;
-			memcpy(&crcValueAsInt, &crcValue, sizeof(crcValue));
-			int header[FLANN_INDEX_HEADER_SIZE] = {
-					RTABMAP_VERSION_MAJOR, RTABMAP_VERSION_MINOR, RTABMAP_VERSION_PATCH, // 0,1,2
-					algorithm_,                                                          // 3,
-					featuresDim_,                                                        // 4,
-					useDistanceL1_?1:0,                                                  // 5,
-					rebalancingFactorAsInt,                                              // 6, Deprecated
-					dataRows,                                                            // 7,
-					dataCols,                                                            // 8,			
-					dataType,                                                            // 9,
-					crcValueAsInt,                                                       // 10         
-					(int)bytes_written};                                                 // 11
-			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], crcValueAsInt, 
-				header[11]);
+			int header[FLANN_INDEX_HEADER_SIZE];
+			fillIndexHeader(header,
+				algorithm_,
+				featuresDim_,
+				useDistanceL1_,
+				rebalancingFactor_,
+				dataRows,
+				dataCols,
+				dataType,
+				computeChecksum?computeCrc(dataset):0,
+				(int)bytes_written);
 			memcpy(indexData.data(), header, headerSizeBytes);
 			return indexData;
 		}
@@ -230,6 +332,10 @@ std::vector<unsigned char> 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<unsigned int> 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<unsigned int> 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<unsigned char> points(features.data, features.rows, features.cols);
 		rtflann::Index<rtflann::Hamming<unsigned char> > * index = (rtflann::Index<rtflann::Hamming<unsigned char> >*)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<unsigned int> FlannIndex::addPoints(const cv::Mat & features)
 			rtflann::Index<rtflann::L1<float> > * index = (rtflann::Index<rtflann::L1<float> >*)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<unsigned int> FlannIndex::addPoints(const cv::Mat & features)
 			rtflann::Index<rtflann::L2_Simple<float> > * index = (rtflann::Index<rtflann::L2_Simple<float> >*)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<unsigned int> FlannIndex::addPoints(const cv::Mat & features)
 			rtflann::Index<rtflann::L2<float> > * index = (rtflann::Index<rtflann::L2<float> >*)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<unsigned int> 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<unsigned int> indexes;
+	indexes.reserve(features.rows);
 	for(int i=0; i<features.rows; ++i)
 	{
-		indexes.push_back(nextIndex_);
-		addedDescriptors_.insert(std::make_pair(nextIndex_++, features.row(i)));
+		indexes.push_back(nextIndex_ + i);
+	}
+
+	// incremental FLANN: we should add all headers separately in case we remove
+	// some indexes (to keep underlying matrix data allocated)
+	if(rebalancingFactor_ > 1.0f)
+	{
+		for(int i=0; i<features.rows; ++i)
+		{
+			addedDescriptors_.insert(std::make_pair(nextIndex_++, features.row(i)));
+		}
+	}
+	else
+	{
+		// tree won't ever be rebalanced, so just keep only one header for the data
+		addedDescriptors_.insert(std::make_pair(nextIndex_, features));
+		nextIndex_ += features.rows;
 	}

 	return indexes;
@@ -688,6 +867,11 @@ std::vector<unsigned int> 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<int>();
-	for(size_t i=0 ; i<indicesBuffer.size(); i+=2)
+	for(size_t i=0 ; i<indicesBuffer.size(); ++i)
 	{
+		// Note: this loop used to write two entries per iteration, which read
+		// and wrote one past the end when query.rows*knn is odd (an odd knn on
+		// an odd number of queries).
 		ptr[i] = indicesBuffer[i] == std::numeric_limits<size_t>::max()?-1:(int)indicesBuffer[i];
-		ptr[i+1] = indicesBuffer[i+1] == std::numeric_limits<size_t>::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!");