samples/cpp/kmeans.cpp

   1 #include "opencv2/highgui/highgui.hpp"
   2 #include "opencv2/core/core.hpp"
   3 #include <iostream>
   4
   5 using namespace cv;
   6 using namespace std;
   7
   8 // static void help()
   9 // {
  10 //     cout << "\nThis program demonstrates kmeans clustering.\n"
  11 //             "It generates an image with random points, then assigns a random number of cluster\n"
  12 //             "centers and uses kmeans to move those cluster centers to their representitive location\n"
  13 //             "Call\n"
  14 //             "./kmeans\n" << endl;
  15 // }
  16
  17 int main( int /*argc*/, char** /*argv*/ )
  18 {
  19     const int MAX_CLUSTERS = 5;
  20     Scalar colorTab[] =
  21     {
  22         Scalar(0, 0, 255),
  23         Scalar(0,255,0),
  24         Scalar(255,100,100),
  25         Scalar(255,0,255),
  26         Scalar(0,255,255)
  27     };
  28
  29     Mat img(500, 500, CV_8UC3);
  30     RNG rng(12345);
  31
  32     for(;;)
  33     {
  34         int k, clusterCount = rng.uniform(2, MAX_CLUSTERS+1);
  35         int i, sampleCount = rng.uniform(1, 1001);
  36         Mat points(sampleCount, 1, CV_32FC2), labels;
  37
  38         clusterCount = MIN(clusterCount, sampleCount);
  39         Mat centers;
  40
  41         /* generate random sample from multigaussian distribution */
  42         for( k = 0; k < clusterCount; k++ )
  43         {
  44             Point center;
  45             center.x = rng.uniform(0, img.cols);
  46             center.y = rng.uniform(0, img.rows);
  47             Mat pointChunk = points.rowRange(k*sampleCount/clusterCount,
  48                                              k == clusterCount - 1 ? sampleCount :
  49                                              (k+1)*sampleCount/clusterCount);
  50             rng.fill(pointChunk, RNG::NORMAL, Scalar(center.x, center.y), Scalar(img.cols*0.05, img.rows*0.05));
  51         }
  52
  53         randShuffle(points, 1, &rng);
  54
  55         kmeans(points, clusterCount, labels,
  56             TermCriteria( TermCriteria::EPS+TermCriteria::COUNT, 10, 1.0),
  57                3, KMEANS_PP_CENTERS, centers);
  58
  59         img = Scalar::all(0);
  60
  61         for( i = 0; i < sampleCount; i++ )
  62         {
  63             int clusterIdx = labels.at<int>(i);
  64             Point ipt = points.at<Point2f>(i);
  65             circle( img, ipt, 2, colorTab[clusterIdx], FILLED, LINE_AA );
  66         }
  67
  68         imshow("clusters", img);
  69
  70         char key = (char)waitKey();
  71         if( key == 27 || key == 'q' || key == 'Q' ) // 'ESC'
  72             break;
  73     }
  74
  75     return 0;
  76 }