-
Notifications
You must be signed in to change notification settings - Fork 43
/
preprocess.py
197 lines (168 loc) · 8.15 KB
/
preprocess.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
"""
Tools for converting, normalizing, and fixing the brats data.
"""
import glob
import os
import warnings
import shutil
import SimpleITK as sitk
import numpy as np
from nipype.interfaces.ants import N4BiasFieldCorrection
from main import config
from multiprocessing import Pool
import time
NUM_Thread = 16
_truth = "seg"
config["all_modalities"] = ["t1", "t1ce", "flair", "t2"]
_out_folder = "data/BraTs2018_data_validation"
_brats_folder = "../../BraTs_Data/BraTs_2018_Data_Validation"
def append_basename(in_file, append):
dirname, basename = os.path.split(in_file)
base, ext = basename.split(".", 1)
return os.path.join(dirname, base + append + "." + ext)
def get_background_mask(in_folder, out_file,truth_name=_truth):
"""
This function computes a common background mask for all of the data in a subject folder.
:param in_folder: a subject folder from the BRATS dataset.
:param out_file: an image containing a mask that is 1 where the image data for that subject contains the background.
:param truth_name: how the truth file is labeled int he subject folder
:return: the path to the out_file
"""
background_image = None
for name in config["all_modalities"] + [truth_name]:
image = sitk.ReadImage(get_image(in_folder, name))
if background_image:
if name == truth_name and not (image.GetOrigin() == background_image.GetOrigin()):
image.SetOrigin(background_image.GetOrigin())
background_image = sitk.And(image == 0, background_image)
else:
background_image = image == 0
sitk.WriteImage(background_image, out_file)
return os.path.abspath(out_file)
def convert_image_format(in_file, out_file):
sitk.WriteImage(sitk.ReadImage(in_file), out_file)
return out_file
def window_intensities(in_file, out_file, min_percent=1, max_percent=99):
image = sitk.ReadImage(in_file)
image_data = sitk.GetArrayFromImage(image)
out_image = sitk.IntensityWindowing(image, np.percentile(image_data, min_percent), np.percentile(image_data,
max_percent))
sitk.WriteImage(out_image, out_file)
return os.path.abspath(out_file)
def correct_bias(in_file, out_file, image_type=sitk.sitkFloat64):
"""
Corrects the bias using ANTs N4BiasFieldCorrection. If this fails, will then attempt to correct bias using SimpleITK
:param in_file: input file path
:param out_file: output file path
:return: file path to the bias corrected image
"""
correct = N4BiasFieldCorrection()
correct.inputs.input_image = in_file
correct.inputs.output_image = out_file
try:
done = correct.run()
return done.outputs.output_image
except IOError:
warnings.warn(RuntimeWarning("ANTs N4BIasFieldCorrection could not be found."
"Will try using SimpleITK for bias field correction"
" which will take much longer. To fix this problem, add N4BiasFieldCorrection"
" to your PATH system variable. (example: EXPORT PATH=${PATH}:/path/to/ants/bin)"))
input_image = sitk.ReadImage(in_file, image_type)
output_image = sitk.N4BiasFieldCorrection(input_image, input_image > 0)
sitk.WriteImage(output_image, out_file)
return os.path.abspath(out_file)
def rescale(in_file, out_file, minimum=0, maximum=20000):
image = sitk.ReadImage(in_file)
sitk.WriteImage(sitk.RescaleIntensity(image, minimum, maximum), out_file)
return os.path.abspath(out_file)
def get_image(subject_folder, name):
file_card = os.path.join(subject_folder, "*" + name + ".nii.gz")
try:
return glob.glob(file_card)[0]
except IndexError:
raise RuntimeError("Could not find file matching {}".format(file_card))
def background_to_zero(in_file, background_file, out_file):
sitk.WriteImage(sitk.Mask(sitk.ReadImage(in_file), sitk.ReadImage(background_file, sitk.sitkUInt8) == 0),
out_file)
return out_file
def check_origin(in_file, in_file2):
import pdb;pdb.set_trace()
image = sitk.ReadImage(in_file)
image2 = sitk.ReadImage(in_file2)
if not image.GetOrigin() == image2.GetOrigin():
image.SetOrigin(image2.GetOrigin())
sitk.WriteImage(image, in_file)
def normalize_image(in_file, out_file, bias_correction=True):
if bias_correction:
correct_bias(in_file, out_file)
else:
shutil.copy(in_file, out_file)
return out_file
def convert_brats_folder(in_folder):
# truth_name="GlistrBoost_ManuallyCorrected"
print("preprocessing "+in_folder)
out_folder = _out_folder
subject = os.path.basename(in_folder)
new_subject_folder = os.path.join(out_folder, os.path.basename(os.path.dirname(in_folder)),subject)
out_folder = new_subject_folder
no_bias_correction_modalities=("flair",)
truth_name = _truth
for name in config["all_modalities"]:
image_file = get_image(in_folder, name)
out_file = os.path.abspath(os.path.join(out_folder, name + ".nii.gz"))
perform_bias_correction = no_bias_correction_modalities and name not in no_bias_correction_modalities
normalize_image(image_file, out_file, bias_correction=perform_bias_correction)
# copy the truth file
# try:
# truth_file = get_image(in_folder, truth_name)
# except RuntimeError:
# truth_file = get_image(in_folder, truth_name.split("_")[0])
# out_file = os.path.abspath(os.path.join(out_folder, "truth.nii.gz"))
# shutil.copy(truth_file, out_file)
# check_origin(out_file, get_image(in_folder, config["all_modalities"][0]))
def convert_brats_data(brats_folder, out_folder, overwrite=False, no_bias_correction_modalities=("flair",)):
"""
Preprocesses the BRATS data and writes it to a given output folder. Assumes the original folder structure.
:param brats_folder: folder containing the original brats data
:param out_folder: output folder to which the preprocessed data will be written
:param overwrite: set to True in order to redo all the preprocessing
:param no_bias_correction_modalities: performing bias correction could reduce the signal of certain modalities. If
concerned about a reduction in signal for a specific modality, specify by including the given modality in a list
or tuple.
:return:
"""
for subject_folder in glob.glob(os.path.join(brats_folder, "*", "*")):
if os.path.isdir(subject_folder):
subject = os.path.basename(subject_folder)
new_subject_folder = os.path.join(out_folder, os.path.basename(os.path.dirname(subject_folder)),
subject)
if not os.path.exists(new_subject_folder) or overwrite:
if not os.path.exists(new_subject_folder):
os.makedirs(new_subject_folder)
convert_brats_folder(subject_folder, new_subject_folder,
no_bias_correction_modalities=no_bias_correction_modalities)
def get_brats_path(brats_folder, out_folder, overwrite=True, no_bias_correction_modalities=("flair",)):
file_list = []
for subject_folder in glob.glob(os.path.join(brats_folder, "*", "*")):
print(subject_folder)
if os.path.isdir(subject_folder):
subject = os.path.basename(subject_folder)
new_subject_folder = os.path.join(out_folder, os.path.basename(os.path.dirname(subject_folder)),
subject)
# if not os.path.exists(new_subject_folder) or overwrite:
# if not os.path.exists(new_subject_folder):
# os.makedirs(new_subject_folder)
file_list.append((subject_folder))
return file_list
if __name__ == "__main__":
brats_folder = _brats_folder
out_folder = _out_folder
e1 = time.time()
file_list = get_brats_path(brats_folder, out_folder)
print("%d files to preprocessing"%(len(file_list)))
pool = Pool(NUM_Thread)
pool.map(convert_brats_folder,file_list)
pool.close()
pool.join()
e2 = time.time()
print('cost %d s'%(int(e2-e1)))