forked from scikit-image/scikit-image
-
Notifications
You must be signed in to change notification settings - Fork 0
/
_watershed.pyx
95 lines (80 loc) · 3.1 KB
/
_watershed.pyx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
"""watershed.pyx - scithon implementation of guts of watershed
Originally part of CellProfiler, code licensed under both GPL and BSD licenses.
Website: http://www.cellprofiler.org
Copyright (c) 2003-2009 Massachusetts Institute of Technology
Copyright (c) 2009-2011 Broad Institute
All rights reserved.
Original author: Lee Kamentsky
"""
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int32_t DTYPE_INT32_t
DTYPE_BOOL = np.bool
ctypedef np.int8_t DTYPE_BOOL_t
include "heap_watershed.pxi"
@cython.boundscheck(False)
def watershed(DTYPE_INT32_t[::1] image,
DTYPE_INT32_t[:, ::1] pq,
Py_ssize_t age,
DTYPE_INT32_t[:, ::1] structure,
DTYPE_BOOL_t[::1] mask,
DTYPE_INT32_t[::1] output):
"""Do heavy lifting of watershed algorithm
Parameters
----------
image - the flattened image pixels, converted to rank-order
pq - the priority queue, starts with the marked pixels
the first element in each row is the image intensity
the second element is the age at entry into the queue
the third element is the index into the flattened image or labels
the remaining elements are the coordinates of the point
age - the next age to assign to a pixel
structure - a numpy int32 array containing the structuring elements
that define nearest neighbors. For each row, the first
element is the stride from the point to its neighbor
in a flattened array. The remaining elements are the
offsets from the point to its neighbor in the various
dimensions
mask - numpy boolean (char) array indicating which pixels to consider
and which to ignore. Also flattened.
output - put the image labels in here
"""
cdef Heapitem elem
cdef Heapitem new_elem
cdef Py_ssize_t nneighbors = structure.shape[0]
cdef Py_ssize_t i = 0
cdef Py_ssize_t index = 0
cdef Py_ssize_t old_index = 0
cdef Py_ssize_t max_index = image.shape[0]
cdef Heap *hp = <Heap *> heap_from_numpy2()
for i in range(pq.shape[0]):
elem.value = pq[i, 0]
elem.age = pq[i, 1]
elem.index = pq[i, 2]
heappush(hp, &elem)
while hp.items > 0:
#
# Pop off an item to work on
#
heappop(hp, &elem)
####################################################
# loop through each of the structuring elements
#
old_index = elem.index
for i in range(nneighbors):
# get the flattened address of the neighbor
index = structure[i, 0] + old_index
if index < 0 or index >= max_index or output[index] or \
not mask[index]:
continue
new_elem.value = image[index]
new_elem.age = elem.age + 1
new_elem.index = index
age += 1
output[index] = output[old_index]
#
# Push the neighbor onto the heap to work on it later
#
heappush(hp, &new_elem)
heap_done(hp)