mirror of
https://github.com/introlab/rtabmap.git
synced 2026-10-08 11:07:46 +08:00
Parameters: Renamed SuperGlue group to PyMatcher group. Added OANet python script.
This commit is contained in:
matlabbe
@@ -0,0 +1,313 @@
|
||||
/**
|
||||
* Python interface for SuperGlue: https://github.com/magicleap/SuperGluePretrainedNetwork
|
||||
*/
|
||||
|
||||
#include <pymatcher/PyMatcher.h>
|
||||
#include <rtabmap/utilite/ULogger.h>
|
||||
#include <rtabmap/utilite/UDirectory.h>
|
||||
#include <rtabmap/utilite/UFile.h>
|
||||
#include <rtabmap/utilite/UStl.h>
|
||||
#include <rtabmap/utilite/UConversion.h>
|
||||
#include <rtabmap/utilite/UTimer.h>
|
||||
|
||||
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
|
||||
#include <numpy/arrayobject.h>
|
||||
|
||||
namespace rtabmap
|
||||
{
|
||||
|
||||
class PythonSingleTon
|
||||
{
|
||||
public:
|
||||
PythonSingleTon() : initialized_(false) {}
|
||||
void init() {UScopeMutex lock(mutex_); if(!initialized_)Py_Initialize(); initialized_=true;}
|
||||
bool initialized() const {return initialized_;}
|
||||
virtual ~PythonSingleTon() {if(initialized_) Py_Finalize();}
|
||||
private:
|
||||
bool initialized_;
|
||||
UMutex mutex_;
|
||||
};
|
||||
|
||||
static PythonSingleTon g_python;
|
||||
|
||||
std::string getTraceback()
|
||||
{
|
||||
// Author: https://stackoverflow.com/questions/41268061/c-c-python-exception-traceback-not-being-generated
|
||||
|
||||
PyObject* type;
|
||||
PyObject* value;
|
||||
PyObject* traceback;
|
||||
|
||||
PyErr_Fetch(&type, &value, &traceback);
|
||||
PyErr_NormalizeException(&type, &value, &traceback);
|
||||
|
||||
std::string fcn = "";
|
||||
fcn += "def get_pretty_traceback(exc_type, exc_value, exc_tb):\n";
|
||||
fcn += " import sys, traceback\n";
|
||||
fcn += " lines = []\n";
|
||||
fcn += " lines = traceback.format_exception(exc_type, exc_value, exc_tb)\n";
|
||||
fcn += " output = '\\n'.join(lines)\n";
|
||||
fcn += " return output\n";
|
||||
|
||||
PyRun_SimpleString(fcn.c_str());
|
||||
PyObject* mod = PyImport_ImportModule("__main__");
|
||||
PyObject* method = PyObject_GetAttrString(mod, "get_pretty_traceback");
|
||||
PyObject* outStr = PyObject_CallObject(method, Py_BuildValue("OOO", type, value, traceback));
|
||||
std::string pretty = PyBytes_AsString(PyUnicode_AsASCIIString(outStr));
|
||||
|
||||
Py_DECREF(method);
|
||||
Py_DECREF(outStr);
|
||||
Py_DECREF(mod);
|
||||
|
||||
return pretty;
|
||||
}
|
||||
|
||||
PyMatcher::PyMatcher(
|
||||
const std::string & pythonMatcherPath,
|
||||
float matchThreshold,
|
||||
int iterations,
|
||||
bool cuda,
|
||||
const std::string & model) :
|
||||
pModule_(0),
|
||||
pFunc_(0),
|
||||
matchThreshold_(matchThreshold),
|
||||
iterations_(iterations),
|
||||
cuda_(cuda)
|
||||
{
|
||||
path_ = uReplaceChar(pythonMatcherPath, '~', UDirectory::homeDir());
|
||||
model_ = uReplaceChar(model, '~', UDirectory::homeDir());
|
||||
UINFO("path = %s", path_.c_str());
|
||||
UINFO("model = %s", model_.c_str());
|
||||
|
||||
if(!UFile::exists(path_) || UFile::getExtension(path_).compare("py") != 0)
|
||||
{
|
||||
UERROR("Cannot initialize Python matcher, the path is not valid: \"%s\"", path_.c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
if(!g_python.initialized())
|
||||
{
|
||||
g_python.init();
|
||||
}
|
||||
|
||||
std::string matcherPythonDir = UDirectory::getDir(path_);
|
||||
if(!matcherPythonDir.empty())
|
||||
{
|
||||
PyRun_SimpleString("import sys");
|
||||
PyRun_SimpleString(uFormat("sys.path.append(\"%s\")", matcherPythonDir.c_str()).c_str());
|
||||
}
|
||||
|
||||
_import_array();
|
||||
|
||||
std::string scriptName = uSplit(UFile::getName(path_), '.').front();
|
||||
PyObject * pName = PyUnicode_FromString(scriptName.c_str());
|
||||
pModule_ = PyImport_Import(pName);
|
||||
Py_DECREF(pName);
|
||||
|
||||
if(!pModule_)
|
||||
{
|
||||
UERROR("Module \"%s\" could not be imported! (File=\"%s\")", scriptName.c_str(), path_.c_str());
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
}
|
||||
}
|
||||
|
||||
PyMatcher::~PyMatcher()
|
||||
{
|
||||
if(pFunc_)
|
||||
{
|
||||
Py_DECREF(pFunc_);
|
||||
}
|
||||
if(pModule_)
|
||||
{
|
||||
Py_DECREF(pModule_);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<cv::DMatch> PyMatcher::match(
|
||||
const cv::Mat & descriptorsQuery,
|
||||
const cv::Mat & descriptorsTrain,
|
||||
const std::vector<cv::KeyPoint> & keypointsQuery,
|
||||
const std::vector<cv::KeyPoint> & keypointsTrain,
|
||||
const cv::Size & imageSize)
|
||||
{
|
||||
UTimer timer;
|
||||
std::vector<cv::DMatch> matches;
|
||||
|
||||
if(!pModule_)
|
||||
{
|
||||
UERROR("Python matcher module not loaded!");
|
||||
return matches;
|
||||
}
|
||||
|
||||
if(!descriptorsQuery.empty() &&
|
||||
descriptorsQuery.cols == descriptorsTrain.cols &&
|
||||
descriptorsQuery.type() == CV_32F &&
|
||||
descriptorsTrain.type() == CV_32F &&
|
||||
descriptorsQuery.rows == (int)keypointsQuery.size() &&
|
||||
descriptorsTrain.rows == (int)keypointsTrain.size() &&
|
||||
imageSize.width>0 && imageSize.height>0)
|
||||
{
|
||||
|
||||
UDEBUG("matchThreshold=%f, iterations=%d, cuda=%d", matchThreshold_, iterations_, cuda_?1:0);
|
||||
|
||||
if(!pFunc_)
|
||||
{
|
||||
PyObject * pFunc = PyObject_GetAttrString(pModule_, "init");
|
||||
if(pFunc)
|
||||
{
|
||||
if(PyCallable_Check(pFunc))
|
||||
{
|
||||
PyObject * result = PyObject_CallFunction(pFunc, "ifiis", descriptorsQuery.cols, matchThreshold_, iterations_, cuda_?1:0, model_.c_str());
|
||||
|
||||
if(result == NULL)
|
||||
{
|
||||
UERROR("Call to \"init(...)\" in \"%s\" failed!", path_.c_str());
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
return matches;
|
||||
}
|
||||
Py_DECREF(result);
|
||||
|
||||
pFunc_ = PyObject_GetAttrString(pModule_, "match");
|
||||
if(pFunc_ && PyCallable_Check(pFunc_))
|
||||
{
|
||||
// we are ready!
|
||||
}
|
||||
else
|
||||
{
|
||||
UERROR("Cannot find method \"match(...)\" in %s", path_.c_str());
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
if(pFunc_)
|
||||
{
|
||||
Py_DECREF(pFunc_);
|
||||
pFunc_ = 0;
|
||||
}
|
||||
return matches;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
UERROR("Cannot call method \"init(...)\" in %s", path_.c_str());
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
return matches;
|
||||
}
|
||||
Py_DECREF(pFunc);
|
||||
}
|
||||
else
|
||||
{
|
||||
UERROR("Cannot find method \"init(...)\"");
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
return matches;
|
||||
}
|
||||
UDEBUG("init time = %fs", timer.ticks());
|
||||
}
|
||||
|
||||
if(pFunc_)
|
||||
{
|
||||
std::vector<float> descriptorsQueryV(descriptorsQuery.rows * descriptorsQuery.cols);
|
||||
memcpy(descriptorsQueryV.data(), descriptorsQuery.data, descriptorsQuery.total()*sizeof(float));
|
||||
npy_intp dimsFrom[2] = {descriptorsQuery.rows, descriptorsQuery.cols};
|
||||
PyObject* pDescriptorsQuery = PyArray_SimpleNewFromData(2, dimsFrom, NPY_FLOAT, (void*)descriptorsQueryV.data());
|
||||
UASSERT(pDescriptorsQuery);
|
||||
|
||||
npy_intp dimsTo[2] = {descriptorsTrain.rows, descriptorsTrain.cols};
|
||||
std::vector<float> descriptorsTrainV(descriptorsTrain.rows * descriptorsTrain.cols);
|
||||
memcpy(descriptorsTrainV.data(), descriptorsTrain.data, descriptorsTrain.total()*sizeof(float));
|
||||
PyObject* pDescriptorsTrain = PyArray_SimpleNewFromData(2, dimsTo, NPY_FLOAT, (void*)descriptorsTrainV.data());
|
||||
UASSERT(pDescriptorsTrain);
|
||||
|
||||
std::vector<float> keypointsQueryV(keypointsQuery.size()*2);
|
||||
std::vector<float> scoresQuery(keypointsQuery.size());
|
||||
for(size_t i=0; i<keypointsQuery.size(); ++i)
|
||||
{
|
||||
keypointsQueryV[i*2] = keypointsQuery[i].pt.x;
|
||||
keypointsQueryV[i*2+1] = keypointsQuery[i].pt.y;
|
||||
scoresQuery[i] = keypointsQuery[i].response;
|
||||
}
|
||||
|
||||
std::vector<float> keypointsTrainV(keypointsTrain.size()*2);
|
||||
std::vector<float> scoresTrain(keypointsTrain.size());
|
||||
for(size_t i=0; i<keypointsTrain.size(); ++i)
|
||||
{
|
||||
keypointsTrainV[i*2] = keypointsTrain[i].pt.x;
|
||||
keypointsTrainV[i*2+1] = keypointsTrain[i].pt.y;
|
||||
scoresTrain[i] = keypointsTrain[i].response;
|
||||
}
|
||||
|
||||
npy_intp dimsKpQuery[2] = {(int)keypointsQuery.size(), 2};
|
||||
PyObject* pKeypointsQuery = PyArray_SimpleNewFromData(2, dimsKpQuery, NPY_FLOAT, (void*)keypointsQueryV.data());
|
||||
UASSERT(pKeypointsQuery);
|
||||
|
||||
npy_intp dimsKpTrain[2] = {(int)keypointsTrain.size(), 2};
|
||||
PyObject* pkeypointsTrain = PyArray_SimpleNewFromData(2, dimsKpTrain, NPY_FLOAT, (void*)keypointsTrainV.data());
|
||||
UASSERT(pkeypointsTrain);
|
||||
|
||||
npy_intp dimsScoresQuery[1] = {(int)keypointsQuery.size()};
|
||||
PyObject* pScoresQuery = PyArray_SimpleNewFromData(1, dimsScoresQuery, NPY_FLOAT, (void*)scoresQuery.data());
|
||||
UASSERT(pScoresQuery);
|
||||
|
||||
npy_intp dimsScoresTrain[1] = {(int)keypointsTrain.size()};
|
||||
PyObject* pScoresTrain = PyArray_SimpleNewFromData(1, dimsScoresTrain, NPY_FLOAT, (void*)scoresTrain.data());
|
||||
UASSERT(pScoresTrain);
|
||||
|
||||
PyObject * pImageWidth = PyLong_FromLong(imageSize.width);
|
||||
PyObject * pImageHeight = PyLong_FromLong(imageSize.height);
|
||||
|
||||
UDEBUG("Preparing data time = %fs", timer.ticks());
|
||||
|
||||
PyObject *pReturn = PyObject_CallFunctionObjArgs(pFunc_, pKeypointsQuery, pkeypointsTrain, pScoresQuery, pScoresTrain, pDescriptorsQuery, pDescriptorsTrain, pImageWidth, pImageHeight, NULL);
|
||||
if(pReturn == NULL)
|
||||
{
|
||||
UERROR("Failed to call match() function!");
|
||||
UERROR("%s", getTraceback().c_str());
|
||||
}
|
||||
else
|
||||
{
|
||||
UDEBUG("Python matching time = %fs", timer.ticks());
|
||||
|
||||
PyArrayObject *np_ret = reinterpret_cast<PyArrayObject*>(pReturn);
|
||||
|
||||
// Convert back to C++ array and print.
|
||||
int len1 = PyArray_SHAPE(np_ret)[0];
|
||||
int len2 = PyArray_SHAPE(np_ret)[1];
|
||||
int type = PyArray_TYPE(np_ret);
|
||||
UDEBUG("Matches array %dx%d (type=%d)", len1, len2, type);
|
||||
UASSERT_MSG(type == NPY_LONG || type == NPY_INT, uFormat("Returned matches should type INT=5 or LONG=7, received type=%d", type).c_str());
|
||||
if(type == NPY_LONG)
|
||||
{
|
||||
long* c_out = reinterpret_cast<long*>(PyArray_DATA(np_ret));
|
||||
for (int i = 0; i < len1*len2; i+=2)
|
||||
{
|
||||
matches.push_back(cv::DMatch(c_out[i], c_out[i+1], 0));
|
||||
}
|
||||
}
|
||||
else // INT
|
||||
{
|
||||
int* c_out = reinterpret_cast<int*>(PyArray_DATA(np_ret));
|
||||
for (int i = 0; i < len1*len2; i+=2)
|
||||
{
|
||||
matches.push_back(cv::DMatch(c_out[i], c_out[i+1], 0));
|
||||
}
|
||||
}
|
||||
Py_DECREF(pReturn);
|
||||
}
|
||||
|
||||
Py_DECREF(pDescriptorsQuery);
|
||||
Py_DECREF(pDescriptorsTrain);
|
||||
Py_DECREF(pKeypointsQuery);
|
||||
Py_DECREF(pkeypointsTrain);
|
||||
Py_DECREF(pScoresQuery);
|
||||
Py_DECREF(pScoresTrain);
|
||||
Py_DECREF(pImageWidth);
|
||||
Py_DECREF(pImageHeight);
|
||||
|
||||
UDEBUG("Fill matches (%d/%d) and cleanup time = %fs", matches.size(), std::min(descriptorsQuery.rows, descriptorsTrain.rows), timer.ticks());
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
UERROR("Invalid inputs! Supported python matchers require float descriptors.");
|
||||
}
|
||||
return matches;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Python interface for python matchers like:
|
||||
* - SuperGlue: https://github.com/magicleap/SuperGluePretrainedNetwork
|
||||
* - OANET https://github.com/zjhthu/OANet
|
||||
*/
|
||||
|
||||
#ifndef PYMATCHER_H
|
||||
#define PYMATCHER_H
|
||||
|
||||
#include <opencv2/core/types.hpp>
|
||||
#include <opencv2/core/mat.hpp>
|
||||
#include <vector>
|
||||
|
||||
#include <Python.h>
|
||||
|
||||
namespace rtabmap
|
||||
{
|
||||
|
||||
class PyMatcher
|
||||
{
|
||||
public:
|
||||
PyMatcher(const std::string & pythonMatcherPath,
|
||||
float matchThreshold = 0.2f,
|
||||
int iterations = 20,
|
||||
bool cuda = true,
|
||||
const std::string & model = "indoor");
|
||||
virtual ~PyMatcher();
|
||||
|
||||
const std::string & path() const {return path_;}
|
||||
float matchThreshold() const {return matchThreshold_;}
|
||||
int iterations() const {return iterations_;}
|
||||
bool cuda() const {return cuda_;}
|
||||
const std::string & model() const {return model_;}
|
||||
|
||||
std::vector<cv::DMatch> match(
|
||||
const cv::Mat & descriptorsQuery,
|
||||
const cv::Mat & descriptorsTrain,
|
||||
const std::vector<cv::KeyPoint> & keypointsQuery,
|
||||
const std::vector<cv::KeyPoint> & keypointsTrain,
|
||||
const cv::Size & imageSize);
|
||||
|
||||
private:
|
||||
PyObject * pModule_;
|
||||
PyObject * pFunc_;
|
||||
std::string path_;
|
||||
float matchThreshold_;
|
||||
int iterations_;
|
||||
bool cuda_;
|
||||
std::string model_;
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,44 @@
|
||||
#! /usr/bin/env python3
|
||||
#
|
||||
# Drop this file in the "demo" folder of OANet git: https://github.com/zjhthu/OANet
|
||||
# To use with rtabmap:
|
||||
# --Vis/CorNNType 6 --PyMatcher/Path ~/OANet/demo/rtabmap_oanet.py --PyMatcher/Model ~/OANet/model/gl3d/sift-4000/model_best.pth
|
||||
#
|
||||
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.path.dirname(os.path.realpath(__file__))+'/../core')
|
||||
if not hasattr(sys, 'argv'):
|
||||
sys.argv = ['']
|
||||
|
||||
#print(os.sys.path)
|
||||
#print(sys.version)
|
||||
|
||||
import numpy as np
|
||||
from learnedmatcher import LearnedMatcher
|
||||
|
||||
lm = None
|
||||
|
||||
def init(descriptorDim, matchThreshold, iterations, cuda, model_path):
|
||||
print("OANet python init()")
|
||||
global lm
|
||||
lm = LearnedMatcher(model_path, inlier_threshold=1, use_ratio=0, use_mutual=0)
|
||||
|
||||
|
||||
def match(kptsFrom, kptsTo, scoresFrom, scoresTo, descriptorsFrom, descriptorsTo, imageWidth, imageHeight):
|
||||
#print("OANet python match()")
|
||||
|
||||
kpt1 = np.asarray(kptsFrom)
|
||||
kpt2 = np.asarray(kptsTo)
|
||||
desc1 = np.asarray(descriptorsFrom)
|
||||
desc2 = np.asarray(descriptorsTo)
|
||||
|
||||
global lm
|
||||
matches, _, _ = lm.infer([kpt1, kpt2], [desc1, desc2])
|
||||
return matches
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
#test
|
||||
init(128, 0.2, 20, False, True)
|
||||
match([[1, 2], [1,3], [4,6]], [[1, 3], [1,2], [16,2]], [1, 3,6], [1,3,5], np.full((3, 128), 1), np.full((3, 128), 1), 640, 480)
|
||||
@@ -0,0 +1,86 @@
|
||||
#! /usr/bin/env python3
|
||||
#
|
||||
# Drop this file in the root folder of SuperGlue git: https://github.com/magicleap/SuperGluePretrainedNetwork
|
||||
# To use with rtabmap:
|
||||
# --Vis/CorNNType 6 --SuperGlue/Path "~/SuperGluePretrainedNetwork/rtabmap_superglue.py"
|
||||
#
|
||||
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
#import sys
|
||||
#import os
|
||||
#print(os.sys.path)
|
||||
#print(sys.version)
|
||||
|
||||
from models.matching import SuperGlue
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
device = 'cpu'
|
||||
superglue = []
|
||||
|
||||
def init(descriptorDim, matchThreshold, iterations, cuda, model):
|
||||
print("SuperGlue python init()")
|
||||
# Load the SuperGlue model.
|
||||
global device
|
||||
device = 'cuda' if torch.cuda.is_available() and cuda else 'cpu'
|
||||
assert model == "indoor" or model == "outdoor", "Available models for SuperGlue are 'indoor' or 'outdoor'"
|
||||
config = {
|
||||
'superglue': {
|
||||
'weights': model,
|
||||
'sinkhorn_iterations': iterations,
|
||||
'match_threshold': matchThreshold,
|
||||
'descriptor_dim' : descriptorDim
|
||||
}
|
||||
}
|
||||
global superglue
|
||||
superglue = SuperGlue(config.get('superglue', {})).eval().to(device)
|
||||
|
||||
|
||||
def match(kptsFrom, kptsTo, scoresFrom, scoresTo, descriptorsFrom, descriptorsTo, imageWidth, imageHeight):
|
||||
#print("SuperGlue python match()")
|
||||
global device
|
||||
kptsFrom = np.asarray(kptsFrom)
|
||||
kptsFrom = kptsFrom[None, :, :]
|
||||
kptsTo = np.asarray(kptsTo)
|
||||
kptsTo = kptsTo[None, :, :]
|
||||
scoresFrom = np.asarray(scoresFrom)
|
||||
scoresFrom = scoresFrom[None, :]
|
||||
scoresTo = np.asarray(scoresTo)
|
||||
scoresTo = scoresTo[None, :]
|
||||
descriptorsFrom = np.transpose(np.asarray(descriptorsFrom))
|
||||
descriptorsFrom = descriptorsFrom[None, :, :]
|
||||
descriptorsTo = np.transpose(np.asarray(descriptorsTo))
|
||||
descriptorsTo = descriptorsTo[None, :, :]
|
||||
|
||||
data = {
|
||||
'image0': torch.rand(1, 1, imageHeight, imageWidth).to(device),
|
||||
'image1': torch.rand(1, 1, imageHeight, imageWidth).to(device),
|
||||
'scores0': torch.from_numpy(scoresFrom).to(device),
|
||||
'scores1': torch.from_numpy(scoresTo).to(device),
|
||||
'keypoints0': torch.from_numpy(kptsFrom).to(device),
|
||||
'keypoints1': torch.from_numpy(kptsTo).to(device),
|
||||
'descriptors0': torch.from_numpy(descriptorsFrom).to(device),
|
||||
'descriptors1': torch.from_numpy(descriptorsTo).to(device),
|
||||
}
|
||||
|
||||
|
||||
global superglue
|
||||
results = superglue(data)
|
||||
|
||||
matches0 = results['matches0'].to('cpu').numpy()
|
||||
|
||||
matchesFrom = np.nonzero(matches0!=-1)[1]
|
||||
matchesTo = matches0[np.nonzero(matches0!=-1)]
|
||||
|
||||
matchesArray = np.stack((matchesFrom, matchesTo), axis=1);
|
||||
|
||||
return matchesArray
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
#test
|
||||
init(256, 0.2, 20, True, 'indoor')
|
||||
match([[1, 2], [1,3]], [[1, 3], [1,2]], [1, 3], [1,3], np.full((2, 256), 1),np.full((2, 256), 1), 640, 480)
|
||||
Reference in New Issue
Block a user