00001 #include <vocabulary_tree/vocabulary_tree.h>
00002 #include <vocabulary_tree/database.h>
00003 #include <boost/lexical_cast.hpp>
00004 #include <map>
00005 #include <cstdio>
00006 #include <fstream>
00007
00008 static const unsigned int DIMENSION = 176;
00009 typedef Eigen::Matrix<float, 1, DIMENSION> Feature;
00010 typedef std::vector<Feature, Eigen::aligned_allocator<Feature> > FeatureVector;
00011
00012 int main(int argc, char** argv)
00013 {
00014 if (argc < 5) {
00015 printf("Usage: %s vocab.tree vocab.weights signatures.dat objects.dat [NUM_SIGS]\n", argv[0]);
00016 return 0;
00017 }
00018
00019
00020 std::ifstream sig_is(argv[3], std::ios::binary);
00021 int length, num_sigs;
00022 if (argc == 5) {
00023 sig_is.seekg(0, std::ios::end);
00024 length = sig_is.tellg();
00025 num_sigs = length / (DIMENSION * sizeof(float));
00026 sig_is.seekg(0, std::ios::beg);
00027 }
00028 else {
00029 num_sigs = boost::lexical_cast<int>(argv[5]);
00030 length = num_sigs * DIMENSION * sizeof(float);
00031 }
00032 printf("Training from %d descriptors\n", num_sigs);
00033
00034
00035 printf("Reading descriptors\n");
00036 FeatureVector features(num_sigs);
00037 sig_is.read((char*)&features[0], length);
00038
00039
00040 printf("Reading document ids\n");
00041 std::ifstream obj_is(argv[4], std::ios::binary);
00042 std::vector<unsigned int> object_ids(num_sigs);
00043 obj_is.read((char*)&object_ids[0], num_sigs * sizeof(unsigned int));
00044
00045
00046 printf("Loading vocabulary tree\n");
00047 vt::VocabularyTree<Feature> tree(argv[1]);
00048
00049
00050 printf("Computing document vectors\n");
00051 typedef std::map<unsigned int, vt::Document> DocumentMap;
00052 DocumentMap documents;
00053 for (int i = 0; i < num_sigs; ++i) {
00054 documents[ object_ids[i] ].push_back( tree.quantize(features[i]) );
00055 }
00056
00057 printf("Loading database weights\n");
00058 vt::Database db(tree.words());
00059 db.loadWeights(argv[2]);
00060
00061
00062 printf("Inserting documents into database\n");
00063 for (DocumentMap::const_iterator i = documents.begin(), ie = documents.end(); i != ie; ++i) {
00064 db.insert(i->second);
00065 }
00066
00067
00068 vt::Matches matches;
00069 for (DocumentMap::const_iterator i = documents.begin(), ie = documents.end(); i != ie; ++i) {
00070 db.find(i->second, 1, matches);
00071 printf("%u -> %u, score = %f\n", i->first, matches[0].id, matches[0].score);
00072 }
00073
00074 return 0;
00075 }