ImProc ass3: Added K-means clustering to recognize Waldo's location.

improc/ass3/.gitignore

 *.txt
+*.dat

improc/ass3/back_projection.py

 #!/usr/bin/env python
-from numpy import zeros
+from math import sqrt
+from numpy import array, zeros
 from matplotlib.pyplot import subplot, imread, imshow, show, plot
-from intersect import col2bin, domainIterator, colHist, histogramIntersect
 from scipy.ndimage import correlate
+from scipy.cluster.vq import kmeans
+
+from intersect import col2bin, domainIterator, colHist, histogramIntersect
 
 def convolution(image, radius):
     """Calculate the convolution of an image with a specified circle radius."""
@@ -49,38 +52,61 @@ def hbp(image, environment, bins, model, radius):
     print 'Creating convolution...'
     return convolution(b, radius)
 
-def loc(image, color):
-    for p in domainIterator(image):
-        if (image[p] == color).all():
-            return p
+def find_peaks(b, threshold, d_sq):
+    """Find the location of the peak value of a back projection histogram."""
+    # Find all pixels with a value higher than a threshold of the maximum
+    # value. Collect K-means estimators on-the-fly.
+    threshold *= b.max()
+    means = []
+    use = []
+
+    for p in domainIterator(b):
+        if b[p] >= threshold:
+            use.append(p)
+            found = False
+
+            for mean in means:
+                if (mean[0] - p[0]) ** 2 + (mean[1] - p[1]) ** 2 < d_sq:
+                    found = True
+                    break
 
-    return (0, 0)
+            if not found:
+                means.append(p)
 
-def find_image(image, environment, bins, model, radius):
-    """Find the location of the peak value of a back projection histogram."""
-    b = hbp(image, environment, bins, model, radius)
+    # Use K-means to identify possible matches
+    m = kmeans(array(use), array(means))[0]
 
-    return loc(b, b.max())
+    return [m[i].tolist() for i in xrange(m.shape[0])]
 
 if __name__ == '__main__':
     print 'Reading images...'
     waldo = imread('waldo.tiff')
     env = imread('waldo_env.tiff')
 
-    b = hbp(waldo, env, [64] * 3, 'rgb', 10)
+    import pickle
+    b = hbp(waldo, env, [64] * 3, 'rgb', 4)
+    pickle.dump(b, open('projection.dat', 'w'))
+    #b = pickle.load(open('projection.dat', 'r'))
+
+    def plt(peaks):
+        """Draw a rectangle around a list of center pixels."""
+        w, h = waldo.shape[:2]
+
+        for x, y in peaks:
+            l = x - w / 2
+            r = x + w / 2
+            b = y - h / 2
+            t = y + h / 2
+            plot([t, t, b, b, t], [l, r, r, l, l], 'r-')
+
+    w, h = waldo.shape[:2]
+
+    peaks = find_peaks(b, .28, w ** 2 + h ** 2)
+
     subplot(121)
+    plt(peaks)
     imshow(env, origin='lower')
     subplot(122)
+    plt(peaks)
     imshow(b, origin='lower')
-
-    # Draw a rectangle around the found center pixel
-    #x, y = find_image(waldo, env, [64] * 3, 'rgb', 10)
-    #w, h = waldo.shape[:2]
-    #l = x - w / 2
-    #r = x + w / 2
-    #b = y - h / 2
-    #t = y + h / 2
-    #plot([t, t, b, b], [l, r, r, l], 'r-')
-    #imshow(env, origin='lower')
-
     show()