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