-
Notifications
You must be signed in to change notification settings - Fork 7
/
process_mito.py
153 lines (136 loc) · 7.46 KB
/
process_mito.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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# Post-processing functions of mitochondria instance segmentation model outputs
# as described in "MitoEM Dataset: Large-scale 3D Mitochondria Instance Segmentation
# from EM Images (MICCAI 2020, https://donglaiw.github.io/page/mitoEM/index.html)".
import numpy as np
from skimage.measure import label
from skimage.transform import resize
from skimage.morphology import dilation
from skimage.segmentation import watershed
from .utils import remove_small_instances
def binary_connected(volume, thres=0.9, thres_small=128, scale_factors=(1.0, 1.0, 1.0),
remove_small_mode='background'):
"""From binary foreground probability map to instance masks via
connected-component labeling.
Args:
volume (numpy.ndarray): foreground probability of shape :math:`(C, Z, Y, X)`.
thres (float): threshold of foreground. Default: 0.8
thres_small (int): size threshold of small objects to remove. Default: 128
scale_factors (tuple): scale factors for resizing in :math:`(Z, Y, X)` order. Default: :math:`(1.0, 1.0, 1.0)`
remove_small_mode (str): ``'background'`` or ``'neighbor'``. Default: ``'background'``
"""
semantic = volume[0]
foreground = (semantic > int(255*thres))
segm = label(foreground)
segm = remove_small_instances(segm, thres_small, remove_small_mode)
if not all(x==1.0 for x in scale_factors):
target_size = (int(semantic.shape[0]*scale_factors[0]),
int(semantic.shape[1]*scale_factors[1]),
int(semantic.shape[2]*scale_factors[2]))
segm = resize(segm, target_size, order=0, anti_aliasing=False, preserve_range=True)
return segm.astype(np.uint32)
def binary_watershed(volume, thres1=0.98, thres2=0.85, thres_small=128, scale_factors=(1.0, 1.0, 1.0),
remove_small_mode='background'):
"""From binary foreground probability map to instance masks via
watershed segmentation algorithm.
Args:
volume (numpy.ndarray): foreground probability of shape :math:`(C, Z, Y, X)`.
thres1 (float): threshold of seeds. Default: 0.98
thres2 (float): threshold of foreground. Default: 0.85
thres_small (int): size threshold of small objects to remove. Default: 128
scale_factors (tuple): scale factors for resizing in :math:`(Z, Y, X)` order. Default: :math:`(1.0, 1.0, 1.0)`
remove_small_mode (str): ``'background'`` or ``'neighbor'``. Default: ``'background'``
"""
semantic = volume[0]
seed_map = semantic > int(255*thres1)
foreground = semantic > int(255*thres2)
seed = label(seed_map)
segm = watershed(-semantic, seed, mask=foreground)
segm = remove_small_instances(segm, thres_small, remove_small_mode)
if not all(x==1.0 for x in scale_factors):
target_size = (int(semantic.shape[0]*scale_factors[0]),
int(semantic.shape[1]*scale_factors[1]),
int(semantic.shape[2]*scale_factors[2]))
segm = resize(segm, target_size, order=0, anti_aliasing=False, preserve_range=True)
return segm.astype(np.uint32)
def bc_connected(volume, thres1=0.8, thres2=0.5, thres_small=128, scale_factors=(1.0, 1.0, 1.0),
dilation_struct=(1,5,5), remove_small_mode='background'):
"""From binary foreground probability map and instance contours to
instance masks via connected-component labeling.
Note:
The instance contour provides additional supervision to distinguish closely touching
objects. However, the decoding algorithm only keep the intersection of foreground and
non-contour regions, which will systematically result in imcomplete instance masks.
Therefore we apply morphological dilation (check :attr:`dilation_struct`) to enlarge
the object masks.
Args:
volume (numpy.ndarray): foreground and contour probability of shape :math:`(C, Z, Y, X)`.
thres1 (float): threshold of foreground. Default: 0.8
thres2 (float): threshold of instance contours. Default: 0.5
thres_small (int): size threshold of small objects to remove. Default: 128
scale_factors (tuple): scale factors for resizing in :math:`(Z, Y, X)` order. Default: :math:`(1.0, 1.0, 1.0)`
dilation_struct (tuple): the shape of the structure for morphological dilation. Default: :math:`(1, 5, 5)`
remove_small_mode (str): ``'background'`` or ``'neighbor'``. Default: ``'background'``
"""
semantic = volume[0]
boundary = volume[1]
foreground = (semantic > int(255*thres1)) * (boundary < int(255*thres2))
segm = label(foreground)
struct = np.ones(dilation_struct)
segm = dilation(segm, struct)
segm = remove_small_instances(segm, thres_small, remove_small_mode)
if not all(x==1.0 for x in scale_factors):
target_size = (int(semantic.shape[0]*scale_factors[0]),
int(semantic.shape[1]*scale_factors[1]),
int(semantic.shape[2]*scale_factors[2]))
segm = resize(segm, target_size, order=0, anti_aliasing=False, preserve_range=True)
return segm.astype(np.uint32)
def bc_watershed(volume, thres1=0.9, thres2=0.8, thres3=0.85, thres_small=128, scale_factors=(1.0, 1.0, 1.0),
remove_small_mode='background'):
"""From binary foreground probability map and instance contours to
instance masks via watershed segmentation algorithm.
Args:
volume (numpy.ndarray): foreground and contour probability of shape :math:`(C, Z, Y, X)`.
thres1 (float): threshold of seeds. Default: 0.9
thres2 (float): threshold of instance contours. Default: 0.8
thres3 (float): threshold of foreground. Default: 0.85
thres_small (int): size threshold of small objects to remove. Default: 128
scale_factors (tuple): scale factors for resizing in :math:`(Z, Y, X)` order. Default: :math:`(1.0, 1.0, 1.0)`
remove_small_mode (str): ``'background'`` or ``'neighbor'``. Default: ``'background'``
"""
semantic = volume[0]
boundary = volume[1]
seed_map = (semantic > int(255*thres1)) * (boundary < int(255*thres2)) # seed , not contours
foreground = (semantic > int(255*thres3))
seed = label(seed_map)
segm = watershed(-semantic, seed, mask=foreground)
segm = remove_small_instances(segm, thres_small, remove_small_mode)
if not all(x==1.0 for x in scale_factors):
target_size = (int(semantic.shape[0]*scale_factors[0]),
int(semantic.shape[1]*scale_factors[1]),
int(semantic.shape[2]*scale_factors[2]))
segm = resize(segm, target_size, order=0, anti_aliasing=False, preserve_range=True)
return segm.astype(np.uint32)
import waterz
import malis
def malis_watershed(seed_map, thres1=0.9, thres2=0.8):
if isinstance(seed_map, list):
semantic = seed_map[0]
boundary = seed_map[1]
seed_map = (semantic > int(255*thres1)) * (boundary < int(255*thres2))
elif isinstance(seed_map, np.ndarray):
seed_map = seed_map
else:
raise RuntimeError("seed map is wrong!")
# generate affinity
output_mixs = seed_map.astype(np.int32)
affs = malis.seg_to_affgraph(output_mixs, malis.mknhood3d())
del output_mixs
affs = affs.astype(np.float32)
# initial watershed + agglomerate
seg = list(waterz.agglomerate(affs, [0.50]))[0]
del affs
seg = seg.astype(np.uint16)
# grow boundary
seg = dilation(seg, np.ones((1,7,7)))
seg = remove_small_instances(seg)
return seg