This repository has been archived by the owner on Mar 19, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathpreprocess.py
164 lines (134 loc) · 6.74 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
"""
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 brats.train import config
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="GlistrBoost_ManuallyCorrected"):
"""
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):
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, out_folder, truth_name='seg', no_bias_correction_modalities=None):
for name in config["all_modalities"]:
try:
image_file = get_image(in_folder, name)
except RuntimeError as error:
if name == 't1ce':
image_file = get_image(in_folder, 't1Gd')
truth_name = "GlistrBoost_ManuallyCorrected"
else:
raise error
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)