ImProc ass3: Added mask option to colHist.

improc/ass3/back_projection.py

@@ -7,6 +7,32 @@ from scipy.cluster.vq import kmeans
 
 from intersect import col2bin, domainIterator, colHist, histogramIntersect
 
+def find_peaks(b, threshold, d_sq):
+    """Find the location of the peak values 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
+
+            if not found:
+                means.append(p)
+
+    # Use K-means to identify possible matches
+    m = kmeans(array(use), array(means))[0]
+
+    return [m[i].tolist() for i in xrange(m.shape[0])]
+
 def convolution(image, radius):
     """Calculate the convolution of an image with a specified circle radius."""
     # Loop to the square that surrounds the circle, and check if the pixel
@@ -20,11 +46,14 @@ def convolution(image, radius):
 
     return correlate(image, mask, mode='nearest')
 
-def hbp(image, environment, bins, model, radius):
+def hbp(image, environment, bins, model, radius, **kwargs):
     """Create the histogram back projection of two images."""
+    options = dict(mask=None)
+    options.update(kwargs)
+
     # Create image histograms
     print 'Creating histograms...'
-    M = colHist(image, bins, model)
+    M = colHist(image, bins, model, mask=options['mask'])
     I = colHist(environment, bins, model)
 
     # Create ratio histogram
@@ -52,39 +81,24 @@ def hbp(image, environment, bins, model, radius):
     print 'Creating convolution...'
     return convolution(b, radius)
 
-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
-
-            if not found:
-                means.append(p)
+def exclude_color(color, image):
+    mask = zeros(image.shape[:2], dtype=int)
+    color = array(color)
 
-    # Use K-means to identify possible matches
-    m = kmeans(array(use), array(means))[0]
+    for p in domainIterator(image):
+        if (image[p] != color).any():
+            mask[p] = 1
 
-    return [m[i].tolist() for i in xrange(m.shape[0])]
+    return mask
 
 if __name__ == '__main__':
     print 'Reading images...'
     waldo = imread('waldo.tiff')
     env = imread('waldo_env.tiff')
+    mask = exclude_color([255] * 3, waldo)
 
     import pickle
-    b = hbp(waldo, env, [64] * 3, 'rgb', 4)
+    b = hbp(waldo, env, [64] * 3, 'rgb', 2, mask=mask)
     pickle.dump(b, open('projection.dat', 'w'))
     #b = pickle.load(open('projection.dat', 'r'))
 
@@ -101,7 +115,9 @@ if __name__ == '__main__':
 
     w, h = waldo.shape[:2]
 
-    peaks = find_peaks(b, .28, w ** 2 + h ** 2)
+    print 'Locating peaks...'
+    peaks = find_peaks(b, .2, w ** 2 + h ** 2)
+    print 'Done'
 
     subplot(121)
     plt(peaks)

improc/ass3/intersect.py

@@ -18,22 +18,23 @@ def domainIterator(image, dim=2):
                 for z in xrange(image.shape[2]):
                     yield x, y, z
 
-def colHist(image, bins, model):
+def colHist(image, bins, model, **kwargs):
     """Create the color histogram of an image."""
     h = zeros(bins, dtype=int)
     use = image.astype(float) * map(lambda x: x - 1, bins)
 
     if model == 'rgb':
         use /= 255
-    elif model == 'rgba':
-        pass
     elif model == 'hsv':
         # TODO: implement HSV color model
         pass
     else:
         raise ValueError('Color model "%s" is not supported.' % model)
 
+    mask = kwargs['mask'] if 'mask' in kwargs else None
+
     for p in domainIterator(image):
+        if mask is None or mask[p].any():
             h[col2bin(use[p])] += 1
 
     return h