Program Listing for File inverted-multi-product-quantization-index.h

↰ Return to documentation for file (algorithms/loopclosure/inverted-multi-index/include/inverted-multi-index/inverted-multi-product-quantization-index.h)

#ifndef INVERTED_MULTI_INDEX_INVERTED_MULTI_PRODUCT_QUANTIZATION_INDEX_H_
#define INVERTED_MULTI_INDEX_INVERTED_MULTI_PRODUCT_QUANTIZATION_INDEX_H_

#include <algorithm>
#include <functional>
#include <limits>
#include <queue>
#include <tuple>
#include <unordered_map>
#include <utility>
#include <vector>

#include <Eigen/Core>
#include <aslam/common/memory.h>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <nabo/nabo.h>
#include <product-quantization/product-quantization.h>

#include <inverted-multi-index/inverted-multi-index-common.h>

DECLARE_double(lc_knn_max_radius);

namespace loop_closure {
namespace inverted_multi_index {
// Variant of the inverted multi-index, which stores the product quantized
// versions of the original descriptors to save memory and accelerate the search
// for large datasets. The index uses one product quantizer for each word from
// each lower dimensional vocabulary.
// The template parameters are the data type used to store the product quantized
// descriptors, the number of components used for product quantization, the
// number of dimension used for each component, and the number of cluster
// centers used for each product quantizer (see also
// product-quantization/include/product-quantization.h for a more detailed
// explanation of the parameters). The inverted multi-index splits descriptors
// into two parts of dimension kNumComponents / 2 * kNumDimPerComp each. Thus,
// kNumComponents needs to be a multiple of 2.
template <typename DataType, int kNumComponents, int kNumDimPerComp,
          int kNumCenters>
class InvertedMultiProductQuantizationIndex {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  // The dimensionality of the original descriptors.
  static constexpr int kOriginalDescDim = kNumDimPerComp * kNumComponents;
  static constexpr int kHalfNumComponents = kNumComponents / 2;
  // The type of the input descriptors.
  typedef Eigen::Matrix<float, kOriginalDescDim, 1> InputDescriptorType;
  typedef Eigen::Matrix<float, kOriginalDescDim, Eigen::Dynamic>
      InputDescriptorMatrixType;
  typedef Eigen::Matrix<float, kOriginalDescDim / 2, 1> HalfInputDescriptorType;
  // The type of the quantized descriptors that is stored.
  typedef Eigen::Matrix<DataType, kNumComponents, 1> StoredDescriptorType;

  typedef common::InvertedFile<DataType, kNumComponents> InvFile;
  typedef product_quantization::ProductQuantization<
      kHalfNumComponents, kNumDimPerComp, kNumCenters, DataType>
      ProductQuantizer;

