1"""watershed.pyx - scithon implementation of guts of watershed 2 3Originally part of CellProfiler, code licensed under both GPL and BSD licenses. 4Website: http://www.cellprofiler.org 5 6Copyright (c) 2003-2009 Massachusetts Institute of Technology 7Copyright (c) 2009-2011 Broad Institute 8All rights reserved. 9 10Original author: Lee Kamentsky 11""" 12import numpy as np 13from libc.math cimport sqrt 14 15cimport numpy as cnp 16cimport cython 17cnp.import_array() 18 19ctypedef cnp.int32_t DTYPE_INT32_t 20ctypedef cnp.int8_t DTYPE_BOOL_t 21 22 23include "heap_watershed.pxi" 24 25 26@cython.wraparound(False) 27@cython.boundscheck(False) 28@cython.cdivision(True) 29@cython.overflowcheck(False) 30@cython.unraisable_tracebacks(False) 31cdef inline double _euclid_dist(Py_ssize_t pt0, Py_ssize_t pt1, 32 cnp.intp_t[::1] strides) nogil: 33 """Return the Euclidean distance between raveled points pt0 and pt1.""" 34 cdef double result = 0 35 cdef double curr = 0 36 for i in range(strides.shape[0]): 37 curr = (pt0 // strides[i]) - (pt1 // strides[i]) 38 result += curr * curr 39 pt0 = pt0 % strides[i] 40 pt1 = pt1 % strides[i] 41 return sqrt(result) 42 43 44@cython.wraparound(False) 45@cython.boundscheck(False) 46@cython.cdivision(True) 47@cython.unraisable_tracebacks(False) 48cdef inline DTYPE_BOOL_t _diff_neighbors(DTYPE_INT32_t[::1] output, 49 cnp.intp_t[::1] structure, 50 DTYPE_BOOL_t[::1] mask, 51 Py_ssize_t index) nogil: 52 """ 53 Return ``True`` and set ``mask[index]`` to ``False`` if the neighbors of 54 ``index`` (as given by the offsets in ``structure``) have more than one 55 distinct nonzero label. 56 """ 57 cdef: 58 Py_ssize_t i, neighbor_index 59 DTYPE_INT32_t neighbor_label0, neighbor_label1 60 Py_ssize_t nneighbors = structure.shape[0] 61 62 if not mask[index]: 63 return True 64 65 neighbor_label0, neighbor_label1 = 0, 0 66 for i in range(nneighbors): 67 neighbor_index = structure[i] + index 68 if mask[neighbor_index]: # neighbor not a watershed line 69 if not neighbor_label0: 70 neighbor_label0 = output[neighbor_index] 71 else: 72 neighbor_label1 = output[neighbor_index] 73 if neighbor_label1 and neighbor_label1 != neighbor_label0: 74 mask[index] = False 75 return True 76 return False 77 78@cython.boundscheck(False) 79@cython.wraparound(False) 80def watershed_raveled(cnp.float64_t[::1] image, 81 cnp.intp_t[::1] marker_locations, 82 cnp.intp_t[::1] structure, 83 DTYPE_BOOL_t[::1] mask, 84 cnp.intp_t[::1] strides, 85 cnp.double_t compactness, 86 DTYPE_INT32_t[::1] output, 87 DTYPE_BOOL_t wsl): 88 """Perform watershed algorithm using a raveled image and neighborhood. 89 90 Parameters 91 ---------- 92 93 image : array of float 94 The flattened image pixels. 95 marker_locations : array of int 96 The raveled coordinates of the initial markers (aka seeds) for the 97 watershed. NOTE: these should *all* point to nonzero entries in the 98 output, or the algorithm will never terminate and blow up your memory! 99 structure : array of int 100 A list of coordinate offsets to compute the raveled coordinates of each 101 neighbor from the raveled coordinates of the current pixel. 102 mask : array of int 103 An array of the same shape as `image` where each pixel contains a 104 nonzero value if it is to be considered for flooding with watershed, 105 zero otherwise. NOTE: it is *essential* that the border pixels (those 106 with neighbors falling outside the volume) are all set to zero, or 107 segfaults could occur. 108 strides : array of int 109 An array representing the number of steps to move along each dimension. 110 This is used in computing the Euclidean distance between raveled 111 indices. 112 compactness : float 113 A value greater than 0 implements the compact watershed algorithm 114 (see .py file). 115 output : array of int 116 The output array, which must already contain nonzero entries at all the 117 seed locations. 118 wsl : bool 119 Parameter indicating whether the watershed line is calculated. 120 If wsl is set to True, the watershed line is calculated. 121 """ 122 cdef Heapitem elem 123 cdef Heapitem new_elem 124 cdef Py_ssize_t nneighbors = structure.shape[0] 125 cdef Py_ssize_t i = 0 126 cdef Py_ssize_t age = 1 127 cdef Py_ssize_t index = 0 128 cdef Py_ssize_t neighbor_index = 0 129 cdef DTYPE_BOOL_t compact = (compactness > 0) 130 131 cdef Heap *hp = <Heap *> heap_from_numpy2() 132 133 with nogil: 134 for i in range(marker_locations.shape[0]): 135 index = marker_locations[i] 136 elem.value = image[index] 137 elem.age = 0 138 elem.index = index 139 elem.source = index 140 heappush(hp, &elem) 141 142 while hp.items > 0: 143 heappop(hp, &elem) 144 145 if compact or wsl: 146 # in the compact case, we need to label pixels as they come off 147 # the heap, because the same pixel can be pushed twice, *and* the 148 # later push can have lower cost because of the compactness. 149 # 150 # In the case of preserving watershed lines, a similar argument 151 # applies: we can only observe that all neighbors have been labeled 152 # as the pixel comes off the heap. Trying to do so at push time 153 # is a bug. 154 if output[elem.index] and elem.index != elem.source: 155 # non-marker, already visited from another neighbor 156 continue 157 if wsl: 158 # if the current element has different-labeled neighbors and we 159 # want to preserve watershed lines, we mask it and move on 160 if _diff_neighbors(output, structure, mask, elem.index): 161 continue 162 output[elem.index] = output[elem.source] 163 164 for i in range(nneighbors): 165 # get the flattened address of the neighbor 166 neighbor_index = structure[i] + elem.index 167 168 if not mask[neighbor_index]: 169 # this branch includes basin boundaries, aka watershed lines 170 # neighbor is not in mask 171 continue 172 173 if output[neighbor_index]: 174 # pre-labeled neighbor is not added to the queue. 175 continue 176 177 age += 1 178 new_elem.value = image[neighbor_index] 179 if compact: 180 new_elem.value += (compactness * 181 _euclid_dist(neighbor_index, elem.source, 182 strides)) 183 elif not wsl: 184 # in the simplest watershed case (no compactness and no 185 # watershed lines), we can label a pixel at the time that 186 # we push it onto the heap, because it can't be reached with 187 # lower cost later. 188 # This results in a very significant performance gain, see: 189 # https://github.com/scikit-image/scikit-image/issues/2636 190 output[neighbor_index] = output[elem.index] 191 new_elem.age = age 192 new_elem.index = neighbor_index 193 new_elem.source = elem.source 194 195 heappush(hp, &new_elem) 196 197 heap_done(hp) 198