inverted_file.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        CloudViewer: www.cloudViewer.org                  -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.cloudViewer.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <Eigen/Core>
#include <algorithm>
#include <bitset>
#include <cstdint>
#include <fstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>
 
#include "retrieval/geometry.h"
#include "retrieval/inverted_file_entry.h"
#include "retrieval/utils.h"
#include "util/alignment.h"
#include "util/logging.h"
#include "util/math.h"
 
namespace colmap {
namespace retrieval {
 
// Implements an inverted file, including the ability to compute image scores
// and matches. The template parameter is the length of the binary vectors
// in the Hamming Embedding.
// This class is based on an original implementation by Torsten Sattler.
template <int kEmbeddingDim>
class InvertedFile {
public:
    typedef Eigen::VectorXf DescType;
    typedef FeatureGeometry GeomType;
    typedef InvertedFileEntry<kEmbeddingDim> EntryType;
 
    enum Status {
        UNUSABLE = 0x00,
        HAS_EMBEDDING = 0x01,
        ENTRIES_SORTED = 0x02,
        USABLE = 0x03,
    };
 
    InvertedFile();
 
    // The number of added entries.
    size_t NumEntries() const;
 
    // Return all entries in the file.
    const std::vector<EntryType>& GetEntries() const;
 
    // Whether the Hamming embedding was computed for this file.
    bool HasHammingEmbedding() const;
 
    // Whether the entries in this file are sorted.
    bool EntriesSorted() const;
 
    // Whether this file is usable for scoring, i.e. the entries are sorted and
    // the Hamming embedding has been computed.
    bool IsUsable() const;
 
    // Adds an inverted file entry given a projected descriptor and its image
    // information stored in an inverted file entry. In particular, this
    // function generates the binary descriptor for the inverted file entry and
    // then stores the entry in the inverted file.
    void AddEntry(const int image_id,
                  typename DescType::Index feature_idx,
                  const DescType& descriptor,
                  const GeomType& geometry);
 
    // Sorts the inverted file entries in ascending order of image ids. This is
    // required for efficient scoring and must be called before ScoreFeature.
    void SortEntries();
 
    // Clear all entries in this file.
    void ClearEntries();
 
    // Reset all computed weights/thresholds and clear all entries.
    void Reset();
 
    // Given a projected descriptor, returns the corresponding binary string.
    void ConvertToBinaryDescriptor(
            const DescType& descriptor,
            std::bitset<kEmbeddingDim>* binary_descriptor) const;
 
    // Compute the idf-weight for this inverted file.
    void ComputeIDFWeight(const int num_total_images);
 
    // Return the idf-weight of this inverted file.
    float IDFWeight() const;
 
    // Given a set of descriptors, learns the thresholds required for the
    // Hamming embedding. Each row in descriptors represents a single descriptor
    // projected into the kEmbeddingDim dimensional space used for Hamming
    // embedding.
    void ComputeHammingEmbedding(
            const Eigen::Matrix<float, Eigen::Dynamic, kEmbeddingDim>&
                    descriptors);
 
    // Given a query feature, performs inverted file scoring.
    void ScoreFeature(const DescType& descriptor,
                      std::vector<ImageScore>* image_scores) const;
 
    // Get the identifiers of all indexed images in this file.
    void GetImageIds(std::unordered_set<int>* ids) const;
 
    // For each image in the inverted file, computes the self-similarity of each
    // image in the file (the part caused by this word) and adds the weight to
    // the entry corresponding to that image. This function is useful to
    // determine the normalization factor for each image that is used during
    // retrieval.
    void ComputeImageSelfSimilarities(
            std::unordered_map<int, double>* self_similarities) const;
 
    // Read/write the inverted file from/to a binary file.
    void Read(std::ifstream* ifs);
    void Write(std::ofstream* ofs) const;
 
private:
    // Whether the inverted file is initialized.
    uint8_t status_;
 
    // The inverse document frequency weight of this inverted file.
    float idf_weight_;
 
    // The entries of the inverted file system.
    std::vector<EntryType> entries_;
 
    // The thresholds used for Hamming embedding.
    DescType thresholds_;
 
