kmeans.py 1.82 KB
Newer Older
wester committed
1 2 3 4 5
#!/usr/bin/python
import urllib2
import cv2.cv as cv
from random import randint
MAX_CLUSTERS = 5
wester committed
6

wester committed
7
if __name__ == "__main__":
wester committed
8

wester committed
9 10 11 12 13 14 15 16
    color_tab = [
        cv.CV_RGB(255, 0,0),
        cv.CV_RGB(0, 255, 0),
        cv.CV_RGB(100, 100, 255),
        cv.CV_RGB(255, 0,255),
        cv.CV_RGB(255, 255, 0)]
    img = cv.CreateImage((500, 500), 8, 3)
    rng = cv.RNG(-1)
wester committed
17

wester committed
18
    cv.NamedWindow("clusters", 1)
wester committed
19

wester committed
20 21 22 23 24
    while True:
        cluster_count = randint(2, MAX_CLUSTERS)
        sample_count = randint(1, 1000)
        points = cv.CreateMat(sample_count, 1, cv.CV_32FC2)
        clusters = cv.CreateMat(sample_count, 1, cv.CV_32SC1)
wester committed
25

wester committed
26 27 28 29 30 31 32
        # generate random sample from multigaussian distribution
        for k in range(cluster_count):
            center = (cv.RandInt(rng)%img.width, cv.RandInt(rng)%img.height)
            first = k*sample_count/cluster_count
            last = sample_count
            if k != cluster_count:
                last = (k+1)*sample_count/cluster_count
wester committed
33

wester committed
34
            point_chunk = cv.GetRows(points, first, last)
wester committed
35

wester committed
36 37 38
            cv.RandArr(rng, point_chunk, cv.CV_RAND_NORMAL,
                       cv.Scalar(center[0], center[1], 0, 0),
                       cv.Scalar(img.width*0.1, img.height*0.1, 0, 0))
wester committed
39 40


wester committed
41 42
        # shuffle samples
        cv.RandShuffle(points, rng)
wester committed
43

wester committed
44 45
        cv.KMeans2(points, cluster_count, clusters,
                   (cv.CV_TERMCRIT_EPS + cv.CV_TERMCRIT_ITER, 10, 1.0))
wester committed
46

wester committed
47
        cv.Zero(img)
wester committed
48

wester committed
49 50 51 52
        for i in range(sample_count):
            cluster_idx = int(clusters[i, 0])
            pt = (cv.Round(points[i, 0][0]), cv.Round(points[i, 0][1]))
            cv.Circle(img, pt, 2, color_tab[cluster_idx], cv.CV_FILLED, cv.CV_AA, 0)
wester committed
53

wester committed
54 55 56 57
        cv.ShowImage("clusters", img)

        key = cv.WaitKey(0) % 0x100
        if key in [27, ord('q'), ord('Q')]:
wester committed
58
            break
wester committed
59 60

    cv.DestroyWindow("clusters")