1 #include "opencv2/highgui.hpp"
2 #include "opencv2/imgproc.hpp"
3 #include "opencv2/ml.hpp"
4
5 using namespace cv;
6 using namespace cv::ml;
7
main(int,char **)8 int main( int /*argc*/, char** /*argv*/ )
9 {
10 const int N = 4;
11 const int N1 = (int)sqrt((double)N);
12 const Scalar colors[] =
13 {
14 Scalar(0,0,255), Scalar(0,255,0),
15 Scalar(0,255,255),Scalar(255,255,0)
16 };
17
18 int i, j;
19 int nsamples = 100;
20 Mat samples( nsamples, 2, CV_32FC1 );
21 Mat labels;
22 Mat img = Mat::zeros( Size( 500, 500 ), CV_8UC3 );
23 Mat sample( 1, 2, CV_32FC1 );
24
25 samples = samples.reshape(2, 0);
26 for( i = 0; i < N; i++ )
27 {
28 // form the training samples
29 Mat samples_part = samples.rowRange(i*nsamples/N, (i+1)*nsamples/N );
30
31 Scalar mean(((i%N1)+1)*img.rows/(N1+1),
32 ((i/N1)+1)*img.rows/(N1+1));
33 Scalar sigma(30,30);
34 randn( samples_part, mean, sigma );
35 }
36 samples = samples.reshape(1, 0);
37
38 // cluster the data
39 Ptr<EM> em_model = EM::create();
40 em_model->setClustersNumber(N);
41 em_model->setCovarianceMatrixType(EM::COV_MAT_SPHERICAL);
42 em_model->setTermCriteria(TermCriteria(TermCriteria::COUNT+TermCriteria::EPS, 300, 0.1));
43 em_model->trainEM( samples, noArray(), labels, noArray() );
44
45 // classify every image pixel
46 for( i = 0; i < img.rows; i++ )
47 {
48 for( j = 0; j < img.cols; j++ )
49 {
50 sample.at<float>(0) = (float)j;
51 sample.at<float>(1) = (float)i;
52 int response = cvRound(em_model->predict2( sample, noArray() )[1]);
53 Scalar c = colors[response];
54
55 circle( img, Point(j, i), 1, c*0.75, FILLED );
56 }
57 }
58
59 //draw the clustered samples
60 for( i = 0; i < nsamples; i++ )
61 {
62 Point pt(cvRound(samples.at<float>(i, 0)), cvRound(samples.at<float>(i, 1)));
63 circle( img, pt, 1, colors[labels.at<int>(i)], FILLED );
64 }
65
66 imshow( "EM-clustering result", img );
67 waitKey(0);
68
69 return 0;
70 }
71