_watershed.py 9.6 KB

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