kdtree_index.h

/***********************************************************************
 * Software License Agreement (BSD License)
 *
 * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
 * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
 *
 * THE BSD LICENSE
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``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 AUTHOR 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.
 *************************************************************************/
 
#ifndef OPENCV_FLANN_KDTREE_INDEX_H_
#define OPENCV_FLANN_KDTREE_INDEX_H_
 
 
#include <algorithm>
#include <map>
#include <cstring>
 
#include "nn_index.h"
#include "dynamic_bitset.h"
#include "matrix.h"
#include "result_set.h"
#include "heap.h"
#include "allocator.h"
#include "random.h"
#include "saving.h"
 
 
namespace cvflann
{
 
struct KDTreeIndexParams : public IndexParams
{
    KDTreeIndexParams(int trees = 4)
    {
        (*this)["algorithm"] = FLANN_INDEX_KDTREE;
        (*this)["trees"] = trees;
    }
};
 
 
template <typename Distance>
class KDTreeIndex : public NNIndex<Distance>
{
public:
    typedef typename Distance::ElementType ElementType;
    typedef typename Distance::ResultType DistanceType;
 
 
    KDTreeIndex(const Matrix<ElementType>& inputData, const IndexParams& params = KDTreeIndexParams(),
                Distance d = Distance() ) :
        dataset_(inputData), index_params_(params), distance_(d)
    {
        size_ = dataset_.rows;
        veclen_ = dataset_.cols;
 
        trees_ = get_param(index_params_,"trees",4);
        tree_roots_ = new NodePtr[trees_];
 
        // Create a permutable array of indices to the input vectors.
        vind_.resize(size_);
        for (size_t i = 0; i < size_; ++i) {
            vind_[i] = int(i);
        }
 
        mean_ = new DistanceType[veclen_];
        var_ = new DistanceType[veclen_];
    }
 
 
    KDTreeIndex(const KDTreeIndex&);
    KDTreeIndex& operator=(const KDTreeIndex&);
 
    ~KDTreeIndex()
    {
        if (tree_roots_!=NULL) {
            delete[] tree_roots_;
        }
        delete[] mean_;
        delete[] var_;
    }
 
    void buildIndex() CV_OVERRIDE
    {
        /* Construct the randomized trees. */
        for (int i = 0; i < trees_; i++) {
            /* Randomize the order of vectors to allow for unbiased sampling. */
#ifndef OPENCV_FLANN_USE_STD_RAND
            cv::randShuffle(vind_);
#else
            std::random_shuffle(vind_.begin(), vind_.end());
#endif
 
            tree_roots_[i] = divideTree(&vind_[0], int(size_) );
        }
    }
 
 
    flann_algorithm_t getType() const CV_OVERRIDE
    {
        return FLANN_INDEX_KDTREE;
    }
 
 
    void saveIndex(FILE* stream) CV_OVERRIDE
    {
        save_value(stream, trees_);
        for (int i=0; i<trees_; ++i) {
            save_tree(stream, tree_roots_[i]);
        }
    }
 
 
 
    void loadIndex(FILE* stream) CV_OVERRIDE
    {
        load_value(stream, trees_);
        if (tree_roots_!=NULL) {
            delete[] tree_roots_;
        }
        tree_roots_ = new NodePtr[trees_];
        for (int i=0; i<trees_; ++i) {
            load_tree(stream,tree_roots_[i]);
        }
 
        index_params_["algorithm"] = getType();
        index_params_["trees"] = tree_roots_;
    }
 
    size_t size() const CV_OVERRIDE
    {
        return size_;
    }
 
    size_t veclen() const CV_OVERRIDE
    {
        return veclen_;
    }
 
    int usedMemory() const CV_OVERRIDE
    {
        return int(pool_.usedMemory+pool_.wastedMemory+dataset_.rows*sizeof(int));  // pool memory and vind array memory
    }
 
    void findNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, const SearchParams& searchParams) CV_OVERRIDE
    {
        const int maxChecks = get_param(searchParams,"checks", 32);
        const float epsError = 1+get_param(searchParams,"eps",0.0f);
        const bool explore_all_trees = get_param(searchParams,"explore_all_trees",false);
 
        if (maxChecks==FLANN_CHECKS_UNLIMITED) {
            getExactNeighbors(result, vec, epsError);
        }
        else {
            getNeighbors(result, vec, maxChecks, epsError, explore_all_trees);
        }
    }
 
    IndexParams getParameters() const CV_OVERRIDE
    {
        return index_params_;
    }
 
private:
 
 
    /*--------------------- Internal Data Structures --------------------------*/
    struct Node
    {
        int divfeat;
        DistanceType divval;
        Node* child1, * child2;
    };
    typedef Node* NodePtr;
    typedef BranchStruct<NodePtr, DistanceType> BranchSt;
    typedef BranchSt* Branch;
 
 
 
