fcw_train.cpp 5.92 KB
Newer Older
Vlad Shakhuro's avatar
Vlad Shakhuro committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
#include <cstdio>
#include <cstring>

#include <string>
using std::string;

#include <vector>
using std::vector;

#include <fstream>
using std::ifstream;
using std::getline;

#include <sstream>
using std::stringstream;

#include <iostream>
using std::cerr;
using std::endl;

#include <opencv2/core.hpp>
using cv::Rect;
using cv::Size;
#include <opencv2/highgui.hpp>
using cv::imread;
#include <opencv2/core/utility.hpp>
using cv::CommandLineParser;
using cv::FileStorage;
manuele's avatar
manuele committed
29
#include <opencv2/core/utility.hpp>
Vlad Shakhuro's avatar
Vlad Shakhuro committed
30

manuele's avatar
manuele committed
31 32 33
#include <ctime>        // std::time
#include <cstdlib>      // std::rand, std::srand

Vlad Shakhuro's avatar
Vlad Shakhuro committed
34 35
#include <opencv2/xobjdetect.hpp>

manuele's avatar
manuele committed
36

Vlad Shakhuro's avatar
Vlad Shakhuro committed
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
using cv::xobjdetect::ICFDetectorParams;
using cv::xobjdetect::ICFDetector;
using cv::xobjdetect::WaldBoost;
using cv::xobjdetect::WaldBoostParams;
using cv::Mat;

static bool read_model_size(const char *str, int *rows, int *cols)
{
    int pos = 0;
    if( sscanf(str, "%dx%d%n", rows, cols, &pos) != 2 || str[pos] != '\0' ||
        *rows <= 0 || *cols <= 0)
    {
        return false;
    }
    return true;
}

manuele's avatar
manuele committed
54 55
static int randomPred (int i) { return std::rand()%i;}

