git.maemo.org Git - opencv/blob - samples/swig_python/kmeans.py

   1 #!/usr/bin/python
   2 from opencv.cv import *
   3 from opencv.highgui import *
   4 from random import randint
   5 MAX_CLUSTERS = 5
   6
   7 if __name__ == "__main__":
   8
   9     color_tab = [
  10         CV_RGB(255,0,0),
  11         CV_RGB(0,255,0),
  12         CV_RGB(100,100,255),
  13         CV_RGB(255,0,255),
  14         CV_RGB(255,255,0)]
  15     img = cvCreateImage( cvSize( 500, 500 ), 8, 3 )
  16     rng = cvRNG(-1)
  17
  18     cvNamedWindow( "clusters", 1 )
  19
  20     while True:
  21         cluster_count = randint(2, MAX_CLUSTERS)
  22         sample_count = randint(1, 1000)
  23         points = cvCreateMat( sample_count, 1, CV_32FC2 )
  24         clusters = cvCreateMat( sample_count, 1, CV_32SC1 )
  25
  26         # generate random sample from multigaussian distribution
  27         for k in range(cluster_count):
  28             center = CvPoint()
  29             center.x = cvRandInt(rng)%img.width
  30             center.y = cvRandInt(rng)%img.height
  31             first = k*sample_count/cluster_count
  32             last = sample_count
  33             if k != cluster_count:
  34                 last = (k+1)*sample_count/cluster_count
  35
  36             point_chunk = cvGetRows(points, first, last)
  37
  38             cvRandArr( rng, point_chunk, CV_RAND_NORMAL,
  39                        cvScalar(center.x,center.y,0,0),
  40                        cvScalar(img.width*0.1,img.height*0.1,0,0))
  41
  42
  43         # shuffle samples
  44         cvRandShuffle( points, rng )
  45
  46         cvKMeans2( points, cluster_count, clusters,
  47                    cvTermCriteria( CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, 10, 1.0 ))
  48
  49         cvZero( img )
  50
  51         for i in range(sample_count):
  52             cluster_idx = clusters[i]
  53             # a multi channel matrix access returns a scalar of
  54             #dimension 4,0, which is not considerate a cvPoint
  55             #we have to create a tuple with the first two elements
  56             pt = (cvRound(points[i][0]), cvRound(points[i][1]))
  57             cvCircle( img, pt, 2, color_tab[cluster_idx], CV_FILLED, CV_AA, 0 )
  58
  59         cvShowImage( "clusters", img )
  60
  61         key = cvWaitKey(0)
  62         if( key == 27 or key == 'q' or key == 'Q' ): # 'ESC'
  63             break
  64
  65
  66     cvDestroyWindow( "clusters" )