colmap · ahojnnes · Aug 26, 2024 · Aug 21, 2024 · Aug 22, 2024 · Aug 22, 2024
diff --git a/src/colmap/controllers/feature_matching_utils.cc b/src/colmap/controllers/feature_matching_utils.cc
@@ -53,11 +53,6 @@ FeatureMatcherWorker::FeatureMatcherWorker(
       output_queue_(output_queue) {
   THROW_CHECK(matching_options_.Check());
 
-  prev_keypoints_image_ids_[0] = kInvalidImageId;
-  prev_keypoints_image_ids_[1] = kInvalidImageId;
-  prev_descriptors_image_ids_[0] = kInvalidImageId;
-  prev_descriptors_image_ids_[1] = kInvalidImageId;
-
   if (matching_options_.use_gpu) {
 #if !defined(COLMAP_CUDA_ENABLED)
     opengl_context_ = std::make_unique<OpenGLContextManager>();
@@ -77,6 +72,10 @@ void FeatureMatcherWorker::Run() {
 #endif
   }
 
+  matching_options_.cpu_descriptor_index_cache =
+      &cache_->GetFeatureDescriptorIndexCache();
+  THROW_CHECK_NOTNULL(matching_options_.cpu_descriptor_index_cache);
+
   std::unique_ptr<FeatureMatcher> matcher =
       CreateSiftFeatureMatcher(matching_options_);
   if (matcher == nullptr) {
@@ -104,48 +103,35 @@ void FeatureMatcherWorker::Run() {
 
       if (matching_options_.guided_matching) {
         matcher->MatchGuided(geometry_options_.ransac_options.max_error,
-                             GetKeypointsPtr(0, data.image_id1),
-                             GetKeypointsPtr(1, data.image_id2),
-                             GetDescriptorsPtr(0, data.image_id1),
-                             GetDescriptorsPtr(1, data.image_id2),
+                             {
+                                 data.image_id1,
+                                 cache_->GetDescriptors(data.image_id1),
+                                 cache_->GetKeypoints(data.image_id1),
+                             },
+                             {
+                                 data.image_id2,
+                                 cache_->GetDescriptors(data.image_id2),
+                                 cache_->GetKeypoints(data.image_id2),
+                             },
                              &data.two_view_geometry);
       } else {
-        matcher->Match(GetDescriptorsPtr(0, data.image_id1),
-                       GetDescriptorsPtr(1, data.image_id2),
-                       &data.matches);
+        matcher->Match(
+            {
+                data.image_id1,
+                cache_->GetDescriptors(data.image_id1),
+            },
+            {
+                data.image_id2,
+                cache_->GetDescriptors(data.image_id2),
+            },
+            &data.matches);
       }
 
       THROW_CHECK(output_queue_->Push(std::move(data)));
     }
   }
 }
 
-std::shared_ptr<FeatureKeypoints> FeatureMatcherWorker::GetKeypointsPtr(
-    const int index, const image_t image_id) {
-  THROW_CHECK_GE(index, 0);
-  THROW_CHECK_LE(index, 1);
-  if (prev_keypoints_image_ids_[index] == image_id) {
-    return nullptr;
-  } else {
-    prev_keypoints_image_ids_[index] = image_id;
-    prev_keypoints_[index] = cache_->GetKeypoints(image_id);
-    return prev_keypoints_[index];
-  }
-}
-
-std::shared_ptr<FeatureDescriptors> FeatureMatcherWorker::GetDescriptorsPtr(
-    const int index, const image_t image_id) {
-  THROW_CHECK_GE(index, 0);
-  THROW_CHECK_LE(index, 1);
-  if (prev_descriptors_image_ids_[index] == image_id) {
-    return nullptr;
-  } else {
-    prev_descriptors_image_ids_[index] = image_id;
-    prev_descriptors_[index] = cache_->GetDescriptors(image_id);
-    return prev_descriptors_[index];
-  }
-}
-
 namespace {
 
 class VerifierWorker : public Thread {

diff --git a/src/colmap/controllers/feature_matching_utils.h b/src/colmap/controllers/feature_matching_utils.h
@@ -66,23 +66,13 @@ class FeatureMatcherWorker : public Thread {
  private:
   void Run() override;
 
-  std::shared_ptr<FeatureKeypoints> GetKeypointsPtr(int index,
-                                                    image_t image_id);
-  std::shared_ptr<FeatureDescriptors> GetDescriptorsPtr(int index,
-                                                        image_t image_id);
-
   SiftMatchingOptions matching_options_;
   TwoViewGeometryOptions geometry_options_;
   FeatureMatcherCache* cache_;
   JobQueue<Input>* input_queue_;
   JobQueue<Output>* output_queue_;
 
   std::unique_ptr<OpenGLContextManager> opengl_context_;
-
-  std::array<image_t, 2> prev_keypoints_image_ids_;
-  std::array<std::shared_ptr<FeatureKeypoints>, 2> prev_keypoints_;
-  std::array<image_t, 2> prev_descriptors_image_ids_;
-  std::array<std::shared_ptr<FeatureDescriptors>, 2> prev_descriptors_;
 };
 
 // Multi-threaded and multi-GPU SIFT feature matcher, which writes the computed

diff --git a/src/colmap/feature/CMakeLists.txt b/src/colmap/feature/CMakeLists.txt
@@ -49,6 +49,7 @@ COLMAP_ADD_LIBRARY(
     NAME colmap_feature
     SRCS
         extractor.h
+        index.h index.cc
         matcher.h matcher.cc
         pairing.h pairing.cc
         sift.h sift.cc
@@ -76,7 +77,12 @@ if(GPU_ENABLED)
 endif()
 
 COLMAP_ADD_TEST(
-    NAME feature_utils_test
+    NAME index_test
+    SRCS index_test.cc
+    LINK_LIBS colmap_feature
+)
+COLMAP_ADD_TEST(
+    NAME utils_test
     SRCS utils_test.cc
     LINK_LIBS colmap_feature
 )

diff --git a/src/colmap/feature/index.cc b/src/colmap/feature/index.cc
@@ -0,0 +1,116 @@
+// Copyright (c) 2023, ETH Zurich and UNC Chapel Hill.
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+//     * Redistributions of source code must retain the above copyright
+//       notice, this list of conditions and the following disclaimer.
+//
+//     * Redistributions in binary form must reproduce the above copyright
+//       notice, this list of conditions and the following disclaimer in the
+//       documentation and/or other materials provided with the distribution.
+//
+//     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
+//       its contributors may be used to endorse or promote products derived
+//       from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+#include "colmap/feature/matcher.h"
+
+#include <flann/flann.hpp>
+
+namespace colmap {
+namespace {
+
+class FlannFeatureDescriptorIndex : public FeatureDescriptorIndex {
+ public:
+  void Build(const FeatureDescriptors& index_descriptors) override {
+    THROW_CHECK_EQ(index_descriptors.cols(), 128);
+    num_index_descriptors_ = index_descriptors.rows();
+    if (num_index_descriptors_ == 0) {
+      // Flann is not happy when the input has no descriptors.
+      index_ = nullptr;
+      return;
+    }
+    const flann::Matrix<uint8_t> descriptors_matrix(
+        const_cast<uint8_t*>(index_descriptors.data()),
+        num_index_descriptors_,
+        index_descriptors.cols());
+    index_ = std::make_unique<FlannIndexType>(
+        descriptors_matrix, flann::KDTreeIndexParams(kNumTreesInForest));
+    index_->buildIndex();
+  }
+
+  void Search(int num_neighbors,
+              const FeatureDescriptors& query_descriptors,
+              Eigen::RowMajorMatrixXi& indices,
+              Eigen::RowMajorMatrixXi& l2_dists) const override {
+    THROW_CHECK_NOTNULL(index_);
+    THROW_CHECK_EQ(query_descriptors.cols(), 128);
+
+    const int num_query_descriptors = query_descriptors.rows();
+    if (num_query_descriptors == 0) {
+      return;
+    }
+
+    const int num_eff_neighbors =
+        std::min(num_neighbors, num_index_descriptors_);
+
+    indices.resize(num_query_descriptors, num_eff_neighbors);
+    l2_dists.resize(num_query_descriptors, num_eff_neighbors);
+    const flann::Matrix<uint8_t> query_matrix(
+        const_cast<uint8_t*>(query_descriptors.data()),
+        num_query_descriptors,
+        query_descriptors.cols());
+
+    flann::Matrix<int> indices_matrix(
+        indices.data(), num_query_descriptors, num_eff_neighbors);
+    std::vector<float> l2_dist_vector(num_query_descriptors *
+                                      num_eff_neighbors);
+    flann::Matrix<float> l2_dist_matrix(
+        l2_dist_vector.data(), num_query_descriptors, num_eff_neighbors);
+    index_->knnSearch(query_matrix,
+                      indices_matrix,
+                      l2_dist_matrix,
+                      num_eff_neighbors,
+                      flann::SearchParams(kNumLeavesToVisit));
+
+    for (int query_idx = 0; query_idx < num_query_descriptors; ++query_idx) {
+      for (int k = 0; k < num_eff_neighbors; ++k) {
+        l2_dists(query_idx, k) = static_cast<int>(
+            std::round(l2_dist_vector[query_idx * num_eff_neighbors + k]));
+      }
+    }
+  }
+
+ private:
+  // Tuned to produce similar results to brute-force matching. If speed is
+  // important, the parameters can be reduced. The biggest speed improvement can
+  // be gained by reducing the number of leaves.
+  constexpr static int kNumTreesInForest = 4;
+  constexpr static int kNumLeavesToVisit = 128;
+
+  using FlannIndexType = flann::Index<flann::L2<uint8_t>>;
+  std::unique_ptr<FlannIndexType> index_;
+  int num_index_descriptors_ = 0;
+};
+
+}  // namespace
+
+std::unique_ptr<FeatureDescriptorIndex> FeatureDescriptorIndex::Create() {
+  return std::make_unique<FlannFeatureDescriptorIndex>();
+}
+
+}  // namespace colmap
diff --git a/src/colmap/feature/index.h b/src/colmap/feature/index.h
@@ -0,0 +1,53 @@
+// Copyright (c) 2023, ETH Zurich and UNC Chapel Hill.
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+//     * Redistributions of source code must retain the above copyright
+//       notice, this list of conditions and the following disclaimer.
+//
+//     * Redistributions in binary form must reproduce the above copyright
+//       notice, this list of conditions and the following disclaimer in the
+//       documentation and/or other materials provided with the distribution.
+//
+//     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
+//       its contributors may be used to endorse or promote products derived
+//       from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include "colmap/feature/types.h"
+#include "colmap/util/types.h"
+
+#include <memory>
+
+namespace colmap {
+
+class FeatureDescriptorIndex {
+ public:
+  virtual ~FeatureDescriptorIndex() = default;
+
+  static std::unique_ptr<FeatureDescriptorIndex> Create();
+
+  virtual void Build(const FeatureDescriptors& descriptors) = 0;
+
+  virtual void Search(int num_neighbors,
+                      const FeatureDescriptors& query_descriptors,
+                      Eigen::RowMajorMatrixXi& indices,
+                      Eigen::RowMajorMatrixXi& l2_dists) const = 0;
+};
+
+}  // namespace colmap