forked from e-koch/ewky_scripts
-
Notifications
You must be signed in to change notification settings - Fork 0
/
iterative_watershed.py
136 lines (109 loc) · 4.67 KB
/
iterative_watershed.py
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import skimage.morphology as mo
import skimage.measure as me
from skimage.feature import peak_local_max
import scipy.ndimage as nd
import numpy as np
import warnings
try:
import cv2
CV2_FLAG = True
except ImportError:
warnings.warn("Cannot import cv2. Computing with scipy.ndimage")
CV2_FLAG = False
def iterative_watershed(array, scale, start_value=5., end_value=3.,
delta_value=1., mask_below=2., min_pix=0):
'''
Iterative Watershed algorithm.
'''
initial_mask = array >= mask_below
if array.ndim == 3:
# Define a footprint to use in the peak finding
footprint = np.tile(mo.disk(scale), (array.shape[0], 1, 1))
use_footprint = True
initial_peaks = peak_local_max(array, footprint=footprint,
threshold_abs=start_value,
exclude_border=False)
initial_markers = np.zeros_like(array, dtype=bool)
initial_markers[:, initial_peaks[:, 1], initial_peaks[:, 2]] = True
elif array.ndim == 2:
initial_peaks = peak_local_max(array, min_distance=scale,
threshold_abs=start_value)
use_footprint = False
initial_markers = np.zeros_like(array, dtype=bool)
initial_markers[initial_peaks[:, 0], initial_peaks[:, 1]] = True
else:
raise Exception("Function only implemented for 2D and 3D arrays.")
initial_markers *= initial_mask
wshed_input = -array.copy()
wshed_input[wshed_input > 0] = 0
labels = mo.watershed(wshed_input, me.label(initial_markers),
mask=initial_mask)
initial_markers *= labels > 0
# Now decrease the local maxima, trying to subdivide
# regions that had a peak at the higher level.
if start_value - end_value < delta_value:
return labels, np.vstack(np.where(initial_markers)).T
peak_levels = \
np.arange(start_value - delta_value,
end_value - delta_value, - 1 * delta_value)
for value in peak_levels:
if use_footprint:
new_peaks = peak_local_max(array, footprint=footprint,
threshold_abs=value,
exclude_border=False)
else:
new_peaks = peak_local_max(array, min_distance=scale,
threshold_abs=value)
markers = initial_markers.copy()
markers[new_peaks[:, 0], new_peaks[:, 1]] = True
# Remove markers not in the last watershed
markers *= labels > 0
# Search for label regions that now have multiple peaks
# and re-run the watershed on them
for lab in range(1, labels.max() + 1):
num_peaks = np.sum(markers * (labels == lab))
if num_peaks == 1:
continue
elif num_peaks == 0:
raise Exception("No peaks found??")
else:
split_marker = me.label(markers * (labels == lab))
split_label = mo.watershed(wshed_input, split_marker,
mask=labels == lab)
orig_used = False
for lab2 in range(1, split_label.max() + 1):
posns = np.where(split_label == lab2)
if array[posns].max() >= start_value:
if not orig_used:
labels[posns] = lab
orig_used = True
else:
labels[posns] = labels.max() + 1
else:
labels[posns] = 0
# Remove small regions
pix = nd.sum(labels > 0, labels, range(1, labels.max() + 1))
for i in xrange(labels.max()):
if pix[i] < min_pix:
labels[labels == i + 1] = 0
markers *= labels > 0
# Return only the peaks that were used.
final_peaks = np.vstack(np.where(markers)).T
return labels, final_peaks
def remove_spurs(mask, min_distance=9):
'''
Remove spurious mask features with reconstruction.
'''
# Distance transform of the mask
dist_trans = nd.distance_transform_edt(mask)
# We don't want to return local maxima within the minimum distance
# Use reconstruction to remove.
seed = dist_trans + min_distance
reconst = mo.reconstruction(seed, dist_trans, method='erosion') - \
min_distance
if CV2_FLAG:
return cv2.morphologyEx((reconst > 0).astype("uint8"),
cv2.MORPH_DILATE,
mo.disk(min_distance).astype("uint8")).astype(bool)
else:
return mo.dilation(reconst > 0, selem=mo.disk(min_distance))