Blame view

src/class_rf.h 2.98 KB
5f3cba02   David Mayerich   initial public co...
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
  //OpenCV
  #include <opencv2/opencv.hpp>
  
  #include "progress_thread.h"
  
  //create a custom classifier to access the number of classes (OpenCV protected variable)
  class RFClass : public CvRTrees{
  public:
  	int get_nclasses(){ return CvRTrees::nclasses;}
  };
  
  /// Perform classification of the ENVI file using the current RF classifier
  std::vector< stim::image<unsigned char> > predict_rf(stim::envi* E, RFClass* RF, unsigned char* MASK = NULL){
  
  	size_t nC = RF->get_nclasses();							//get the number of classes
  	size_t X = E->header.samples;
  	size_t Y = E->header.lines;
  	size_t B = E->header.bands;
  	size_t XY = E->header.samples * E->header.lines;
  
  	size_t tP = 0;											//calculate the total number of pixels
  	if(MASK){
  		for(size_t xy = 0; xy < XY; xy++){
  			if(MASK[xy]) tP++;
  		}
  	}
  	else
  		tP = X * Y;
  
  	std::vector< stim::image<unsigned char> > C;					//create an array of mask images
  	C.resize(nC);
  
  	for(unsigned long long c = 0; c < nC; c++){						//for each class mask
  		C[c] = stim::image<unsigned char>(X, Y, 1);					//allocate space for the mask
  		memset(C[c].data(), 0, X * Y * sizeof(unsigned char));		//initialize all of the pixels to zero
  	}
  
  	cv::Mat classF(1, (int)B, CV_32FC1);							//allocate space for a single feature vector
  
  
  	double progress = 0;											//initialize the progress bar variable
  	std::thread t1(progressbar_thread, &progress);					//start the progress bar thread
  
  	unsigned long long t = 0;
  	for(unsigned long long p = 0; p < XY; p++){						//for each pixel
  		if(!MASK || MASK[p] > 0){
  			E->spectrum<float>((float*)classF.data, (int)p);		//fill the classF matrix with a single spectrum
  			float c = RF->predict(classF);							//classify the feature vector
  			C[(size_t)c].data()[p] = 255;							//write a white pixel to the appropriate class image
  			t++;
  			progress = (double)(t+1) / (double)(tP) * 100.0;		//update the progress bar variable
  		}
  	}
  	t1.join();														//finish the progress bar thread
  
  	return C;
  }
  
  //function for training a random forest classifier
  void train_rf(RFClass* RF, cv::Mat &trainF, cv::Mat &trainR, int tree_depth, int num_trees, bool VERBOSE = false){
  
  	//unsigned tP = trainF.cols;								//calculate the number of individual measurements
  
      CvRTParams  params( tree_depth, // max_depth,
                          1, // min_sample_count,
                          0.f, // regression_accuracy,
                          false, // use_surrogates,
                          16, // max_categories,
                          0, // priors,
                          false, // calc_var_importance,
                          0, // nactive_vars,
                          num_trees, // max_num_of_trees_in_the_forest,
                          0, // forest_accuracy,
                          CV_TERMCRIT_ITER // termcrit_type
                         );
  	
  	if(VERBOSE) std::cout<<"Starting OpenCV CvRT training algorithm...";
  	RF->train(trainF, CV_ROW_SAMPLE, trainR, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), params );		//train the classifier
  	if(VERBOSE) std::cout<<"done"<<std::endl;
  }