├── ex.png ├── ws.png ├── README.md ├── LICENSE └── Watershed.py /ex.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mzur/watershed/HEAD/ex.png -------------------------------------------------------------------------------- /ws.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mzur/watershed/HEAD/ws.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Watershed 2 | 3 | A simple (but not very fast) Python implementation of [Determining watersheds in digital pictures via flooding simulations](http://dx.doi.org/10.1117/12.24211). 4 | 5 | ![source image](ex.png) ![labelled image](ws.png) 6 | 7 | In contrast to [`skimage.morphology.watershed`](http://scikit-image.org/docs/dev/api/skimage.morphology.html#skimage.morphology.watershed) and [`cv2.watershed`](https://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_watershed/py_watershed.html) this implementation does not use marker seeds. 8 | 9 | ## Usage 10 | 11 | ```python 12 | import numpy as np 13 | from Watershed import Watershed 14 | from PIL import Image 15 | import matplotlib.pyplot as plt 16 | 17 | w = Watershed() 18 | image = np.array(Image.open('ex.png')) 19 | labels = w.apply(image) 20 | plt.imshow(labels, cmap='Paired', interpolation='nearest') 21 | plt.show() 22 | ``` 23 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Martin Zurowietz 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Watershed.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from collections import deque 3 | 4 | # Implementation of: 5 | # Pierre Soille, Luc M. Vincent, "Determining watersheds in digital pictures via 6 | # flooding simulations", Proc. SPIE 1360, Visual Communications and Image Processing 7 | # '90: Fifth in a Series, (1 September 1990); doi: 10.1117/12.24211; 8 | # http://dx.doi.org/10.1117/12.24211 9 | class Watershed(object): 10 | MASK = -2 11 | WSHD = 0 12 | INIT = -1 13 | INQE = -3 14 | 15 | def __init__(self, levels = 256): 16 | self.levels = levels 17 | 18 | # Neighbour (coordinates of) pixels, including the given pixel. 19 | def _get_neighbors(self, height, width, pixel): 20 | return np.mgrid[ 21 | max(0, pixel[0] - 1):min(height, pixel[0] + 2), 22 | max(0, pixel[1] - 1):min(width, pixel[1] + 2) 23 | ].reshape(2, -1).T 24 | 25 | def apply(self, image): 26 | current_label = 0 27 | flag = False 28 | fifo = deque() 29 | 30 | height, width = image.shape 31 | total = height * width 32 | labels = np.full((height, width), self.INIT, np.int32) 33 | 34 | reshaped_image = image.reshape(total) 35 | # [y, x] pairs of pixel coordinates of the flattened image. 36 | pixels = np.mgrid[0:height, 0:width].reshape(2, -1).T 37 | # Coordinates of neighbour pixels for each pixel. 38 | neighbours = np.array([self._get_neighbors(height, width, p) for p in pixels]) 39 | if len(neighbours.shape) == 3: 40 | # Case where all pixels have the same number of neighbours. 41 | neighbours = neighbours.reshape(height, width, -1, 2) 42 | else: 43 | # Case where pixels may have a different number of pixels. 44 | neighbours = neighbours.reshape(height, width) 45 | 46 | indices = np.argsort(reshaped_image) 47 | sorted_image = reshaped_image[indices] 48 | sorted_pixels = pixels[indices] 49 | 50 | # self.levels evenly spaced steps from minimum to maximum. 51 | levels = np.linspace(sorted_image[0], sorted_image[-1], self.levels) 52 | level_indices = [] 53 | current_level = 0 54 | 55 | # Get the indices that deleimit pixels with different values. 56 | for i in xrange(total): 57 | if sorted_image[i] > levels[current_level]: 58 | # Skip levels until the next highest one is reached. 59 | while sorted_image[i] > levels[current_level]: current_level += 1 60 | level_indices.append(i) 61 | level_indices.append(total) 62 | 63 | start_index = 0 64 | for stop_index in level_indices: 65 | # Mask all pixels at the current level. 66 | for p in sorted_pixels[start_index:stop_index]: 67 | labels[p[0], p[1]] = self.MASK 68 | # Initialize queue with neighbours of existing basins at the current level. 69 | for q in neighbours[p[0], p[1]]: 70 | # p == q is ignored here because labels[p] < WSHD 71 | if labels[q[0], q[1]] >= self.WSHD: 72 | labels[p[0], p[1]] = self.INQE 73 | fifo.append(p) 74 | break 75 | 76 | # Extend basins. 77 | while fifo: 78 | p = fifo.popleft() 79 | # Label p by inspecting neighbours. 80 | for q in neighbours[p[0], p[1]]: 81 | # Don't set lab_p in the outer loop because it may change. 82 | lab_p = labels[p[0], p[1]] 83 | lab_q = labels[q[0], q[1]] 84 | if lab_q > 0: 85 | if lab_p == self.INQE or (lab_p == self.WSHD and flag): 86 | labels[p[0], p[1]] = lab_q 87 | elif lab_p > 0 and lab_p != lab_q: 88 | labels[p[0], p[1]] = self.WSHD 89 | flag = False 90 | elif lab_q == self.WSHD: 91 | if lab_p == self.INQE: 92 | labels[p[0], p[1]] = self.WSHD 93 | flag = True 94 | elif lab_q == self.MASK: 95 | labels[q[0], q[1]] = self.INQE 96 | fifo.append(q) 97 | 98 | # Detect and process new minima at the current level. 99 | for p in sorted_pixels[start_index:stop_index]: 100 | # p is inside a new minimum. Create a new label. 101 | if labels[p[0], p[1]] == self.MASK: 102 | current_label += 1 103 | fifo.append(p) 104 | labels[p[0], p[1]] = current_label 105 | while fifo: 106 | q = fifo.popleft() 107 | for r in neighbours[q[0], q[1]]: 108 | if labels[r[0], r[1]] == self.MASK: 109 | fifo.append(r) 110 | labels[r[0], r[1]] = current_label 111 | 112 | start_index = stop_index 113 | 114 | return labels 115 | 116 | if __name__ == "__main__": 117 | import sys 118 | from PIL import Image 119 | import matplotlib.pyplot as plt 120 | from scipy.misc import imsave 121 | 122 | w = Watershed() 123 | image = np.array(Image.open(sys.argv[1])) 124 | labels = w.apply(image) 125 | imsave('ws.png', labels) 126 | # plt.imshow(labels, cmap='Paired', interpolation='nearest') 127 | # plt.show() 128 | 129 | --------------------------------------------------------------------------------