Merge pull request #2211 from jni/compact-watershed

Add compact watershed and clean up existing watershed

doc/examples/segmentation/plot_compact_watershed.py

@@ -0,0 +1,43 @@
+"""
+=============================================
+Find Regular Segments Using Compact Watershed
+=============================================
+
+The watershed transform is commonly used as a starting point for many
+segmentation algorithms. However, without a judicious choice of seeds, it
+can produce very uneven fragment sizes, which can be difficult to deal with
+in downstream analyses.
+
+The *compact* watershed transform remedies this by favoring seeds that are
+close to the pixel being considered.
+
+Both algorithms are implemented in the :py:func:`skimage.morphology.watershed`
+function. To use the compact form, simply pass a ``compactness`` value greater
+than 0.
+"""
+
+import numpy as np
+from skimage import data, util, filters, color
+from skimage.morphology import watershed
+import matplotlib.pyplot as plt
+
+coins = data.coins()
+edges = filters.sobel(coins)
+
+grid = util.regular_grid(coins.shape, n_points=468)
+
+seeds = np.zeros(coins.shape, dtype=int)
+seeds[grid] = np.arange(seeds[grid].size).reshape(seeds[grid].shape) + 1
+
+w0 = watershed(edges, seeds)
+w1 = watershed(edges, seeds, compactness=0.01)
+
+fig, (ax0, ax1) = plt.subplots(1, 2)
+
+ax0.imshow(color.label2rgb(w0, coins))
+ax0.set_title('Classical watershed')
+
+ax1.imshow(color.label2rgb(w1, coins))
+ax1.set_title('Compact watershed')
+
+plt.show()

doc/examples/segmentation/plot_segmentations.py

@@ -51,46 +51,74 @@
 image location and is therefore closely related to quickshift. As the
 clustering method is simpler, it is very efficient. It is essential for this
 algorithm to work in Lab color space to obtain good results.  The algorithm
-quickly gained momentum and is now widely used. See [3] for details.  The
+quickly gained momentum and is now widely used. See [3]_ for details.  The
 ``compactness`` parameter trades off color-similarity and proximity, as in the
 case of Quickshift, while ``n_segments`` chooses the number of centers for
 kmeans.
 
 .. [3] Radhakrishna Achanta, Appu Shaji, Kevin Smith, Aurelien Lucchi,
     Pascal Fua, and Sabine Suesstrunk, SLIC Superpixels Compared to
     State-of-the-art Superpixel Methods, TPAMI, May 2012.
+
+
+Compact watershed segmentation of gradient images
+-------------------------------------------------
+
+Instead of taking a color image as input, watershed requires a grayscale
+*gradient* image, where bright pixels denote a boundary between regions.
+The algorithm views the image as a landscape, with bright pixels forming high
+peaks. This landscape is then flooded from the given *markers*, until separate
+flood basins meet at the peaks. Each distinct basin then forms a different
+image segment. [4]_
+
+As with SLIC, there is an additional *compactness* argument that makes it
+harder for markers to flood faraway pixels. This makes the watershed regions
+more regularly shaped. [5]_
+
+.. [4] http://en.wikipedia.org/wiki/Watershed_%28image_processing%29
+
+.. [5] Peer Neubert & Peter Protzel (2014). Compact Watershed and
+       Preemptive SLIC: On Improving Trade-offs of Superpixel Segmentation
+       Algorithms. ICPR 2014, pp 996-1001. DOI:10.1109/ICPR.2014.181
+       https://www.tu-chemnitz.de/etit/proaut/forschung/rsrc/cws_pSLIC_ICPR.pdf
 """
 from __future__ import print_function
 
 import matplotlib.pyplot as plt
 import numpy as np
 
 from skimage.data import astronaut
-from skimage.segmentation import felzenszwalb, slic, quickshift
+from skimage.color import rgb2gray
+from skimage.filters import sobel
+from skimage.segmentation import felzenszwalb, slic, quickshift, watershed
 from skimage.segmentation import mark_boundaries
 from skimage.util import img_as_float
 
 img = img_as_float(astronaut()[::2, ::2])
 segments_fz = felzenszwalb(img, scale=100, sigma=0.5, min_size=50)
 segments_slic = slic(img, n_segments=250, compactness=10, sigma=1)
 segments_quick = quickshift(img, kernel_size=3, max_dist=6, ratio=0.5)
+gradient = sobel(rgb2gray(img))
+segments_watershed = watershed(gradient, markers=250, compactness=0.001)
 
 print("Felzenszwalb's number of segments: %d" % len(np.unique(segments_fz)))
-print("Slic number of segments: %d" % len(np.unique(segments_slic)))
-print("Quickshift number of segments: %d" % len(np.unique(segments_quick)))
+print('SLIC number of segments: %d' % len(np.unique(segments_slic)))
+print('Quickshift number of segments: %d' % len(np.unique(segments_quick)))
 
-fig, ax = plt.subplots(1, 3, sharex=True, sharey=True,
+fig, ax = plt.subplots(2, 2, sharex=True, sharey=True,
                        subplot_kw={'adjustable': 'box-forced'})
 fig.set_size_inches(8, 3, forward=True)
 fig.tight_layout()
 
-ax[0].imshow(mark_boundaries(img, segments_fz))
-ax[0].set_title("Felzenszwalbs's method")
-ax[1].imshow(mark_boundaries(img, segments_slic))
-ax[1].set_title("SLIC")
-ax[2].imshow(mark_boundaries(img, segments_quick))
-ax[2].set_title("Quickshift")
-for a in ax:
+ax[0, 0].imshow(mark_boundaries(img, segments_fz))
+ax[0, 0].set_title("Felzenszwalbs's method")
+ax[0, 1].imshow(mark_boundaries(img, segments_slic))
+ax[0, 1].set_title('SLIC')
+ax[1, 0].imshow(mark_boundaries(img, segments_quick))
+ax[1, 0].set_title('Quickshift')
+ax[1, 1].imshow(mark_boundaries(img, segments_watershed))
+ax[1, 1].set_title('Compact watershed')
+for a in ax.ravel():
     a.set_xticks(())
     a.set_yticks(())
 plt.show()

skimage/morphology/_watershed.pyx

@@ -10,87 +10,114 @@ All rights reserved.
 Original author: Lee Kamentsky
 """
 import numpy as np
