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