Added iterator_ij to iterate over pixel coordinates.

spectral/algorithms/algorithms.py

@@ -149,6 +149,55 @@ def iterator(image, mask=None, index=None):
         return ImageIterator(image)
 
 
+def iterator_ij(image, mask=None, index=None):
+    '''
+    Returns an iterator over image pixel coordinates for a given mask.
+
+    Arguments:
+
+        `image` (ndarray or :class:`spectral.Image`):
+
+            An image over whose pixels will be iterated.
+
+        `mask` (ndarray) [default None]:
+
+            An array of integers that specify over which pixels in `image`
+            iteration should be performed.
+
+        `index` (int) [default None]:
+
+            Specifies which value in `mask` should be used for iteration.
+
+    Returns:
+
+        An iterator over image pixel coordinates. Each returned item is a
+        2-tuple of the form (row, col).
+
+    If neither `mask` nor `index` are defined, iteration is performed over all
+    pixels.  If `mask` (but not `index`) is defined, iteration is performed
+    over all pixels for which `mask` is nonzero.  If both `mask` and `index`
+    are defined, iteration is performed over all pixels `image[i,j]` for which
+    `mask[i,j] == index`.
+    '''
+
+    if mask is None and index is None:
+        # Iterate over all pixels
+        (nrows, ncols) = image.shape[:2]
+        for r in range(nrows):
+            for c in range(ncols):
+                yield (r, c)
+    else:
+        if mask.shape != image.shape[:len(mask.shape)]:
+            raise ValueError('Mask shape does not match image.')
+
+        if index is None:
+            mask = mask != 0
+        else:
+            mask = mask == index
+        for rc in np.argwhere(mask):
+            yield tuple(rc)
+
+
 def mean_cov(image, mask=None, index=None):
     '''
     Return the mean and covariance of the set of vectors.

spectral/tests/iterators.py

@@ -78,6 +78,28 @@ def test_iterator_index(self):
         itsum = np.sum(np.array([x for x in iterator(data, self.gt, cls)]), 0)
         assert_allclose(sum, itsum)
 
+    def test_iterator_ij_nonzero(self):
+        '''Iteration over all non-background pixels.'''
+        from spectral.algorithms.algorithms import iterator_ij
+        data = self.image.load()
+        classes = self.gt.ravel()
+        pixels = data.reshape((-1, data.shape[-1]))
+        sum = np.sum(pixels[classes > 0], 0)
+        itsum = np.sum(np.array([data[ij] for ij in iterator_ij(data,
+                                                                self.gt)]), 0)
+        assert_allclose(sum, itsum)
+
+    def test_iterator_ij_index(self):
+        '''Iteration over single ground truth index'''
+        from spectral.algorithms.algorithms import iterator_ij
+        cls = 5
+        data = self.image.load()
+        classes = self.gt.ravel()
+        pixels = data.reshape((-1, data.shape[-1]))
+        sum = np.sum(pixels[classes == cls], 0)
+        itsum = np.sum(np.array([data[ij] for ij in iterator_ij(data, self.gt, cls)]), 0)
+        assert_allclose(sum, itsum)
+
     def test_iterator_spyfile(self):
         '''Iteration over SpyFile object for single ground truth index'''
         from spectral.algorithms.algorithms import iterator