vocab_tree.h

Go to the documentation of this file.

/*
Copyright (c) 2015, Tianwei Shen
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 libvot 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 HOLDER 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.

*/

#ifndef VOCAB_TREE_HEADER
#define VOCAB_TREE_HEADER

#include <vector>
#include <cstdlib>
#include <mutex>

#include "utils/global_params.h"
#include "utils/data_types.h"

namespace vot
{
    typedef enum
    {
        L1 = 1,
        L2 = 2 
    } DistanceType;

    class ImageCount 
    {
    public: 
        ImageCount() : index(0), count(0.0) { }
        ImageCount(size_t index_, float count_) : index(index_), count(count_) { }

        size_t index; 
        float count; 
    };

    class TreeNode
    {
        public:
            TreeNode():des(NULL)  {}
            virtual ~TreeNode()
            {
                if(des != NULL) 
                {
                    delete [] des;
                    des = NULL;
                }
            }
            // pure virtual function that will be implemented in derived classes
            virtual bool RecursiveBuild(size_t num_keys, int dim, int depth, int depth_curr, int bf, DTYPE **p , double *means, int *assign, int thread_num = 1) = 0;
            virtual bool ClearNode(int bf) = 0;
            virtual bool WriteNode(FILE *f, int branch_num, int dim) const = 0;
            virtual bool ReadNode(FILE *f, int branch_num, int dim) = 0;
            virtual size_t CountNodes(int branch_num) const = 0;
            virtual size_t CountLeaves(int branch_num) const = 0;
            virtual bool Compare(TreeNode *in, int branch_num, int dim) const = 0;
            virtual bool ClearScores(int bf) = 0;             
            // function for build image database
            virtual size_t DescendFeature(float *q, DTYPE *v, size_t image_index, int branch_num, int dim, bool add = true) = 0;
            virtual double ComputeImageVectorMagnitude(int bf, DistanceType dt) = 0;
            virtual bool SetConstantWeight(int bf) = 0;   
            virtual bool ComputeTFIDFWeight(int bf, size_t n) = 0;  
            virtual bool ComputeDatabaseMagnitude(int bf, DistanceType dis_type, size_t start_id, std::vector<float> &database_mag) = 0; 
            virtual bool NormalizeDatabase(int bf, size_t start_id, std::vector<float> &database_mag) = 0;   
            // member functions for querying database
            virtual bool IndexLeaves(int branch_num) = 0;
            virtual bool FillQueryVector(float *q, int branch_num, float normalize_factor) = 0;     
            virtual bool ScoreQuery(float *q, int branch_num, DistanceType dt, float *scores) = 0;       

            DTYPE *des; 
            size_t id; 
    };

    class TreeInNode: public TreeNode
    {
        public:
            TreeInNode():children(NULL) {}
            virtual ~TreeInNode();
            virtual bool RecursiveBuild(size_t num_keys, int dim, int depth, int depth_curr, int bf, DTYPE **p, double *means, int *assign, int thread_num = 1);
            virtual bool ClearNode(int bf);
            virtual bool WriteNode(FILE *f, int branch_num, int dim) const;
            virtual bool ReadNode(FILE *f, int branch_num, int dim);
            virtual size_t CountNodes(int branch_num) const ;
            virtual size_t CountLeaves(int branch_num) const ;
            virtual bool Compare(TreeNode *in, int branch_num, int dim) const;
            virtual bool ClearScores(int bf);             
            // function for build image database
            virtual size_t DescendFeature(float *q, DTYPE *v, size_t image_index, int branch_num, int dim, bool add = true);
            virtual double ComputeImageVectorMagnitude(int bf, DistanceType dt);
            virtual bool SetConstantWeight(int bf);   
            virtual bool ComputeTFIDFWeight(int bf, size_t n);  
            virtual bool ComputeDatabaseMagnitude(int bf, DistanceType dis_type, size_t start_id, std::vector<float> &database_mag); 
            virtual bool NormalizeDatabase(int bf, size_t start_id, std::vector<float> &database_mag);   
            // member functions for querying database
            virtual bool IndexLeaves(int branch_num);
            virtual bool FillQueryVector(float *q, int branch_num, float normalize_factor);     
            virtual bool ScoreQuery(float *q, int branch_num, DistanceType dt, float *scores);       

            TreeNode **children;
    };

    class TreeLeafNode: public TreeNode
    {
        public:
            TreeLeafNode():score(0.0), weight(0.0) {}
            virtual ~TreeLeafNode();
            virtual bool RecursiveBuild(size_t num_keys, int dim, int depth, int depth_curr, int bf, DTYPE **p , double *means, int *assign, int thread_num = 1);
            virtual bool ClearNode(int bf);
            virtual bool WriteNode(FILE *f, int branch_num, int dim) const;
            virtual bool ReadNode(FILE *f, int branch_num, int dim);
            virtual size_t CountNodes(int branch_num) const;
            virtual size_t CountLeaves(int branch_num) const;
            virtual bool Compare(TreeNode *leaf, int branch_num, int dim) const;
            virtual bool ClearScores(int bf);             
            // function for build image database
            virtual size_t DescendFeature(float *q, DTYPE *v, size_t image_index, int branch_num, int dim, bool add = true);
            virtual double ComputeImageVectorMagnitude(int bf, DistanceType dt);
            virtual bool SetConstantWeight(int bf);   
            virtual bool ComputeTFIDFWeight(int bf, size_t n);  
            virtual bool ComputeDatabaseMagnitude(int bf, DistanceType dis_type, size_t start_id, std::vector<float> &database_mag); 
            virtual bool NormalizeDatabase(int bf, size_t start_id, std::vector<float> &database_mag);   
            // member functions for querying database
            virtual bool IndexLeaves(int branch_num);
            virtual bool FillQueryVector(float *q, int branch_num, float normalize_factor);     
            virtual bool ScoreQuery(float *q, int branch_num, DistanceType dt, float *scores);       

            float score;            
            float weight;           
            std::vector<ImageCount> inv_list;   
            std::mutex add_lock;        
    };

    class VocabTree
    {
        public:
            // constructors and destructors
            VocabTree();
            VocabTree(int depth_, int branch_num_, int dim_ = 128, DistanceType dis_type_ = L2);
            ~VocabTree();

            // member function for building vocabulary tree
            bool BuildTree(size_t num_keys, int dim, int depth, int bf, DTYPE **p, int thread_num = 1);      
            bool WriteTree(const char *filename) const;       
            bool ReadTree(const char *filename);            
            bool ClearTree();                   
            bool Compare(VocabTree &v) const; 
            // member function for building image database
            double AddImage2Tree(size_t image_index, tw::SiftData &sift, int thread_num);   
            void Show() const;      
            bool SetConstantWeight();   
            bool ComputeTFIDFWeight(size_t image_num);  
            bool NormalizeDatabase(size_t start_id, size_t image_num);    
            // member function for querying database
            bool Query(tw::SiftData &sift, float *scores);   
            size_t IndexLeaves(); 

            // public data member 
            int branch_num;             
            int depth;                  
            int dim;                    
            size_t database_image_num;  
            size_t num_nodes;           
            size_t num_leaves;          
            DistanceType dis_type;      
            TreeNode *root;             
    };

}   // end of namespace vot

#endif  //VOCAB_TREE_HEADER