    void save_tree(FILE* stream, NodePtr tree)
    {
        save_value(stream, *tree);
        if (tree->child1!=NULL) {
            save_tree(stream, tree->child1);
        }
        if (tree->child2!=NULL) {
            save_tree(stream, tree->child2);
        }
    }
 
 
    void load_tree(FILE* stream, NodePtr& tree)
    {
        tree = pool_.allocate<Node>();
        load_value(stream, *tree);
        if (tree->child1!=NULL) {
            load_tree(stream, tree->child1);
        }
        if (tree->child2!=NULL) {
            load_tree(stream, tree->child2);
        }
    }
 
 
    NodePtr divideTree(int* ind, int count)
    {
        NodePtr node = pool_.allocate<Node>(); // allocate memory
 
        /* If too few exemplars remain, then make this a leaf node. */
        if ( count == 1) {
            node->child1 = node->child2 = NULL;    /* Mark as leaf node. */
            node->divfeat = *ind;    /* Store index of this vec. */
        }
        else {
            int idx;
            int cutfeat;
            DistanceType cutval;
            meanSplit(ind, count, idx, cutfeat, cutval);
 
            node->divfeat = cutfeat;
            node->divval = cutval;
            node->child1 = divideTree(ind, idx);
            node->child2 = divideTree(ind+idx, count-idx);
        }
 
        return node;
    }
 
 
    void meanSplit(int* ind, int count, int& index, int& cutfeat, DistanceType& cutval)
    {
        memset(mean_,0,veclen_*sizeof(DistanceType));
        memset(var_,0,veclen_*sizeof(DistanceType));
 
        /* Compute mean values.  Only the first SAMPLE_MEAN values need to be
            sampled to get a good estimate.
         */
        int cnt = std::min((int)SAMPLE_MEAN+1, count);
        for (int j = 0; j < cnt; ++j) {
            ElementType* v = dataset_[ind[j]];
            for (size_t k=0; k<veclen_; ++k) {
                mean_[k] += v[k];
            }
        }
        for (size_t k=0; k<veclen_; ++k) {
            mean_[k] /= cnt;
        }
 
        /* Compute variances (no need to divide by count). */
        for (int j = 0; j < cnt; ++j) {
            ElementType* v = dataset_[ind[j]];
            for (size_t k=0; k<veclen_; ++k) {
                DistanceType dist = v[k] - mean_[k];
                var_[k] += dist * dist;
            }
        }
        /* Select one of the highest variance indices at random. */
        cutfeat = selectDivision(var_);
        cutval = mean_[cutfeat];
 
        int lim1, lim2;
        planeSplit(ind, count, cutfeat, cutval, lim1, lim2);
 
        if (lim1>count/2) index = lim1;
        else if (lim2<count/2) index = lim2;
        else index = count/2;
 
        /* If either list is empty, it means that all remaining features
         * are identical. Split in the middle to maintain a balanced tree.
         */
        if ((lim1==count)||(lim2==0)) index = count/2;
    }
 
 
    int selectDivision(DistanceType* v)
    {
        int num = 0;
        size_t topind[RAND_DIM];
 
        /* Create a list of the indices of the top RAND_DIM values. */
        for (size_t i = 0; i < veclen_; ++i) {
            if ((num < RAND_DIM)||(v[i] > v[topind[num-1]])) {
                /* Put this element at end of topind. */
                if (num < RAND_DIM) {
                    topind[num++] = i;            /* Add to list. */
                }
                else {
                    topind[num-1] = i;         /* Replace last element. */
                }
                /* Bubble end value down to right location by repeated swapping. */
                int j = num - 1;
                while (j > 0  &&  v[topind[j]] > v[topind[j-1]]) {
                    std::swap(topind[j], topind[j-1]);
                    --j;
                }
            }
        }
        /* Select a random integer in range [0,num-1], and return that index. */
        int rnd = rand_int(num);
        return (int)topind[rnd];
    }
 
 
    void planeSplit(int* ind, int count, int cutfeat, DistanceType cutval, int& lim1, int& lim2)
    {
        /* Move vector indices for left subtree to front of list. */
        int left = 0;
        int right = count-1;
        for (;; ) {
            while (left<=right && dataset_[ind[left]][cutfeat]<cutval) ++left;
            while (left<=right && dataset_[ind[right]][cutfeat]>=cutval) --right;
            if (left>right) break;
            std::swap(ind[left], ind[right]); ++left; --right;
        }
        lim1 = left;
        right = count-1;
        for (;; ) {
            while (left<=right && dataset_[ind[left]][cutfeat]<=cutval) ++left;
            while (left<=right && dataset_[ind[right]][cutfeat]>cutval) --right;
            if (left>right) break;
            std::swap(ind[left], ind[right]); ++left; --right;
        }
        lim2 = left;
    }
 
