super point nms

Author: xmba15

examples/SuperPointApp.cpp

@@ -7,14 +7,33 @@
 
 #include "SuperPoint.hpp"
 #include "Utility.hpp"
+#include <algorithm>
+#include <iterator>
 
 namespace
 {
 using KeyPointAndDesc = std::pair<std::vector<cv::KeyPoint>, cv::Mat>;
 
 KeyPointAndDesc processOneFrame(const Ort::SuperPoint& osh, const cv::Mat& inputImg, float* dst, int borderRemove = 4,
-                                float confidenceThresh = 0.015, bool alignCorners = true);
+                                float confidenceThresh = 0.015, bool alignCorners = true, int distThresh = 2);
 
+/**
+ *  @brief detect super point
+ *
+ *  @return vector of detected key points
+ */
+std::vector<cv::KeyPoint> getKeyPoints(const std::vector<Ort::OrtSessionHandler::DataOutputType>& inferenceOutput,
+                                       int borderRemove, float confidenceThresh);
+
+/**
+ *  @brief estimate super point's keypoint descriptor
+ *
+ *  @return keypoint Mat of shape [num key point x 256]
+ */
+cv::Mat getDescriptors(const cv::Mat& coarseDescriptors, const std::vector<cv::KeyPoint>& keyPoints, int height,
+                       int width, bool alignCorners);
+
+std::vector<int> nmsFast(const std::vector<cv::KeyPoint>& keyPoints, int height, int width, int distThresh = 2);
 }  // namespace
 
 int main(int argc, char* argv[])
@@ -66,18 +85,13 @@ int main(int argc, char* argv[])
     cv::drawMatches(images[0], results[0].first, images[1], results[1].first, goodMatches, matchesImage,
                     cv::Scalar::all(-1), cv::Scalar::all(-1), std::vector<char>(),
                     cv::DrawMatchesFlags::NOT_DRAW_SINGLE_POINTS);
-    cv::imwrite("good_matches.jpg", matchesImage);
+    cv::imwrite("super_point_good_matches.jpg", matchesImage);
 
     return EXIT_SUCCESS;
 }
 
 namespace
 {
-/**
- *  @brief detect super point
- *
- *  @return vector of detected key points
- */
 std::vector<cv::KeyPoint> getKeyPoints(const std::vector<Ort::OrtSessionHandler::DataOutputType>& inferenceOutput,
                                        int borderRemove, float confidenceThresh)
 {
@@ -119,12 +133,6 @@ std::vector<cv::KeyPoint> getKeyPoints(const std::vector<Ort::OrtSessionHandler:
 
     return keyPoints;
 }
-
-/**
- *  @brief estimate super point's keypoint descriptor
- *
- *  @return keypoint Mat of shape [num key point x 256]
- */
 cv::Mat getDescriptors(const cv::Mat& coarseDescriptors, const std::vector<cv::KeyPoint>& keyPoints, int height,
                        int width, bool alignCorners)
 {
@@ -145,8 +153,50 @@ cv::Mat getDescriptors(const cv::Mat& coarseDescriptors, const std::vector<cv::K
     return buffer;
 }
 
+std::vector<int> nmsFast(const std::vector<cv::KeyPoint>& keyPoints, int height, int width, int distThresh)
+{
+    static const int TO_PROCESS = 1;
+    static const int EMPTY_OR_SUPPRESSED = 0;
+    static const int KEPT = -1;
+
+    std::vector<int> sortedIndices(keyPoints.size());
+    std::iota(sortedIndices.begin(), sortedIndices.end(), 0);
+
+    // sort in descending order base on confidence
+    std::stable_sort(sortedIndices.begin(), sortedIndices.end(),
+                     [&keyPoints](int lidx, int ridx) { return keyPoints[lidx].response > keyPoints[ridx].response; });
+
+    cv::Mat grid = cv::Mat(height, width, CV_8S, TO_PROCESS);
+    std::vector<int> keepIndices;
+
+    for (int idx : sortedIndices) {
+        int x = keyPoints[idx].pt.x;
+        int y = keyPoints[idx].pt.y;
+
+        if (grid.at<schar>(y, x) == TO_PROCESS) {
+            for (int i = y - distThresh; i < y + distThresh; ++i) {
+                if (i < 0 || i >= height) {
+                    continue;
+                }
+
+                for (int j = x - distThresh; j < x + distThresh; ++j) {
+                    if (j < 0 || j >= width) {
+                        continue;
+                    }
+                    grid.at<int>(i, j) = EMPTY_OR_SUPPRESSED;
+                }
+            }
+
+            grid.at<int>(y, x) = KEPT;
+            keepIndices.emplace_back(idx);
+        }
+    }
+
+    return keepIndices;
+}
+
 KeyPointAndDesc processOneFrame(const Ort::SuperPoint& osh, const cv::Mat& inputImg, float* dst, int borderRemove,
-                                float confidenceThresh, bool alignCorners)
+                                float confidenceThresh, bool alignCorners, int distThresh)
 {
     int origW = inputImg.cols, origH = inputImg.rows;
     std::vector<float> originImageSize{static_cast<float>(origH), static_cast<float>(origW)};
@@ -161,6 +211,14 @@ KeyPointAndDesc processOneFrame(const Ort::SuperPoint& osh, const cv::Mat& input
     cv::Mat coarseDescriptorMat(descriptorShape.size(), descriptorShape.data(), CV_32F,
                                 inferenceOutput[1].first);  // 1 x 256 x H/8 x W/8
 
+    std::vector<int> keepIndices = nmsFast(keyPoints, Ort::SuperPoint::IMG_H, Ort::SuperPoint::IMG_W, distThresh);
+
+    std::vector<cv::KeyPoint> keepKeyPoints;
+    keepKeyPoints.reserve(keepIndices.size());
+    std::transform(keepIndices.begin(), keepIndices.end(), std::back_inserter(keepKeyPoints),
+                   [&keyPoints](int idx) { return keyPoints[idx]; });
+    keyPoints = std::move(keepKeyPoints);
+
     cv::Mat descriptors =
         getDescriptors(coarseDescriptorMat, keyPoints, Ort::SuperPoint::IMG_H, Ort::SuperPoint::IMG_W, alignCorners);