    // The functor to derive a voting weight from a Hamming distance.
    static const HammingDistWeightFunctor<kEmbeddingDim>
            hamming_dist_weight_functor_;
};
 
// Implementation
 
template <int kEmbeddingDim>
const HammingDistWeightFunctor<kEmbeddingDim>
        InvertedFile<kEmbeddingDim>::hamming_dist_weight_functor_;
 
template <int kEmbeddingDim>
InvertedFile<kEmbeddingDim>::InvertedFile()
    : status_(UNUSABLE), idf_weight_(0.0f) {
    static_assert(kEmbeddingDim % 8 == 0,
                  "Dimensionality of projected space needs to"
                  " be a multiple of 8.");
    static_assert(kEmbeddingDim > 0,
                  "Dimensionality of projected space needs to be > 0.");
 
    thresholds_.resize(kEmbeddingDim);
    thresholds_.setZero();
}
 
template <int kEmbeddingDim>
size_t InvertedFile<kEmbeddingDim>::NumEntries() const {
    return entries_.size();
}
 
template <int kEmbeddingDim>
const std::vector<typename InvertedFile<kEmbeddingDim>::EntryType>&
InvertedFile<kEmbeddingDim>::GetEntries() const {
    return entries_;
}
 
template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::HasHammingEmbedding() const {
    return status_ & HAS_EMBEDDING;
}
 
template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::EntriesSorted() const {
    return status_ & ENTRIES_SORTED;
}
 
template <int kEmbeddingDim>
bool InvertedFile<kEmbeddingDim>::IsUsable() const {
    return status_ & USABLE;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::AddEntry(const int image_id,
                                           typename DescType::Index feature_idx,
                                           const DescType& descriptor,
                                           const GeomType& geometry) {
    CHECK_GE(image_id, 0);
    CHECK_EQ(descriptor.size(), kEmbeddingDim);
    EntryType entry;
    entry.image_id = image_id;
    entry.feature_idx = feature_idx;
    entry.geometry = geometry;
    ConvertToBinaryDescriptor(descriptor, &entry.descriptor);
    entries_.push_back(entry);
    status_ &= ~ENTRIES_SORTED;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::SortEntries() {
    std::sort(entries_.begin(), entries_.end(),
              [](const EntryType& entry1, const EntryType& entry2) {
                  return entry1.image_id < entry2.image_id;
              });
    status_ |= ENTRIES_SORTED;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ClearEntries() {
    entries_.clear();
    status_ &= ~ENTRIES_SORTED;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Reset() {
    status_ = UNUSABLE;
    idf_weight_ = 0.0f;
    entries_.clear();
    thresholds_.setZero();
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ConvertToBinaryDescriptor(
        const DescType& descriptor,
        std::bitset<kEmbeddingDim>* binary_descriptor) const {
    CHECK_EQ(descriptor.size(), kEmbeddingDim);
    for (int i = 0; i < kEmbeddingDim; ++i) {
        (*binary_descriptor)[i] = descriptor[i] > thresholds_[i];
    }
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeIDFWeight(const int num_total_images) {
    if (entries_.empty()) {
        return;
    }
 
    std::unordered_set<int> image_ids;
    GetImageIds(&image_ids);
 
    idf_weight_ = std::log(static_cast<double>(num_total_images) /
                           static_cast<double>(image_ids.size()));
}
 
template <int kEmbeddingDim>
float InvertedFile<kEmbeddingDim>::IDFWeight() const {
    return idf_weight_;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeHammingEmbedding(
        const Eigen::Matrix<float, Eigen::Dynamic, kEmbeddingDim>&
                descriptors) {
    const int num_descriptors = static_cast<int>(descriptors.rows());
    if (num_descriptors < 2) {
        return;
    }
 
    std::vector<float> elements(num_descriptors);
    for (int n = 0; n < kEmbeddingDim; ++n) {
        for (int i = 0; i < num_descriptors; ++i) {
            elements[i] = descriptors(i, n);
        }
        thresholds_[n] = Median(elements);
    }
 
    status_ |= HAS_EMBEDDING;
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ScoreFeature(
        const DescType& descriptor,
        std::vector<ImageScore>* image_scores) const {
    CHECK_EQ(descriptor.size(), kEmbeddingDim);
 
    image_scores->clear();
 
    if (!IsUsable()) {
        return;
    }
 
    if (entries_.size() == 0) {
        return;
    }
 
    const float squared_idf_weight = idf_weight_ * idf_weight_;
 
    std::bitset<kEmbeddingDim> bin_descriptor;
    ConvertToBinaryDescriptor(descriptor, &bin_descriptor);
 
    ImageScore image_score;
    image_score.image_id = entries_.front().image_id;
    image_score.score = 0.0f;
    int num_image_votes = 0;
 
    // Note that this assumes that the entries are sorted using SortEntries
    // according to their image identifiers.
    for (const auto& entry : entries_) {
        if (image_score.image_id < entry.image_id) {
            if (num_image_votes > 0) {
                // Finalizes the voting since we now know how many features from
                // the database image match the current image feature. This is
                // required to perform burstiness normalization (cf. Eqn. 2 in
                // Arandjelovic, Zisserman: Scalable descriptor
                // distinctiveness for location recognition. ACCV 2014).
                // Notice that the weight from the descriptor matching is
                // already accumulated in image_score.score, i.e., we only need
                // to apply the burstiness weighting.
                image_score.score /=
                        std::sqrt(static_cast<float>(num_image_votes));
                image_score.score *= squared_idf_weight;
                image_scores->push_back(image_score);
            }
 
            image_score.image_id = entry.image_id;
            image_score.score = 0.0f;
            num_image_votes = 0;
        }
 
        const size_t hamming_dist = (bin_descriptor ^ entry.descriptor).count();
 
        if (hamming_dist <= hamming_dist_weight_functor_.kMaxHammingDistance) {
            image_score.score += hamming_dist_weight_functor_(hamming_dist);
            num_image_votes += 1;
        }
    }
 
    // Add the voting for the largest image_id in the entries.
    if (num_image_votes > 0) {
        image_score.score /= std::sqrt(static_cast<float>(num_image_votes));
        image_score.score *= squared_idf_weight;
        image_scores->push_back(image_score);
    }
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::GetImageIds(
        std::unordered_set<int>* ids) const {
    for (const EntryType& entry : entries_) {
        ids->insert(entry.image_id);
    }
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::ComputeImageSelfSimilarities(
        std::unordered_map<int, double>* self_similarities) const {
    const double squared_idf_weight = idf_weight_ * idf_weight_;
    for (const auto& entry : entries_) {
        (*self_similarities)[entry.image_id] += squared_idf_weight;
    }
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Read(std::ifstream* ifs) {
    CHECK(ifs->is_open());
 
    ifs->read(reinterpret_cast<char*>(&status_), sizeof(uint8_t));
    ifs->read(reinterpret_cast<char*>(&idf_weight_), sizeof(float));
 
    for (int i = 0; i < kEmbeddingDim; ++i) {
        ifs->read(reinterpret_cast<char*>(&thresholds_[i]), sizeof(float));
    }
 
    uint32_t num_entries = 0;
    ifs->read(reinterpret_cast<char*>(&num_entries), sizeof(uint32_t));
    entries_.resize(num_entries);
 
    for (uint32_t i = 0; i < num_entries; ++i) {
        entries_[i].Read(ifs);
    }
}
 
template <int kEmbeddingDim>
void InvertedFile<kEmbeddingDim>::Write(std::ofstream* ofs) const {
    CHECK(ofs->is_open());
 
    ofs->write(reinterpret_cast<const char*>(&status_), sizeof(uint8_t));
    ofs->write(reinterpret_cast<const char*>(&idf_weight_), sizeof(float));
 
    for (int i = 0; i < kEmbeddingDim; ++i) {
        ofs->write(reinterpret_cast<const char*>(&thresholds_[i]),
                   sizeof(float));
    }
 
    const uint32_t num_entries = static_cast<uint32_t>(entries_.size());
    ofs->write(reinterpret_cast<const char*>(&num_entries), sizeof(uint32_t));
 
    for (uint32_t i = 0; i < num_entries; ++i) {
        entries_[i].Write(ofs);
    }
}
 
}  // namespace retrieval
}  // namespace colmap