Vlad Shakhuro's avatar
Vlad Shakhuro committed
56 57
int main(int argc, char *argv[])
{
manuele's avatar
manuele committed
58
    
Vlad Shakhuro's avatar
Vlad Shakhuro committed
59 60 61 62
    const string keys =
        "{help           |           | print this message}"
        "{pos_path       |       pos | path to training object samples}"
        "{bg_path        |        bg | path to background images}"
Vlad Shakhuro's avatar
Vlad Shakhuro committed
63
        "{bg_per_image   |         5 | number of windows to sample per bg image}"
Vlad Shakhuro's avatar
Vlad Shakhuro committed
64 65 66 67
        "{feature_count  |     10000 | number of features to generate}"
        "{weak_count     |       100 | number of weak classifiers in cascade}"
        "{model_size     |     40x40 | model size in pixels}"
        "{model_filename | model.xml | filename for saving model}"
manuele's avatar
manuele committed
68 69 70 71 72 73
        "{features_type  |       icf | features type, \"icf\" or \"acf\"}"
        "{alpha          |      0.02 | alpha value}"
        "{is_grayscale   |     false | read the image as grayscale}"
        "{use_fast_log   |     false | use fast log function}"
        "{limit_ps       |        -1 | limit to positive samples (-1 means all)}"
        "{limit_bg       |        -1 | limit to negative samples (-1 means all)}"
Vlad Shakhuro's avatar
Vlad Shakhuro committed
74 75
        ;

manuele's avatar
manuele committed
76

Vlad Shakhuro's avatar
Vlad Shakhuro committed
77 78 79 80 81 82 83 84 85 86 87 88 89 90
    CommandLineParser parser(argc, argv, keys);
    parser.about("FCW trainer");

    if( parser.has("help") || argc == 1)
    {
        parser.printMessage();
        return 0;
    }

    string pos_path = parser.get<string>("pos_path");
    string bg_path = parser.get<string>("bg_path");
    string model_filename = parser.get<string>("model_filename");

    ICFDetectorParams params;
Vlad Shakhuro's avatar
Vlad Shakhuro committed
91 92 93
    params.feature_count = parser.get<int>("feature_count");
    params.weak_count = parser.get<int>("weak_count");
    params.bg_per_image = parser.get<int>("bg_per_image");
manuele's avatar
manuele committed
94 95 96 97 98 99 100 101
    params.features_type = parser.get<string>("features_type");
    params.alpha = parser.get<float>("alpha");
    params.is_grayscale = parser.get<bool>("is_grayscale");
    params.use_fast_log = parser.get<bool>("use_fast_log");
    
    int limit_ps = parser.get<int>("limit_ps");
    int limit_bg = parser.get<int>("limit_bg");    
    
Vlad Shakhuro's avatar
Vlad Shakhuro committed
102 103 104 105 106 107 108
    string model_size = parser.get<string>("model_size");
    if( !read_model_size(model_size.c_str(), &params.model_n_rows,
        &params.model_n_cols) )
    {
        cerr << "Error reading model size from `" << model_size << "`" << endl;
        return 1;
    }
manuele's avatar
manuele committed
109
    
Vlad Shakhuro's avatar
Vlad Shakhuro committed
110 111 112 113 114 115 116 117 118 119 120 121
    if( params.feature_count <= 0 )
    {
        cerr << "feature_count must be positive number" << endl;
        return 1;
    }

    if( params.weak_count <= 0 )
    {
        cerr << "weak_count must be positive number" << endl;
        return 1;
    }

manuele's avatar
manuele committed
122
    if( params.features_type != "icf" &&  params.features_type != "acf" )
Vlad Shakhuro's avatar
Vlad Shakhuro committed
123
    {
manuele's avatar
manuele committed
124 125 126 127 128 129
        cerr << "features_type must be \"icf\" or \"acf\"" << endl;
        return 1;
    }
    if( params.alpha <= 0 )
    {
        cerr << "alpha must be positive float number" << endl;
Vlad Shakhuro's avatar
Vlad Shakhuro committed
130 131
        return 1;
    }
Vlad Shakhuro's avatar
Vlad Shakhuro committed
132 133 134 135 136
    if( !parser.check() )
    {
        parser.printErrors();
        return 1;
    }
manuele's avatar
manuele committed
137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161
    
    std::vector<cv::String> pos_filenames;
    glob(pos_path, pos_filenames);

    std::vector<cv::String> bg_filenames;
    glob(bg_path, bg_filenames);
        
    if(limit_ps != -1 && (int)pos_filenames.size() > limit_ps)
      pos_filenames.erase(pos_filenames.begin()+limit_ps, pos_filenames.end());
    if(limit_bg != -1 && (int)bg_filenames.size() > limit_bg)
      bg_filenames.erase(bg_filenames.begin()+limit_bg, bg_filenames.end());
    
    //random pick input images
    bool random_shuffle = false;
    if(random_shuffle)
    {
      std::srand ( unsigned ( std::time(0) ) );
      std::random_shuffle ( pos_filenames.begin(), pos_filenames.end(), randomPred );
      std::random_shuffle ( bg_filenames.begin(), bg_filenames.end(), randomPred );
    }
    
    int samples_size = (int)((params.bg_per_image * bg_filenames.size()) + pos_filenames.size());
    int features_size = params.feature_count;
    int max_features_allowed = (int)(INT_MAX/(sizeof(int)* samples_size));
    int max_samples_allowed = (int)(INT_MAX/(sizeof(int)* features_size));
manuele's avatar
manuele committed
162 163
    int total_samples = (int)((params.bg_per_image * bg_filenames.size()) + pos_filenames.size());
    
manuele's avatar
manuele committed
164
    
manuele's avatar
manuele committed
165
    if(total_samples >max_samples_allowed)
manuele's avatar
manuele committed
166
    {
manuele's avatar
manuele committed
167
      CV_Error_(1, ("exceeded maximum number of samples. Maximum number of samples with %d features is %d, you have %d (%d positive samples + (%d bg * %d bg_per_image))\n",features_size,max_samples_allowed,total_samples,pos_filenames.size(),bg_filenames.size(),params.bg_per_image ));
manuele's avatar
manuele committed
168 169 170 171
    }
    
    if(params.feature_count >max_features_allowed)
    {
manuele's avatar
manuele committed
172
      CV_Error_(1, ("exceeded maximum number of features. Maximum number of features with %d samples is %d, you have %d\n",samples_size,max_features_allowed, features_size ));
manuele's avatar
manuele committed
173 174 175 176 177 178 179 180 181
    }
    
    std::cout<<pos_filenames.size()<<std::endl;
    std::cout<<bg_filenames.size()<<std::endl;

    ICFDetector detector;    

    
    detector.train(pos_filenames, bg_filenames, params);
Vlad Shakhuro's avatar
Vlad Shakhuro committed
182 183

    FileStorage fs(model_filename, FileStorage::WRITE);
Vlad Shakhuro's avatar
Vlad Shakhuro committed
184 185
    fs << "icfdetector";
    detector.write(fs);
Vlad Shakhuro's avatar
Vlad Shakhuro committed
186 187
    fs.release();
}