  // Creates the index from a given set of visual words. Each column in words_i
  // specifies a cluster center coordinate. quantizer_centers_i define the
  // cluster centers used by the product quantizer for the i-th lower
  // dimensional vocabulary. quantizer_centers_i is required to have size
  // kNumDimPerComp / 2 x (kNumComponents * kNumCenters * words_i.cols()).
  // The first kNumComponents * kNumCenters columns of quantizer_centers_i
  // define the cluster centers of the product quantizer used for the 1st word
  // in the i-th vocabulary. The next kNumComponents * kNumCenters columns give
  // the centers for the second quantizer, etc.
  InvertedMultiProductQuantizationIndex(
      const Eigen::MatrixXf& words_1, const Eigen::MatrixXf& words_2,
      const Eigen::MatrixXf& quantizer_centers_1,
      const Eigen::MatrixXf& quantizer_centers_2,
      int num_closest_words_for_nn_search)
      : words_1_(words_1),
        words_2_(words_2),
        words_1_index_(
            common::NNSearch::createKDTreeLinearHeap(
                words_1_, kOriginalDescDim / 2,
                common::kCollectTouchStatistics)),
        words_2_index_(
            common::NNSearch::createKDTreeLinearHeap(
                words_2_, kOriginalDescDim / 2,
                common::kCollectTouchStatistics)),
        num_closest_words_for_nn_search_(num_closest_words_for_nn_search),
        max_db_descriptor_index_(0) {
    static_assert(
        kNumComponents % 2 == 0,
        "The number of components needs to be a multiple of 2.");

    CHECK_EQ(words_1.rows(), kOriginalDescDim / 2);
    CHECK_GT(words_1.cols(), 0);
    CHECK_EQ(words_2.rows(), kOriginalDescDim / 2);
    CHECK_GT(words_2.cols(), 0);

    // The number of columns corresponding to the cluster centers of a product
    // quantizer (pq).
    const int num_cols_per_pq = kHalfNumComponents * kNumCenters;

    // The number of components per product quantizer is half the total number
    // of components since each product quantizer is used to quantize one half
    // of the original descriptor.
    CHECK_EQ(quantizer_centers_1.rows(), kNumDimPerComp);
    CHECK_EQ(quantizer_centers_2.rows(), kNumDimPerComp);
    CHECK_EQ(quantizer_centers_1.cols(), num_cols_per_pq * words_1.cols());
    CHECK_EQ(quantizer_centers_1.cols(), num_cols_per_pq * words_1.cols());

    CHECK_GT(num_closest_words_for_nn_search_, 0);

    quantizers_words_1_.resize(words_1.cols());
    int index = 0;
    for (int i = 0; i < words_1.cols(); ++i, index += num_cols_per_pq) {
      quantizers_words_1_[i].SetClusterCenters(
          quantizer_centers_1.block<kNumDimPerComp, num_cols_per_pq>(0, index));
    }

    quantizers_words_2_.resize(words_2.cols());
    index = 0;
    for (int i = 0; i < words_1.cols(); ++i, index += num_cols_per_pq) {
      quantizers_words_2_[i].SetClusterCenters(
          quantizer_centers_2.block<kNumDimPerComp, num_cols_per_pq>(0, index));
    }
  }

  void SetNumClosestWordsForNNSearch(int num_closest_words_for_nn_search) {
    CHECK_GT(num_closest_words_for_nn_search, 0);
    num_closest_words_for_nn_search_ = num_closest_words_for_nn_search;
  }

  inline int GetNumDescriptorsInIndex() const {
    return max_db_descriptor_index_;
  }

  // Clears the inverted multi-index by removing all references to the database
  // descriptors stored in it. Does NOT remove the underlying quantization.
  inline void Clear() {
    inverted_files_.clear();
    word_index_map_.clear();
    max_db_descriptor_index_ = 0;
  }

  // Adds a set of database descriptors to the inverted multi-index.
  // Each column defines a database descriptor.
  void AddDescriptors(const InputDescriptorMatrixType& descriptors) {
    const int num_descriptors = descriptors.cols();
    std::vector<std::pair<int, int> > closest_word;

    for (int i = 0; i < num_descriptors; ++i) {
      common::FindClosestWords<kOriginalDescDim / 2>(
          descriptors.col(i), 1, *words_1_index_, *words_2_index_,
          words_1_.cols(), words_2_.cols(), &closest_word);
      CHECK(!closest_word.empty());
      const int word1 = closest_word[0].first;
      const int word2 = closest_word[0].second;
      const int word_index = word1 * words_2_.cols() + word2;

      // Quantizes the descriptor using the product quantizers belonging to the
      // product word: We first compute the residual between the descriptor and
      // the cluster center before quantizing the residual.
      HalfInputDescriptorType residual_part_1;
      ComputeResidual(
          descriptors.col(i).template head<kOriginalDescDim / 2>(), words_1_,
          word1, &residual_part_1);
      HalfInputDescriptorType residual_part_2;
      ComputeResidual(
          descriptors.col(i).template tail<kOriginalDescDim / 2>(), words_2_,
          word2, &residual_part_2);

      StoredDescriptorType quantized_residual;
      Eigen::Matrix<DataType, kHalfNumComponents, 1> quantized_part;
      quantizers_words_1_[closest_word[0].first].Quantize(
          residual_part_1, &quantized_part);
      quantized_residual.template head<kHalfNumComponents>() = quantized_part;
      quantizers_words_2_[closest_word[0].second].Quantize(
          residual_part_2, &quantized_part);
      quantized_residual.template tail<kHalfNumComponents>() = quantized_part;

      common::AddDescriptor<DataType, kNumComponents>(
          quantized_residual, max_db_descriptor_index_, word_index,
          &word_index_map_, &inverted_files_);
      ++max_db_descriptor_index_;
    }
  }

