-
Notifications
You must be signed in to change notification settings - Fork 2
/
common_test.py
337 lines (269 loc) · 14.3 KB
/
common_test.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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
import os
import nibabel as nib
import torch
import numpy as np
from scipy.ndimage.filters import gaussian_filter
from utilities import ComputMetric
from typing import Tuple, List
def testlitstumor(model, saveresults, name, trainval = False, ImgsegmentSize = [80, 160, 160], deepsupervision=False, DatafileValFold=None, tta=False, ttalist = [0], ttalistprob=[1], NumsClass = 2):
batch_size = 1
NumsInputChannel = 1
if trainval == False:
DatafileFold = DatafileValFold
DatafileImgc1 = DatafileFold + 'Imgpre-eval.txt'
DatafileLabel = DatafileFold + 'seg-eval.txt'
DatafileMask = DatafileFold + 'mask-eval.txt'
else:
DatafileFold = DatafileValFold
DatafileImgc1 = DatafileFold + 'Imgpre-train.txt'
DatafileLabel = DatafileFold + 'seg-train.txt'
DatafileMask = DatafileFold + 'mask-train.txt'
Imgfilec1 = open(DatafileImgc1)
Imgreadc1 = Imgfilec1.read().splitlines()
Maskfile = open(DatafileMask)
Maskread = Maskfile.read().splitlines()
Labelfile = open(DatafileLabel)
Labelread = Labelfile.read().splitlines()
DSClist = []
SENSlist = []
PREClist = []
PredSumlist = []
for numr in range(len(Imgreadc1)):
# for numr in range(50, 51):
Imgnamec1 = Imgreadc1[numr]
Imgloadc1 = nib.load(Imgnamec1)
Imgc1 = Imgloadc1.get_fdata()
Maskname = Maskread[numr]
Maskload = nib.load(Maskname)
roi_mask = Maskload.get_fdata()
Labelname = Labelread[numr]
Labelload = nib.load(Labelname)
gtlabel = Labelload.get_fdata()
Imgc1 = np.float32(Imgc1)
knamelist = Imgnamec1.split("/")
kname = knamelist[-1][0:-7]
channels = Imgc1[None, ...]
hp_results = tta_rolling(model, channels, batch_size, ImgsegmentSize, NumsInputChannel, NumsClass, tta, ttalist, ttalistprob, deepsupervision)
predSegmentation = np.argmax(hp_results, axis=0)
'''This is debuggin code for mask softmax'''
# y_eye = np.eye(2)
# y_onehot = y_eye[gtlabel.astype(int)]
# y_onehot = np.transpose(y_onehot, (3,0,1,2))
# y_onehot_repeat = np.repeat(y_onehot, 2, axis=0)
# print(y_onehot_repeat[:, 0, 0, 0])
# hp_results = np.exp(hp_results) / np.sum(np.exp(hp_results), axis=0)
# hp_results = hp_results * y_onehot_repeat
# hp_results[hp_results==0] = 100
# predSegmentation = np.argmin(hp_results, axis=0)
# hp_results = np.exp(hp_results) / np.sum(np.exp(hp_results), axis=0)
# predSegmentation = hp_results[0, :, :, :]
# print(np.sum(predSegmentation))
## use the mask to constratin the results
PredSegmentationWithinRoi = predSegmentation * roi_mask
# PredSegmentationWithinRoi = predSegmentation
# sio.savemat('./result.mat', {'results': PredSegmentationWithinRoi})
imgToSave = PredSegmentationWithinRoi
if saveresults:
npDtype = np.dtype(np.float32)
proxy_origin = nib.load(Imgnamec1)
hdr_origin = proxy_origin.header
affine_origin = proxy_origin.affine
proxy_origin.uncache()
newLabelImg = nib.Nifti1Image(imgToSave, affine_origin)
newLabelImg.set_data_dtype(npDtype)
dimsImgToSave = len(imgToSave.shape)
newZooms = list(hdr_origin.get_zooms()[:dimsImgToSave])
if len(newZooms) < dimsImgToSave: # Eg if original image was 3D, but I need to save a multi-channel image.
newZooms = newZooms + [1.0] * (dimsImgToSave - len(newZooms))
newLabelImg.header.set_zooms(newZooms)
directory = "./output/litstumor/%s/" % (name)
if not os.path.exists(directory):
os.makedirs(directory)
savename = directory + 'pred_' + kname + '_Segm.nii.gz'
nib.save(newLabelImg, savename)
labelwt = gtlabel == 1
predwt = imgToSave == 1
DSCwt, SENSwt, PRECwt = ComputMetric(labelwt, predwt)
# print(DSCwt)
DSClist.append([DSCwt])
SENSlist.append([SENSwt])
PREClist.append([PRECwt])
PredSumlist.append(np.sum(labelwt))
print('case ' + str(numr) + ' done')
sel = [n for n, i in enumerate(PredSumlist) if i > 0]
DSClist = np.array(DSClist)
DSCmean = DSClist[sel, :].mean(axis=0)
SENSlist = np.array(SENSlist)
SENSmean = SENSlist[sel, :].mean(axis=0)
PREClist = np.array(PREClist)
PRECmean = PREClist[sel, :].mean(axis=0)
return DSCmean, SENSmean, PRECmean
def tta_rolling(model, channels, batch_size, ImgsegmentSize, NumsInputChannel, NumsClass, tta, ttalist, ttalistprob, deepsupervision):
hp_results = 0
for ttaindex, ttaindexprob in zip(ttalist, ttalistprob):
channels_per_path = []
channels_per_path.append(channels.copy())
channels_augment = channels_per_path[0]
offset = [0, 0, 0]
pad_border_mode = 'constant'
pad_kwargs = dict()
pad_kwargs['constant_values'] = 0
data, slicer = pad_nd_image(channels_augment, ImgsegmentSize, pad_border_mode, pad_kwargs, True, None)
data_shape = data.shape
step_size = 0.5
steps = _compute_steps_for_sliding_window(ImgsegmentSize, data_shape[1:], step_size)
hp = np.zeros([NumsClass] + list(data.shape[1:]), dtype=np.float32)
aggregated_nb_of_predictions = np.zeros([NumsClass] + list(data.shape[1:]), dtype=np.float32)
xpixels = []
ypixels = []
zpixels = []
for jx in steps[0]:
for jy in steps[1]:
for jz in steps[2]:
xpixels.append(jx)
ypixels.append(jy)
zpixels.append(jz)
inputxnor = getallbatch(data, ImgsegmentSize, xpixels, ypixels, zpixels, offset)
## gaussian filter
patch_size = [ImgsegmentSize[0], ImgsegmentSize[1], ImgsegmentSize[2]]
tmp = np.zeros(patch_size)
center_coords = [i // 2 for i in patch_size]
sigmas = [i * 1. / 8 for i in patch_size]
tmp[tuple(center_coords)] = 1
gaussian_importance_map = gaussian_filter(tmp, sigmas, 0, mode='constant', cval=0)
gaussian_importance_map = gaussian_importance_map / np.max(gaussian_importance_map) * 1
gaussian_importance_map = gaussian_importance_map.astype(np.float32)
# gaussian_importance_map cannot be 0, otherwise we may end up with nans!
gaussian_importance_map[gaussian_importance_map == 0] = np.min(
gaussian_importance_map[gaussian_importance_map != 0])
inputxnor = torch.tensor(np.array(inputxnor))
inputxnor = inputxnor.float().cuda()
for xlist in range(0, len(inputxnor), batch_size):
batchxnor = inputxnor[xlist: xlist + batch_size, :, :, :, :]
xstarts = xpixels[xlist: xlist + batch_size]
ystarts = ypixels[xlist: xlist + batch_size]
zstarts = zpixels[xlist: xlist + batch_size]
with torch.no_grad():
pred = model(batchxnor)
## in case it is multi-task model.
if len(pred) == 2:
pred = pred[0]
if deepsupervision:
output = pred[0]
else:
output = pred
output = output.data.cpu().numpy()
kbatch = 0
for xstart, ystart, zstart in zip(xstarts, ystarts, zstarts):
# only crop the center parts.
# maybe I should use gaussain? to do ..
# hp[:, xstart + offset:xstart + offset + PredSizetest, ystart + offset:ystart + offset + PredSizetest,
# zstart + offset:zstart + offset + PredSizetest] = output[kbatch, :, offset:offset + PredSizetest,
# offset:offset + PredSizetest,
# offset:offset + PredSizetest]
hp[:, xstart:xstart + ImgsegmentSize[0], ystart:ystart + ImgsegmentSize[1],
zstart:zstart + ImgsegmentSize[2]] += output[kbatch, :, :, :, :] * gaussian_importance_map
aggregated_nb_of_predictions[:, xstart:xstart + ImgsegmentSize[0], ystart:ystart + ImgsegmentSize[1],
zstart:zstart + ImgsegmentSize[2]] += gaussian_importance_map
kbatch = kbatch + 1
slicer = tuple(
[slice(0, hp.shape[i]) for i in
range(len(hp.shape) - (len(slicer) - 1))] + slicer[1:])
hp = hp[slicer]
aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer]
hp = hp / aggregated_nb_of_predictions
## to see if the probability map needs revert spatial transformations
hp_revert = hp
## hp_revert with shape: 1, D, W, H
## Here I want to try to ensemble with prediction score.
# hp_revert = np.exp(hp_revert) / np.sum(np.exp(hp_revert), axis=0)
hp_results += hp_revert / sum(ttalistprob) * ttaindexprob
return hp_results
def getallbatch(Imgenlarge, ImgsegmentSize, xpixels, ypixels, zpixels, offset):
inputxnor = []
# normal pathway
for (selindex_x, selindex_y, selindex_z) in zip(xpixels, ypixels, zpixels):
coord_center = np.zeros(3, dtype=int)
coord_center[0] = selindex_x + ImgsegmentSize[0] // 2
coord_center[1] = selindex_y + ImgsegmentSize[1] // 2
coord_center[2] = selindex_z + ImgsegmentSize[2] // 2
samplekernal_primary = 1
channs_of_sample_per_path_normal = Imgenlarge[:,
coord_center[0] - ImgsegmentSize[0] // 2: coord_center[0] + ImgsegmentSize[0] // 2,
coord_center[1] - ImgsegmentSize[1] // 2: coord_center[1] + ImgsegmentSize[1] // 2,
coord_center[2] - ImgsegmentSize[2] // 2: coord_center[2] + ImgsegmentSize[2] // 2]
inputxnor.append(channs_of_sample_per_path_normal)
return inputxnor
def pad_nd_image(image, new_shape=None, mode="constant", kwargs=None, return_slicer=False, shape_must_be_divisible_by=None):
"""
one padder to pad them all. Documentation? Well okay. A little bit
:param image: nd image. can be anything
:param new_shape: what shape do you want? new_shape does not have to have the same dimensionality as image. If
len(new_shape) < len(image.shape) then the last axes of image will be padded. If new_shape < image.shape in any of
the axes then we will not pad that axis, but also not crop! (interpret new_shape as new_min_shape)
Example:
image.shape = (10, 1, 512, 512); new_shape = (768, 768) -> result: (10, 1, 768, 768). Cool, huh?
image.shape = (10, 1, 512, 512); new_shape = (364, 768) -> result: (10, 1, 512, 768).
:param mode: see np.pad for documentation
:param return_slicer: if True then this function will also return what coords you will need to use when cropping back
to original shape
:param shape_must_be_divisible_by: for network prediction. After applying new_shape, make sure the new shape is
divisibly by that number (can also be a list with an entry for each axis). Whatever is missing to match that will
be padded (so the result may be larger than new_shape if shape_must_be_divisible_by is not None)
:param kwargs: see np.pad for documentation
"""
if kwargs is None:
kwargs = {'constant_values': 0}
if new_shape is not None:
old_shape = np.array(image.shape[-len(new_shape):])
else:
assert shape_must_be_divisible_by is not None
assert isinstance(shape_must_be_divisible_by, (list, tuple, np.ndarray))
new_shape = image.shape[-len(shape_must_be_divisible_by):]
old_shape = new_shape
num_axes_nopad = len(image.shape) - len(new_shape)
new_shape = [max(new_shape[i], old_shape[i]) for i in range(len(new_shape))]
if not isinstance(new_shape, np.ndarray):
new_shape = np.array(new_shape)
if shape_must_be_divisible_by is not None:
if not isinstance(shape_must_be_divisible_by, (list, tuple, np.ndarray)):
shape_must_be_divisible_by = [shape_must_be_divisible_by] * len(new_shape)
else:
assert len(shape_must_be_divisible_by) == len(new_shape)
for i in range(len(new_shape)):
if new_shape[i] % shape_must_be_divisible_by[i] == 0:
new_shape[i] -= shape_must_be_divisible_by[i]
new_shape = np.array([new_shape[i] + shape_must_be_divisible_by[i] - new_shape[i] % shape_must_be_divisible_by[i] for i in range(len(new_shape))])
difference = new_shape - old_shape
pad_below = difference // 2
pad_above = difference // 2 + difference % 2
pad_list = [[0, 0]]*num_axes_nopad + list([list(i) for i in zip(pad_below, pad_above)])
if not ((all([i == 0 for i in pad_below])) and (all([i == 0 for i in pad_above]))):
res = np.pad(image, pad_list, mode, **kwargs)
else:
res = image
if not return_slicer:
return res
else:
pad_list = np.array(pad_list)
pad_list[:, 1] = np.array(res.shape) - pad_list[:, 1]
slicer = list(slice(*i) for i in pad_list)
return res, slicer
def _compute_steps_for_sliding_window(patch_size: Tuple[int, ...], image_size: Tuple[int, ...], step_size: float) -> List[List[int]]:
assert [i >= j for i, j in zip(image_size, patch_size)], "image size must be as large or larger than patch_size"
assert 0 < step_size <= 1, 'step_size must be larger than 0 and smaller or equal to 1'
# our step width is patch_size*step_size at most, but can be narrower. For example if we have image size of
# 110, patch size of 32 and step_size of 0.5, then we want to make 4 steps starting at coordinate 0, 27, 55, 78
target_step_sizes_in_voxels = [i * step_size for i in patch_size]
num_steps = [int(np.ceil((i - k) / j)) + 1 for i, j, k in zip(image_size, target_step_sizes_in_voxels, patch_size)]
steps = []
for dim in range(len(patch_size)):
# the highest step value for this dimension is
max_step_value = image_size[dim] - patch_size[dim]
if num_steps[dim] > 1:
actual_step_size = max_step_value / (num_steps[dim] - 1)
else:
actual_step_size = 99999999999 # does not matter because there is only one step at 0
steps_here = [int(np.round(actual_step_size * i)) for i in range(num_steps[dim])]
steps.append(steps_here)
return steps