-cimport numpy as np
+cimport numpy as cnp
 cimport cython
 
 
-ctypedef np.int32_t DTYPE_INT32_t
-DTYPE_BOOL = np.bool
-ctypedef np.int8_t DTYPE_BOOL_t
+ctypedef cnp.int32_t DTYPE_INT32_t
+ctypedef cnp.int8_t DTYPE_BOOL_t
 
 
 include "heap_watershed.pxi"
 
 
+from libc.math cimport sqrt
+
+
+@cython.boundscheck(False)
+@cython.cdivision(True)
+@cython.overflowcheck(False)
+@cython.unraisable_tracebacks(False)
+cdef inline double _euclid_dist(cnp.int32_t pt0, cnp.int32_t pt1,
+                                cnp.int32_t[::1] strides):
+    """Return the Euclidean distance between raveled points pt0 and pt1."""
+    cdef double result = 0
+    cdef double curr = 0
+    for i in range(strides.shape[0]):
+        curr = (pt0 // strides[i]) - (pt1 // strides[i])
+        result += curr * curr
+        pt0 = pt0 % strides[i]
+        pt1 = pt1 % strides[i]
+    return sqrt(result)
+
+
 @cython.boundscheck(False)
-def watershed(DTYPE_INT32_t[::1] image,
-              DTYPE_INT32_t[:, ::1] pq,
-              Py_ssize_t age,
-              DTYPE_INT32_t[:, ::1] structure,
-              DTYPE_BOOL_t[::1] mask,
-              DTYPE_INT32_t[::1] output):
-    """Do heavy lifting of watershed algorithm
+def watershed_raveled(cnp.float64_t[::1] image,
+                      DTYPE_INT32_t[::1] marker_locations,
+                      DTYPE_INT32_t[::1] structure,
+                      DTYPE_BOOL_t[::1] mask,
+                      cnp.int32_t[::1] strides,
+                      cnp.double_t compactness,
+                      DTYPE_INT32_t[::1] output):
+    """Perform watershed algorithm using a raveled image and neighborhood.
 
     Parameters
     ----------
 
-    image - the flattened image pixels, converted to rank-order
-    pq    - the priority queue, starts with the marked pixels
-            the first element in each row is the image intensity
-            the second element is the age at entry into the queue
-            the third element is the index into the flattened image or labels
-            the remaining elements are the coordinates of the point
-    age   - the next age to assign to a pixel
-    structure - a numpy int32 array containing the structuring elements
-                that define nearest neighbors. For each row, the first
-                element is the stride from the point to its neighbor
-                in a flattened array. The remaining elements are the
-                offsets from the point to its neighbor in the various
-                dimensions
-    mask  - numpy boolean (char) array indicating which pixels to consider
-            and which to ignore. Also flattened.
-    output - put the image labels in here
+    image : array of float
+        The flattened image pixels.
+    marker_locations : array of int
+        The raveled coordinates of the initial markers (aka seeds) for the
+        watershed. NOTE: these should *all* point to nonzero entries in the
+        output, or the algorithm will never terminate and blow up your memory!
+    structure : array of int
+        A list of coordinate offsets to compute the raveled coordinates of each
+        neighbor from the raveled coordinates of the current pixel.
+    mask : array of int
+        An array of the same shape as `image` where each pixel contains a
+        nonzero value if it is to be considered for flooding with watershed,
+        zero otherwise. NOTE: it is *essential* that the border pixels (those
+        with neighbors falling outside the volume) are all set to zero, or
+        segfaults could occur.
+    strides : array of int
+        An array representing the number of steps to move along each dimension.
+        This is used in computing the Euclidean distance between raveled
+        indices.
+    compactness : float
+        A value greater than 0 implements the compact watershed algorithm
+        (see .py file).
+    output : array of int
+        The output array, which must already contain nonzero entries at all the
+        seed locations.
     """
     cdef Heapitem elem
     cdef Heapitem new_elem
     cdef Py_ssize_t nneighbors = structure.shape[0]
     cdef Py_ssize_t i = 0
+    cdef Py_ssize_t age = 1
     cdef Py_ssize_t index = 0
-    cdef Py_ssize_t old_index = 0
-    cdef Py_ssize_t max_index = image.shape[0]
 
     cdef Heap *hp = <Heap *> heap_from_numpy2()
 
-    for i in range(pq.shape[0]):
-        elem.value = pq[i, 0]
-        elem.age = pq[i, 1]
-        elem.index = pq[i, 2]
+    for i in range(marker_locations.shape[0]):
+        index = marker_locations[i]
+        elem.value = image[index]
+        elem.age = 0
+        elem.index = index
+        elem.source = index
         heappush(hp, &elem)
 
     while hp.items > 0:
-        #
-        # Pop off an item to work on
-        #
         heappop(hp, &elem)
-        ####################################################
-        # loop through each of the structuring elements
-        #
-        old_index = elem.index
+        if output[elem.index] and elem.index != elem.source:
+            # non-marker, already visited from another neighbor
+            continue
+        output[elem.index] = output[elem.source]
         for i in range(nneighbors):
             # get the flattened address of the neighbor
-            index = structure[i, 0] + old_index
-            if index < 0 or index >= max_index or output[index] or \
-                    not mask[index]:
+            index = structure[i] + elem.index
+            if output[index] or not mask[index]:
+                # previously visited, masked, or border pixel
                 continue
 
             age += 1
             new_elem.value = image[index]
+            if compactness > 0:
+                new_elem.value += (compactness *
+                                   _euclid_dist(index, elem.source, strides))
             new_elem.age = age
             new_elem.index = index
+            new_elem.source = elem.source
 
-            output[index] = output[old_index]
-            #
-            # Push the neighbor onto the heap to work on it later
-            #
             heappush(hp, &new_elem)
     heap_done(hp)

skimage/morphology/heap_watershed.pxi

@@ -13,13 +13,14 @@ cimport numpy as cnp
 
 
 cdef struct Heapitem:
-    cnp.int32_t value
+    cnp.float64_t value
     cnp.int32_t age
     Py_ssize_t index
+    Py_ssize_t source
 
 
 cdef inline int smaller(Heapitem *a, Heapitem *b):
-    if a.value <> b.value:
+    if a.value != b.value:
         return a.value < b.value
     return a.age < b.age
 

skimage/morphology/tests/test_watershed.py

@@ -48,7 +48,7 @@
 import numpy as np
 from scipy import ndimage as ndi
 
-from skimage.morphology.watershed import watershed, _slow_watershed
+from skimage.morphology.watershed import watershed
 
 eps = 1e-12
 
@@ -114,9 +114,6 @@ def test_watershed01(self):
                       [-1, -1, -1, -1, -1, -1, -1]])
         error = diff(expected, out)
         assert error < eps
-        out = _slow_watershed(data, markers, 8)
-        error = diff(expected, out)
-        assert error < eps
 
     def test_watershed02(self):
         "watershed 2"
@@ -426,5 +423,38 @@ def test_watershed11(self):
         dmin = np.min(d, 2)
         self.assertTrue(np.all(d[i, j, out[i, j]-1] == dmin))
 
+
+def test_compact_watershed():
+    image = np.zeros((5, 6))
+    image[:, 3:] = 1
+    seeds = np.zeros((5, 6), dtype=int)
+    seeds[2, 0] = 1
+    seeds[2, 3] = 2
+    compact = watershed(image, seeds, compactness=0.01)
+    expected = np.array([[1, 1, 1, 2, 2, 2],
+                         [1, 1, 1, 2, 2, 2],
+                         [1, 1, 1, 2, 2, 2],
+                         [1, 1, 1, 2, 2, 2],
+                         [1, 1, 1, 2, 2, 2]], dtype=int)
+    np.testing.assert_equal(compact, expected)
+    normal = watershed(image, seeds)
+    expected = np.ones(image.shape, dtype=int)
+    expected[2, 3:] = 2
+    np.testing.assert_equal(normal, expected)
+
+
+def test_numeric_seed_watershed():
+    """Test that passing just the number of seeds to watershed works."""
+    image = np.zeros((5, 6))
+    image[:, 3:] = 1
+    compact = watershed(image, 2, compactness=0.01)
+    expected = np.array([[1, 1, 1, 1, 2, 2],
+                         [1, 1, 1, 1, 2, 2],
+                         [1, 1, 1, 1, 2, 2],
+                         [1, 1, 1, 1, 2, 2],
+                         [1, 1, 1, 1, 2, 2]], dtype=np.int32)
+    np.testing.assert_equal(compact, expected)
+
+
 if __name__ == "__main__":
     np.testing.run_module_suite()