← Back to branch summary

~paparazzi-uav/paparazzi/v5.0-manual

« back to all changes in this revision

Viewing changes to sw/ext/opencv_bebop/opencv/samples/python/gaussian_mix.py

  • Committer: Paparazzi buildbot
  • Date: 2016-05-18 15:00:29 UTC
  • Revision ID: felix.ruess+docbot@gmail.com-20160518150029-e8lgzi5kvb4p7un9
Manual import commit 4b8bbb730080dac23cf816b98908dacfabe2a8ec from v5.0 branch.

Show diffs side-by-side

added added

removed removed

Lines of Context:
sw/ext/opencv_bebop/opencv/samples/python/gaussian_mix.py
 
1
#!/usr/bin/env python
 
2
 
 
3
# Python 2/3 compatibility
 
4
from __future__ import print_function
 
5
import sys
 
6
PY3 = sys.version_info[0] == 3
 
7
 
 
8
if PY3:
 
9
    xrange = range
 
10
 
 
11
import numpy as np
 
12
from numpy import random
 
13
import cv2
 
14
 
 
15
def make_gaussians(cluster_n, img_size):
 
16
    points = []
 
17
    ref_distrs = []
 
18
    for i in xrange(cluster_n):
 
19
        mean = (0.1 + 0.8*random.rand(2)) * img_size
 
20
        a = (random.rand(2, 2)-0.5)*img_size*0.1
 
21
        cov = np.dot(a.T, a) + img_size*0.05*np.eye(2)
 
22
        n = 100 + random.randint(900)
 
23
        pts = random.multivariate_normal(mean, cov, n)
 
24
        points.append( pts )
 
25
        ref_distrs.append( (mean, cov) )
 
26
    points = np.float32( np.vstack(points) )
 
27
    return points, ref_distrs
 
28
 
 
29
def draw_gaussain(img, mean, cov, color):
 
30
    x, y = np.int32(mean)
 
31
    w, u, vt = cv2.SVDecomp(cov)
 
32
    ang = np.arctan2(u[1, 0], u[0, 0])*(180/np.pi)
 
33
    s1, s2 = np.sqrt(w)*3.0
 
34
    cv2.ellipse(img, (x, y), (s1, s2), ang, 0, 360, color, 1, cv2.LINE_AA)
 
35
 
 
36
 
 
37
if __name__ == '__main__':
 
38
    cluster_n = 5
 
39
    img_size = 512
 
40
 
 
41
    print('press any key to update distributions, ESC - exit\n')
 
42
 
 
43
    while True:
 
44
        print('sampling distributions...')
 
45
        points, ref_distrs = make_gaussians(cluster_n, img_size)
 
46
 
 
47
        print('EM (opencv) ...')
 
48
        em = cv2.ml.EM_create()
 
49
        em.setClustersNumber(cluster_n)
 
50
        em.setCovarianceMatrixType(cv2.ml.EM_COV_MAT_GENERIC)
 
51
        em.trainEM(points)
 
52
        means = em.getMeans()
 
53
        covs = em.getCovs()  # Known bug: https://github.com/Itseez/opencv/pull/4232
 
54
        found_distrs = zip(means, covs)
 
55
        print('ready!\n')
 
56
 
 
57
        img = np.zeros((img_size, img_size, 3), np.uint8)
 
58
        for x, y in np.int32(points):
 
59
            cv2.circle(img, (x, y), 1, (255, 255, 255), -1)
 
60
        for m, cov in ref_distrs:
 
61
            draw_gaussain(img, m, cov, (0, 255, 0))
 
62
        for m, cov in found_distrs:
 
63
            draw_gaussain(img, m, cov, (0, 0, 255))
 
64
 
 
65
        cv2.imshow('gaussian mixture', img)
 
66
        ch = 0xFF & cv2.waitKey(0)
 
67
        if ch == 27:
 
68
            break
 
69
    cv2.destroyAllWindows()