feat: init FeatureMatcher

MaaXYZ · Oct 10, 2023 · 825bc52 · 825bc52
1 parent 8f4a4bf
commit 825bc52
Show file tree

Hide file tree

Showing 6 changed files with 200 additions and 4 deletions.
diff --git a/source/MaaFramework/MaaFramework.vcxproj b/source/MaaFramework/MaaFramework.vcxproj
@@ -71,6 +71,7 @@
     <ClInclude Include="Utils\StringMisc.hpp" />
     <ClInclude Include="Utils\TempPath.hpp" />
     <ClInclude Include="Utils\Time.hpp" />
+    <ClInclude Include="Vision\FeatureMatcher.h" />
     <ClInclude Include="Vision\NeuralNetworkClassifier.h" />
     <ClInclude Include="Vision\ColorMatcher.h" />
     <ClInclude Include="Vision\TemplateComparator.h" />
@@ -107,6 +108,7 @@
     <ClCompile Include="Task\Recognizer.cpp" />
     <ClCompile Include="Task\SyncContext.cpp" />
     <ClCompile Include="Task\PipelineTask.cpp" />
+    <ClCompile Include="Vision\FeatureMatcher.cpp" />
     <ClCompile Include="Vision\NeuralNetworkClassifier.cpp" />
     <ClCompile Include="Vision\ColorMatcher.cpp" />
     <ClCompile Include="Vision\TemplateComparator.cpp" />

diff --git a/source/MaaFramework/Vision/FeatureMatcher.cpp b/source/MaaFramework/Vision/FeatureMatcher.cpp
@@ -0,0 +1,112 @@
+#include "FeatureMatcher.h"
+
+MAA_SUPPRESS_CV_WARNINGS_BEGIN
+#include <opencv2/opencv.hpp>
+#include <opencv2/xfeatures2d.hpp>
+MAA_SUPPRESS_CV_WARNINGS_END
+
+#include "Utils/Logger.h"
+
+MAA_VISION_NS_BEGIN
+
+FeatureMatcher::ResultsVec FeatureMatcher::analyze() const
+{
+    if (!template_) {
+        LogError << name_ << "template_ is empty" << VAR(param_.template_path);
+        return {};
+    }
+
+    const cv::Mat& templ = *template_;
+
+    auto start_time = std::chrono::steady_clock::now();
+    ResultsVec results = foreach_rois(templ);
+
+    auto costs = duration_since(start_time);
+    LogDebug << name_ << "Raw:" << VAR(results) << VAR(param_.template_path) << VAR(costs);
+
+    int count = param_.count;
+    filter(results, count);
+
+    costs = duration_since(start_time);
+    LogDebug << name_ << "Filter:" << VAR(results) << VAR(path) << VAR(count) << VAR(costs);
+
+    return results;
+}
+
+FeatureMatcher::ResultsVec FeatureMatcher::foreach_rois(const cv::Mat& templ) const
+{
+    if (templ.empty()) {
+        LogWarn << name_ << "template is empty" << VAR(param_.template_path);
+        return {};
+    }
+
+    auto matcher = create_matcher(templ);
+
+    if (param_.roi.empty()) {
+        return match(cv::Rect(0, 0, image_.cols, image_.rows), matcher);
+    }
+
+    ResultsVec results;
+    for (const cv::Rect& roi : param_.roi) {
+        ResultsVec res = match(roi, matcher);
+        results.insert(results.end(), std::make_move_iterator(res.begin()), std::make_move_iterator(res.end()));
+    }
+
+    return results;
+}
+
+cv::FlannBasedMatcher FeatureMatcher::create_matcher(const cv::Mat& templ) const
+{
+    std::vector<cv::KeyPoint> keypoints_1;
+    cv::Mat descriptors_1;
+    detect(templ, param_.green_mask, keypoints_1, descriptors_1);
+
+    std::vector<cv::Mat> train_desc(1, descriptors_1);
+    cv::FlannBasedMatcher matcher;
+    matcher.add(train_desc);
+    matcher.train();
+
+    return matcher;
+}
+
+void FeatureMatcher::detect(const cv::Mat& image, bool green_mask, std::vector<cv::KeyPoint>& keypoints,
+                            cv::Mat& descriptors) const
+{
+    auto detector = cv::xfeatures2d::SURF::create(param_.hessian);
+
+    cv::Mat mask = cv::Mat::ones(image.size(), CV_8UC1);
+    if (green_mask) {
+        cv::inRange(image, cv::Scalar(0, 255, 0), cv::Scalar(0, 255, 0), mask);
+        mask = ~mask;
+    }
+
+    detector->detectAndCompute(image, mask, keypoints, descriptors);
+}
+
+FeatureMatcher::ResultsVec FeatureMatcher::match(const cv::Rect& roi, cv::FlannBasedMatcher& matcher) const
+{
+    cv::Mat image = image_with_roi(roi);
+    std::vector<cv::KeyPoint> keypoints_2;
+    cv::Mat descriptors_2;
+    detect(image, false, keypoints_2, descriptors_2);
+
+    std::vector<std::vector<cv::DMatch>> match_points;
+    matcher.knnMatch(descriptors_2, match_points, 2);
+
+    ResultsVec results;
+    for (const auto& point : match_points) {
+        if (point.size() != 2) {
+            continue;
+        }
+
+        if (point[0].distance < param_.distance_ratio * point[1].distance) {
+            // TODO
+        }
+    }
+}
+
+void FeatureMatcher::draw_result(const cv::Rect& roi, const ResultsVec& results) const {}
+
+void FeatureMatcher::filter(ResultsVec& results, int count) const {}
+
+MAA_VISION_NS_END
diff --git a/source/MaaFramework/Vision/FeatureMatcher.h b/source/MaaFramework/Vision/FeatureMatcher.h
@@ -0,0 +1,51 @@
+#pragma once
+
+#include <ostream>
+#include <vector>
+
+#include <opencv2/features2d.hpp>
+
+#include "VisionBase.h"
+#include "VisionTypes.h"
+
+MAA_VISION_NS_BEGIN
+
+class FeatureMatcher : public VisionBase
+{
+public:
+    struct Result
+    {
+        cv::Rect box {};
+        double score = 0.0;
+
+        json::value to_json() const
+        {
+            json::value root;
+            root["box"] = json::array({ box.x, box.y, box.width, box.height });
+            root["score"] = score;
+            return root;
+        }
+    };
+
+    using ResultsVec = std::vector<Result>;
+
+public:
+    void set_template(std::shared_ptr<cv::Mat> templ) { template_ = std::move(templ); }
+    void set_param(FeatureMatcherParam param) { param_ = std::move(param); }
+
+    ResultsVec analyze() const;
+
+private:
+    ResultsVec foreach_rois(const cv::Mat& templ) const;
+    cv::FlannBasedMatcher create_matcher(const cv::Mat& templ) const;
+    void detect(const cv::Mat& image, bool green_mask, std::vector<cv::KeyPoint>& keypoints,
+                cv::Mat& descriptors) const;
+    ResultsVec match(const cv::Rect& roi, cv::FlannBasedMatcher& matcher) const;
+    void draw_result(const cv::Rect& roi, const ResultsVec& results) const;
+    void filter(ResultsVec& results, int count) const;
+
+    FeatureMatcherParam param_;
+    std::shared_ptr<cv::Mat> template_;
+};
+
+MAA_VISION_NS_END
diff --git a/source/MaaFramework/Vision/TemplateMatcher.cpp b/source/MaaFramework/Vision/TemplateMatcher.cpp
@@ -58,19 +58,19 @@ TemplateMatcher::ResultsVec TemplateMatcher::foreach_rois(const cv::Mat& templ)
     }
 
     if (param_.roi.empty()) {
-        return match_and_postproc(cv::Rect(0, 0, image_.cols, image_.rows), templ);
+        return match(cv::Rect(0, 0, image_.cols, image_.rows), templ);
     }
 
     ResultsVec results;
     for (const cv::Rect& roi : param_.roi) {
-        ResultsVec res = match_and_postproc(roi, templ);
+        ResultsVec res = match(roi, templ);
         results.insert(results.end(), std::make_move_iterator(res.begin()), std::make_move_iterator(res.end()));
     }
 
     return results;
 }
 