  // Finds the n nearest neighbors for a given query feature.
  // This function is thread-safe.
  template <typename DerivedQuery, typename DerivedIndices,
            typename DerivedDistances>
  inline void GetNNearestNeighbors(
      const Eigen::MatrixBase<DerivedQuery>& query_feature, int num_neighbors,
      const Eigen::MatrixBase<DerivedIndices>& out_indices,
      const Eigen::MatrixBase<DerivedDistances>& out_distances) const {
    EIGEN_STATIC_ASSERT_MATRIX_SPECIFIC_SIZE(
        DerivedQuery, static_cast<int>(kOriginalDescDim), 1);
    CHECK_EQ(out_indices.cols(), 1);
    CHECK_EQ(out_distances.cols(), 1);
    CHECK_GT(num_neighbors, 0);

    CHECK_EQ(out_indices.rows(), num_neighbors)
        << "The indices parameter must be pre-allocated to hold all results.";
    CHECK_EQ(out_distances.rows(), num_neighbors)
        << "The distances parameter must be pre-allocated to hold all results.";

    // Zero copy passing of Eigen-Block Expressions.
    // http://eigen.tuxfamily.org/dox/TopicFunctionTakingEigenTypes.html
    Eigen::MatrixBase<DerivedIndices>& indices =
        const_cast<Eigen::MatrixBase<DerivedIndices>&>(out_indices);
    Eigen::MatrixBase<DerivedDistances>& distances =
        const_cast<Eigen::MatrixBase<DerivedDistances>&>(out_distances);

    // Finds the closest visual words.
    std::vector<std::pair<int, int> > closest_words;
    common::FindClosestWords<kOriginalDescDim / 2>(
        query_feature, num_closest_words_for_nn_search_, *words_1_index_,
        *words_2_index_, words_1_.cols(), words_2_.cols(), &closest_words);

    // Performs exhaustive search through all descriptors assigned to the
    // closest words, using product quantization to compute the distances.
    std::vector<std::pair<float, int> > nearest_neighbors;
    nearest_neighbors.reserve(num_neighbors + 1);

    const int num_words_to_use = static_cast<int>(closest_words.size());
    std::unordered_map<int, int>::const_iterator word_index_map_it;

    // In order to avoid re-computing the distances between the descriptor and
    // the cluster centers of a product quantizer (which are stored in look-up
    // tables), we cache them independently for each of the two lower
    // dimensional vocabularies.
    typedef Eigen::Matrix<float, kHalfNumComponents, kNumCenters> LookUpTable;
    AlignedUnorderedMap<int, LookUpTable> table_cache_words_1;
    AlignedUnorderedMap<int, LookUpTable> table_cache_words_2;

    for (int i = 0; i < num_words_to_use; ++i) {
      const int word1 = closest_words[i].first;
      const int word2 = closest_words[i].second;
      const int word_index = word1 * words_2_.cols() + word2;
      word_index_map_it = word_index_map_.find(word_index);
      if (word_index_map_it == word_index_map_.end())
        continue;

      // Computes the look-up table for the product quantizer corresponding to
      // the word in the first vocabulary.
      typename AlignedUnorderedMap<int, LookUpTable>::iterator table_it =
          table_cache_words_1.find(word1);
      if (table_it == table_cache_words_1.end()) {
        HalfInputDescriptorType residual_part;
        ComputeResidual(
            query_feature.template head<kOriginalDescDim / 2>(), words_1_,
            word1, &residual_part);
        LookUpTable lut;
        quantizers_words_1_[word1].FillLUT(residual_part, &lut);
        table_it = table_cache_words_1.insert(std::make_pair(word1, lut)).first;
      }
      const LookUpTable& lut1 = table_it->second;

      // Computes the look-up table for the product quantizer corresponding to
      // the word in the first vocabulary.
      table_it = table_cache_words_2.find(word2);
      if (table_it == table_cache_words_2.end()) {
        HalfInputDescriptorType residual_part;
        ComputeResidual(
            query_feature.template tail<kOriginalDescDim / 2>(), words_2_,
            word2, &residual_part);
        LookUpTable lut;
        quantizers_words_2_[word2].FillLUT(residual_part, &lut);
        table_it = table_cache_words_2.insert(std::make_pair(word2, lut)).first;
      }
      const LookUpTable& lut2 = table_it->second;

      const InvFile& inverted_file = inverted_files_[word_index_map_it->second];
      const int num_descriptors =
          static_cast<int>(inverted_file.descriptors_.size());
      for (int j = 0; j < num_descriptors; ++j) {
        // Computes the distance between the query and the quantized descriptors
        // using the look-up tables.
        float distance = quantizers_words_1_[word1].ComputeDistance(
            lut1,
            inverted_file.descriptors_[j].template head<kHalfNumComponents>());
        distance += quantizers_words_2_[word2].ComputeDistance(
            lut2,
            inverted_file.descriptors_[j].template tail<kHalfNumComponents>());
        common::InsertNeighbor(
            inverted_file.indices_[j], distance, num_neighbors,
            &nearest_neighbors);
      }
    }

    for (size_t i = 0; i < nearest_neighbors.size(); ++i) {
      indices(i, 0) = nearest_neighbors[i].second;
      distances(i, 0) = nearest_neighbors[i].first;
    }
    for (int i = nearest_neighbors.size(); i < num_neighbors; ++i) {
      indices(i, 0) = -1;
      distances(i, 0) = std::numeric_limits<float>::infinity();
    }
  }

