Automated dictionary recovery/repairing (#1669)

* Automated dictionary recovery/repairing

* typo

* updated comment

* updating dictionary after fixing it

* Fixing missing ref

* Adding helping assert msg in case of repairing failure

corelib/include/rtabmap/core/VWDictionary.h

@@ -89,7 +89,7 @@ public:
 	std::vector<int> findNN(const std::list<VisualWord *> & vws) const;
 	std::vector<int> findNN(const cv::Mat & descriptors) const;
 
-	void addWordRef(int wordId, int signatureId);
+	bool addWordRef(int wordId, int signatureId);
 	void removeAllWordRef(int wordId, int signatureId);
 	const VisualWord * getWord(int id) const;
 	VisualWord * getUnusedWord(int id) const;

corelib/src/Memory.cpp

@@ -437,7 +437,8 @@ void Memory::loadDataFromDb(bool postInitClosingEvents)
 		UTimer timer;
 		// Enable loaded signatures
 		const std::map<int, Signature *> & signatures = this->getSignatures();
-		for(std::map<int, Signature *>::const_iterator i=signatures.begin(); i!=signatures.end(); ++i)
+		bool corruptedDictionary = false;
+		for(std::map<int, Signature *>::const_iterator i=signatures.begin(); i!=signatures.end() && !corruptedDictionary; ++i)
 		{
 			Signature * s = this->_getSignature(i->first);
 			UASSERT(s != 0);
@@ -450,12 +451,114 @@ void Memory::loadDataFromDb(bool postInitClosingEvents)
 				{
 					if(iter->first > 0)
 					{
-						_vwd->addWordRef(iter->first, i->first);
+						if(!_vwd->addWordRef(iter->first, s->id()))
+						{
+							corruptedDictionary = true;
+							break;
+						}
+					}
+				}
+				s->setEnabled(!corruptedDictionary);
+				if(corruptedDictionary)
+				{
+					//revert all changes from that signature till it broke above
+					for(std::multimap<int, int>::const_iterator iter = words.begin(); iter!=words.end(); ++iter)
+					{
+						if(iter->first > 0)
+						{
+							_vwd->removeAllWordRef(iter->first, s->id());
+						}
 					}
 				}
-				s->setEnabled(true);
 			}
 		}
+		if(corruptedDictionary)
+		{
+			if(!_vwd->isIncremental())
+			{
+				UERROR("The dictionary is empty or missing some words from nodes in WM, "
+					"we cannot repair it because it is a fixed dictionary. Make sure you "
+					"are using the right fixed dictionary that was used to generate the map.");
+			}
+			else
+			{
+				std::string msg = uFormat(
+					"The dictionary is empty or missing some words from nodes in WM, "
+					"we will try to repair it. This can be caused by rtabmap closing before it has time "
+					"to save the dictionary. Re-creating the dictionary from %ld nodes...",
+					signatures.size());
+				UWARN("%s", msg.c_str());
+				if(postInitClosingEvents) UEventsManager::post(new RtabmapEventInit(msg));
+
+				//remove all words ref
+
+				const std::map<int, VisualWord *> & addedWords = _vwd->getVisualWords();
+				int nodesRepaired = 0;
+				size_t oldSize = addedWords.size();
+				std::string assertMsg = 
+					"If we assert here, the problem is maybe deeper. Try "
+					"to use rtabmap-recovery tool instead to fix the database.";
+				for(std::map<int, Signature *>::const_iterator i=signatures.begin(); i!=signatures.end(); ++i)
+				{
+					Signature * s = this->_getSignature(i->first);
+					UASSERT_MSG(s != 0, assertMsg.c_str());
+
+					if(s->isEnabled())
+					{
+						// Words already in dictionary and references added
+						continue;
+					}
+
+					const std::multimap<int, int> * words = &s->getWords();
+					if(words->size())
+					{
+						cv::Mat descriptors = s->getWordsDescriptors();
+						std::multimap<int, int> loadedWords;
+						if(descriptors.empty())
+						{
+							// We may have started rtabmap without loading features, check in the database
+							std::multimap<int, int> w;
+							std::vector<cv::KeyPoint> k;
+							std::vector<cv::Point3f> p;
+							_dbDriver->getLocalFeatures(s->id(), loadedWords, k, p, descriptors);
+							UASSERT_MSG(loadedWords.size() == words->size(), assertMsg.c_str()); // Just doublecheck
+							words = &loadedWords; // The index will be set
+							UASSERT_MSG(!descriptors.empty(), assertMsg.c_str());
+						}
+						bool repaired = false;
+						for(std::multimap<int, int>::const_iterator iter = words->begin(); iter!=words->end(); ++iter)
+						{
+							if(iter->first > 0)
+							{
+								if(addedWords.find(iter->first) == addedWords.end())
+								{
+									UASSERT_MSG(iter->second >= 0 && iter->second < descriptors.rows, 
+										uFormat("iter->second=%d descriptors.rows=%d (signature=%d word=%d). %s",
+										iter->second, descriptors.rows, s->id(), iter->first, assertMsg.c_str()).c_str());
+									_vwd->addWord(new VisualWord(iter->first, descriptors.row(iter->second).clone()));
+									repaired = true;
+								}
+								UASSERT_MSG(_vwd->addWordRef(iter->first, s->id()), assertMsg.c_str());
+							}
+						}
+						nodesRepaired += (repaired?1:0);
+						s->setEnabled(true);
+					}
+				}
+
+				msg = uFormat(
+					"Regenerated the dictionary with %ld missing words (%ld -> %ld) from %d nodes.",
+					addedWords.size() - oldSize,
+					oldSize,
+					addedWords.size(),
+					nodesRepaired);
+				UWARN("%s", msg.c_str());
+				if(postInitClosingEvents) UEventsManager::post(new RtabmapEventInit(msg));
+				_memoryChanged = true; // This will force rtabmap to save back the dictionary even if we don't process any new data
+				_vwd->update();
+			}
+		}
+
 		if(postInitClosingEvents) UEventsManager::post(new RtabmapEventInit(uFormat("Adding word references, done! (%d)", _vwd->getTotalActiveReferences())));
 
 		if(_vwd->getUnusedWordsSize() && _vwd->isIncremental())
@@ -6599,7 +6702,10 @@ void Memory::enableWordsRef(const std::list<int> & signatureIds)
 			{
 				if(keys.at(i)>0)
 				{
-					_vwd->addWordRef(keys.at(i), (*j)->id());
+					if(_vwd->addWordRef(keys.at(i), (*j)->id()))
+					{
+						UERROR("Could not add word ref %d to node %d!?", keys.at(i), (*j)->id());
+					}
 				}
 			}
 			(*j)->setEnabled(true);

corelib/src/VWDictionary.cpp

@@ -868,7 +868,7 @@ int VWDictionary::getNextId()
 	return ++_lastWordId;
 }
 
-void VWDictionary::addWordRef(int wordId, int signatureId)
+bool VWDictionary::addWordRef(int wordId, int signatureId)
 {
 	VisualWord * vw = 0;
 	vw = uValue(_visualWords, wordId, vw);
@@ -878,10 +878,12 @@ void VWDictionary::addWordRef(int wordId, int signatureId)
 		_totalActiveReferences += 1;
 
 		_unusedWords.erase(vw->id());
+		return true;
 	}
 	else
 	{
-		UERROR("Not found word %d (dict size=%d)", wordId, (int)_visualWords.size());
+		UWARN("Not found word %d (dict size=%d)", wordId, (int)_visualWords.size());
+		return false;
 	}
 }