i/rtabmap
1
0
Fork 0
mirror of https://github.com/introlab/rtabmap.git synced 2026-09-02 09:30:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

0.19.7: added SuperPoint Torch feature support. RegVis: keep Feature2D detectors as class members instead of recreating them at each registration.

Browse Source
This commit is contained in:
matlabbe
2020-04-16 17:59:45 -04:00
parent 2ff582f06f
commit f575652456
16 changed files with 939 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
corelib/src/CMakeLists.txt
View File

@@ -173,6 +173,22 @@ ELSE()
)
ENDIF()
IF(TORCH_FOUND)
SET(LIBRARIES
${LIBRARIES}
${TORCH_LIBRARIES}
)
SET(SRC_FILES
${SRC_FILES}
superpoint_torch/SuperPoint.cc
)
SET(INCLUDE_DIRS
${TORCH_INCLUDE_DIRS}
${CMAKE_CURRENT_SOURCE_DIR}/superpoint_torch
${INCLUDE_DIRS}
)
ENDIF(TORCH_FOUND)
IF(Freenect_FOUND)
IF(Freenect_DASH_INCLUDES)
ADD_DEFINITIONS("-DFREENECT_DASH_INCLUDES")

85
corelib/src/Features2d.cpp
View File

@@ -44,6 +44,10 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#include "opencv/ORBextractor.h"
#endif
#ifdef RTABMAP_SP_TORCH
#include "superpoint_torch/SuperPoint.h"
#endif
#if CV_MAJOR_VERSION < 3
#include "opencv/Orb.h"
#ifdef HAVE_OPENCV_GPU
@@ -461,6 +465,14 @@ Feature2D * Feature2D::create(Feature2D::Type type, const ParametersMap & parame
}
#endif
#ifndef RTABMAP_SP_TORCH
if(type == Feature2D::kFeatureSuperPointTorch)
{
UWARN("SupertPoint Torch feature cannot be used as RTAB-Map is not built with the option enabled. GFTT/ORB is used instead.");
type = Feature2D::kFeatureGfttOrb;
}
#endif
Feature2D * feature2D = 0;
switch(type)
{
@@ -497,6 +509,11 @@ Feature2D * Feature2D::create(Feature2D::Type type, const ParametersMap & parame
case Feature2D::kFeatureOrbOctree:
feature2D = new ORBOctree(parameters);
break;
#ifdef RTABMAP_SP_TORCH
case Feature2D::kFeatureSuperPointTorch:
feature2D = new SuperPointTorch(parameters);
break;
#endif
#ifdef RTABMAP_NONFREE
default:
feature2D = new SURF(parameters);
@@ -1795,4 +1812,72 @@ cv::Mat ORBOctree::generateDescriptorsImpl(const cv::Mat & image, std::vector<cv
return descriptors_;
}
//////////////////////////
//SuperPointTorch
//////////////////////////
SuperPointTorch::SuperPointTorch(const ParametersMap & parameters) :
path_(Parameters::defaultSPTorchModelPath()),
threshold_(Parameters::defaultSPTorchThreshold()),
nms_(Parameters::defaultSPTorchNMS()),
minDistance_(Parameters::defaultSPTorchMinDistance()),
cuda_(Parameters::defaultSPTorchCuda())
{
parseParameters(parameters);
}
SuperPointTorch::~SuperPointTorch()
{
}
void SuperPointTorch::parseParameters(const ParametersMap & parameters)
{
Feature2D::parseParameters(parameters);
std::string previousPath = path_;
bool previousCuda = cuda_;
Parameters::parse(parameters, Parameters::kSPTorchModelPath(), path_);
Parameters::parse(parameters, Parameters::kSPTorchThreshold(), threshold_);
Parameters::parse(parameters, Parameters::kSPTorchNMS(), nms_);
Parameters::parse(parameters, Parameters::kSPTorchMinDistance(), minDistance_);
Parameters::parse(parameters, Parameters::kSPTorchCuda(), cuda_);
#ifdef RTABMAP_SP_TORCH
if(superPoint_.get() == 0 || path_.compare(previousPath) != 0 || previousCuda != cuda_)
{
superPoint_ = cv::Ptr<SPDetector>(new SPDetector(path_, threshold_, nms_, minDistance_, cuda_));
}
else
{
superPoint_->setThreshold(threshold_);
superPoint_->SetNMS(nms_);
superPoint_->setMinDistance(minDistance_);
}
#else
UWARN("RTAB-Map is not built with SuperPoint Torch support so SuperPoint Torch feature cannot be used!");
#endif
}
std::vector<cv::KeyPoint> SuperPointTorch::generateKeypointsImpl(const cv::Mat & image, const cv::Rect & roi, const cv::Mat & mask)
{
#ifdef RTABMAP_SP_TORCH
UASSERT(!image.empty() && image.channels() == 1 && image.depth() == CV_8U);
UASSERT_MSG(roi.x==0 && roi.y ==0, "Not supporting ROI");
return superPoint_->detect(image);
#else
UWARN("RTAB-Map is not built with SuperPoint Torch support so SuperPoint Torch feature cannot be used!");
return std::vector<cv::KeyPoint>();
#endif
}
cv::Mat SuperPointTorch::generateDescriptorsImpl(const cv::Mat & image, std::vector<cv::KeyPoint> & keypoints) const
{
#ifdef RTABMAP_SP_TORCH
UASSERT(!image.empty() && image.channels() == 1 && image.depth() == CV_8U);
return superPoint_->compute(keypoints);
#else
UWARN("RTAB-Map is not built with SuperPoint Torch support so SuperPoint Torch feature cannot be used!");
return cv::Mat();
#endif
}
}

9
corelib/src/Parameters.cpp
View File

@@ -165,7 +165,8 @@ bool Parameters::isFeatureParameter(const std::string & parameter)
group.compare("BRIEF") == 0 ||
group.compare("GFTT") == 0 ||
group.compare("BRISK") == 0 ||
group.compare("KAZE") == 0;
group.compare("KAZE") == 0 ||
group.compare("SPTorch") == 0;
}
rtabmap::ParametersMap Parameters::getDefaultOdometryParameters(bool stereo, bool vis, bool icp)
@@ -605,6 +606,12 @@ ParametersMap Parameters::parseArguments(int argc, char * argv[], bool onlyParam
std::cout << str << std::setw(spacing - str.size()) << "true" << std::endl;
#else
std::cout << str << std::setw(spacing - str.size()) << "false" << std::endl;
#endif
str = "With SuperPoint Torch:";
#ifdef RTABMAP_SP_TORCH
std::cout << str << std::setw(spacing - str.size()) << "true" << std::endl;
#else
std::cout << str << std::setw(spacing - str.size()) << "false" << std::endl;
#endif
str = "With FastCV:";
#ifdef RTABMAP_FASTCV

40
corelib/src/RegistrationVis.cpp
View File

@@ -71,7 +71,9 @@ RegistrationVis::RegistrationVis(const ParametersMap & parameters, Registration
_bundleAdjustment(Parameters::defaultVisBundleAdjustment()),
_depthAsMask(Parameters::defaultVisDepthAsMask()),
_minInliersDistributionThr(Parameters::defaultVisMinInliersDistribution()),
_maxInliersMeanDistance(Parameters::defaultVisMeanInliersDistance())
_maxInliersMeanDistance(Parameters::defaultVisMeanInliersDistance()),
_detectorFrom(0),
_detectorTo(0)
{
_featureParameters = Parameters::getDefaultParameters();
uInsert(_featureParameters, ParametersPair(Parameters::kKpNNStrategy(), _featureParameters.at(Parameters::kVisCorNNType())));
@@ -185,15 +187,17 @@ void RegistrationVis::parseParameters(const ParametersMap & parameters)
{
uInsert(_featureParameters, ParametersPair(Parameters::kKpGridCols(), parameters.at(Parameters::kVisGridCols())));
}
delete _detectorFrom;
delete _detectorTo;
_detectorFrom = Feature2D::create(_featureParameters);
_detectorTo = Feature2D::create(_featureParameters);
}
RegistrationVis::~RegistrationVis()
{
}
Feature2D * RegistrationVis::createFeatureDetector() const
{
return Feature2D::create(_featureParameters);
delete _detectorFrom;
delete _detectorTo;
}
Transform RegistrationVis::computeTransformationImpl(
@@ -280,8 +284,6 @@ Transform RegistrationVis::computeTransformationImpl(
toSignature.sensorData().imageRaw().type() == CV_8UC1 ||
toSignature.sensorData().imageRaw().type() == CV_8UC3);
Feature2D * detectorFrom = createFeatureDetector();
Feature2D * detectorTo = createFeatureDetector();
std::vector<cv::KeyPoint> kptsFrom;
cv::Mat imageFrom = fromSignature.sensorData().imageRaw();
cv::Mat imageTo = toSignature.sensorData().imageRaw();
@@ -311,7 +313,7 @@ Transform RegistrationVis::computeTransformationImpl(
}
}
kptsFrom = detectorFrom->generateKeypoints(
kptsFrom = _detectorFrom->generateKeypoints(
imageFrom,
depthMask);
}
@@ -378,7 +380,7 @@ Transform RegistrationVis::computeTransformationImpl(
}
else
{
kptsFrom3D = detectorFrom->generateKeypoints3D(fromSignature.sensorData(), kptsFrom);
kptsFrom3D = _detectorFrom->generateKeypoints3D(fromSignature.sensorData(), kptsFrom);
}
if(!imageFrom.empty() && !imageTo.empty())
@@ -452,7 +454,7 @@ Transform RegistrationVis::computeTransformationImpl(
std::vector<cv::Point3f> kptsTo3D;
if(_estimationType == 0 || _estimationType == 1 || !_forwardEstimateOnly)
{
kptsTo3D = detectorTo->generateKeypoints3D(toSignature.sensorData(), kptsTo);
kptsTo3D = _detectorTo->generateKeypoints3D(toSignature.sensorData(), kptsTo);
}
UASSERT(kptsFrom.size() == kptsFrom3DKept.size());
@@ -519,7 +521,7 @@ Transform RegistrationVis::computeTransformationImpl(
}
}
kptsTo = detectorTo->generateKeypoints(
kptsTo = _detectorTo->generateKeypoints(
imageTo,
depthMask);
}
@@ -566,7 +568,7 @@ Transform RegistrationVis::computeTransformationImpl(
}
UDEBUG("cleared orignalWordsFromIds");
orignalWordsFromIds.clear();
descriptorsFrom = detectorFrom->generateDescriptors(imageFrom, kptsFrom);
descriptorsFrom = _detectorFrom->generateDescriptors(imageFrom, kptsFrom);
}
cv::Mat descriptorsTo;
@@ -598,7 +600,7 @@ Transform RegistrationVis::computeTransformationImpl(
imageTo = tmp;
}
descriptorsTo = detectorTo->generateDescriptors(imageTo, kptsTo);
descriptorsTo = _detectorTo->generateDescriptors(imageTo, kptsTo);
}
}
@@ -629,9 +631,9 @@ Transform RegistrationVis::computeTransformationImpl(
kptsFrom.size(),
fromSignature.sensorData().keypoints3D().size());
}
kptsFrom3D = detectorFrom->generateKeypoints3D(fromSignature.sensorData(), kptsFrom);
kptsFrom3D = _detectorFrom->generateKeypoints3D(fromSignature.sensorData(), kptsFrom);
UDEBUG("generated kptsFrom3D=%d", (int)kptsFrom3D.size());
if(!kptsFrom3D.empty() && (detectorFrom->getMinDepth() > 0.0f || detectorFrom->getMaxDepth() > 0.0f))
if(!kptsFrom3D.empty() && (_detectorFrom->getMinDepth() > 0.0f || _detectorFrom->getMaxDepth() > 0.0f))
{
//remove all keypoints/descriptors with no valid 3D points
UASSERT_MSG((int)kptsFrom.size() == descriptorsFrom.rows &&
@@ -701,8 +703,8 @@ Transform RegistrationVis::computeTransformationImpl(
(int)kptsTo.size(),
(int)toSignature.sensorData().keypoints3D().size());
}
kptsTo3D = detectorTo->generateKeypoints3D(toSignature.sensorData(), kptsTo);
if(kptsTo3D.size() && (detectorTo->getMinDepth() > 0.0f || detectorTo->getMaxDepth() > 0.0f))
kptsTo3D = _detectorTo->generateKeypoints3D(toSignature.sensorData(), kptsTo);
if(kptsTo3D.size() && (_detectorTo->getMinDepth() > 0.0f || _detectorTo->getMaxDepth() > 0.0f))
{
UDEBUG("");
//remove all keypoints/descriptors with no valid 3D points
@@ -1174,8 +1176,6 @@ Transform RegistrationVis::computeTransformationImpl(
toSignature.setWords(wordsTo);
toSignature.setWords3(words3To);
toSignature.setWordsDescriptors(wordsDescTo);
delete detectorFrom;
delete detectorTo;
}
/////////////////////

336
corelib/src/superpoint_torch/SuperPoint.cc Normal file
View File

@@ -0,0 +1,336 @@
/**
* Original code from https://github.com/KinglittleQ/SuperPoint_SLAM
*/
#include <superpoint_torch/SuperPoint.h>
#include <rtabmap/utilite/ULogger.h>
namespace rtabmap
{
const int c1 = 64;
const int c2 = 64;
const int c3 = 128;
const int c4 = 128;
const int c5 = 256;
const int d1 = 256;
SuperPoint::SuperPoint()
: conv1a(torch::nn::Conv2dOptions( 1, c1, 3).stride(1).padding(1)),
conv1b(torch::nn::Conv2dOptions(c1, c1, 3).stride(1).padding(1)),
conv2a(torch::nn::Conv2dOptions(c1, c2, 3).stride(1).padding(1)),
conv2b(torch::nn::Conv2dOptions(c2, c2, 3).stride(1).padding(1)),
conv3a(torch::nn::Conv2dOptions(c2, c3, 3).stride(1).padding(1)),
conv3b(torch::nn::Conv2dOptions(c3, c3, 3).stride(1).padding(1)),
conv4a(torch::nn::Conv2dOptions(c3, c4, 3).stride(1).padding(1)),
conv4b(torch::nn::Conv2dOptions(c4, c4, 3).stride(1).padding(1)),
convPa(torch::nn::Conv2dOptions(c4, c5, 3).stride(1).padding(1)),
convPb(torch::nn::Conv2dOptions(c5, 65, 1).stride(1).padding(0)),
convDa(torch::nn::Conv2dOptions(c4, c5, 3).stride(1).padding(1)),
convDb(torch::nn::Conv2dOptions(c5, d1, 1).stride(1).padding(0))
{
register_module("conv1a", conv1a);
register_module("conv1b", conv1b);
register_module("conv2a", conv2a);
register_module("conv2b", conv2b);
register_module("conv3a", conv3a);
register_module("conv3b", conv3b);
register_module("conv4a", conv4a);
register_module("conv4b", conv4b);
register_module("convPa", convPa);
register_module("convPb", convPb);
register_module("convDa", convDa);
register_module("convDb", convDb);
}
std::vector<torch::Tensor> SuperPoint::forward(torch::Tensor x) {
x = torch::relu(conv1a->forward(x));
x = torch::relu(conv1b->forward(x));
x = torch::max_pool2d(x, 2, 2);
x = torch::relu(conv2a->forward(x));
x = torch::relu(conv2b->forward(x));
x = torch::max_pool2d(x, 2, 2);
x = torch::relu(conv3a->forward(x));
x = torch::relu(conv3b->forward(x));
x = torch::max_pool2d(x, 2, 2);
x = torch::relu(conv4a->forward(x));
x = torch::relu(conv4b->forward(x));
auto cPa = torch::relu(convPa->forward(x));
auto semi = convPb->forward(cPa); // [B, 65, H/8, W/8]
auto cDa = torch::relu(convDa->forward(x));
auto desc = convDb->forward(cDa); // [B, d1, H/8, W/8]
auto dn = torch::norm(desc, 2, 1);
desc = desc.div(torch::unsqueeze(dn, 1));
semi = torch::softmax(semi, 1);
semi = semi.slice(1, 0, 64);
semi = semi.permute({0, 2, 3, 1}); // [B, H/8, W/8, 64]
int Hc = semi.size(1);
int Wc = semi.size(2);
semi = semi.contiguous().view({-1, Hc, Wc, 8, 8});
semi = semi.permute({0, 1, 3, 2, 4});
semi = semi.contiguous().view({-1, Hc * 8, Wc * 8}); // [B, H, W]
std::vector<torch::Tensor> ret;
ret.push_back(semi);
ret.push_back(desc);
return ret;
}
void NMS(const std::vector<cv::KeyPoint> & ptsIn,
const cv::Mat & conf,
const cv::Mat & descriptorsIn,
std::vector<cv::KeyPoint> & ptsOut,
cv::Mat & descriptorsOut,
int border, int dist_thresh, int img_width, int img_height);
SPDetector::SPDetector(const std::string & modelPath, float threshold, bool nms, int minDistance, bool cuda) :
threshold_(threshold),
nms_(nms),
minDistance_(minDistance),
detected_(false)
{
UDEBUG("modelPath=%s thr=%f nms=%d cuda=%d", modelPath.c_str(), threshold, nms?1:0, cuda?1:0);
if(modelPath.empty())
{
return;
}
model_ = std::make_shared<SuperPoint>();
torch::load(model_, modelPath);
if(cuda && !torch::cuda::is_available())
{
UWARN("Cuda option is enabled but torch doesn't have cuda support on this platform, using CPU instead.");
}
cuda_ = cuda && torch::cuda::is_available();
torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
model_->to(device);
}
SPDetector::~SPDetector()
{
}
std::vector<cv::KeyPoint> SPDetector::detect(const cv::Mat &img)
{
detected_ = false;
if(model_)
{
torch::NoGradGuard no_grad_guard;
auto x = torch::from_blob(img.data, {1, 1, img.rows, img.cols}, torch::kByte);
x = x.to(torch::kFloat) / 255;
torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
x = x.set_requires_grad(false);
auto out = model_->forward(x.to(device));
prob_ = out[0].squeeze(0); // [H, W]
desc_ = out[1]; // [1, 256, H/8, W/8]
auto kpts = (prob_ > threshold_);
kpts = torch::nonzero(kpts); // [n_keypoints, 2] (y, x)
std::vector<cv::KeyPoint> keypoints_no_nms;
for (int i = 0; i < kpts.size(0); i++) {
float response = prob_[kpts[i][0]][kpts[i][1]].item<float>();
keypoints_no_nms.push_back(cv::KeyPoint(kpts[i][1].item<float>(), kpts[i][0].item<float>(), 8, -1, response));
}
detected_ = true;
if (nms_ && !keypoints_no_nms.empty()) {
cv::Mat conf(keypoints_no_nms.size(), 1, CV_32F);
for (size_t i = 0; i < keypoints_no_nms.size(); i++) {
int x = keypoints_no_nms[i].pt.x;
int y = keypoints_no_nms[i].pt.y;
conf.at<float>(i, 0) = prob_[y][x].item<float>();
}
int border = 0;
int dist_thresh = minDistance_;
int height = img.rows;
int width = img.cols;
std::vector<cv::KeyPoint> keypoints;
cv::Mat descEmpty;
NMS(keypoints_no_nms, conf, descEmpty, keypoints, descEmpty, border, dist_thresh, width, height);
return keypoints;
}
else {
return keypoints_no_nms;
}
}
else
{
UERROR("No model is loaded!");
return std::vector<cv::KeyPoint>();
}
}
cv::Mat SPDetector::compute(const std::vector<cv::KeyPoint> &keypoints)
{
if(!detected_)
{
UERROR("SPDetector has been reset before extracting the descriptors! detect() should be called before compute().");
return cv::Mat();
}
if(model_.get())
{
cv::Mat kpt_mat(keypoints.size(), 2, CV_32F); // [n_keypoints, 2] (y, x)
for (size_t i = 0; i < keypoints.size(); i++) {
kpt_mat.at<float>(i, 0) = (float)keypoints[i].pt.y;
kpt_mat.at<float>(i, 1) = (float)keypoints[i].pt.x;
}
auto fkpts = torch::from_blob(kpt_mat.data, {(long int)keypoints.size(), 2}, torch::kFloat);
torch::Device device(cuda_?torch::kCUDA:torch::kCPU);
auto grid = torch::zeros({1, 1, fkpts.size(0), 2}).to(device); // [1, 1, n_keypoints, 2]
grid[0][0].slice(1, 0, 1) = 2.0 * fkpts.slice(1, 1, 2) / prob_.size(1) - 1; // x
grid[0][0].slice(1, 1, 2) = 2.0 * fkpts.slice(1, 0, 1) / prob_.size(0) - 1; // y
auto desc = torch::grid_sampler(desc_, grid, 0, 0); // [1, 256, 1, n_keypoints]
desc = desc.squeeze(0).squeeze(1); // [256, n_keypoints]
// normalize to 1
auto dn = torch::norm(desc, 2, 1);
desc = desc.div(torch::unsqueeze(dn, 1));
desc = desc.transpose(0, 1).contiguous(); // [n_keypoints, 256]
if(cuda_)
desc = desc.to(torch::kCPU);
cv::Mat desc_mat(cv::Size(desc.size(1), desc.size(0)), CV_32FC1, desc.data<float>());
return desc_mat.clone();
}
else
{
UERROR("No model is loaded!");
return cv::Mat();
}
}
void NMS(const std::vector<cv::KeyPoint> & ptsIn,
const cv::Mat & conf,
const cv::Mat & descriptorsIn,
std::vector<cv::KeyPoint> & ptsOut,
cv::Mat & descriptorsOut,
int border, int dist_thresh, int img_width, int img_height)
{
std::vector<cv::Point2f> pts_raw;
for (size_t i = 0; i < ptsIn.size(); i++)
{
int u = (int) ptsIn[i].pt.x;
int v = (int) ptsIn[i].pt.y;
pts_raw.push_back(cv::Point2f(u, v));
}
cv::Mat grid = cv::Mat(cv::Size(img_width, img_height), CV_8UC1);
cv::Mat inds = cv::Mat(cv::Size(img_width, img_height), CV_16UC1);
cv::Mat confidence = cv::Mat(cv::Size(img_width, img_height), CV_32FC1);
grid.setTo(0);
inds.setTo(0);
confidence.setTo(0);
for (size_t i = 0; i < pts_raw.size(); i++)
{
int uu = (int) pts_raw[i].x;
int vv = (int) pts_raw[i].y;
grid.at<char>(vv, uu) = 1;
inds.at<unsigned short>(vv, uu) = i;
confidence.at<float>(vv, uu) = conf.at<float>(i, 0);
}
cv::copyMakeBorder(grid, grid, dist_thresh, dist_thresh, dist_thresh, dist_thresh, cv::BORDER_CONSTANT, 0);
for (size_t i = 0; i < pts_raw.size(); i++)
{
int uu = (int) pts_raw[i].x + dist_thresh;
int vv = (int) pts_raw[i].y + dist_thresh;
if (grid.at<char>(vv, uu) != 1)
continue;
for(int k = -dist_thresh; k < (dist_thresh+1); k++)
for(int j = -dist_thresh; j < (dist_thresh+1); j++)
{
if(j==0 && k==0) continue;
if ( conf.at<float>(vv + k, uu + j) < conf.at<float>(vv, uu) )
grid.at<char>(vv + k, uu + j) = 0;
}
grid.at<char>(vv, uu) = 2;
}
size_t valid_cnt = 0;
std::vector<int> select_indice;
for (int v = 0; v < (img_height + dist_thresh); v++){
for (int u = 0; u < (img_width + dist_thresh); u++)
{
if (u -dist_thresh>= (img_width - border) || u-dist_thresh < border || v-dist_thresh >= (img_height - border) || v-dist_thresh < border)
continue;
if (grid.at<char>(v,u) == 2)
{
int select_ind = (int) inds.at<unsigned short>(v-dist_thresh, u-dist_thresh);
float response = conf.at<float>(select_ind, 0);
ptsOut.push_back(cv::KeyPoint(pts_raw[select_ind], 8.0f, -1, response));
select_indice.push_back(select_ind);
valid_cnt++;
}
}
}
if(!descriptorsIn.empty())
{
UASSERT(descriptorsIn.rows == (int)ptsIn.size());
descriptorsOut.create(select_indice.size(), 256, CV_32F);
for (size_t i=0; i<select_indice.size(); i++)
{
for (int j=0; j < 256; j++)
{
descriptorsOut.at<float>(i, j) = descriptorsIn.at<float>(select_indice[i], j);
}
}
}
}
}

75
corelib/src/superpoint_torch/SuperPoint.h Normal file
View File

@@ -0,0 +1,75 @@
/**
* Original code from https://github.com/KinglittleQ/SuperPoint_SLAM
*/
#ifndef SUPERPOINT_H
#define SUPERPOINT_H
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <vector>
#ifdef EIGEN_MPL2_ONLY
#undef EIGEN_MPL2_ONLY
#endif
namespace rtabmap
{
struct SuperPoint : torch::nn::Module {
SuperPoint();
std::vector<torch::Tensor> forward(torch::Tensor x);
torch::nn::Conv2d conv1a;
torch::nn::Conv2d conv1b;
torch::nn::Conv2d conv2a;
torch::nn::Conv2d conv2b;
torch::nn::Conv2d conv3a;
torch::nn::Conv2d conv3b;
torch::nn::Conv2d conv4a;
torch::nn::Conv2d conv4b;
torch::nn::Conv2d convPa;
torch::nn::Conv2d convPb;
// descriptor
torch::nn::Conv2d convDa;
torch::nn::Conv2d convDb;
};
class SPDetector {
public:
SPDetector(const std::string & modelPath, float threshold = 0.2f, bool nms = true, int minDistance = 4, bool cuda = false);
virtual ~SPDetector();
std::vector<cv::KeyPoint> detect(const cv::Mat &img);
cv::Mat compute(const std::vector<cv::KeyPoint> &keypoints);
void setThreshold(float threshold) {threshold_ = threshold;}
void SetNMS(bool enabled) {nms_ = enabled;}
void setMinDistance(float minDistance) {minDistance_ = minDistance;}
private:
std::shared_ptr<SuperPoint> model_;
torch::Tensor prob_;
torch::Tensor desc_;
float threshold_;
bool nms_;
int minDistance_;
bool cuda_;
bool detected_;
};
}
#endif