 protected:
  // Given a half of a original descriptor, a vocabulary, and the word from this
  // vocabulary that is closest to the half, computes between residual the
  // descriptor and the cluster center of the word.
  // For performance reasons, we do not CHECK_NOTNULL as both AddDescriptors and
  // GetNNearestNeighbors guarantee that the pointer is not null.
  // Similar, both functions guarantee that word is valid.
  inline void ComputeResidual(
      const HalfInputDescriptorType& descriptor_half,
      const Eigen::MatrixXf& words, int word,
      HalfInputDescriptorType* residual) const {
    *residual = descriptor_half - words.col(word);
  }

  // The two sets of cluster centers defining the quantization of the descriptor
  // space as the Cartesian product of the two sets of words.
  Eigen::MatrixXf words_1_;
  Eigen::MatrixXf words_2_;
  std::shared_ptr<common::NNSearch> words_1_index_;
  std::shared_ptr<common::NNSearch> words_2_index_;

  // A set of product quantizers, one for each word from the lower dimensional
  // vocabularies.
  Aligned<std::vector, ProductQuantizer> quantizers_words_1_;
  Aligned<std::vector, ProductQuantizer> quantizers_words_2_;

  // The number of closest words from the product vocabulary that should be used
  // during nearest neighbor search.
  int num_closest_words_for_nn_search_;
  // Hashmap storing for each combined visual word the index in inverted_files_
  // in which all database descriptors assigned to that word can be found.
  // This allows us to easily add descriptors assigned to words that have not
  // been used previously without having to re-order large amounts of memory.
  std::unordered_map<int, int> word_index_map_;
  // Vector containing the inverted files, one for each visual word in the
  // product vocabulary. Each inverted file holds all descriptors assigned to
  // the corresponding word and their indices.
  Aligned<std::vector, InvFile> inverted_files_;
  // The maximum index of the descriptor indices.
  int max_db_descriptor_index_;
};
}  // namespace inverted_multi_index
}  // namespace loop_closure

#endif  // INVERTED_MULTI_INDEX_INVERTED_MULTI_PRODUCT_QUANTIZATION_INDEX_H_