kmeans.cpp 2.53 KB
Newer Older
1 2 3
#include "opencv2/highgui.hpp"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
Gary Bradski's avatar
Gary Bradski committed
4
#include <iostream>
5 6

using namespace cv;
Gary Bradski's avatar
Gary Bradski committed
7 8
using namespace std;

9 10 11 12 13 14 15 16
// static void help()
// {
//     cout << "\nThis program demonstrates kmeans clustering.\n"
//             "It generates an image with random points, then assigns a random number of cluster\n"
//             "centers and uses kmeans to move those cluster centers to their representitive location\n"
//             "Call\n"
//             "./kmeans\n" << endl;
// }
17

18
int main( int /*argc*/, char** /*argv*/ )
19 20 21 22 23 24 25 26 27 28
{
    const int MAX_CLUSTERS = 5;
    Scalar colorTab[] =
    {
        Scalar(0, 0, 255),
        Scalar(0,255,0),
        Scalar(255,100,100),
        Scalar(255,0,255),
        Scalar(0,255,255)
    };
29

30 31 32 33 34 35 36
    Mat img(500, 500, CV_8UC3);
    RNG rng(12345);

    for(;;)
    {
        int k, clusterCount = rng.uniform(2, MAX_CLUSTERS+1);
        int i, sampleCount = rng.uniform(1, 1001);
37
        Mat points(sampleCount, 1, CV_32FC2), labels;
38

39
        clusterCount = MIN(clusterCount, sampleCount);
40
        std::vector<Point2f> centers;
41 42 43 44 45 46 47 48 49 50

        /* generate random sample from multigaussian distribution */
        for( k = 0; k < clusterCount; k++ )
        {
            Point center;
            center.x = rng.uniform(0, img.cols);
            center.y = rng.uniform(0, img.rows);
            Mat pointChunk = points.rowRange(k*sampleCount/clusterCount,
                                             k == clusterCount - 1 ? sampleCount :
                                             (k+1)*sampleCount/clusterCount);
51
            rng.fill(pointChunk, RNG::NORMAL, Scalar(center.x, center.y), Scalar(img.cols*0.05, img.rows*0.05));
52 53 54 55
        }

        randShuffle(points, 1, &rng);

Maksim Shabunin's avatar
Maksim Shabunin committed
56
        double compactness = kmeans(points, clusterCount, labels,
57
            TermCriteria( TermCriteria::EPS+TermCriteria::COUNT, 10, 1.0),
58
               3, KMEANS_PP_CENTERS, centers);
59 60 61 62 63 64 65

        img = Scalar::all(0);

        for( i = 0; i < sampleCount; i++ )
        {
            int clusterIdx = labels.at<int>(i);
            Point ipt = points.at<Point2f>(i);
66
            circle( img, ipt, 2, colorTab[clusterIdx], FILLED, LINE_AA );
67
        }
68
        for (i = 0; i < (int)centers.size(); ++i)
Maksim Shabunin's avatar
Maksim Shabunin committed
69
        {
70
            Point2f c = centers[i];
Maksim Shabunin's avatar
Maksim Shabunin committed
71 72 73
            circle( img, c, 40, colorTab[i], 1, LINE_AA );
        }
        cout << "Compactness: " << compactness << endl;
74 75 76 77 78 79 80 81 82 83

        imshow("clusters", img);

        char key = (char)waitKey();
        if( key == 27 || key == 'q' || key == 'Q' ) // 'ESC'
            break;
    }

    return 0;
}