1
|
|
|
import os |
2
|
|
|
import random |
3
|
|
|
import shutil |
4
|
|
|
import zipfile |
5
|
|
|
|
6
|
|
|
import nibabel as nib |
7
|
|
|
import numpy as np |
8
|
|
|
from tensorflow.keras.utils import get_file |
9
|
|
|
from tqdm import tqdm |
10
|
|
|
|
11
|
|
|
# if wget is installed remove the following line from comment |
12
|
|
|
# import wget |
13
|
|
|
|
14
|
|
|
# if already in the abc/DeepReg directory then do nothing, otherwise |
15
|
|
|
# use os.chdir(r'abc/DeepReg') before this line |
16
|
|
|
main_path = os.getcwd() |
17
|
|
|
os.chdir(main_path) |
18
|
|
|
|
19
|
|
|
######## DOWNLOADING AND UNZIPPING ALL FILES INTO CORRECT PATH ######## |
20
|
|
|
|
21
|
|
|
project_dir = "demos/paired_ct_lung" |
22
|
|
|
data_folder_name = "dataset" |
23
|
|
|
os.chdir(project_dir) |
24
|
|
|
|
25
|
|
|
url = "https://zenodo.org/record/3835682/files/training.zip" |
26
|
|
|
|
27
|
|
|
# if wget is installed remove following line from comments and comment |
28
|
|
|
# out the fname = 'training.zip' line |
29
|
|
|
# fname = wget.download(url) |
30
|
|
|
fname = "training.zip" |
31
|
|
|
|
32
|
|
|
# os.system("wget " + url) |
33
|
|
|
|
34
|
|
|
get_file(os.path.join(os.getcwd(), fname), url) |
35
|
|
|
|
36
|
|
|
print("The file ", fname, " has successfully been downloaded!") |
37
|
|
|
path_to_data_folder = os.path.join(main_path, project_dir, data_folder_name) |
38
|
|
|
if os.path.exists(path_to_data_folder): |
39
|
|
|
shutil.rmtree(path_to_data_folder) |
40
|
|
|
os.mkdir(path_to_data_folder) |
41
|
|
|
|
42
|
|
|
with zipfile.ZipFile(fname, "r") as zip_ref: |
43
|
|
|
zip_ref.extractall(data_folder_name) |
44
|
|
|
|
45
|
|
|
print("Files unzipped!") |
46
|
|
|
|
47
|
|
|
os.remove(fname) |
48
|
|
|
os.chdir(main_path) |
49
|
|
|
|
50
|
|
|
######## MOVING FILES INTO TRAIN DIRECTORY ######## |
51
|
|
|
|
52
|
|
|
path_to_train = os.path.join(main_path, project_dir, data_folder_name, "train") |
53
|
|
|
path_to_test = os.path.join(main_path, project_dir, data_folder_name, "test") |
54
|
|
|
path_to_images_and_labels = os.path.join( |
55
|
|
|
main_path, project_dir, data_folder_name, "training" |
56
|
|
|
) |
57
|
|
|
|
58
|
|
|
labels_fnames = os.listdir(os.path.join(path_to_images_and_labels, "lungMasks")) |
59
|
|
|
images_fnames = os.listdir(os.path.join(path_to_images_and_labels, "scans")) |
60
|
|
|
|
61
|
|
|
if os.path.exists(path_to_train) is not True: |
62
|
|
|
os.mkdir(path_to_train) |
63
|
|
|
os.mkdir(os.path.join(path_to_train, "fixed_images")) |
64
|
|
|
os.mkdir(os.path.join(path_to_train, "fixed_labels")) |
65
|
|
|
os.mkdir(os.path.join(path_to_train, "moving_images")) |
66
|
|
|
os.mkdir(os.path.join(path_to_train, "moving_labels")) |
67
|
|
|
|
68
|
|
|
|
69
|
|
View Code Duplication |
def move_files_into_correct_path( |
|
|
|
|
70
|
|
|
fnames, path_to_images_and_labels, new_path, suffix, sub_folder_name |
71
|
|
|
): |
72
|
|
|
os.chdir(os.path.join(path_to_images_and_labels, sub_folder_name)) |
73
|
|
|
for file in fnames: |
74
|
|
|
if "insp" in file: |
75
|
|
|
source = file |
76
|
|
|
destination = os.path.join(path_to_train, "fixed_" + suffix) |
77
|
|
|
shutil.move(source, destination) |
78
|
|
|
if "exp" in file: |
79
|
|
|
source = file |
80
|
|
|
destination = os.path.join(path_to_train, "moving_" + suffix) |
81
|
|
|
shutil.move(source, destination) |
82
|
|
|
|
83
|
|
|
|
84
|
|
|
if os.path.exists(path_to_images_and_labels): |
85
|
|
|
move_files_into_correct_path( |
86
|
|
|
images_fnames, path_to_images_and_labels, path_to_train, "images", "scans" |
87
|
|
|
) |
88
|
|
|
move_files_into_correct_path( |
89
|
|
|
labels_fnames, path_to_images_and_labels, path_to_train, "labels", "lungMasks" |
90
|
|
|
) |
91
|
|
|
|
92
|
|
|
os.chdir(main_path) |
93
|
|
|
|
94
|
|
|
######## MOVING FILES INTO TEST AND VALID DIRECTORY ######## |
95
|
|
|
path_to_test = os.path.join(path_to_data_folder, "test") |
96
|
|
|
path_to_valid = os.path.join(path_to_data_folder, "valid") |
97
|
|
|
|
98
|
|
View Code Duplication |
if os.path.exists(path_to_test) is not True: |
|
|
|
|
99
|
|
|
|
100
|
|
|
os.mkdir(path_to_test) |
101
|
|
|
os.mkdir(os.path.join(path_to_test, "fixed_images")) |
102
|
|
|
os.mkdir(os.path.join(path_to_test, "fixed_labels")) |
103
|
|
|
os.mkdir(os.path.join(path_to_test, "moving_images")) |
104
|
|
|
os.mkdir(os.path.join(path_to_test, "moving_labels")) |
105
|
|
|
|
106
|
|
|
ratio_of_test_and_valid_samples = 0.2 |
107
|
|
|
|
108
|
|
|
unique_case_names = [] |
109
|
|
|
for file in images_fnames: |
110
|
|
|
case_name_as_list = file.split("_")[0:2] |
111
|
|
|
case_name = case_name_as_list[0] + "_" + case_name_as_list[1] |
112
|
|
|
unique_case_names.append(case_name) |
113
|
|
|
unique_case_names = np.unique(unique_case_names) |
114
|
|
|
|
115
|
|
|
test_and_valid_cases = random.sample( |
116
|
|
|
list(unique_case_names), |
117
|
|
|
int(ratio_of_test_and_valid_samples * len(unique_case_names)), |
118
|
|
|
) |
119
|
|
|
test_cases = test_and_valid_cases[ |
120
|
|
|
0 : int(int(ratio_of_test_and_valid_samples * len(unique_case_names) / 2)) |
121
|
|
|
] |
122
|
|
|
valid_cases = test_and_valid_cases[ |
123
|
|
|
int(int(ratio_of_test_and_valid_samples * len(unique_case_names) / 2)) + 1 : |
124
|
|
|
] |
125
|
|
|
|
126
|
|
|
def move_test_cases_into_correct_path(test_cases, path_to_train, path_to_test): |
127
|
|
|
folder_names = os.listdir(path_to_train) |
128
|
|
|
os.chdir(path_to_train) |
129
|
|
|
for case in test_cases: |
130
|
|
|
for folder in folder_names: |
131
|
|
|
file_names = os.listdir(os.path.join(path_to_train, folder)) |
132
|
|
|
for file in file_names: |
133
|
|
|
if case in file: |
134
|
|
|
os.chdir(os.path.join(path_to_train, folder)) |
135
|
|
|
source = file |
136
|
|
|
destination = os.path.join(path_to_test, folder) |
137
|
|
|
shutil.move(source, destination) |
138
|
|
|
|
139
|
|
|
move_test_cases_into_correct_path(test_cases, path_to_train, path_to_test) |
140
|
|
|
|
141
|
|
|
os.mkdir(path_to_valid) |
142
|
|
|
os.mkdir(os.path.join(path_to_valid, "fixed_images")) |
143
|
|
|
os.mkdir(os.path.join(path_to_valid, "fixed_labels")) |
144
|
|
|
os.mkdir(os.path.join(path_to_valid, "moving_images")) |
145
|
|
|
os.mkdir(os.path.join(path_to_valid, "moving_labels")) |
146
|
|
|
|
147
|
|
|
move_test_cases_into_correct_path(valid_cases, path_to_train, path_to_valid) |
148
|
|
|
|
149
|
|
|
######## NAMING FILES SUCH THAT THEIR NAMES MATCH FOR PAIRING ######## |
150
|
|
|
|
151
|
|
|
# name all files such that names match exactly for training |
152
|
|
|
for folder in os.listdir(path_to_train): |
153
|
|
|
path_to_folder = os.path.join(path_to_train, folder) |
154
|
|
|
os.chdir(path_to_folder) |
155
|
|
|
for file in os.listdir(path_to_folder): |
156
|
|
|
if "_insp" in file: |
157
|
|
|
new_name = file.replace("_insp", "") |
158
|
|
|
elif "_exp" in file: |
159
|
|
|
new_name = file.replace("_exp", "") |
160
|
|
|
else: |
161
|
|
|
continue |
162
|
|
|
source = file |
163
|
|
|
destination = new_name |
164
|
|
|
os.rename(source, destination) |
165
|
|
|
|
166
|
|
|
# name all files such that names match exactly for testing |
167
|
|
|
|
168
|
|
|
for folder in os.listdir(path_to_test): |
169
|
|
|
path_to_folder = os.path.join(path_to_test, folder) |
170
|
|
|
os.chdir(path_to_folder) |
171
|
|
|
for file in os.listdir(path_to_folder): |
172
|
|
|
if "_insp" in file: |
173
|
|
|
new_name = file.replace("_insp", "") |
174
|
|
|
elif "_exp" in file: |
175
|
|
|
new_name = file.replace("_exp", "") |
176
|
|
|
else: |
177
|
|
|
continue |
178
|
|
|
source = file |
179
|
|
|
destination = new_name |
180
|
|
|
os.rename(source, destination) |
181
|
|
|
|
182
|
|
|
# name all files such that names match exactly for validation |
183
|
|
|
|
184
|
|
|
for folder in os.listdir(path_to_valid): |
185
|
|
|
path_to_folder = os.path.join(path_to_valid, folder) |
186
|
|
|
os.chdir(path_to_folder) |
187
|
|
|
for file in os.listdir(path_to_folder): |
188
|
|
|
if "_insp" in file: |
189
|
|
|
new_name = file.replace("_insp", "") |
190
|
|
|
elif "_exp" in file: |
191
|
|
|
new_name = file.replace("_exp", "") |
192
|
|
|
else: |
193
|
|
|
continue |
194
|
|
|
source = file |
195
|
|
|
destination = new_name |
196
|
|
|
os.rename(source, destination) |
197
|
|
|
|
198
|
|
|
print("All files moved and restructured") |
199
|
|
|
|
200
|
|
|
shutil.rmtree(os.path.join(path_to_images_and_labels)) |
201
|
|
|
os.chdir(main_path) |
202
|
|
|
|
203
|
|
|
######## NOW WE NEED TO RESCALE EACH IMAGE ######## |
204
|
|
|
|
205
|
|
|
data_dir = f"demos/paired_ct_lung/{data_folder_name}" |
206
|
|
|
folders = os.listdir(data_dir) |
207
|
|
|
|
208
|
|
|
for folder in folders: |
209
|
|
|
subfolders = os.listdir(os.path.join(data_dir, folder)) |
210
|
|
|
print("\n Working on ", folder, ", progress:") |
211
|
|
|
for subfolder in tqdm(subfolders): |
212
|
|
|
files = os.listdir(os.path.join(data_dir, folder, subfolder)) |
213
|
|
|
for file in files: |
214
|
|
View Code Duplication |
if file.startswith("case_020"): # this case did not laod correctly |
|
|
|
|
215
|
|
|
os.remove(os.path.join(data_dir, folder, subfolder, file)) |
216
|
|
|
else: |
217
|
|
|
im_data = np.asarray( |
218
|
|
|
nib.load(os.path.join(data_dir, folder, subfolder, file)).dataobj, |
219
|
|
|
dtype=np.float32, |
220
|
|
|
) |
221
|
|
|
if np.max(im_data) > 255.0: |
222
|
|
|
im_data = ((im_data + 285) / (3770 + 285)) * 255.0 # rescale image |
223
|
|
|
img = nib.Nifti1Image(im_data, affine=None) |
224
|
|
|
nib.save(img, os.path.join(data_dir, folder, subfolder, file)) |
225
|
|
|
if np.max(img.dataobj) > 255.0: |
226
|
|
|
print( |
227
|
|
|
"Recheck the following file: ", |
228
|
|
|
os.path.join(data_dir, folder, subfolder, file), |
229
|
|
|
) |
230
|
|
|
nib.save(img, os.path.join(data_dir, folder, subfolder, file)) |
231
|
|
|
|
232
|
|
|
######## DOWNLOAD MODEL CKPT FROM MODEL ZOO ######## |
233
|
|
|
|
234
|
|
|
url = "https://github.com/DeepRegNet/deepreg-model-zoo/raw/master/demo/paired_ct_lung/20210110.zip" |
235
|
|
|
fname = "pretrained.zip" |
236
|
|
|
os.chdir(os.path.join(main_path, project_dir)) |
237
|
|
|
|
238
|
|
|
# download and unzip into pretrained subfolder |
239
|
|
|
get_file(os.path.join(os.getcwd(), fname), url) |
240
|
|
|
with zipfile.ZipFile(fname, "r") as zip_ref: |
241
|
|
|
zip_ref.extractall(os.path.join(data_folder_name, "pretrained")) |
242
|
|
|
|
243
|
|
|
# remove pretrained.zip |
244
|
|
|
os.remove(fname) |
245
|
|
|
print( |
246
|
|
|
"Pretrained model downloaded: %s" |
247
|
|
|
% os.path.abspath(os.path.join(data_folder_name, "pretrained")) |
248
|
|
|
) |
249
|
|
|
|