    void getExactNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, float epsError)
    {
        //      checkID -= 1;  /* Set a different unique ID for each search. */
 
        if (trees_ > 1) {
            fprintf(stderr,"It doesn't make any sense to use more than one tree for exact search");
        }
        if (trees_>0) {
            searchLevelExact(result, vec, tree_roots_[0], 0.0, epsError);
        }
        CV_Assert(result.full());
    }
 
    void getNeighbors(ResultSet<DistanceType>& result, const ElementType* vec,
                      int maxCheck, float epsError, bool explore_all_trees = false)
    {
        int i;
        BranchSt branch;
        int checkCount = 0;
        DynamicBitset checked(size_);
 
        // Priority queue storing intermediate branches in the best-bin-first search
        const cv::Ptr<Heap<BranchSt>>& heap = Heap<BranchSt>::getPooledInstance(cv::utils::getThreadID(), (int)size_);
 
        /* Search once through each tree down to root. */
        for (i = 0; i < trees_; ++i) {
            searchLevel(result, vec, tree_roots_[i], 0, checkCount, maxCheck,
                        epsError, heap, checked, explore_all_trees);
            if (!explore_all_trees && (checkCount >= maxCheck) && result.full())
                break;
        }
 
        /* Keep searching other branches from heap until finished. */
        while ( heap->popMin(branch) && (checkCount < maxCheck || !result.full() )) {
            searchLevel(result, vec, branch.node, branch.mindist, checkCount, maxCheck,
                        epsError, heap, checked, false);
        }
 
        CV_Assert(result.full());
    }
 
 
    void searchLevel(ResultSet<DistanceType>& result_set, const ElementType* vec, NodePtr node, DistanceType mindist, int& checkCount, int maxCheck,
                     float epsError, const cv::Ptr<Heap<BranchSt>>& heap, DynamicBitset& checked, bool explore_all_trees = false)
    {
        if (result_set.worstDist()<mindist) {
            //          printf("Ignoring branch, too far\n");
            return;
        }
 
        /* If this is a leaf node, then do check and return. */
        if ((node->child1 == NULL)&&(node->child2 == NULL)) {
            /*  Do not check same node more than once when searching multiple trees.
                Once a vector is checked, we set its location in vind to the
                current checkID.
             */
            int index = node->divfeat;
            if ( checked.test(index) ||
                 (!explore_all_trees && (checkCount>=maxCheck) && result_set.full()) ) {
                return;
            }
            checked.set(index);
            checkCount++;
 
            DistanceType dist = distance_(dataset_[index], vec, veclen_);
            result_set.addPoint(dist,index);
 
            return;
        }
 
        /* Which child branch should be taken first? */
        ElementType val = vec[node->divfeat];
        DistanceType diff = val - node->divval;
        NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
        NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;
 
        /* Create a branch record for the branch not taken.  Add distance
            of this feature boundary (we don't attempt to correct for any
            use of this feature in a parent node, which is unlikely to
            happen and would have only a small effect).  Don't bother
            adding more branches to heap after halfway point, as cost of
            adding exceeds their value.
         */
 
        DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
        //      if (2 * checkCount < maxCheck  ||  !result.full()) {
        if ((new_distsq*epsError < result_set.worstDist())||  !result_set.full()) {
            heap->insert( BranchSt(otherChild, new_distsq) );
        }
 
        /* Call recursively to search next level down. */
        searchLevel(result_set, vec, bestChild, mindist, checkCount, maxCheck, epsError, heap, checked);
    }
 
    void searchLevelExact(ResultSet<DistanceType>& result_set, const ElementType* vec, const NodePtr node, DistanceType mindist, const float epsError)
    {
        /* If this is a leaf node, then do check and return. */
        if ((node->child1 == NULL)&&(node->child2 == NULL)) {
            int index = node->divfeat;
            DistanceType dist = distance_(dataset_[index], vec, veclen_);
            result_set.addPoint(dist,index);
            return;
        }
 
        /* Which child branch should be taken first? */
        ElementType val = vec[node->divfeat];
        DistanceType diff = val - node->divval;
        NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
        NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;
 
        /* Create a branch record for the branch not taken.  Add distance
            of this feature boundary (we don't attempt to correct for any
            use of this feature in a parent node, which is unlikely to
            happen and would have only a small effect).  Don't bother
            adding more branches to heap after halfway point, as cost of
            adding exceeds their value.
         */
 
        DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
 
        /* Call recursively to search next level down. */
        searchLevelExact(result_set, vec, bestChild, mindist, epsError);
 
        if (new_distsq*epsError<=result_set.worstDist()) {
            searchLevelExact(result_set, vec, otherChild, new_distsq, epsError);
        }
    }
 
 
private:
 
    enum
    {
        SAMPLE_MEAN = 100,
        RAND_DIM=5
    };
 
 
    int trees_;
 
    std::vector<int> vind_;
 
    const Matrix<ElementType> dataset_;
 
    IndexParams index_params_;
 
    size_t size_;
    size_t veclen_;
 
 
    DistanceType* mean_;
    DistanceType* var_;
 
 
    NodePtr* tree_roots_;
 
    PooledAllocator pool_;
 
    Distance distance_;
 
 
};   // class KDTreeForest
 
}
 
 
#endif //OPENCV_FLANN_KDTREE_INDEX_H_