diff python/ml.py @ 787:0a428b449b80 dev

improved script to display over world image
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Thu, 17 Mar 2016 16:01:19 -0400
parents 1f2b2d1f4fbf
children 5b970a5bc233
line wrap: on
line diff
--- a/python/ml.py	Fri Mar 11 17:38:48 2016 -0500
+++ b/python/ml.py	Thu Mar 17 16:01:19 2016 -0400
@@ -195,23 +195,26 @@
     return clusterSizes
 
 # Gaussian Mixture Models
-def plotGMMClusters(model, dataset = None, colors = utils.colors):
+def plotGMMClusters(model, dataset = None, fig = None, colors = utils.colors, nPixelsPerUnit = 1., alpha = 0.3):
     '''plot the ellipse corresponding to the Gaussians
     and the predicted classes of the instances in the dataset'''
-    fig = plt.figure()
+    if fig is None:
+        fig = plt.figure()
     labels = model.predict(dataset)
+    tmpDataset = nPixelsPerUnit*dataset
     for i in xrange(model.n_components):
-        mean = model.means_[i]
+        mean = nPixelsPerUnit*model.means_[i]
+        covariance = nPixelsPerUnit*model.covars_[i]
         if dataset is not None:
-            plt.scatter(dataset[labels == i, 0], dataset[labels == i, 1], .8, color=colors[i])
+            plt.scatter(tmpDataset[labels == i, 0], tmpDataset[labels == i, 1], .8, color=colors[i])
         plt.annotate(str(i), xy=(mean[0]+1, mean[1]+1))
 
         # Plot an ellipse to show the Gaussian component                                                  
-        v, w = np.linalg.eigh(model.covars_[i])
+        v, w = np.linalg.eigh(covariance)
         angle = np.arctan2(w[0][1], w[0][0])
         angle = 180*angle/np.pi  # convert to degrees                                             
 	v *= 4
         ell = mpl.patches.Ellipse(mean, v[0], v[1], 180+angle, color=colors[i])
         ell.set_clip_box(fig.bbox)
-        ell.set_alpha(.5)
+        ell.set_alpha(alpha)
         fig.axes[0].add_artist(ell)