/
nic_infer_segmentation_batch.py
180 lines (151 loc) · 6.34 KB
/
nic_infer_segmentation_batch.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
# ---------------------------------------------------------------------------------------
# MS lesion segmentation pipeline
# ---------------------------------
# - incorporates:
# - MRI identification
# - registration
# - skull stripping
# - MS lesion segmentation using the CNN Valverde et al (NI2017)
#
# Sergi Valverde 2017
# svalverde@eia.udg.edu
# ---------------------------------------------------------------------------------------
import os
import argparse
import sys
import platform
import time
import ConfigParser
from utils.load_options import load_options, print_options
from utils.preprocess import preprocess_scan
from utils.postprocess import invert_registration
from shutil import copyfile
os.system('cls' if platform.system() == 'Windows' else 'clear')
print "##################################################"
print "# MS WM lesion segmentation #"
print "# #"
print "# ------------------------------- #"
print "# (c) Sergi Valverde 2019 #"
print "# Neuroimage Computing Group #"
print "# ------------------------------- #"
print "##################################################\n"
# link related libraries
CURRENT_PATH = os.path.split(os.path.realpath(__file__))[0]
sys.path.append(os.path.join(CURRENT_PATH, 'libs'))
# load options from input
parser = argparse.ArgumentParser()
parser.add_argument('--docker',
dest='docker',
action='store_true')
parser.set_defaults(docker=False)
args = parser.parse_args()
container = args.docker
# --------------------------------------------------
# load default options and update them with user information
# from utils.load_options import *
# --------------------------------------------------
default_config = ConfigParser.SafeConfigParser()
default_config.read(os.path.join(CURRENT_PATH, 'config', 'default.cfg'))
user_config = ConfigParser.RawConfigParser()
user_config.read(os.path.join(CURRENT_PATH, 'config', 'configuration.cfg'))
# read user's configuration file
options = load_options(default_config, user_config)
if options['debug']:
print_options(options)
# tensorflow backend
device = str(options['gpu_number'])
print "DEBUG: ", device
os.environ['KERAS_BACKEND'] = 'tensorflow'
os.environ["CUDA_VISIBLE_DEVICES"] = device
# set paths taking into account the host OS
host_os = platform.system()
if host_os == 'Linux':
options['niftyreg_path'] = CURRENT_PATH + '/libs/linux/niftyreg'
options['robex_path'] = CURRENT_PATH + '/libs/linux/ROBEX/runROBEX.sh'
options['test_slices'] = 256
elif host_os == 'Windows':
options['niftyreg_path'] = os.path.normpath(
os.path.join(CURRENT_PATH, 'libs', 'win', 'niftyreg'))
options['robex_path'] = os.path.normpath(
os.path.join(CURRENT_PATH, 'libs', 'win', 'ROBEX', 'runROBEX.bat'))
options['test_slices'] = 256
else:
"> ERROR: The OS system", host_os, "is not currently supported"
from CNN.base import test_cascaded_model
from CNN.build_model import cascade_model
# --------------------------------------------------
# net configuration
# take into account if the pretrained models have to be used
# all images share the same network model
# --------------------------------------------------
options['full_train'] = True
options['load_weights'] = True
options['weight_paths'] = os.path.join(CURRENT_PATH, 'nets')
options['net_verbose'] = 0
model = cascade_model(options)
# --------------------------------------------------
# process each of the scans
# - image identification
# - image registration
# - skull-stripping
# - WM segmentation
# --------------------------------------------------
if container:
options['test_folder'] = os.path.normpath('/data' + options['test_folder'])
else:
options['test_folder'] = os.path.normpath(options['test_folder'])
# set task to train
options['task'] = 'inference'
# list scans
scan_list = os.listdir(options['test_folder'])
scan_list.sort()
for scan in scan_list:
total_time = time.time()
options['tmp_scan'] = scan
# --------------------------------------------------
# move things to a tmp folder before starting
# --------------------------------------------------
current_folder = os.path.join(options['test_folder'], scan)
options['tmp_folder'] = os.path.normpath(
os.path.join(current_folder, 'tmp'))
# --------------------------------------------------
# preprocess scans
# --------------------------------------------------
preprocess_scan(current_folder, options)
# --------------------------------------------------
# WM MS lesion inference
# --------------------------------------------------
seg_time = time.time()
"> CNN:", scan, "running WM lesion segmentation"
sys.stdout.flush()
options['test_scan'] = scan
test_x_data = {scan: {m: os.path.join(options['tmp_folder'], n)
for m, n in zip(options['modalities'],
options['x_names'])}}
out_seg = test_cascaded_model(model, test_x_data, options)
print "> INFO:", scan, "CNN Segmentation time: ",\
round(time.time() - seg_time), "sec"
# If input images have been registered before segmentation -> T1w space,
# then resample the segmentation back to the original space
if options['register_modalities']:
print "> INFO:", scan, "Inverting lesion segmentation masks"
invert_registration(current_folder, options)
print "> INFO:", scan, "total pipeline time: ",\
round(time.time() - total_time), "sec"
# remove tmps if not set
if options['save_tmp'] is False:
try:
copyfile(os.path.join(current_folder,
options['experiment'],
options['experiment'] +
'_out_CNN.nii.gz'),
os.path.join(current_folder,
'out_seg_' +
options['experiment'] +
'.nii.gz'))
os.rmdir(options['tmp_folder'])
os.rmdir(os.path.join(options['current_folder'],
options['experiment']))
except:
pass
print "> INFO: All processes have been finished. Have a good day!"