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

Limit Max Features (#1614)

Browse Source 
* limit max features internally inside pydetector and superpoint_rpautrat to  limit keypoint/desc data size and improve performance

* avoid full reset when max features is changed to work around the memory temprorary param update logic

* remove leftover comment

* try using image roi instead

* revert and regenerate

* unintended change

* fixing issue with pydetector, adding mask filtering

---------

Co-authored-by: Felix Toft <[email protected]>
This commit is contained in:
Felix Toft
2025-11-16 11:37:03 -08:00
committed by GitHub
co-authored by Felix Toft
parent e612d103bf
commit d1b5d62d21
5 changed files with 284 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files
corelib/src/Features2d.cpp
+9 -3
View File
@@ -2661,14 +2661,14 @@ void SuperPointRpautrat::parseParameters(const ParametersMap & parameters)
Parameters::parse(parameters, Parameters::kSuperPointRpautratNMSRadius(), minDistance_); Parameters::parse(parameters, Parameters::kSuperPointRpautratNMSRadius(), minDistance_);
Parameters::parse(parameters, Parameters::kSuperPointRpautratCuda(), cuda_); Parameters::parse(parameters, Parameters::kSuperPointRpautratCuda(), cuda_);
Parameters::parse(parameters, Parameters::kRtabmapWorkingDirectory(), outputDir_); Parameters::parse(parameters, Parameters::kRtabmapWorkingDirectory(), outputDir_);
// If working directory is not set, use the default // If working directory is not set, use the default
if(outputDir_.empty()) if(outputDir_.empty())
{ {
outputDir_ = Parameters::createDefaultWorkingDirectory(); outputDir_ = Parameters::createDefaultWorkingDirectory();
} }
// Delete the detector to force re-initialization on next frame if any parameter changed // Reinitialize detector if model-affecting parameters changed
if(superPoint_.get() == 0 || if(superPoint_.get() == 0 ||
superpointWeightsPath_.compare(previousWeightsPath) != 0 || superpointWeightsPath_.compare(previousWeightsPath) != 0 ||
superpointModelPath_.compare(previousModelPath) != 0 || superpointModelPath_.compare(previousModelPath) != 0 ||
@@ -2677,7 +2677,13 @@ void SuperPointRpautrat::parseParameters(const ParametersMap & parameters)
previousNms != nms_ || previousNms != nms_ ||
previousMinDistance != minDistance_) previousMinDistance != minDistance_)
{ {
superPoint_ = cv::Ptr<SPDetectorRpautrat>(new SPDetectorRpautrat(superpointWeightsPath_, superpointModelPath_, outputDir_, threshold_, nms_, minDistance_, cuda_)); superPoint_ = cv::Ptr<SPDetectorRpautrat>(new SPDetectorRpautrat(superpointWeightsPath_, superpointModelPath_, outputDir_, threshold_, nms_, minDistance_, cuda_, this->getMaxFeatures(), this->getSSC()));
}
else if(superPoint_.get() != 0)
{
// Update post-processing parameters without reinitializing
superPoint_->setMaxFeatures(this->getMaxFeatures());
superPoint_->setSSC(this->getSSC());
} }
#else #else
UWARN("RTAB-Map is not built with Torch support so SuperPoint Rpautrat feature cannot be used!"); UWARN("RTAB-Map is not built with Torch support so SuperPoint Rpautrat feature cannot be used!");
corelib/src/python/PyDetector.cpp
+242 -230
View File
@@ -1,230 +1,242 @@
/** /**
* Python interface for SuperGlue: https://github.com/magicleap/SuperGluePretrainedNetwork * Python interface for SuperGlue: https://github.com/magicleap/SuperGluePretrainedNetwork
*/ */
#include "PyDetector.h" #include "PyDetector.h"
#include <rtabmap/utilite/ULogger.h> #include <rtabmap/utilite/ULogger.h>
#include <rtabmap/utilite/UDirectory.h> #include <rtabmap/utilite/UDirectory.h>
#include <rtabmap/utilite/UFile.h> #include <rtabmap/utilite/UFile.h>
#include <rtabmap/utilite/UStl.h> #include <rtabmap/utilite/UStl.h>
#include <rtabmap/utilite/UConversion.h> #include <rtabmap/utilite/UConversion.h>
#include <rtabmap/utilite/UTimer.h> #include <rtabmap/utilite/UTimer.h>
#include <pybind11/embed.h> #include <pybind11/embed.h>
#define NPY_NO_DEPRECATED_API NPY_API_VERSION #define NPY_NO_DEPRECATED_API NPY_API_VERSION
#include <numpy/arrayobject.h> #include <numpy/arrayobject.h>
namespace rtabmap namespace rtabmap
{ {
PyDetector::PyDetector(const ParametersMap & parameters) : PyDetector::PyDetector(const ParametersMap & parameters) :
pModule_(0), pModule_(0),
pFunc_(0), pFunc_(0),
path_(Parameters::defaultPyDetectorPath()), path_(Parameters::defaultPyDetectorPath()),
cuda_(Parameters::defaultPyDetectorCuda()) cuda_(Parameters::defaultPyDetectorCuda())
{ {
this->parseParameters(parameters); this->parseParameters(parameters);
UDEBUG("path = %s", path_.c_str()); UDEBUG("path = %s", path_.c_str());
if(!UFile::exists(path_) || UFile::getExtension(path_).compare("py") != 0) if(!UFile::exists(path_) || UFile::getExtension(path_).compare("py") != 0)
{ {
UERROR("Cannot initialize Python detector, the path is not valid: \"%s\"=\"%s\"", UERROR("Cannot initialize Python detector, the path is not valid: \"%s\"=\"%s\"",
Parameters::kPyDetectorPath().c_str(), path_.c_str()); Parameters::kPyDetectorPath().c_str(), path_.c_str());
return; return;
} }
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
std::string matcherPythonDir = UDirectory::getDir(path_); std::string matcherPythonDir = UDirectory::getDir(path_);
if(!matcherPythonDir.empty()) if(!matcherPythonDir.empty())
{ {
PyRun_SimpleString("import sys"); PyRun_SimpleString("import sys");
PyRun_SimpleString(uFormat("sys.path.append(\"%s\")", matcherPythonDir.c_str()).c_str()); PyRun_SimpleString(uFormat("sys.path.append(\"%s\")", matcherPythonDir.c_str()).c_str());
} }
_import_array(); _import_array();
std::string scriptName = uSplit(UFile::getName(path_), '.').front(); std::string scriptName = uSplit(UFile::getName(path_), '.').front();
PyObject * pName = PyUnicode_FromString(scriptName.c_str()); PyObject * pName = PyUnicode_FromString(scriptName.c_str());
UDEBUG("PyImport_Import() beg"); UDEBUG("PyImport_Import() beg");
pModule_ = PyImport_Import(pName); pModule_ = PyImport_Import(pName);
UDEBUG("PyImport_Import() end"); UDEBUG("PyImport_Import() end");
Py_DECREF(pName); Py_DECREF(pName);
if(!pModule_) if(!pModule_)
{ {
UERROR("Module \"%s\" could not be imported! (File=\"%s\")", scriptName.c_str(), path_.c_str()); UERROR("Module \"%s\" could not be imported! (File=\"%s\")", scriptName.c_str(), path_.c_str());
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
} }
} }
PyDetector::~PyDetector() PyDetector::~PyDetector()
{ {
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
if(pFunc_) if(pFunc_)
{ {
Py_DECREF(pFunc_); Py_DECREF(pFunc_);
} }
if(pModule_) if(pModule_)
{ {
Py_DECREF(pModule_); Py_DECREF(pModule_);
} }
} }
void PyDetector::parseParameters(const ParametersMap & parameters) void PyDetector::parseParameters(const ParametersMap & parameters)
{ {
Feature2D::parseParameters(parameters); Feature2D::parseParameters(parameters);
Parameters::parse(parameters, Parameters::kPyDetectorPath(), path_); Parameters::parse(parameters, Parameters::kPyDetectorPath(), path_);
Parameters::parse(parameters, Parameters::kPyDetectorCuda(), cuda_); Parameters::parse(parameters, Parameters::kPyDetectorCuda(), cuda_);
path_ = uReplaceChar(path_, '~', UDirectory::homeDir()); path_ = uReplaceChar(path_, '~', UDirectory::homeDir());
} }
std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & image, const cv::Rect & roi, const cv::Mat & mask) std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & image, const cv::Rect & roi, const cv::Mat & mask)
{ {
UDEBUG(""); UDEBUG("");
descriptors_ = cv::Mat(); descriptors_ = cv::Mat();
UASSERT(!image.empty() && image.channels() == 1 && image.depth() == CV_8U); UASSERT(!image.empty() && image.channels() == 1 && image.depth() == CV_8U);
std::vector<cv::KeyPoint> keypoints; std::vector<cv::KeyPoint> keypoints;
cv::Mat imgRoi(image, roi); cv::Mat imgRoi(image, roi);
UTimer timer; UTimer timer;
if(!pModule_) if(!pModule_)
{ {
UERROR("Python detector module not loaded!"); UERROR("Python detector module not loaded!");
return keypoints; return keypoints;
} }
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
if(!pFunc_) if(!pFunc_)
{ {
PyObject * pFunc = PyObject_GetAttrString(pModule_, "init"); PyObject * pFunc = PyObject_GetAttrString(pModule_, "init");
if(pFunc) if(pFunc)
{ {
if(PyCallable_Check(pFunc)) if(PyCallable_Check(pFunc))
{ {
PyObject * result = PyObject_CallFunction(pFunc, "i", cuda_?1:0); PyObject * result = PyObject_CallFunction(pFunc, "i", cuda_?1:0);
if(result == NULL) if(result == NULL)
{ {
UERROR("Call to \"init(...)\" in \"%s\" failed!", path_.c_str()); UERROR("Call to \"init(...)\" in \"%s\" failed!", path_.c_str());
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
return keypoints; return keypoints;
} }
Py_DECREF(result); Py_DECREF(result);
pFunc_ = PyObject_GetAttrString(pModule_, "detect"); pFunc_ = PyObject_GetAttrString(pModule_, "detect");
if(pFunc_ && PyCallable_Check(pFunc_)) if(pFunc_ && PyCallable_Check(pFunc_))
{ {
// we are ready! // we are ready!
} }
else else
{ {
UERROR("Cannot find method \"detect(...)\" in %s", path_.c_str()); UERROR("Cannot find method \"detect(...)\" in %s", path_.c_str());
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
if(pFunc_) if(pFunc_)
{ {
Py_DECREF(pFunc_); Py_DECREF(pFunc_);
pFunc_ = 0; pFunc_ = 0;
} }
return keypoints; return keypoints;
} }
} }
else else
{ {
UERROR("Cannot call method \"init(...)\" in %s", path_.c_str()); UERROR("Cannot call method \"init(...)\" in %s", path_.c_str());
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
return keypoints; return keypoints;
} }
Py_DECREF(pFunc); Py_DECREF(pFunc);
} }
else else
{ {
UERROR("Cannot find method \"init(...)\""); UERROR("Cannot find method \"init(...)\"");
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
return keypoints; return keypoints;
} }
UDEBUG("init time = %fs", timer.ticks()); UDEBUG("init time = %fs", timer.ticks());
} }
if(pFunc_) if(pFunc_)
{ {
npy_intp dims[2] = {imgRoi.rows, imgRoi.cols}; npy_intp dims[2] = {imgRoi.rows, imgRoi.cols};
PyObject* pImageBuffer = PyArray_SimpleNewFromData(2, dims, NPY_UBYTE, (void*)imgRoi.data); PyObject * pImageBuffer = PyArray_SimpleNewFromData(2, dims, NPY_UBYTE, (void*)imgRoi.data);
UASSERT(pImageBuffer); UASSERT(pImageBuffer);
UDEBUG("Preparing data time = %fs", timer.ticks()); UDEBUG("Preparing data time = %fs", timer.ticks());
PyObject *pReturn = PyObject_CallFunctionObjArgs(pFunc_, pImageBuffer, NULL); PyObject * pReturn = PyObject_CallFunctionObjArgs(pFunc_, pImageBuffer, NULL);
if(pReturn == NULL) if(pReturn == NULL)
{ {
UERROR("Failed to call match() function!"); UERROR("Failed to call match() function!");
UERROR("%s", getPythonTraceback().c_str()); UERROR("%s", getPythonTraceback().c_str());
} }
else else
{ {
UDEBUG("Python detector time = %fs", timer.ticks()); UDEBUG("Python detector time = %fs", timer.ticks());
if (PyTuple_Check(pReturn) && PyTuple_GET_SIZE(pReturn) == 2) if (PyTuple_Check(pReturn) && PyTuple_GET_SIZE(pReturn) == 2)
{ {
PyObject *kptsPtr = PyTuple_GET_ITEM(pReturn, 0); PyObject * kptsPtr = PyTuple_GET_ITEM(pReturn, 0);
PyObject *descPtr = PyTuple_GET_ITEM(pReturn, 1); PyObject * descPtr = PyTuple_GET_ITEM(pReturn, 1);
if(PyArray_Check(kptsPtr) && PyArray_Check(descPtr)) if(PyArray_Check(kptsPtr) && PyArray_Check(descPtr))
{ {
PyArrayObject *arrayPtr = reinterpret_cast<PyArrayObject*>(kptsPtr); PyArrayObject *arrayPtr = reinterpret_cast<PyArrayObject*>(kptsPtr);
int nKpts = PyArray_SHAPE(arrayPtr)[0]; int nKpts = PyArray_SHAPE(arrayPtr)[0];
int kptSize = PyArray_SHAPE(arrayPtr)[1]; int kptSize = PyArray_SHAPE(arrayPtr)[1];
int type = PyArray_TYPE(arrayPtr); int type = PyArray_TYPE(arrayPtr);
UDEBUG("Kpts array %dx%d (type=%d)", nKpts, kptSize, type); UDEBUG("Kpts array %dx%d (type=%d)", nKpts, kptSize, type);
UASSERT(kptSize == 3); UASSERT(kptSize == 3);
UASSERT_MSG(type == NPY_FLOAT, uFormat("Returned matches should type FLOAT=11, received type=%d", type).c_str()); UASSERT_MSG(type == NPY_FLOAT, uFormat("Returned matches should type FLOAT=11, received type=%d", type).c_str());
float* c_out = reinterpret_cast<float*>(PyArray_DATA(arrayPtr)); float* c_out = reinterpret_cast<float*>(PyArray_DATA(arrayPtr));
keypoints.reserve(nKpts); std::vector<bool> keep_kpt(nKpts);
for (int i = 0; i < nKpts*kptSize; i+=kptSize) keypoints.reserve(nKpts);
{ for (int i = 0, kpt_idx = 0; i < nKpts*kptSize; i+=kptSize, kpt_idx++)
cv::KeyPoint kpt(c_out[i], c_out[i+1], 8, -1, c_out[i+2]); {
keypoints.push_back(kpt); // x,y in full image coordinates. Mask is in full image coordinates too.
} int full_x = (int)(c_out[i] + roi.x);
int full_y = (int)(c_out[i+1] + roi.y);
arrayPtr = reinterpret_cast<PyArrayObject*>(descPtr); keep_kpt[kpt_idx] = mask.empty() || (full_x >= 0 && full_x < mask.cols && full_y >= 0 && full_y < mask.rows && mask.at<unsigned char>(full_y, full_x) != 0);
int nDesc = PyArray_SHAPE(arrayPtr)[0]; if(keep_kpt[kpt_idx]) {
UASSERT(nDesc = nKpts); cv::KeyPoint kpt(c_out[i], c_out[i+1], 8, -1, c_out[i+2]);
int dim = PyArray_SHAPE(arrayPtr)[1]; keypoints.push_back(kpt);
type = PyArray_TYPE(arrayPtr); }
UDEBUG("Desc array %dx%d (type=%d)", nDesc, dim, type); }
UASSERT_MSG(type == NPY_FLOAT, uFormat("Returned matches should type FLOAT=11, received type=%d", type).c_str());
arrayPtr = reinterpret_cast<PyArrayObject*>(descPtr);
c_out = reinterpret_cast<float*>(PyArray_DATA(arrayPtr)); int nDesc = PyArray_SHAPE(arrayPtr)[0];
for (int i = 0; i < nDesc*dim; i+=dim) UASSERT(nDesc = nKpts);
{ int dim = PyArray_SHAPE(arrayPtr)[1];
cv::Mat descriptor = cv::Mat(1, dim, CV_32FC1, &c_out[i]).clone(); type = PyArray_TYPE(arrayPtr);
descriptors_.push_back(descriptor); UDEBUG("Desc array %dx%d (type=%d)", nDesc, dim, type);
} UASSERT_MSG(type == NPY_FLOAT, uFormat("Returned matches should type FLOAT=11, received type=%d", type).c_str());
}
} c_out = reinterpret_cast<float*>(PyArray_DATA(arrayPtr));
else for (int i = 0, kpt_idx = 0; i < nDesc*dim; i+=dim, kpt_idx++)
{ {
UWARN("Expected tuple (Kpts 3 x N, Descriptors dim x N), returning empty features."); if(keep_kpt[kpt_idx]) {
} cv::Mat descriptor = cv::Mat(1, dim, CV_32FC1, &c_out[i]).clone();
Py_DECREF(pReturn); descriptors_.push_back(descriptor);
} }
Py_DECREF(pImageBuffer); }
} }
}
return keypoints; else
} {
UWARN("Expected tuple (Kpts 3 x N, Descriptors dim x N), returning empty features.");
cv::Mat PyDetector::generateDescriptorsImpl(const cv::Mat & image, std::vector<cv::KeyPoint> & keypoints) const }
{ Py_DECREF(pReturn);
UASSERT((int)keypoints.size() == descriptors_.rows); }
return descriptors_; Py_DECREF(pImageBuffer);
} }
} // Apply limitKeypoints to enforce maxFeatures and SSC
this->limitKeypoints(keypoints, descriptors_, this->getMaxFeatures(), cv::Size(roi.width, roi.height), this->getSSC());
return keypoints;
}
cv::Mat PyDetector::generateDescriptorsImpl(const cv::Mat & image, std::vector<cv::KeyPoint> & keypoints) const
{
UASSERT((int)keypoints.size() == descriptors_.rows);
return descriptors_;
}
}
corelib/src/superpoint_rpautrat/SuperpointRpautrat.cpp
+22 -18
View File
@@ -4,6 +4,7 @@
*/ */
#include "SuperpointRpautrat.h" #include "SuperpointRpautrat.h"
#include <rtabmap/core/Features2d.h>
#include <rtabmap/utilite/ULogger.h> #include <rtabmap/utilite/ULogger.h>
#include <rtabmap/utilite/UDirectory.h> #include <rtabmap/utilite/UDirectory.h>
#include <rtabmap/utilite/UFile.h> #include <rtabmap/utilite/UFile.h>
@@ -95,7 +96,7 @@ static std::string exportSuperPointTorchScript(
return output; return output;
} }
SPDetectorRpautrat::SPDetectorRpautrat(std::string superpointWeightsPath, std::string superpointModelPath, std::string outputDir, float threshold, bool nms, int minDistance, bool cuda) : SPDetectorRpautrat::SPDetectorRpautrat(std::string superpointWeightsPath, std::string superpointModelPath, std::string outputDir, float threshold, bool nms, int minDistance, bool cuda, int maxFeatures, bool ssc) :
device_(torch::kCPU), device_(torch::kCPU),
superpointWeightsPath_(superpointWeightsPath), superpointWeightsPath_(superpointWeightsPath),
superpointModelPath_(superpointModelPath), superpointModelPath_(superpointModelPath),
@@ -103,6 +104,8 @@ SPDetectorRpautrat::SPDetectorRpautrat(std::string superpointWeightsPath, std::s
threshold_(threshold), threshold_(threshold),
nms_(nms), nms_(nms),
minDistance_(minDistance), minDistance_(minDistance),
maxFeatures_(maxFeatures),
ssc_(ssc),
detected_(false) detected_(false)
{ {
if(cuda && !torch::cuda::is_available()) if(cuda && !torch::cuda::is_available())
@@ -136,12 +139,9 @@ cv::Mat SPDetectorRpautrat::compute(const std::vector<cv::KeyPoint> &keypoints)
} }
// These should have the same size // These should have the same size
UASSERT(static_cast<size_t>(desc_.size(0)) == keypoints.size()); UASSERT(static_cast<size_t>(desc_.rows) == keypoints.size());
// Move to CPU and return descriptors computed in the forward pass return desc_;
torch::Tensor desc_cpu = desc_.to(torch::kCPU);
cv::Mat desc_mat(cv::Size(desc_cpu.size(1), desc_cpu.size(0)), CV_32FC1, desc_cpu.data_ptr<float>());
return desc_mat.clone();
} }
std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const cv::Mat & mask) std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const cv::Mat & mask)
@@ -189,12 +189,12 @@ std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const c
x = x.set_requires_grad(false).to(device_); x = x.set_requires_grad(false).to(device_);
auto outputs = model_.forward({x}).toTuple(); auto outputs = model_.forward({x}).toTuple();
keypoints_tensor_ = outputs->elements()[0].toTensor(); // [N, 2] keypoint coordinates auto kpts_tensor = outputs->elements()[0].toTensor(); // [N, 2] keypoint coordinates
auto scores_tensor = outputs->elements()[1].toTensor(); // [N] keypoint scores auto scores_tensor = outputs->elements()[1].toTensor(); // [N] keypoint scores
desc_ = outputs->elements()[2].toTensor(); // [N, 256] descriptors torch::Tensor desc_tensor = outputs->elements()[2].toTensor(); // [N, 256] descriptors
// Convert to CPU for processing // Convert to CPU for processing
auto keypoints_cpu = keypoints_tensor_.to(torch::kCPU); auto keypoints_cpu = kpts_tensor.to(torch::kCPU);
auto scores_cpu = scores_tensor.to(torch::kCPU); auto scores_cpu = scores_tensor.to(torch::kCPU);
std::vector<cv::KeyPoint> filtered_keypoints; std::vector<cv::KeyPoint> filtered_keypoints;
@@ -213,16 +213,20 @@ std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const c
} }
} }
// Update the stored tensors to maintain correspondence // Filter descriptors based on mask
// This way if keypoints are re-ordered, we can still match kpts->descs in the compute step
auto keep_indices = torch::from_blob(keep_indices_vec.data(), {(long int)keep_indices_vec.size()}, torch::kLong); auto keep_indices = torch::from_blob(keep_indices_vec.data(), {(long int)keep_indices_vec.size()}, torch::kLong);
keep_indices = keep_indices.to(keypoints_tensor_.device()); keep_indices = keep_indices.to(desc_tensor.device());
auto filtered_keypoints_tensor = keypoints_tensor_.index_select(0, keep_indices); auto filtered_descriptors = desc_tensor.index_select(0, keep_indices);
auto filtered_descriptors = desc_.index_select(0, keep_indices);
keypoints_tensor_ = filtered_keypoints_tensor; // Convert descriptors to cv::Mat
desc_ = filtered_descriptors; auto filtered_descriptors_cpu = filtered_descriptors.to(torch::kCPU);
cv::Mat descriptors_mat(filtered_descriptors_cpu.size(0), filtered_descriptors_cpu.size(1), CV_32FC1, filtered_descriptors_cpu.data_ptr<float>());
cv::Mat descriptors_clone = descriptors_mat.clone(); // Clone to own the memory
// Apply limitKeypoints to enforce maxFeatures and SSC
Feature2D::limitKeypoints(filtered_keypoints, descriptors_clone, maxFeatures_, cv::Size(img.cols, img.rows), ssc_);
desc_ = descriptors_clone;
detected_ = true; detected_ = true;
return filtered_keypoints; return filtered_keypoints;
} }
corelib/src/superpoint_rpautrat/SuperpointRpautrat.h
+10 -3
View File
@@ -23,17 +23,22 @@ class SPDetectorRpautrat {
float threshold = 0.005f, float threshold = 0.005f,
bool nms = true, bool nms = true,
int nmsRadius = 4, int nmsRadius = 4,
bool cuda = false bool cuda = false,
int maxFeatures = 1000,
bool ssc = false
); );
virtual ~SPDetectorRpautrat(); virtual ~SPDetectorRpautrat();
std::vector<cv::KeyPoint> detect(const cv::Mat &img, const cv::Mat & mask = cv::Mat()); std::vector<cv::KeyPoint> detect(const cv::Mat &img, const cv::Mat & mask = cv::Mat());
cv::Mat compute(const std::vector<cv::KeyPoint> &keypoints); cv::Mat compute(const std::vector<cv::KeyPoint> &keypoints);
// Setters for post-processing parameters that don't require model reinitialization
void setMaxFeatures(int maxFeatures) { maxFeatures_ = maxFeatures; }
void setSSC(bool ssc) { ssc_ = ssc; }
private: private:
torch::jit::script::Module model_; torch::jit::script::Module model_;
torch::Device device_; torch::Device device_;
torch::Tensor desc_; cv::Mat desc_;
torch::Tensor keypoints_tensor_;
std::string superpointWeightsPath_; std::string superpointWeightsPath_;
std::string superpointModelPath_; std::string superpointModelPath_;
@@ -42,6 +47,8 @@ class SPDetectorRpautrat {
bool nms_; bool nms_;
int minDistance_; int minDistance_;
bool cuda_; bool cuda_;
int maxFeatures_;
bool ssc_;
bool detected_; bool detected_;
}; };
corelib/src/superpoint_rpautrat/superpoint_to_torchscript.py
+1 -1
View File
@@ -59,7 +59,7 @@ def generate_model(
# Load SuperPoint model and weights # Load SuperPoint model and weights
model = SuperPoint( model = SuperPoint(
nms_radius=nms_radius, nms_radius=nms_radius,
threshold=threshold, detection_threshold=threshold,
).eval().to(device) ).eval().to(device)
# Load weights without forcing CPU location to allow CUDA usage # Load weights without forcing CPU location to allow CUDA usage