-TemplateMatcher::ResultsVec TemplateMatcher::match_and_postproc(const cv::Rect& roi, const cv::Mat& templ) const
+TemplateMatcher::ResultsVec TemplateMatcher::match(const cv::Rect& roi, const cv::Mat& templ) const
 {
     cv::Mat image = image_with_roi(roi);
     cv::Mat matched = match_template(image, templ, param_.method, param_.green_mask);

diff --git a/source/MaaFramework/Vision/TemplateMatcher.h b/source/MaaFramework/Vision/TemplateMatcher.h
@@ -33,7 +33,7 @@ class TemplateMatcher : public VisionBase
 
 private:
     ResultsVec foreach_rois(const cv::Mat& templ) const;
-    ResultsVec match_and_postproc(const cv::Rect& roi, const cv::Mat& templ) const;
+    ResultsVec match(const cv::Rect& roi, const cv::Mat& templ) const;
     void draw_result(const cv::Rect& roi, const cv::Mat& templ, const ResultsVec& results) const;
 
     void filter(ResultsVec& results, double threshold) const;

diff --git a/source/MaaFramework/Vision/VisionTypes.h b/source/MaaFramework/Vision/VisionTypes.h
@@ -97,4 +97,35 @@ struct ColorMatcherParam
     bool connected = false; // 是否计算连通域
 };
 
+struct FeatureMatcherParam
+{
+    enum Detector
+    {
+        SURF,
+        // TODO: ORB, SIFT, ...
+    };
+    enum Matcher
+    {
+        KNN,
+    };
+
+    inline static constexpr Detector kDefaultDetector = Detector::SURF;
+    inline static constexpr Matcher kDefaultMatcher = Matcher::KNN;
+    inline static constexpr int kDefaultHessianThreshold = 100;
+    inline static constexpr double kDefaultDistanceRatio = 0.6;
+    inline static constexpr int kDefaultCount = 4;
+
+    std::vector<cv::Rect> roi;
+    std::string template_path;
+    bool green_mask = false;
+
+    Detector detector = kDefaultDetector;
+    int hessian = kDefaultHessianThreshold;
+
+    Matcher matcher = kDefaultMatcher;
+
+    double distance_ratio = kDefaultDistanceRatio;
+    int count = kDefaultCount;
+};
+
 MAA_VISION_NS_END