python cv2 sample: GMM expectation-maximization

diff --git a/samples/python2/gaussian_mix.py b/samples/python2/gaussian_mix.py
new file mode 100644 (file)
index 0000000..6f08e88
--- /dev/null
+++ b/samples/python2/gaussian_mix.py
@@ -0,0 +1,57 @@
+import numpy as np\r
+from numpy import random\r
+import cv2, cv\r
+\r
+\r
+def make_gaussians(cluster_n, img_size):\r
+    points = []\r
+    ref_distrs = []\r
+    for i in xrange(cluster_n):\r
+        mean = (0.1 + 0.8*random.rand(2)) * img_size\r
+        a = (random.rand(2, 2)-0.5)*img_size*0.1\r
+        cov = np.dot(a.T, a) + img_size*0.05*np.eye(2)\r
+        n = 100 + random.randint(900)\r
+        pts = random.multivariate_normal(mean, cov, n)\r
+        points.append( pts )\r
+        ref_distrs.append( (mean, cov) )\r
+    points = np.float32( np.vstack(points) )\r
+    return points, ref_distrs\r
+\r
+def draw_gaussain(img, mean, cov, color):\r
+    x, y = np.int32(mean)\r
+    w, u, vt = cv2.SVDecomp(cov)\r
+    ang = np.rad2deg( np.arctan2(u[1, 0], u[0, 0]) )\r
+    s1, s2 = np.sqrt(w)*3.0\r
+    cv2.ellipse(img, (x, y), (s1, s2), ang, 0, 360, color, 1, cv.CV_AA)\r
+\r
+\r
+if __name__ == '__main__':\r
+    cluster_n = 5\r
+    img_size = 512\r
+\r
+    print 'press any key to update distributions, ESC - exit\n'\r
+\r
+    while True:\r
+        print 'sampling distributions...'\r
+        points, ref_distrs = make_gaussians(cluster_n, img_size)\r
+\r
+        print 'EM (opencv) ...'\r
+        em = cv2.EM(points, params = dict( nclusters = cluster_n, cov_mat_type = cv2.EM_COV_MAT_GENERIC) )\r
+        means = em.getMeans()\r
+        covs = np.zeros((cluster_n, 2, 2), np.float32) \r
+        covs = em.getCovs(covs)         # FIXME\r
+        found_distrs = zip(means, covs)\r
+        print 'ready!\n'\r
+\r
+        img = np.zeros((img_size, img_size, 3), np.uint8)\r
+        for x, y in np.int32(points):\r
+            cv2.circle(img, (x, y), 1, (255, 255, 255), -1)\r
+        for m, cov in ref_distrs:\r
+            draw_gaussain(img, m, cov, (0, 255, 0))\r
+        for m, cov in found_distrs:\r
+            draw_gaussain(img, m, cov, (0, 0, 255))\r
+\r
+        cv2.imshow('gaussian mixture', img)\r
+        ch = cv2.waitKey(0)\r
+        if ch == 27:\r
+            break\r