1"""watershed.py - watershed algorithm
2
3This module implements a watershed algorithm that apportions pixels into
4marked basins. The algorithm uses a priority queue to hold the pixels
5with the metric for the priority queue being pixel value, then the time
6of entry into the queue - this settles ties in favor of the closest marker.
7
8Some ideas taken from
9Soille, "Automated Basin Delineation from Digital Elevation Models Using
10Mathematical Morphology", Signal Processing 20 (1990) 171-182.
11
12The most important insight in the paper is that entry time onto the queue
13solves two problems: a pixel should be assigned to the neighbor with the
14largest gradient or, if there is no gradient, pixels on a plateau should
15be split between markers on opposite sides.
16
17Originally part of CellProfiler, code licensed under both GPL and BSD licenses.
18Website: http://www.cellprofiler.org
19
20Copyright (c) 2003-2009 Massachusetts Institute of Technology
21Copyright (c) 2009-2011 Broad Institute
22All rights reserved.
23
24Original author: Lee Kamentsky
25"""
26
27import numpy as np
28from scipy import ndimage as ndi
29
30from . import _watershed_cy
31from ..morphology.extrema import local_minima
32from ..morphology._util import (_validate_connectivity,
33                                _offsets_to_raveled_neighbors)
34from ..util import crop, regular_seeds
35
36
37def _validate_inputs(image, markers, mask, connectivity):
38    """Ensure that all inputs to watershed have matching shapes and types.
39
40    Parameters
41    ----------
42    image : array
43        The input image.
44    markers : int or array of int
45        The marker image.
46    mask : array, or None
47        A boolean mask, True where we want to compute the watershed.
48    connectivity : int in {1, ..., image.ndim}
49        The connectivity of the neighborhood of a pixel.
50
51    Returns
52    -------
53    image, markers, mask : arrays
54        The validated and formatted arrays. Image will have dtype float64,
55        markers int32, and mask int8. If ``None`` was given for the mask,
56        it is a volume of all 1s.
57
58    Raises
59    ------
60    ValueError
61        If the shapes of the given arrays don't match.
62    """
63    n_pixels = image.size
64    if mask is None:
65        # Use a complete `True` mask if none is provided
66        mask = np.ones(image.shape, bool)
67    else:
68        mask = np.asanyarray(mask, dtype=bool)
69        n_pixels = np.sum(mask)
70        if mask.shape != image.shape:
71            message = (f'`mask` (shape {mask.shape}) must have same shape '
72                       f'as `image` (shape {image.shape})')
73            raise ValueError(message)
74    if markers is None:
75        markers_bool = local_minima(image, connectivity=connectivity) * mask
76        footprint = ndi.generate_binary_structure(markers_bool.ndim, connectivity)
77        markers = ndi.label(markers_bool, structure=footprint)[0]
78    elif not isinstance(markers, (np.ndarray, list, tuple)):
79        # not array-like, assume int
80        # given int, assume that number of markers *within mask*.
81        markers = regular_seeds(image.shape,
82                                int(markers / (n_pixels / image.size)))
83        markers *= mask
84    else:
85        markers = np.asanyarray(markers) * mask
86        if markers.shape != image.shape:
87            message = (f'`markers` (shape {markers.shape}) must have same '
88                       f'shape as `image` (shape {image.shape})')
89            raise ValueError(message)
90    return (image.astype(np.float64),
91            markers.astype(np.int32),
92            mask.astype(np.int8))
93
94
95def watershed(image, markers=None, connectivity=1, offset=None, mask=None,
96              compactness=0, watershed_line=False):
97    """Find watershed basins in `image` flooded from given `markers`.
98
99    Parameters
100    ----------
101    image : ndarray (2-D, 3-D, ...)
102        Data array where the lowest value points are labeled first.
103    markers : int, or ndarray of int, same shape as `image`, optional
104        The desired number of markers, or an array marking the basins with the
105        values to be assigned in the label matrix. Zero means not a marker. If
106        ``None`` (no markers given), the local minima of the image are used as
107        markers.
108    connectivity : ndarray, optional
109        An array with the same number of dimensions as `image` whose
110        non-zero elements indicate neighbors for connection.
111        Following the scipy convention, default is a one-connected array of
112        the dimension of the image.
113    offset : array_like of shape image.ndim, optional
114        offset of the connectivity (one offset per dimension)
115    mask : ndarray of bools or 0s and 1s, optional
116        Array of same shape as `image`. Only points at which mask == True
117        will be labeled.
118    compactness : float, optional
119        Use compact watershed [3]_ with given compactness parameter.
120        Higher values result in more regularly-shaped watershed basins.
121    watershed_line : bool, optional
122        If watershed_line is True, a one-pixel wide line separates the regions
123        obtained by the watershed algorithm. The line has the label 0.
124
125    Returns
126    -------
127    out : ndarray
128        A labeled matrix of the same type and shape as markers
129
130    See Also
131    --------
132    skimage.segmentation.random_walker : random walker segmentation
133        A segmentation algorithm based on anisotropic diffusion, usually
134        slower than the watershed but with good results on noisy data and
135        boundaries with holes.
136
137    Notes
138    -----
139    This function implements a watershed algorithm [1]_ [2]_ that apportions
140    pixels into marked basins. The algorithm uses a priority queue to hold
141    the pixels with the metric for the priority queue being pixel value, then
142    the time of entry into the queue - this settles ties in favor of the
143    closest marker.
144
145    Some ideas taken from
146    Soille, "Automated Basin Delineation from Digital Elevation Models Using
147    Mathematical Morphology", Signal Processing 20 (1990) 171-182
148
149    The most important insight in the paper is that entry time onto the queue
150    solves two problems: a pixel should be assigned to the neighbor with the
151    largest gradient or, if there is no gradient, pixels on a plateau should
152    be split between markers on opposite sides.
153
154    This implementation converts all arguments to specific, lowest common
155    denominator types, then passes these to a C algorithm.
156
157    Markers can be determined manually, or automatically using for example
158    the local minima of the gradient of the image, or the local maxima of the
159    distance function to the background for separating overlapping objects
160    (see example).
161
162    References
163    ----------
164    .. [1] https://en.wikipedia.org/wiki/Watershed_%28image_processing%29
165
166    .. [2] http://cmm.ensmp.fr/~beucher/wtshed.html
167
168    .. [3] Peer Neubert & Peter Protzel (2014). Compact Watershed and
169           Preemptive SLIC: On Improving Trade-offs of Superpixel Segmentation
170           Algorithms. ICPR 2014, pp 996-1001. :DOI:`10.1109/ICPR.2014.181`
171           https://www.tu-chemnitz.de/etit/proaut/publications/cws_pSLIC_ICPR.pdf
172
173    Examples
174    --------
175    The watershed algorithm is useful to separate overlapping objects.
176
177    We first generate an initial image with two overlapping circles:
178
179    >>> x, y = np.indices((80, 80))
180    >>> x1, y1, x2, y2 = 28, 28, 44, 52
181    >>> r1, r2 = 16, 20
182    >>> mask_circle1 = (x - x1)**2 + (y - y1)**2 < r1**2
183    >>> mask_circle2 = (x - x2)**2 + (y - y2)**2 < r2**2
184    >>> image = np.logical_or(mask_circle1, mask_circle2)
185
186    Next, we want to separate the two circles. We generate markers at the
187    maxima of the distance to the background:
188
189    >>> from scipy import ndimage as ndi
190    >>> distance = ndi.distance_transform_edt(image)
191    >>> from skimage.feature import peak_local_max
192    >>> max_coords = peak_local_max(distance, labels=image,
193    ...                             footprint=np.ones((3, 3)))
194    >>> local_maxima = np.zeros_like(image, dtype=bool)
195    >>> local_maxima[tuple(max_coords.T)] = True
196    >>> markers = ndi.label(local_maxima)[0]
197
198    Finally, we run the watershed on the image and markers:
199
200    >>> labels = watershed(-distance, markers, mask=image)
201
202    The algorithm works also for 3-D images, and can be used for example to
203    separate overlapping spheres.
204    """
205    image, markers, mask = _validate_inputs(image, markers, mask, connectivity)
206    connectivity, offset = _validate_connectivity(image.ndim, connectivity,
207                                                  offset)
208
209    # pad the image, markers, and mask so that we can use the mask to
210    # keep from running off the edges
211    pad_width = [(p, p) for p in offset]
212    image = np.pad(image, pad_width, mode='constant')
213    mask = np.pad(mask, pad_width, mode='constant').ravel()
214    output = np.pad(markers, pad_width, mode='constant')
215
216    flat_neighborhood = _offsets_to_raveled_neighbors(
217        image.shape, connectivity, center=offset)
218    marker_locations = np.flatnonzero(output)
219    image_strides = np.array(image.strides, dtype=np.intp) // image.itemsize
220
221    _watershed_cy.watershed_raveled(image.ravel(),
222                                    marker_locations, flat_neighborhood,
223                                    mask, image_strides, compactness,
224                                    output.ravel(),
225                                    watershed_line)
226
227    output = crop(output, pad_width, copy=True)
228
229    return output
230