Limit Max Features (#1614)

* 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]>

corelib/src/Features2d.cpp

@@ -2668,7 +2668,7 @@ void SuperPointRpautrat::parseParameters(const ParametersMap & parameters)
 		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 || 
 	   superpointWeightsPath_.compare(previousWeightsPath) != 0 || 
 	   superpointModelPath_.compare(previousModelPath) != 0 ||
@@ -2677,7 +2677,13 @@ void SuperPointRpautrat::parseParameters(const ParametersMap & parameters)
 	   previousNms != nms_ ||
 	   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
 	UWARN("RTAB-Map is not built with Torch support so SuperPoint Rpautrat feature cannot be used!");

corelib/src/python/PyDetector.cpp

@@ -156,12 +156,12 @@ std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & imag
 	if(pFunc_)
 	{
 		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);
 
 		UDEBUG("Preparing data time = %fs", timer.ticks());
 
-		PyObject *pReturn = PyObject_CallFunctionObjArgs(pFunc_, pImageBuffer, NULL);
+		PyObject * pReturn = PyObject_CallFunctionObjArgs(pFunc_, pImageBuffer, NULL);
 		if(pReturn == NULL)
 		{
 			UERROR("Failed to call match() function!");
@@ -173,8 +173,8 @@ std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & imag
 
 			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))
 				{
 					PyArrayObject *arrayPtr = reinterpret_cast<PyArrayObject*>(kptsPtr);
@@ -186,11 +186,18 @@ std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & imag
 					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));
+					std::vector<bool> keep_kpt(nKpts);
 					keypoints.reserve(nKpts);
-					for (int i = 0; i < nKpts*kptSize; i+=kptSize)
+					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]);
+						// x,y in full image coordinates. Mask is in full image coordinates too.
-						keypoints.push_back(kpt);
+						int full_x = (int)(c_out[i] + roi.x);
+						int full_y = (int)(c_out[i+1] + roi.y);
+						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);
+						if(keep_kpt[kpt_idx]) {
+							cv::KeyPoint kpt(c_out[i], c_out[i+1], 8, -1, c_out[i+2]);
+							keypoints.push_back(kpt);
+						}
 					}
 
 					arrayPtr = reinterpret_cast<PyArrayObject*>(descPtr);
@@ -202,10 +209,12 @@ std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & imag
 					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));
-					for (int i = 0; i < nDesc*dim; i+=dim)
+					for (int i = 0, kpt_idx = 0; i < nDesc*dim; i+=dim, kpt_idx++)
 					{
-						cv::Mat descriptor = cv::Mat(1, dim, CV_32FC1, &c_out[i]).clone();
+						if(keep_kpt[kpt_idx]) {
-						descriptors_.push_back(descriptor);
+							cv::Mat descriptor = cv::Mat(1, dim, CV_32FC1, &c_out[i]).clone();
+							descriptors_.push_back(descriptor);
+						}
 					}
 				}
 			}
@@ -218,6 +227,9 @@ std::vector<cv::KeyPoint> PyDetector::generateKeypointsImpl(const cv::Mat & imag
 		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;
 }
 

corelib/src/superpoint_rpautrat/SuperpointRpautrat.cpp

@@ -4,6 +4,7 @@
  */
 
 #include "SuperpointRpautrat.h"
+#include <rtabmap/core/Features2d.h>
 #include <rtabmap/utilite/ULogger.h>
 #include <rtabmap/utilite/UDirectory.h>
 #include <rtabmap/utilite/UFile.h>
@@ -95,7 +96,7 @@ static std::string exportSuperPointTorchScript(
 	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),
 		superpointWeightsPath_(superpointWeightsPath),
         superpointModelPath_(superpointModelPath),
@@ -103,6 +104,8 @@ SPDetectorRpautrat::SPDetectorRpautrat(std::string superpointWeightsPath, std::s
         threshold_(threshold),
 		nms_(nms),
 		minDistance_(minDistance),
+        maxFeatures_(maxFeatures),
+        ssc_(ssc),
 		detected_(false)
 {
 	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
-    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)
@@ -189,12 +189,12 @@ std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const c
     x = x.set_requires_grad(false).to(device_);
     
     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
-    auto keypoints_cpu = keypoints_tensor_.to(torch::kCPU);
+    auto keypoints_cpu = kpts_tensor.to(torch::kCPU);
     auto scores_cpu = scores_tensor.to(torch::kCPU);
 
     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);
-    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;
     return filtered_keypoints;
 }

corelib/src/superpoint_rpautrat/SuperpointRpautrat.h

@@ -23,17 +23,22 @@ class SPDetectorRpautrat {
             float threshold = 0.005f, 
             bool nms = true, 
             int nmsRadius = 4, 
-            bool cuda = false
+            bool cuda = false,
+            int maxFeatures = 1000,
+            bool ssc = false
         );
         virtual ~SPDetectorRpautrat();
         std::vector<cv::KeyPoint> detect(const cv::Mat &img, const cv::Mat & mask = cv::Mat());
         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:
         torch::jit::script::Module model_;
         torch::Device device_;
-        torch::Tensor desc_;
+        cv::Mat desc_;
-        torch::Tensor keypoints_tensor_;
         
         std::string superpointWeightsPath_;
         std::string superpointModelPath_;
@@ -42,6 +47,8 @@ class SPDetectorRpautrat {
         bool nms_;
         int minDistance_;
         bool cuda_;
+        int maxFeatures_;
+        bool ssc_;
         
         bool detected_;
     };    

corelib/src/superpoint_rpautrat/superpoint_to_torchscript.py

@@ -59,7 +59,7 @@ def generate_model(
     # Load SuperPoint model and weights
     model = SuperPoint(
         nms_radius=nms_radius,
-        threshold=threshold,
+        detection_threshold=threshold,
     ).eval().to(device)
     
     # Load weights without forcing CPU location to allow CUDA usage