Skip to content

Commit

Permalink
Merge pull request #2211 from jni/compact-watershed
Browse files Browse the repository at this point in the history
Add compact watershed and clean up existing watershed
  • Loading branch information
sciunto committed Aug 27, 2016
2 parents c369eac + d34f7ab commit 74c6b8d
Show file tree
Hide file tree
Showing 9 changed files with 377 additions and 293 deletions.
43 changes: 43 additions & 0 deletions doc/examples/segmentation/plot_compact_watershed.py
Original file line number Diff line number Diff line change
@@ -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()
52 changes: 40 additions & 12 deletions doc/examples/segmentation/plot_segmentations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
121 changes: 74 additions & 47 deletions skimage/morphology/_watershed.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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)
5 changes: 3 additions & 2 deletions skimage/morphology/heap_watershed.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
38 changes: 34 additions & 4 deletions skimage/morphology/tests/test_watershed.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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()

0 comments on commit 74c6b8d

Please sign in to comment.