00001 #include <vocabulary_tree/tree_builder.h>
00002 #include <boost/lexical_cast.hpp>
00003 #include <cstdio>
00004 #include <fstream>
00005
00006 static const unsigned int DIMENSION = 176;
00007 static const uint32_t K = 10;
00008 static const uint32_t LEVELS = 5;
00009
00010
00011
00012 int main(int argc, char** argv)
00013 {
00014 if (argc < 3) {
00015 printf("Usage: %s signatures.dat output.tree [NUM_SIGS]\n", argv[0]);
00016 return 0;
00017 }
00018
00019 std::ifstream sig_is(argv[1], std::ios::binary);
00020 std::string tree_file = argv[2];
00021
00022
00023 int length, num_sigs;
00024 if (argc == 3) {
00025 sig_is.seekg(0, std::ios::end);
00026 length = sig_is.tellg();
00027 num_sigs = length / (DIMENSION * sizeof(float));
00028 sig_is.seekg(0, std::ios::beg);
00029 }
00030 else {
00031 num_sigs = boost::lexical_cast<int>(argv[3]);
00032 length = num_sigs * DIMENSION * sizeof(float);
00033 }
00034 printf("Training from %d descriptors\n", num_sigs);
00035 printf("Data length = %d\n", length);
00036
00037
00038 typedef Eigen::Matrix<float, 1, DIMENSION> Feature;
00039 typedef std::vector<Feature, Eigen::aligned_allocator<Feature> > FeatureVector;
00040 FeatureVector features(num_sigs);
00041 sig_is.read((char*)&features[0], length);
00042 printf("Done reading in descriptors\n");
00043
00044
00045 vt::TreeBuilder<Feature> builder(Feature::Zero());
00046 builder.kmeans().setRestarts(5);
00047 builder.build(features, K, LEVELS);
00048 printf("%u centers\n", builder.tree().centers().size());
00049 builder.tree().save(tree_file);
00050
00051 return 0;
00052 }