def train_net(
net,
x_train,
y_train,
images,
b_name='\033[30mbaseline_%s\033[0m',
f_name='\033[30mfollow_%s\033[0m',
d_name='\033[30mdeformation_%s\033[0m'
):
defo = False
d_inputs = []
c = color_codes()
n_images = len(images)
# We try to get the last weights to keep improving the net over and over
if isinstance(x_train, tuple):
defo = True
x_train, defo_train = x_train
defo_train = np.split(defo_train, len(images), axis=1)
d_inputs = [(d_name % im, np.squeeze(d_im)) for im, d_im in zip(images, defo_train)]
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Training' + c['nc'])
n_channels = x_train.shape[1]
x_train = np.split(x_train, n_channels, axis=1)
b_inputs = [(b_name % im, x_im) for im, x_im in zip(images, x_train[:n_images])]
f_inputs = [(f_name % im, x_im) for im, x_im in zip(images, x_train[n_images:])]
inputs = dict(b_inputs + f_inputs) if not defo else dict(b_inputs + f_inputs + d_inputs)
net.fit(inputs, y_train)
def test_net(net, p, outputname):
c = color_codes()
options = parse_inputs()
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['test_size']
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname_path = os.path.join(patient_path, outputname + '.nii.gz')
pr_outputname_path = os.path.join(patient_path, outputname + '.pr.nii.gz')
try:
image = load_nii(outputname_path).get_data()
except IOError:
print('%s[%s] %sTesting the network%s' % (c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
nii = load_nii(p[0])
roi = nii.get_data().astype(dtype=np.bool)
centers = get_mask_voxels(roi)
test_samples = np.count_nonzero(roi)
image = np.zeros_like(roi).astype(dtype=np.uint8)
pr = np.zeros_like(roi).astype(dtype=np.float32)
print('%s[%s] %s<Creating the probability map %s%s%s%s - %s%s%s%s (%d samples)>%s' % (
c['c'], strftime("%H:%M:%S"),
c['g'], c['b'], p_name, c['nc'],
c['g'], c['b'], outputname, c['nc'],
c['g'], test_samples, c['nc']
))
n_centers = len(centers)
image_list = [load_norm_list(p)]
for i in range(0, n_centers, batch_size):
print(
'%f%% tested (step %d/%d)' % (100.0 * i / n_centers, (i / batch_size) + 1, -(-n_centers/batch_size)),
end='\r'
)
sys.stdout.flush()
centers_i = [centers[i:i + batch_size]]
x = get_patches_list(image_list, centers_i, patch_size, True)
x = np.concatenate(x).astype(dtype=np.float32)
y_pr_pred = net.predict(x, batch_size=options['batch_size'])
[x, y, z] = np.stack(centers_i[0], axis=1)
# We store the results
image[x, y, z] = np.argmax(y_pr_pred, axis=1).astype(dtype=np.int8)
pr[x, y, z] = y_pr_pred[:, 1].astype(dtype=np.float32)
print(' '.join([''] * 50), end='\r')
sys.stdout.flush()
# Post-processing (Basically keep the biggest connected region)
# image = get_biggest_region(image)
print('%s -- Saving image %s%s%s' % (c['g'], c['b'], outputname_path, c['nc']))
nii.get_data()[:] = image
nii.to_filename(outputname_path)
nii.get_data()[:] = pr
nii.to_filename(pr_outputname_path)
return image
def main():
# Init
options = parse_inputs()
c = color_codes()
# Data loading (or preparation)
d_path = options['val_dir']
gt_names = sorted(filter(
lambda xi: not os.path.isdir(xi) and re.search(options['lab_tag'], xi),
os.listdir(d_path)),
key=find_number)
cases_pre = [str(find_number(r)) for r in gt_names]
gt_names = [
gt for c, gt in zip(cases_pre, gt_names)
if find_file('Z{:}.jpg'.format(c), d_path)
]
cases = [c for c in cases_pre if find_file('Z{:}.jpg'.format(c), d_path)]
print('%s[%s] %s<Tree detection pipeline>%s' %
(c['c'], time.strftime("%H:%M:%S"), c['y'], c['nc']))
''' <Detection task> '''
net_name = 'tree-detection.nDEM.unet'
train(cases, gt_names, net_name, 'nDEM')
net_name = 'tree-detection.DEM.unet'
train(cases, gt_names, net_name, 'DEM')
eval(cases, gt_names)
def main():
options = parse_inputs()
c = color_codes()
experimental = options['experimental']
exp_s = c['b'] + '(experimental %d)' % experimental if experimental else c[
'b'] + '(baseline)'
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
'Starting cross-validation ' + exp_s + c['nc'])
# N-fold cross validation main loop (we'll do 2 training iterations with testing for each patient)
data_names, label_names = get_names_from_path(options)
folds = len(data_names)
fold_generator = izip(
nfold_cross_validation(data_names, label_names, n=folds),
xrange(folds))
dsc_results = list()
for (train_data, train_labels, test_data,
test_labels), i in fold_generator:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] +
'Fold %d/%d: ' % (i + 1, folds) + c['g'] +
'Number of training/testing images (%d=%d/%d)' %
(len(train_data), len(train_labels), len(test_data)) + c['nc'])
# Prepare the data relevant to the leave-one-out (subtract the patient from the dataset and set the path)
# Also, prepare the network
if not check_image_list(test_data, options):
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] +
c['g'] + 'Training' + c['nc'])
net = train_net(i, train_data, train_labels, options)
else:
net = None
# Then we test the net.
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] +
c['g'] + 'Testing' + c['nc'])
for p, gt_name in zip(test_data, test_labels):
image, gt = test_net(net, p, gt_name, options)
p_name = '-'.join(
p[0].rsplit('/')[-1].rsplit('.')[0].rsplit('-')[:-1])
vals = np.unique(gt.flatten())
gt_mask = np.sum(map(
lambda (l, val): np.array(gt == val, dtype=np.uint8) * l,
enumerate(vals)),
axis=0)
results = (str.capitalize(p_name), dsc_seg(gt_mask == 1,
image == 1),
dsc_seg(gt_mask == 2,
image == 2), dsc_seg(gt_mask == 3, image == 3))
dsc_results.append(results)
print('%s DSC: %f/%f/%f' % results)
dsc_results = sorted(dsc_results,
cmp=lambda x, y: int(x[0][8:]) - int(y[0][8:]))
for results in dsc_results:
print(c['c'] + '%s DSC: \033[32;1m%f/%f/%f' % results + c['nc'])
f_dsc = tuple(
np.asarray([results[1:] for results in dsc_results]).mean(axis=0))
print(c['c'] + 'Final results DSC: \033[32;1m%f/%f/%f' % f_dsc + c['nc'])
def get_norm_patch_vectors(image_names, positive_masks, negative_masks, size, balanced=True, random_state=42):
# Get all the centers for each image
c = color_codes()
print(c['lgy'] + ' ' + image_names[0].rsplit('/')[-1] + c['nc'])
# Get all the patches for each image
positive_centers, negative_centers = get_centers_from_masks(positive_masks, negative_masks, balanced, random_state)
return get_patch_vectors(norm_image_generator(image_names), positive_centers, negative_centers, size)
def permute(x, seed, datatype=np.float32):
c = color_codes()
print(c['g'] + ' Vector shape ='
' (' + ','.join([c['bg'] + str(length) + c['nc'] + c['g'] for length in x.shape]) + ')' + c['nc'])
np.random.seed(seed)
x_permuted = np.random.permutation(x.astype(dtype=datatype))
return x_permuted
def train_net(net, net_name, nlabels):
options = parse_inputs()
c = color_codes()
# Data stuff
train_data, train_labels = get_names_from_path(options)
# Prepare the net architecture parameters
dfactor = options['dfactor']
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
conv_blocks = options['conv_blocks']
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
val_rate = options['val_rate']
preload = options['preload']
fc_width = patch_width - sum(kernel_size_list) + conv_blocks
fc_shape = (fc_width, ) * 3
try:
net = load_model(net_name + '.md')
except IOError:
centers = np.random.permutation(
get_cnn_centers(train_data[:, 0], train_labels, balanced=balanced))
print(' '.join([''] * 15) + c['g'] + 'Total number of centers = ' +
c['b'] + '(%d centers)' % (len(centers)) + c['nc'])
for i in range(dfactor):
print(' '.join([''] * 16) + c['g'] + 'Round ' + c['b'] + '%d' %
(i + 1) + c['nc'] + c['g'] + '/%d' % dfactor + c['nc'])
batch_centers = centers[i::dfactor]
print(' '.join([''] * 16) + c['g'] + 'Loading data ' + c['b'] +
'(%d centers)' % (len(batch_centers)) + c['nc'])
x, y = load_patches_train(
image_names=train_data,
label_names=train_labels,
batch_centers=batch_centers,
size=patch_size,
fc_shape=fc_shape,
nlabels=nlabels,
preload=preload,
)
print(' '.join([''] * 16) + c['g'] + 'Training the model for ' +
c['b'] +
'(%d parameters)' % net.count_trainable_parameters() +
c['nc'])
net.fit(x,
y,
batch_size=batch_size,
validation_split=val_rate,
epochs=epochs)
net.save(net_name + '.mod')
def main():
c = color_codes()
options = parse_inputs()
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting the registration training' + c['nc'])
dir_name = options['dir_name']
baseline_name = options['b_folder']
followup_name = options['f_folder']
image_name = options['im_name']
input_size = options['input_size']
n_filters = options['number_filters']
pool_size = options['pool_size']
conv_width = options['conv_width']
conv_blocks = options['conv_blocks']
augment_p = options['augment_p']
seed = np.random.randint(np.iinfo(np.int32).max)
input_size_s = 'x'.join([str(length) for length in input_size])
sufix = 's%s.c%s.n%s.a%f' % (input_size_s, conv_width, n_filters, augment_p)
net_name = os.path.join(dir_name, 'deep-exp_registration.' + sufix + '.')
net = create_cnn3d_register(
input_shape=input_size,
convo_size=conv_width,
convo_blocks=conv_blocks,
pool_size=pool_size,
number_filters=n_filters,
data_augment_p=augment_p,
patience=100,
name=net_name,
epochs=2000
)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + 'Loading the data for ' + c['b'] + 'iteration 1' + c['nc'])
names = get_names_from_path(dir_name, baseline_name, followup_name, image_name)
x_train, y_train = load_register_data(
names=names,
image_size=input_size,
seed=seed,
)
# We try to get the last weights to keep improving the net over and over
try:
net.load_params_from(net_name + 'model_weights.pkl')
except IOError:
pass
train_net(net, x_train, y_train)
test_net(net, x_train)
def test_survival(net, image_names, features, slices, n_slices):
c = color_codes()
x_vol = get_reshaped_data(image_names,
slices, (224, 224),
n_slices=n_slices)
x_vol = np.stack(x_vol, axis=0)
x = [x_vol.astype(np.float32), np.array(features, dtype=np.float32)]
print('%s[%s] %sTesting the survival network %s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
survival = net.predict(x, batch_size=1)[0]
return np.fabs(np.squeeze(survival))
def test_net(
net,
x,
b_name='\033[30mbaseline\033[0m',
f_name='\033[30mfollow\033[0m',
loc_name='\033[33mloc_net\033[0m'
):
import theano
import lasagne
c = color_codes()
# We try to get the last weights to keep improving the net over and over
x_test = np.split(x, 2, axis=1)
b_inputs = (b_name, x_test[0])
f_inputs = (f_name, x_test[1])
inputs = dict([b_inputs, f_inputs])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Predicting (' + c['b'] + 'images' + c['nc'] + c['g'] + ')' + c['nc'])
print(' Testing vector shape ='
' (' + ','.join([str(length) for length in x.shape]) + ')')
y_test = net.predict(inputs)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Predicting (' + c['b'] + 'transformations' + c['nc'] + c['g'] + ')' + c['nc'])
f = theano.function(
[
net.layers_[b_name].input_var,
net.layers_[f_name].input_var
],
lasagne.layers.get_output(
net.layers_[loc_name],
{
b_name: net.layers_[b_name].input_var,
f_name: net.layers_[f_name].input_var
}
),
name='registration'
)
rand_transf = [random_affine3d_matrix(x_range=1/36.0, y_range=1/36.0, z_range=1/36.0) for x_t in x_test[0]]
x_test_random = np.stack([affine_transform(x_t, t) for x_t, t in zip(x_test[0], rand_transf)])
transforms = f(x_test_random, x_test[0])
for tt, t in zip(transforms, rand_transf):
print('%s %s' % (','.join(['%f' % ts for ts in t[0, :]]), ','.join(['%f' % tts for tts in tt[:4]])))
print('%s = %s' % (','.join(['%f' % ts for ts in t[1, :]]), ','.join(['%f' % tts for tts in tt[4:8]])))
print('%s %s' % (','.join(['%f' % ts for ts in t[2, :]]), ','.join(['%f' % tts for tts in tt[8:12]])))
return y_test, transforms
def main():
# Init
c = color_codes()
options = parse_inputs()
patch_size = options['patch_size']
filters = options['filters']
depth = options['blocks']
# Prepare the sufix that will be added to the results for the net and images
n_folds = 5
filters_s = '-filt%d' % filters
patch_s = '-ps%d' % patch_size if patch_size is not None else ''
depth_s = '-d%d' % depth
print(
'%s[%s] %s<BRATS 2019 pipeline testing>%s' % (
c['c'], strftime("%H:%M:%S"), c['y'], c['nc']
)
)
''' <Survival task> '''
net_name = 'brats2019-survival%s%s%s' % (
filters_s, depth_s, patch_s
)
print(
'%s[%s] %sStarting cross-validation (survival) - %d folds%s' % (
c['c'], strftime("%H:%M:%S"), c['g'], n_folds, c['nc']
)
)
# train_test_survival(net_name, n_folds)
''' <Segmentation task> '''
print(
'%s[%s] %sStarting cross-validation (segmentation) - %d folds%s' % (
c['c'], strftime("%H:%M:%S"), c['g'], n_folds, c['nc']
)
)
net_name = 'brats2019-seg-nodrop%s%s' % (
filters_s, depth_s
)
train_test_seg(net_name, n_folds)
test_seg_validation(net_name)
def main():
c = color_codes()
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
try:
net = load_model('/home/mariano/Desktop/test.tf')
except IOError:
x = Input([784])
x_image = Reshape([28, 28, 1])(x)
x_conv1 = Conv(filters=32,
kernel_size=(5, 5),
activation='relu',
padding='same')(x_image)
h_pool1 = MaxPool((2, 2), padding='same')(x_conv1)
h_conv2 = Conv(filters=64,
kernel_size=(5, 5),
activation='relu',
padding='same')(h_pool1)
h_pool2 = MaxPool((2, 2), padding='same')(h_conv2)
h_fc1 = Dense(1024, activation='relu')(h_pool2)
h_drop = Dropout(0.5)(h_fc1)
y_conv = Dense(10)(h_drop)
net = Model(x,
y_conv,
optimizer='adam',
loss='categorical_cross_entropy',
metrics='accuracy')
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + c['b'] +
'Original (MNIST)' + c['nc'] + c['g'] + ' net ' + c['nc'] + c['b'] +
'(%d parameters)' % net.count_trainable_parameters() + c['nc'])
net.fit(mnist.train.images,
mnist.train.labels,
val_data=mnist.test.images,
val_labels=mnist.test.labels,
patience=10,
epochs=200,
batch_size=1024)
save_model(net, '/home/mariano/Desktop/test.tf')
def train_net(
net,
x,
y,
b_name='\033[30mbaseline\033[0m',
f_name='\033[30mfollow\033[0m'
):
c = color_codes()
print(' Training vector shape ='
' (' + ','.join([str(length) for length in x.shape]) + ')')
print(' Training labels shape ='
' (' + ','.join([str(length) for length in y.shape]) + ')')
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Training (' + c['b'] + 'registration' + c['nc'] + c['g'] + ')' + c['nc'])
x_train = np.split(x, 2, axis=1)
b_inputs = (b_name, x_train[0])
f_inputs = (f_name, x_train[1])
inputs = dict([b_inputs, f_inputs])
net.fit(inputs, y)
def get_defo_patch_vectors(image_names, masks, size=(5, 5, 5), balanced=True, random_state=42):
# Get all the centers for each image
c = color_codes()
print(c['lgy'] + ' ' + image_names[0].rsplit('/')[-1] + c['nc'])
defo = norm_defo_generator(image_names)
# We divide the 4D deformation image into 3D images for each component (x, y, z)
defo_xyz = [np.squeeze(d) for d in [np.split(d, 3, axis=4) for d in defo]]
positive_masks, negative_masks = masks
positive_centers, negative_centers = get_centers_from_masks(positive_masks, negative_masks, balanced, random_state)
patches = np.stack(
[np.concatenate(get_patch_vectors(list(d), positive_centers, negative_centers, size))
for d in izip(*defo_xyz)],
axis=1
)
# Get all the patches for each image
return patches
def train_greenspan(
net,
x_train,
y_train,
images,
b_name='\033[30mbaseline_%s\033[0m',
f_name='\033[30mfollow_%s\033[0m'
):
c = color_codes()
n_axis = x_train.shape[1]
n_images = x_train.shape[2]/2
print(' Training vector shape ='
' (' + ','.join([str(length) for length in x_train.shape]) + ')')
print(' Training labels shape ='
' (' + ','.join([str(length) for length in y_train.shape]) + ')')
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Training' + c['nc'])
# We try to get the last weights to keep improving the net over and over
x_train = np.split(x_train, n_axis, axis=1)
b_inputs = [(b_name % im, np.squeeze(x_im[:, :, :n_images, :, :])) for im, x_im in zip(images, x_train)]
f_inputs = [(f_name % im, np.squeeze(x_im[:, :, n_images:, :, :])) for im, x_im in zip(images, x_train)]
inputs = dict(b_inputs + f_inputs)
net.fit(inputs, y_train)
def train_all_nets(
names_lou,
defo_names_lou,
mask_names,
roi_names,
net_names,
conv_blocks,
patch_sizes,
defo_sizes,
conv_sizes,
n_filters,
images,
pool_size,
dense_sizes,
epochs,
seed,
):
# We need to prepare the name list to load the leave-one-out data.
# Since these names are common for all the nets, we can define them before looping.
c = color_codes()
net_combos = itertools.product(zip(conv_blocks, patch_sizes, conv_sizes, defo_sizes), n_filters, dense_sizes)
nets = list()
# Data loading. We load the largest possible patch size, and then we resize everything
max_patch = max(patch_sizes)
max_defo = max(defo_sizes)
x_train, y_train = load_lesion_cnn_data(
names=names_lou,
mask_names=mask_names,
defo_names=defo_names_lou,
roi_names=roi_names,
pr_names=None,
patch_size=max_patch,
defo_size=max_defo,
random_state=seed
)
for ((blocks, patch, convo, defo), filters, dense), net_name in zip(net_combos, net_names):
net = create_cnn3d_longitudinal(
convo_blocks=blocks,
input_shape=(None, 6) + patch,
images=images,
convo_size=convo,
pool_size=pool_size,
dense_size=dense,
number_filters=filters,
padding='valid',
drop=0.5,
register=False,
defo=True,
patience=epochs,
name=net_name,
epochs=epochs
)
# First we check that we did not train that patient, in order to save time
try:
net.load_params_from(net_name + 'model_weights.pkl')
except IOError:
combo_s = ' [c%d.k%d.d%d.n%d]' % (
blocks,
convo[0],
dense,
filters
)
# Afterwards we train. Check the relevant training function.
# NOTE: We have to resize the training patches if necessary.
patch_ratio = (1, 1) + tuple(itertools.imap(lambda x, y: float(x) / y, patch, max_patch))
defo_ratio = (1, 1, 1) + tuple(itertools.imap(lambda x, y: float(x) / y, defo, max_defo))
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Patch shape = (' +
','.join([c['bg'] + str(length) + c['nc'] + c['g'] for length in patch]) + ')' +
combo_s + c['nc'])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Deformation shape = (' +
','.join([c['bg'] + str(length) + c['nc'] + c['g'] for length in defo]) + ')' +
combo_s + c['nc'])
train_net(
net=net,
x_train=(zoom(x_train[0], patch_ratio), zoom(x_train[1], defo_ratio)),
y_train=y_train,
images=images
) if max_patch != patch else train_net(
net=net,
x_train=x_train,
y_train=y_train,
images=images
)
for callback in net.on_epoch_finished:
if isinstance(callback, WeightsLogger):
callback.save()
with open(net_name + 'layers.pkl', 'wb') as fnet:
pickle.dump(net.layers, fnet, -1)
nets.append(net)
return nets
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
# Data loading parameters
preload = options['preload']
# Prepare the sufix that will be added to the results for the net and images
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
ub_s = '.ub' if not balanced else ''
params_s = (ub_s, patch_width, conv_s, filters_s, dense_size, epochs)
sufix = '%s.p%d.c%s.n%s.d%d.e%d' % params_s
preload_s = ' (with ' + c['b'] + 'preloading' + c['nc'] + c['c'] + ')' if preload else ''
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting training' + preload_s + c['nc'])
train_data, _ = get_names_from_path(options)
test_data, test_labels = get_names_from_path(options, False)
input_shape = (train_data.shape[1],) + patch_size
dsc_results_gan = list()
dsc_results_cnn = list()
dsc_results_caps = list()
train_data, train_labels = get_names_from_path(options)
centers_s = np.random.permutation(
get_cnn_centers(train_data[:, 0], train_labels, balanced=balanced)
)[::options['down_sampling']]
x_seg, y_seg = load_patches_ganseg_by_batches(
image_names=train_data,
label_names=train_labels,
source_centers=centers_s,
size=patch_size,
nlabels=5,
preload=preload,
batch_size=51200
)
y_seg_roi = np.empty((len(y_seg), 2), dtype=np.bool)
y_seg_roi[:, 0] = y_seg[:, 0]
y_seg_roi[:, 1] = np.sum(y_seg[:, 1:], axis=1)
for i, (p, gt_name) in enumerate(zip(test_data, test_labels)):
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
print('%s[%s] %sCase %s%s%s%s%s (%d/%d):%s' % (
c['c'],
strftime("%H:%M:%S"),
c['nc'],
c['c'],
c['b'],
p_name,
c['nc'],
c['c'],
i + 1,
len(test_data),
c['nc']
))
# ROI segmentation
adversarial_w = K.variable(0)
roi_cnn = get_brats_fc(input_shape, filters_list, kernel_size_list, dense_size, 2)
roi_caps = get_brats_caps(input_shape, filters_list, kernel_size_list, 8, 2)
roi_gan, _ = get_brats_gan_fc(
input_shape,
filters_list,
kernel_size_list,
dense_size,
2,
lambda_var=adversarial_w
)
train_nets(
x=x_seg,
y=y_seg_roi,
gan=roi_gan,
cnn=roi_cnn,
caps=roi_caps,
p=p,
name='brats2017-roi' + sufix,
adversarial_w=adversarial_w
)
# Tumor substructures net
adversarial_w = K.variable(0)
seg_cnn = get_brats_fc(input_shape, filters_list, kernel_size_list, dense_size, 5)
seg_caps = get_brats_caps(input_shape, filters_list, kernel_size_list, 8, 5)
seg_gan_tr, seg_gan_tst = get_brats_gan_fc(
input_shape,
filters_list,
kernel_size_list,
dense_size,
5,
lambda_var=adversarial_w
)
roi_net_conv_layers = [l for l in roi_gan.layers if 'conv' in l.name]
seg_net_conv_layers = [l for l in seg_gan_tr.layers if 'conv' in l.name]
for lr, ls in zip(roi_net_conv_layers[:conv_blocks], seg_net_conv_layers[:conv_blocks]):
ls.set_weights(lr.get_weights())
train_nets(
x=x_seg,
y=y_seg,
gan=seg_gan_tr,
cnn=seg_cnn,
caps=seg_caps,
p=p,
name='brats2017-full' + sufix,
adversarial_w=adversarial_w
)
image_cnn_name = os.path.join(patient_path, p_name + '.cnn.test')
try:
image_cnn = load_nii(image_cnn_name + '.nii.gz').get_data()
except IOError:
image_cnn = test_net(seg_cnn, p, image_cnn_name)
image_caps_name = os.path.join(patient_path, p_name + '.caps.test')
try:
image_caps = load_nii(image_caps_name + '.nii.gz').get_data()
except IOError:
image_caps = test_net(seg_caps, p, image_caps_name)
image_gan_name = os.path.join(patient_path, p_name + '.gan.test')
try:
image_gan = load_nii(image_gan_name + '.nii.gz').get_data()
except IOError:
image_gan = test_net(seg_gan_tst, p, image_gan_name)
results_cnn = check_dsc(gt_name, image_cnn)
dsc_string = c['g'] + '/'.join(['%f'] * len(results_cnn)) + c['nc']
print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] + ' CNN DSC: ' +
dsc_string % tuple(results_cnn))
results_caps = check_dsc(gt_name, image_caps)
dsc_string = c['g'] + '/'.join(['%f'] * len(results_caps)) + c['nc']
print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] + ' CAPS DSC: ' +
dsc_string % tuple(results_caps))
results_gan = check_dsc(gt_name, image_gan)
dsc_string = c['g'] + '/'.join(['%f'] * len(results_gan)) + c['nc']
print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] + ' GAN DSC: ' +
dsc_string % tuple(results_gan))
dsc_results_cnn.append(results_cnn)
dsc_results_caps.append(results_caps)
dsc_results_gan.append(results_gan)
f_dsc = tuple(
[np.array([dsc[i] for dsc in dsc_results_cnn if len(dsc) > i]).mean() for i in range(3)]
) + tuple(
[np.array([dsc[i] for dsc in dsc_results_caps if len(dsc) > i]).mean() for i in range(3)]
) + tuple(
[np.array([dsc[i] for dsc in dsc_results_gan if len(dsc) > i]).mean() for i in range(3)]
)
print('Final results DSC: (%f/%f/%f) vs (%f/%f/%f) vs (%f/%f/%f)' % f_dsc)
def test_all_nets(
path,
names_test,
defo_names_test,
roi_name,
nets,
case,
batch_size,
patch_sizes,
defo_sizes,
dense_sizes,
n_filters,
sufixes,
iter_name,
train_case=False
):
c = color_codes()
net_combos = itertools.product(zip(patch_sizes, defo_sizes), n_filters, dense_sizes)
image_nii = load_nii(names_test[0])
mask_nii = load_nii(roi_name)
images = list()
for net, sufix, ((patch_size, defo_size), _, _) in zip(nets, sufixes, net_combos):
outputname = os.path.join(path, 't' + case + sufix + iter_name + '.nii.gz')
# We save time by checking if we already tested that patient.
try:
image_nii = load_nii(outputname)
image = image_nii.get_data()
if train_case:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + ' Patient ' + names_test[0].rsplit('/')[-4] + ' already done' + c['nc'])
except IOError:
if train_case:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + ' Testing with patient ' + c['b'] + names_test[0].rsplit('/')[-4] + c['nc'])
else:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map for net: ' +
c['b'] + sufix + c['nc'] + c['g'] + '>' + c['nc'])
image = test_net(
net=net,
names=names_test,
mask=mask_nii.get_data(),
batch_size=batch_size,
patch_size=patch_size,
defo_size=defo_size,
image_size=image_nii.get_data().shape,
images=['flair', 'pd', 't2'],
d_names=defo_names_test
)
if train_case:
print(c['g'] + ' -- Saving image ' + c['b'] + outputname + c['nc'])
image_nii.get_data()[:] = image
image_nii.to_filename(outputname)
images.append(image)
return images
def main():
options = parse_inputs()
options['use_flair'] = True
options['use_pd'] = True
options['use_t2'] = True
conv_blocks = [1, 2, 3]
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
conv_width = options['conv_width']
conv_sizes = [[conv_width]*blocks for blocks in conv_blocks]
last_width = options['last_width']
patch_widths = [blocks * (conv_width - 1) + last_width for blocks in conv_blocks]
patch_sizes = [(width, width, width) for width in patch_widths]
pool_size = 1
batch_size = 100000
dense_sizes = [16, 64, 128, 256]
n_filters = [32, 64]
# Prepare the sufix that will be added to the results for the net and images
flair_name = 'flair'
pd_name = 'pd'
t2_name = 't2'
images = [name for name in [flair_name, pd_name, t2_name] if name is not None]
parameter_combos = itertools.product(patch_widths, conv_sizes, n_filters, dense_sizes)
name_parameters = [(
combo[0],
'c'.join([str(cs) for cs in combo[1]]),
combo[2],
combo[3],
epochs
) for combo in parameter_combos]
sufixes = ['.p%d.c%s.n%s.d%d.e%d' % p for p in name_parameters]
# Prepare the data names
mask_name = options['mask']
wm_name = options['wm_mask']
dir_name = options['dir_name']
patients = [f for f in sorted(os.listdir(dir_name))
if os.path.isdir(os.path.join(dir_name, f))]
n_patients = len(patients)
names = get_names_from_path(dir_name, options, patients)
defo_names = get_defonames_from_path(dir_name, options, patients)
defo_widths = [blocks * 2 + 1 for blocks in conv_blocks]
defo_sizes = [(width, width, width) for width in defo_widths]
# Random initialisation
seed = np.random.randint(np.iinfo(np.int32).max)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting leave-one-out' + c['nc'])
# Leave-one-out main loop (we'll do 2 training iterations with testing for each patient)
for i in range(0, n_patients):
# Prepare the data relevant to the leave-one-out (subtract the patient from the dataset and set the path)
# Also, prepare the network
case = patients[i]
path = os.path.join(dir_name, case)
names_lou = np.concatenate([names[:, :i], names[:, i + 1:]], axis=1)
defo_names_lou = np.concatenate([defo_names[:, :i], defo_names[:, i + 1:]], axis=1)
paths = [os.path.join(dir_name, p) for p in np.concatenate([patients[:i], patients[i + 1:]])]
mask_names = [os.path.join(p_path, mask_name) for p_path in paths]
wm_names = [os.path.join(p_path, wm_name) for p_path in paths]
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + 'Patient ' + c['b'] + case + c['nc'] +
c['g'] + ' (%d/%d)' % (i+1, n_patients))
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Running iteration ' + c['b'] + '1' + c['nc'] + c['g'] + '>' + c['nc'])
net_names = [os.path.join(path, 'deep-generalise.init' + sufix + '.') for sufix in sufixes]
nets = train_all_nets(
names_lou=names_lou,
defo_names_lou=defo_names_lou,
mask_names=mask_names,
roi_names=wm_names,
net_names=net_names,
conv_blocks=conv_blocks,
patch_sizes=patch_sizes,
defo_sizes=defo_sizes,
conv_sizes=conv_sizes,
n_filters=n_filters,
images=images,
pool_size=pool_size,
dense_sizes=dense_sizes,
epochs=epochs,
seed=seed
)
# Then we test the net.
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + '<Testing iteration ' + c['b'] + '1' +
c['nc'] + c['g'] + '>' + c['nc'])
names_test = get_names_from_path(path, options)
defo_names_test = get_defonames_from_path(path, options)
roi_name = os.path.join(path, wm_name)
test_all_nets(
path=path,
names_test=names_test,
defo_names_test=defo_names_test,
roi_name=roi_name,
nets=nets,
case=case,
batch_size=batch_size,
patch_sizes=patch_sizes,
defo_sizes=defo_sizes,
dense_sizes=dense_sizes,
n_filters=n_filters,
sufixes=sufixes,
iter_name='.generalise'
)
def test_net(net, p, gt_name, options):
# Testing hyperparameters
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
# Data loading parameters
preload = options['preload']
queue = options['queue']
sufix = get_sufix(options)
c = color_codes()
p_name = '-'.join(p[0].rsplit('/')[-1].rsplit('.')[0].rsplit('-')[:-1])
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname = os.path.join(patient_path,
'deep-' + p_name + sufix + 'brain.hdr')
gt_nii = load_nii(gt_name)
gt = np.copy(np.squeeze(gt_nii.get_data()))
vals = np.unique(gt.flatten())
try:
image = np.squeeze(load_nii(outputname).get_data())
except IOError:
roi = np.squeeze(load_nii(p[0]).get_data())
centers = get_mask_voxels(roi.astype(dtype=np.bool))
test_samples = np.count_nonzero(roi)
image = np.zeros_like(roi).astype(dtype=np.uint8)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + p_name + c['nc'] +
c['g'] + ' (%d samples)>' % test_samples + c['nc'])
test_steps_per_epoch = -(-test_samples / batch_size)
y_pr_pred = net.predict_generator(
generator=load_patch_batch_generator_test(
image_names=p,
centers=centers,
batch_size=batch_size,
size=patch_size,
preload=preload,
),
steps=test_steps_per_epoch,
max_q_size=queue)
[x, y, z] = np.stack(centers, axis=1)
if options['experimental'] >= 3:
y_pr_pred = y_pr_pred[:-1]
for num, results in enumerate(y_pr_pred):
brain = np.argmax(results, axis=1)
image[x, y, z] = brain
if num is 0:
im = sufix + 'csf.'
gt_nii.get_data()[:] = np.expand_dims(image, axis=3)
elif num is 1:
im = sufix + 'gm.'
gt_nii.get_data()[:] = np.expand_dims(image, axis=3)
elif num is 2:
im = sufix + 'wm.'
gt_nii.get_data()[:] = np.expand_dims(image, axis=3)
elif num is 3:
im = sufix + 'brain.'
gt_nii.get_data()[:] = np.expand_dims(vals[image], axis=3)
elif num is 4:
im = sufix + 'rf.'
gt_nii.get_data()[:] = np.expand_dims(vals[image], axis=3)
else:
im = sufix + 'merge.'
gt_nii.get_data()[:] = np.expand_dims(vals[image], axis=3)
roiname = os.path.join(patient_path,
'deep-' + p_name + im + 'roi.hdr')
print(c['g'] + ' -- Saving image ' + c['b'] +
roiname + c['nc'])
save_nii(gt_nii, roiname)
y_pred = np.argmax(y_pr_pred[-1], axis=1)
image[x, y, z] = y_pred
gt_nii.get_data()[:] = np.expand_dims(image, axis=3)
save_nii(gt_nii, outputname)
return image, gt
def train_net(fold_n, train_data, train_labels, options):
# Prepare the net architecture parameters
dfactor = options['dfactor']
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, ) * 3
batch_size = options['batch_size']
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
nfilters = options['n_filters']
filters_list = nfilters if len(nfilters) > 1 else nfilters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
experimental = options['experimental']
fc_width = patch_width - sum(kernel_size_list) + conv_blocks
fc_shape = (fc_width, ) * 3
# Data loading parameters
preload = options['preload']
# Prepare the sufix that will be added to the results for the net and images
path = options['dir_name']
sufix = get_sufix(options)
net_name = os.path.join(path, 'iseg2017.fold%d' % fold_n + sufix + 'mdl')
checkpoint = 'iseg2017.fold%d' % fold_n + sufix + '{epoch:02d}.{val_brain_acc:.2f}.hdf5'
c = color_codes()
try:
net = load_model(net_name)
net.load_weights(os.path.join(path, checkpoint))
except IOError:
# Data loading
train_centers = get_cnn_centers(train_data[:, 0], train_labels)
train_samples = len(train_centers) / dfactor
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Loading data ' + c['b'] + '(%d centers)' % len(train_centers) +
c['nc'])
x, y = load_patches_train(image_names=train_data,
label_names=train_labels,
centers=train_centers,
size=patch_size,
fc_shape=fc_shape,
nlabels=4,
dfactor=dfactor,
preload=preload,
split=True,
iseg=True,
experimental=experimental,
datatype=np.float32)
# NET definition using Keras
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Creating and compiling the model ' + c['b'] +
'(%d samples)' % train_samples + c['nc'])
input_shape = (2, ) + patch_size
# This architecture is based on the functional Keras API to introduce 3 output paths:
# - Whole tumor segmentation
# - Core segmentation (including whole tumor)
# - Whole segmentation (tumor, core and enhancing parts)
# The idea is to let the network work on the three parts to improve the multiclass segmentation.
network_func = [
get_iseg_baseline, get_iseg_experimental1, get_iseg_experimental2,
get_iseg_experimental3, get_iseg_experimental4
]
net = network_func[experimental](input_shape, filters_list,
kernel_size_list, dense_size)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Training the model ' + c['b'] +
'(%d parameters)' % net.count_params() + c['nc'])
print(net.summary())
callbacks = [
EarlyStopping(monitor='val_brain_loss',
patience=options['patience']),
ModelCheckpoint(os.path.join(path, checkpoint),
monitor='val_brain_loss',
save_best_only=True)
]
net.save(net_name)
net.fit(x,
y,
batch_size=batch_size,
validation_split=0.25,
epochs=epochs,
callbacks=callbacks)
net.load_weights(os.path.join(path, checkpoint))
return net
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net architecture parameters
dfactor = options['dfactor']
# Prepare the net hyperparameters
num_classes = 4
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
val_rate = options['val_rate']
# Data loading parameters
preload = options['preload']
queue = options['queue']
# Prepare the sufix that will be added to the results for the net and images
path = options['dir_name']
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
ub_s = '.ub' if not balanced else ''
params_s = (ub_s, dfactor, patch_width, conv_s, filters_s, dense_size,
epochs)
sufix = '%s.D%d.p%d.c%s.n%s.d%d.e%d.' % params_s
n_channels = 4
preload_s = ' (with ' + c['b'] + 'preloading' + c['nc'] + c[
'c'] + ')' if preload else ''
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting training' +
preload_s + c['nc'])
# N-fold cross validation main loop (we'll do 2 training iterations with testing for each patient)
train_data, train_labels = get_names_from_path(options)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + c['g'] +
'Number of training images (%d=%d)' %
(len(train_data), len(train_labels)) + c['nc'])
# Also, prepare the network
net_name = os.path.join(path, 'CBICA-brats2017' + sufix)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Creating and compiling the model ' + c['nc'])
input_shape = (train_data.shape[1], ) + patch_size
# Sequential model that merges all 4 images. This architecture is just a set of convolutional blocks
# that end in a dense layer. This is supposed to be an original baseline.
inputs = Input(shape=input_shape, name='merged_inputs')
conv = inputs
for filters, kernel_size in zip(filters_list, kernel_size_list):
conv = Conv3D(filters,
kernel_size=kernel_size,
activation='relu',
data_format='channels_first')(conv)
conv = Dropout(0.5)(conv)
full = Conv3D(dense_size,
kernel_size=(1, 1, 1),
data_format='channels_first')(conv)
full = PReLU()(full)
full = Conv3D(2, kernel_size=(1, 1, 1), data_format='channels_first')(full)
rf = concatenate([conv, full], axis=1)
while np.product(K.int_shape(rf)[2:]) > 1:
rf = Conv3D(dense_size,
kernel_size=(3, 3, 3),
data_format='channels_first')(rf)
rf = Dropout(0.5)(rf)
full = Reshape((2, -1))(full)
full = Permute((2, 1))(full)
full_out = Activation('softmax', name='fc_out')(full)
tumor = Dense(2, activation='softmax', name='tumor')(rf)
outputs = [tumor, full_out]
net = Model(inputs=inputs, outputs=outputs)
net.compile(optimizer='adadelta',
loss='categorical_crossentropy',
loss_weights=[0.8, 1.0],
metrics=['accuracy'])
fc_width = patch_width - sum(kernel_size_list) + conv_blocks
fc_shape = (fc_width, ) * 3
checkpoint = net_name + '{epoch:02d}.{val_tumor_acc:.2f}.hdf5'
callbacks = [
EarlyStopping(monitor='val_tumor_loss', patience=options['patience']),
ModelCheckpoint(os.path.join(path, checkpoint),
monitor='val_tumor_loss',
save_best_only=True)
]
for i in range(options['r_epochs']):
try:
net = load_model(net_name + ('e%d.' % i) + 'mdl')
except IOError:
train_centers = get_cnn_centers(train_data[:, 0],
train_labels,
balanced=balanced)
train_samples = len(train_centers) / dfactor
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Loading data ' + c['b'] + '(%d centers)' %
(len(train_centers) / dfactor) + c['nc'])
x, y = load_patches_train(image_names=train_data,
label_names=train_labels,
centers=train_centers,
size=patch_size,
fc_shape=fc_shape,
nlabels=2,
dfactor=dfactor,
preload=preload,
split=True,
iseg=False,
experimental=1,
datatype=np.float32)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Training the model for ' + c['b'] +
'(%d parameters)' % net.count_params() + c['nc'])
print(net.summary())
net.fit(x,
y,
batch_size=batch_size,
validation_split=val_rate,
epochs=epochs,
callbacks=callbacks)
net.save(net_name + ('e%d.' % i) + 'mdl')
def train(cases, gt_names, net_name, dem_name, ratio=10, verbose=1):
# Init
options = parse_inputs()
d_path = options['val_dir']
c = color_codes()
n_folds = len(gt_names)
cases = [
c_i for c_i in cases
if find_file('Z{:}.jpg'.format(c_i + dem_name), d_path)
]
print('{:}[{:}]{:} Loading all mosaics and DEMs{:}'.format(
c['c'], time.strftime("%H:%M:%S"), c['g'], c['nc']))
for c_i in cases:
dem_file = os.path.join(d_path, 'Z{:}.jpg'.format(c_i + dem_name))
dem = cv2.imread(dem_file)
mosaic_file = os.path.join(d_path, 'Z{:}.jpg'.format(c_i))
mosaic = cv2.imread(mosaic_file)
print(dem_file, dem.shape, mosaic_file, mosaic.shape)
print('{:}[{:}]{:} Ground truth'.format(c['c'], time.strftime("%H:%M:%S"),
c['nc']))
y = [(np.mean(cv2.imread(os.path.join(d_path, im)), axis=-1) < 50).astype(
np.uint8) for im in gt_names]
print('{:}[{:}]{:} DEM'.format(c['c'], time.strftime("%H:%M:%S"), c['nc']))
dems = [
cv2.imread(os.path.join(d_path, 'Z{:}.jpg'.format(c_i + dem_name)))
for c_i in cases
]
print('{:}[{:}]{:} Mosaics'.format(c['c'], time.strftime("%H:%M:%S"),
c['nc']))
mosaics = [
cv2.imread(os.path.join(d_path, 'Z{:}.jpg'.format(c_i)))
for c_i in cases
]
print('{:}[{:}]{:} Normalising data'.format(c['c'],
time.strftime("%H:%M:%S"),
c['nc']))
x = [
np.moveaxis(
np.concatenate([mosaic, np.expand_dims(dem[..., 0], -1)], -1), -1,
0) for mosaic, dem in zip(mosaics, dems)
]
mean_x = [np.mean(xi.reshape((len(xi), -1)), axis=-1) for xi in x]
std_x = [np.std(xi.reshape((len(xi), -1)), axis=-1) for xi in x]
norm_x = [(xi - meani.reshape((-1, 1, 1))) / stdi.reshape((-1, 1, 1))
for xi, meani, stdi in zip(x, mean_x, std_x)]
print('%s[%s] %sStarting cross-validation (leave-one-mosaic-out)'
' - %d mosaics%s' %
(c['c'], time.strftime("%H:%M:%S"), c['g'], n_folds, c['nc']))
training_start = time.time()
for i, case in enumerate(cases):
if verbose > 0:
print('%s[%s]%s Starting training for mosaic %s %s(%d/%d)%s' %
(c['c'], time.strftime("%H:%M:%S"), c['g'], case, c['c'],
i + 1, len(cases), c['nc']))
test_x = norm_x[i]
train_y = y[:i] + y[i + 1:]
train_x = norm_x[:i] + norm_x[i + 1:]
val_split = 0.1
batch_size = 32
# patch_size = (256, 256)
patch_size = (64, 64)
# overlap = (64, 64)
overlap = (32, 32)
num_workers = 1
model_name = '{:}.d{:}.unc.mosaic{:}.mdl'.format(net_name, ratio, case)
net = Unet2D(n_inputs=len(norm_x[0]))
try:
net.load_model(os.path.join(d_path, model_name))
except IOError:
# Dataloader creation
if verbose > 0:
n_params = sum(p.numel() for p in net.parameters()
if p.requires_grad)
print('%sStarting training with a Unet 2D%s (%d parameters)' %
(c['c'], c['nc'], n_params))
if val_split > 0:
n_samples = len(train_x)
n_t_samples = int(n_samples * (1 - val_split))
d_train = train_x[:n_t_samples]
d_val = train_x[n_t_samples:]
l_train = train_y[:n_t_samples]
l_val = train_y[n_t_samples:]
print('Training dataset (with validation)')
# train_dataset = Cropping2DDataset(
# d_train, l_train, patch_size=patch_size, overlap=overlap,
# filtered=True
# )
train_dataset = CroppingDown2DDataset(d_train,
l_train,
patch_size=patch_size,
overlap=overlap,
filtered=True)
print('Validation dataset (with validation)')
# val_dataset = Cropping2DDataset(
# d_val, l_val, patch_size=patch_size, overlap=overlap,
# filtered=True
# )
val_dataset = CroppingDown2DDataset(d_val,
l_val,
patch_size=patch_size,
overlap=overlap,
filtered=True)
else:
print('Training dataset')
train_dataset = Cropping2DDataset(train_x,
train_y,
patch_size=patch_size,
overlap=overlap,
filtered=True)
print('Validation dataset')
val_dataset = Cropping2DDataset(train_x,
train_y,
patch_size=patch_size,
overlap=overlap)
train_dataloader = DataLoader(train_dataset,
batch_size,
True,
num_workers=num_workers)
val_dataloader = DataLoader(val_dataset,
batch_size,
num_workers=num_workers)
epochs = parse_inputs()['epochs']
patience = parse_inputs()['patience']
net.fit(
train_dataloader,
val_dataloader,
epochs=epochs,
patience=patience,
)
net.save_model(os.path.join(d_path, model_name))
if verbose > 0:
print('%s[%s]%s Starting testing with mosaic %s %s(%d/%d)%s' %
(c['c'], time.strftime("%H:%M:%S"), c['g'], case, c['c'],
i + 1, len(cases), c['nc']))
downtest_x = imresize(
test_x, (test_x.shape[0], ) +
tuple([length // ratio for length in test_x.shape[1:]]))
yi, unci = net.test([downtest_x], patch_size=None)
upyi = imresize(yi[0], test_x.shape[1:])
upunci = imresize(unci[0], test_x.shape[1:])
cv2.imwrite(
os.path.join(
d_path,
'pred.ds{:}.{:}_trees{:}.jpg'.format(ratio, dem_name, case)),
(yi[0] * 255).astype(np.uint8))
cv2.imwrite(
os.path.join(
d_path,
'pred.d{:}.{:}_trees{:}.jpg'.format(ratio, dem_name, case)),
(upyi * 255).astype(np.uint8))
cv2.imwrite(
os.path.join(
d_path,
'unc.ds{:}.{:}_trees{:}.jpg'.format(ratio, dem_name, case)),
(unci[0] * 255).astype(np.uint8))
cv2.imwrite(
os.path.join(
d_path,
'unc.d{:}.{:}_trees{:}.jpg'.format(ratio, dem_name, case)),
(upunci * 255).astype(np.uint8))
if verbose > 0:
time_str = time.strftime('%H hours %M minutes %S seconds',
time.gmtime(time.time() - training_start))
print('%sTraining finished%s (total time %s)\n' %
(c['r'], c['nc'], time_str))
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
# Data loading parameters
# Prepare the sufix that will be added to the results for the net and images
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
flair_s = '.noflair' if not options['use_flair'] else ''
t1_s = '.not1' if not options['use_t1'] else ''
t1ce_s = '.not1ce' if not options['use_t1ce'] else ''
t2_s = '.not2' if not options['use_t2'] else ''
images_s = flair_s + t1_s + t1ce_s + t2_s
params_s = (options['netname'], images_s, options['nlabels'], patch_width,
conv_s, filters_s, epochs)
sufix = '.%s%s.l%d.p%d.c%s.n%s.e%d' % params_s
train_data, _ = get_names_from_path()
unet_seg_results = list()
unet_roi_results = list()
ensemble_seg_results = list()
ensemble_roi_results = list()
image_names, label_names = get_names_from_path()
print('%s[%s] %s<BRATS 2018 pipeline testing>%s' %
(c['c'], strftime("%H:%M:%S"), c['y'], c['nc']))
print('%s[%s] %sCenter computation%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
# Block center computation
overlap = 0 if options['netname'] != 'roinet' else patch_width / 4
brain_centers = get_bounding_centers(image_names, patch_width, overlap)
test_dir = options['train_dir'][1] if options[
'train_dir'] is not None else None
if test_dir is None:
''' <Segmentation task> '''
#
# print('%s[%s] %sStarting leave-one-out (segmentation)%s' % (c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
# for i in range(len(image_names)):
# ''' Training '''
# # Data separation
# p = image_names[i]
# train_images = np.asarray(image_names.tolist()[:i] + image_names.tolist()[i + 1:])
# train_labels = np.asarray(label_names.tolist()[:i] + label_names.tolist()[i + 1:])
# train_centers = brain_centers[:i] + brain_centers[i + 1:]
# # Patient stuff
# p_name = p[0].rsplit('/')[-2]
# patient_path = '/'.join(p[0].rsplit('/')[:-1])
#
# # Data stuff
# print('%s[%s] %sPatient %s%s%s %s(%s%d%s%s/%d)%s' % (
# c['c'], strftime("%H:%M:%S"),
# c['g'], c['b'], p_name, c['nc'],
# c['c'], c['b'], i+1, c['nc'], c['c'], len(brain_centers), c['nc']
# ))
#
# # > Training for everything
#
# net, ensemble = train_seg_function(train_images, train_labels, train_centers, save_path=patient_path)
#
# ''' Testing '''
# # > Testing for the tumor ROI
# #
# # We first test with the ROI segmentation net.
#
# image_unet = test_seg(net, p, p_name, options['nlabels'])
# # image_unet = test_seg(net, p, p_name + '.unet.test' + sufix, options['nlabels'])
#
# seg_dsc = check_dsc(label_names[i], image_unet.get_data(), options['nlabels'])
# roi_dsc = check_dsc(label_names[i], image_unet.get_data().astype(np.bool), 2)
#
# dsc_string = c['g'] + '/'.join(['%f'] * len(seg_dsc)) + c['nc'] + ' (%f)'
# print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] + ' Unet DSC: ' +
# dsc_string % tuple(seg_dsc + roi_dsc))
#
# unet_seg_results.append(seg_dsc)
# unet_roi_results.append(roi_dsc)
#
# # > Testing for the tumor inside the ROI
# #
# # All we need to do now is test with the ensemble.
# image_cnn = test_seg(
# ensemble,
# p,
# p_name,
# options['nlabels'],
# mask=image_unet.get_data().astype(np.bool),
# verbose=False
# )
#
# seg_dsc = check_dsc(label_names[i], image_cnn.get_data(), options['nlabels'])
# roi_dsc = check_dsc(label_names[i], image_cnn.get_data().astype(np.bool), 2)
#
# dsc_string = c['g'] + '/'.join(['%f'] * len(seg_dsc)) + c['nc'] + ' (%f)'
# print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] + ' CNN DSC: ' +
# dsc_string % tuple(seg_dsc + roi_dsc))
#
# ensemble_seg_results.append(seg_dsc)
# ensemble_roi_results.append(roi_dsc)
#
# unet_r_dsc = np.mean(unet_roi_results)
# print('Final ROI results DSC: %f' % unet_r_dsc)
# unet_f_dsc = tuple(map(
# lambda k: np.mean(
# [dsc[k] for dsc in unet_seg_results if len(dsc) > k]
# ),
# range(3)
# ))
# print('Final Unet results DSC: (%f/%f/%f)' % unet_f_dsc)
#
# cnn_r_dsc = np.mean(ensemble_roi_results)
# print('Final ROI results DSC: %f' % cnn_r_dsc)
# cnn_f_dsc = tuple(map(
# lambda k: np.mean(
# [dsc[k] for dsc in ensemble_seg_results if len(dsc) > k]
# ),
# range(3)
# ))
# print('Final CNN results DSC: (%f/%f/%f)' % cnn_f_dsc)
''' <Survival task> '''
tst_simage_names, tst_survival, tst_features, tst_slices = get_survival_data(
options, recession=['GTR'])
max_survival = np.max(tst_survival)
n_folds = len(tst_simage_names)
print('%s[%s] %sStarting leave-one-out (survival)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
with open(os.path.join(options['loo_dir'], 'survival_results.csv'),
'w') as csvfile:
csvwriter = csv.writer(csvfile, delimiter=',')
for i in range(n_folds):
''' Training '''
ini_p = len(tst_simage_names) * i / n_folds
end_p = len(tst_simage_names) * (i + 1) / n_folds
# Validation data
p = tst_simage_names[ini_p:end_p]
p_features = tst_features[ini_p:end_p]
p_slices = tst_slices[ini_p:end_p]
p_survival = tst_survival[ini_p:end_p]
# Training data
train_images = np.concatenate(
[
tst_simage_names[:ini_p, :],
tst_simage_names[end_p:, :],
# simage_names
],
axis=0)
train_survival = np.asarray(
tst_survival.tolist()[:ini_p] +
tst_survival.tolist()[end_p:] # + survival.tolist()
)
train_features = np.asarray(
tst_features.tolist()[:ini_p] +
tst_features.tolist()[end_p:] # + features.tolist()
)
train_slices = tst_slices[:ini_p] + tst_slices[
end_p:] # + slices
# Patient info
p_name = map(lambda pi: pi[0].rsplit('/')[-2], p)
# Data stuff
print('%s[%s] %sFold %s(%s%d%s%s/%d)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['c'], c['b'],
i + 1, c['nc'], c['c'], n_folds, c['nc']))
print('%s[%s] %sStarting training (%ssurvival%s)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'],
c['nc'] + c['g'], c['nc']))
snet = train_survival_function(
train_images,
train_survival / max_survival,
train_features,
train_slices,
thresholds=[300 / max_survival, 450 / max_survival],
save_path=options['loo_dir'],
sufix='-fold%d' % i)
''' Testing '''
survival_out = test_survival(
snet, p, p_features, p_slices,
options['n_slices']) * max_survival
print(
'%s[%s] %sPatient %s%s%s predicted survival = %s%f (%f)%s'
% (c['c'], strftime("%H:%M:%S"), c['g'], c['b'], p_name[0],
c['nc'], c['g'], survival_out, p_survival, c['nc']))
csvwriter.writerow([p_name[0], '%f' % float(survival_out)])
else:
train_dir = options['train_dir'][0]
''' <Segmentation task> '''
print('%s[%s] %sStarting training (%ssegmentation%s)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], c['nc'] + c['g'],
c['nc']))
# net, ensemble = train_seg_function(image_names, label_names, brain_centers, save_path=test_dir)
net, ensemble = train_seg_function(image_names,
label_names,
brain_centers,
save_path=train_dir)
''' Testing '''
print('%s[%s] %sStarting testing (segmentation)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
test_image_names, _ = get_names_from_path(path=test_dir)
for i in range(len(test_image_names)):
# Patient stuff
p = test_image_names[i]
p_name = p[0].rsplit('/')[-2]
print('%s[%s] %sPatient %s%s%s %s(%s%d%s%s/%d)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], p_name,
c['nc'], c['c'], c['b'], i + 1, c['nc'], c['c'],
len(test_image_names), c['nc']))
# > Testing for the tumor ROI
#
# We first test with the ROI segmentation net.
image_unet = test_seg(net, p, p_name + '.unet.test' + sufix,
options['nlabels'])
# > Testing for the tumor inside the ROI
#
# All we need to do now is test with the ensemble.
test_seg(ensemble,
p,
p_name,
options['nlabels'],
mask=image_unet.get_data().astype(np.bool),
verbose=False)
''' <Survival task> '''
''' Training'''
print('%s[%s] %sStarting training (%ssurvival%s)%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], c['nc'] + c['g'],
c['nc']))
simage_names, survival, features, slices = get_survival_data(options)
max_survival = np.max(survival)
snet = train_survival_function(
image_names,
survival / max_survival,
features,
slices,
thresholds=[300 / max_survival, 450 / max_survival],
save_path=train_dir)
''' Testing '''
with open(os.path.join(test_dir, 'survival_results.csv'),
'w') as csvfile:
csvwriter = csv.writer(csvfile, delimiter=',')
names, simage_names, features, slices = get_survival_data(
options, test=True)
survival_out = test_survival(snet, simage_names, features, slices,
options['n_slices']) * max_survival
for i, (p_name_i,
pred_survival_i) in enumerate(zip(names, survival_out)):
print(
'%s[%s] %sPatient %s%s%s %s(%s%d%s%s/%d)%s predicted survival = %s%f%s'
% (c['c'], strftime("%H:%M:%S"), c['g'], c['b'], p_name_i,
c['nc'], c['c'], c['b'], i + 1, c['nc'], c['c'],
len(names), c['nc'], c['g'], pred_survival_i, c['nc']))
csvwriter.writerow([p_name_i, '%f' % pred_survival_i])
def transfer_learning(
net_domain,
net,
data,
train_image,
train_labels,
train_roi,
train_centers,
options
):
c = color_codes()
# Network hyperparameters
epochs = options['epochs']
net_epochs = options['net_epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
d_factor = options['down_factor']
# We prepare the layers for transfer learning
net_domain_conv_layers = [l for l in net_domain.layers if 'conv' in l.name]
net_conv_layers = sorted(
[l for l in net.layers if 'conv' in l.name],
cmp=lambda l1, l2: int(l1.name[7:]) - int(l2.name[7:])
)
# We freeze the convolutional layers for the final net
for layer in net.layers:
if not isinstance(layer, Dense):
layer.trainable = False
net_domain_params = np.sum([K.count_params(p) for p in set(net_domain.trainable_weights)])
# We start retraining.
# First we retrain the convolutional so the tumor rois appear similar after convolution, and then we
# retrain the classifier with the new convolutional weights.
# Data loading
centers = [tuple(center) for center in np.random.permutation(train_centers)[::d_factor]]
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Preparing ' + c['b'] + 'net' + c['nc'] +
c['g'] + ' data (' + c['b'] + '%d' % len(centers) + c['nc'] + c['g'] + ' samples)' + c['nc'])
x = [get_patches(image, centers, patch_size)
for image in train_image]
x = np.stack(x, axis=1).astype(np.float32)
y = np.array([train_labels[center] for center in centers])
y = [
keras.utils.to_categorical(
np.copy(y).astype(dtype=np.bool),
num_classes=2
),
keras.utils.to_categorical(
np.array(y > 0).astype(dtype=np.int8) + np.array(y > 1).astype(dtype=np.int8),
num_classes=3
),
keras.utils.to_categorical(
y,
num_classes=5
)
]
# We start retraining.
# First we retrain the convolutional so the tumor rois appear similar after convolution, and then we
# retrain the classifier with the new convolutional weights.
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Training the models ' + c['nc'] +
c['b'] + '(%d patches)' % len(centers) + c['nc'])
conv_data = net_domain.predict(np.expand_dims(train_roi, axis=0), batch_size=1)
for e in range(epochs):
print(c['b'] + 'Epoch %d/%d ' % (e+1, epochs) + c['nc'])
print(''.join([' ']*14) + c['g'] + c['b'] + 'Domain' + c['nc'] + c['g'] + ' net ' + c['nc'] +
c['b'] + '(%d parameters)' % net_domain_params + c['nc'])
net_domain.fit(np.expand_dims(data, axis=0), conv_data, epochs=1, batch_size=1)
for layer in net.layers:
if isinstance(layer, Dense):
if layer.name in ['core', 'tumor', 'enhancing']:
layer.trainable = False
else:
layer.trainable = True
net.compile(optimizer='adadelta', loss='categorical_crossentropy', metrics=['accuracy'])
net_params = np.sum([K.count_params(p) for p in set(net.trainable_weights)])
print(''.join([' ']*14) + c['g'] + c['b'] + 'Original (dense)' + c['nc'] + c['g'] + ' net ' + c['nc'] +
c['b'] + '(%d parameters)' % net_params + c['nc'])
net.fit(x, y, epochs=net_epochs, batch_size=batch_size)
for layer in net.layers:
if isinstance(layer, Dense):
if layer.name in ['core', 'tumor', 'enhancing']:
layer.trainable = True
else:
layer.trainable = False
net.compile(optimizer='adadelta', loss='categorical_crossentropy', metrics=['accuracy'])
net_params = np.sum([K.count_params(p) for p in set(net.trainable_weights)])
print(''.join([' ']*14) + c['g'] + c['b'] + 'Original (out)' + c['nc'] + c['g'] + ' net ' + c['nc'] +
c['b'] + '(%d parameters)' % net_params + c['nc'])
net.fit(x, y, epochs=net_epochs, batch_size=batch_size)
# We transfer the convolutional weights after retraining the net
for l_new, l_orig in zip(net_domain_conv_layers, net_conv_layers):
l_orig.set_weights(l_new.get_weights())
def train_seg(net,
image_names,
label_names,
train_centers,
save_path,
sufix,
nlabels,
net_type='unet'):
options = parse_inputs()
conv_blocks = options['conv_blocks_seg']
patch_width = options['patch_width']
patch_size = (patch_width, ) * 3 if net_type == 'unet' else (
conv_blocks * 2 + 3, ) * 3
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
batch_size = options['batch_size']
net_name = os.path.join(save_path, 'brats2018%s.mdl' % sufix)
net.save(net_name)
checkpoint = 'brats2018%s.hdf5' % sufix
callbacks = [
EarlyStopping(monitor='val_loss', patience=options['patience']),
ModelCheckpoint(os.path.join(save_path, checkpoint),
monitor='val_loss',
save_best_only=True)
]
try:
net.load_weights(os.path.join(save_path, checkpoint))
except IOError:
trainable_params = int(
np.sum([K.count_params(w) for w in set(net.trainable_weights)]))
print('%s[%s] %sTraining the network %s%s %s(%s%d %sparameters)' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], net_type, c['nc'],
c['b'], trainable_params, c['nc']))
# net.summary()
x = get_data(
image_names=image_names,
list_of_centers=train_centers,
patch_size=patch_size,
verbose=True,
)
print('%s- Concatenating the data' % ' '.join([''] * 12))
x = np.concatenate(x)
get_labels_dict = {
'unet':
lambda: get_fcnn_labels(train_centers, label_names, nlabels),
'ensemble':
lambda: get_cnn_labels(train_centers, label_names, nlabels),
'nets':
lambda: get_cluster_labels(train_centers, label_names, nlabels),
}
y = get_labels_dict[net_type]()
print('%s-- Using %d blocks of data' % (' '.join([''] * 12), len(x)))
print('%s-- X shape: (%s)' %
(' '.join([''] * 12), ', '.join(map(str, x.shape))))
if type(y) is not list:
print('%s-- Y shape: (%s)' %
(' '.join([''] * 12), ', '.join(map(str, y.shape))))
else:
y_message = '%s-- Y shape: (%s)'
for yi in y:
print(y_message %
(' '.join([''] * 12), ', '.join(map(str, yi.shape))))
print('%s- Randomising the training data' % ' '.join([''] * 12))
idx = np.random.permutation(range(len(x)))
x = x[idx]
y = y[idx] if type(y) is not list else map(lambda yi: yi[idx], y)
print('%s%sStarting the training process (%s%s%s%s) %s' % (' '.join(
[''] * 12), c['g'], c['b'], net_type, c['nc'], c['g'], c['nc']))
net.fit(x,
y,
batch_size=batch_size,
validation_split=options['val_rate'],
epochs=epochs,
callbacks=callbacks)
net.load_weights(os.path.join(save_path, checkpoint))
def test_net(net, p, outputname):
c = color_codes()
options = parse_inputs()
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['test_size']
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname_path = os.path.join(patient_path, outputname + '.nii.gz')
roiname = os.path.join(patient_path, outputname + '.roi.nii.gz')
try:
image = load_nii(outputname_path).get_data()
load_nii(roiname)
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Testing the network' + c['nc'])
roi_nii = load_nii(p[0])
roi = roi_nii.get_data().astype(dtype=np.bool)
centers = get_mask_voxels(roi)
test_samples = np.count_nonzero(roi)
image = np.zeros_like(roi).astype(dtype=np.uint8)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + p_name + c['nc'] +
c['g'] + ' - ' + c['b'] + outputname + c['nc'] + c['g'] +
' (%d samples)>' % test_samples + c['nc'])
n_centers = len(centers)
image_list = [load_norm_list(p)]
is_roi = True
roi = np.zeros_like(roi).astype(dtype=np.uint8)
fcn_out = np.zeros_like(roi).astype(dtype=np.uint8)
out = np.zeros_like(roi).astype(dtype=np.uint8)
for i in range(0, n_centers, batch_size):
print('%f%% tested (step %d/%d)' %
(100.0 * i / n_centers,
(i / batch_size) + 1, -(-n_centers / batch_size)),
end='\r')
sys.stdout.flush()
centers_i = [centers[i:i + batch_size]]
x = get_patches_list(image_list, centers_i, patch_size, True)
x = np.concatenate(x).astype(dtype=np.float32)
y_pr_pred = net.predict(x, batch_size=options['batch_size'])
[x, y, z] = np.stack(centers_i[0], axis=1)
out[x, y, z] = np.argmax(y_pr_pred[0], axis=1)
y_pr_pred = y_pr_pred[1][:, y_pr_pred[1].shape[1] / 2 + 1, :]
y_pr_pred = np.squeeze(y_pr_pred)
fcn_out[x, y, z] = np.argmax(y_pr_pred, axis=1)
tumor = np.argmax(y_pr_pred, axis=1)
# We store the ROI
roi[x, y, z] = tumor.astype(dtype=np.bool)
# We store the results
y_pred = np.argmax(y_pr_pred, axis=1)
image[x, y, z] = tumor if is_roi else y_pred
print(' '.join([''] * 50), end='\r')
sys.stdout.flush()
# Post-processing (Basically keep the biggest connected region)
image = get_biggest_region(image, is_roi)
print(c['g'] + ' -- Saving image ' + c['b'] +
outputname_path + c['nc'])
roi_nii.get_data()[:] = roi
roi_nii.to_filename(roiname)
roi_nii.get_data()[:] = get_biggest_region(fcn_out, is_roi)
roi_nii.to_filename(
os.path.join(patient_path, outputname + '.fcn.nii.gz'))
roi_nii.get_data()[:] = get_biggest_region(out, is_roi)
roi_nii.to_filename(
os.path.join(patient_path, outputname + '.dense.nii.gz'))
roi_nii.get_data()[:] = image
roi_nii.to_filename(outputname_path)
return image
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net architecture parameters
register = options['register']
multi = options['multi']
defo = options['deformation']
layers = ''.join(options['layers'])
greenspan = options['greenspan']
freeze = options['freeze']
balanced = options['balanced'] if not freeze else False
# Prepare the net hyperparameters
epochs = options['epochs']
padding = options['padding']
patch_width = options['patch_width']
patch_size = (32, 32) if greenspan else (patch_width, patch_width, patch_width)
pool_size = options['pool_size']
batch_size = options['batch_size']
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['number_filters']
n_filters = n_filters if len(n_filters) > 1 else n_filters*conv_blocks
conv_width = options['conv_width']
conv_size = conv_width if isinstance(conv_width, list) else [conv_width]*conv_blocks
# Prepare the sufix that will be added to the results for the net and images
use_flair = options['use_flair']
use_pd = options['use_pd']
use_t2 = options['use_t2']
flair_name = 'flair' if use_flair else None
pd_name = 'pd' if use_pd else None
t2_name = 't2' if use_t2 else None
images = filter(None, [flair_name, pd_name, t2_name])
reg_s = '.reg' if register else ''
filters_s = 'n'.join(['%d' % nf for nf in n_filters])
conv_s = 'c'.join(['%d' % cs for cs in conv_size])
im_s = '.'.join(images)
mc_s = '.mc' if multi else ''
d_s = 'd%d.' % (conv_blocks*2+defo) if defo else ''
sufix = '.greenspan' if greenspan else '%s.%s%s%s.p%d.c%s.n%s.d%d.e%d.pad_%s' %\
(mc_s, d_s, im_s, reg_s, patch_width, conv_s, filters_s, dense_size, epochs, padding)
# Prepare the data names
mask_name = options['mask']
wm_name = options['wm_mask']
sub_folder = options['sub_folder']
sub_name = options['flair_sub']
dir_name = options['dir_name']
patients = [f for f in sorted(os.listdir(dir_name))
if os.path.isdir(os.path.join(dir_name, f))]
n_patients = len(patients)
names = get_names_from_path(dir_name, options, patients)
defo_names = get_defonames_from_path(dir_name, options, patients) if defo else None
defo_width = conv_blocks*2+defo if defo else None
defo_size = (defo_width, defo_width, defo_width)
# Random initialisation
seed = np.random.randint(np.iinfo(np.int32).max)
# Metrics output
metrics_file = os.path.join(dir_name, 'metrics' + sufix)
with open(metrics_file, 'w') as f:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting leave-one-out' + c['nc'])
# Leave-one-out main loop (we'll do 2 training iterations with testing for each patient)
for i in range(0, n_patients):
# Prepare the data relevant to the leave-one-out (subtract the patient from the dataset and set the path)
# Also, prepare the network
case = patients[i]
path = os.path.join(dir_name, case)
names_lou = np.concatenate([names[:, :i], names[:, i + 1:]], axis=1)
defo_names_lou = np.concatenate([defo_names[:, :i], defo_names[:, i + 1:]], axis=1) if defo else None
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + 'Patient ' + c['b'] + case + c['nc'] +
c['g'] + ' (%d/%d)' % (i+1, n_patients))
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Running iteration ' + c['b'] + '1' + c['nc'] + c['g'] + '>' + c['nc'])
net_name = os.path.join(path, 'deep-longitudinal.init' + sufix + '.')
if greenspan:
net = create_cnn_greenspan(
input_channels=names.shape[0]/2,
patience=25,
name=net_name,
epochs=500
)
images = ['axial', 'coronal', 'sagital']
else:
if multi:
net = create_cnn3d_det_string(
cnn_path=layers,
input_shape=(None, names.shape[0], patch_width, patch_width, patch_width),
convo_size=conv_size,
padding=padding,
dense_size=dense_size,
pool_size=2,
number_filters=n_filters,
patience=10,
multichannel=True,
name=net_name,
epochs=100
)
else:
net = create_cnn3d_longitudinal(
convo_blocks=conv_blocks,
input_shape=(None, names.shape[0], patch_width, patch_width, patch_width),
images=images,
convo_size=conv_size,
pool_size=pool_size,
dense_size=dense_size,
number_filters=n_filters,
padding=padding,
drop=0.5,
register=register,
defo=defo,
patience=10,
name=net_name,
epochs=100
)
names_test = get_names_from_path(path, options)
defo_names_test = get_defonames_from_path(path, options) if defo else None
outputname1 = os.path.join(path, 't' + case + sufix + '.iter1.nii.gz') if not greenspan else os.path.join(
path, 't' + case + sufix + '.nii.gz')
# First we check that we did not train for that patient, in order to save time
try:
net.load_params_from(net_name + 'model_weights.pkl')
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + 'Loading the data for ' + c['b'] + 'iteration 1' + c['nc'])
# Data loading. Most of it is based on functions from data_creation that load the data.
# But we also need to prepare the name list to load the leave-one-out data.
paths = [os.path.join(dir_name, p) for p in np.concatenate([patients[:i], patients[i+1:]])]
mask_names = [os.path.join(p_path, mask_name) for p_path in paths]
wm_names = [os.path.join(p_path, wm_name) for p_path in paths]
pr_names = [os.path.join(p_path, sub_folder, sub_name) for p_path in paths]
x_train, y_train = load_lesion_cnn_data(
names=names_lou,
mask_names=mask_names,
defo_names=defo_names_lou,
roi_names=wm_names,
pr_names=pr_names,
patch_size=patch_size,
defo_size=defo_size,
random_state=seed
)
# Afterwards we train. Check the relevant training function.
if greenspan:
x_train = np.swapaxes(x_train, 1, 2)
train_greenspan(net, x_train, y_train, images)
else:
train_net(net, x_train, y_train, images)
with open(net_name + 'layers.pkl', 'wb') as fnet:
pickle.dump(net.layers, fnet, -1)
# Then we test the net. Again we save time by checking if we already tested that patient.
try:
image_nii = load_nii(outputname1)
image1 = image_nii.get_data()
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + '1' + c['nc'] + c['g'] + '>' + c['nc'])
image_nii = load_nii(os.path.join(path, options['image_folder'], options['flair_f']))
mask_nii = load_nii(os.path.join(path, wm_name))
if greenspan:
image1 = test_greenspan(
net,
names_test,
mask_nii.get_data(),
batch_size,
patch_size,
image_nii.get_data().shape,
images
)
else:
image1 = test_net(
net,
names_test,
mask_nii.get_data(),
batch_size,
patch_size,
defo_size,
image_nii.get_data().shape,
images,
defo_names_test
)
image_nii.get_data()[:] = image1
image_nii.to_filename(outputname1)
if greenspan:
# Since Greenspan did not use two iterations, we must get the final mask here.
outputname_final = os.path.join(path, 't' + case + sufix + '.final.nii.gz')
mask_nii.get_data()[:] = (image1 > 0.5).astype(dtype=np.int8)
mask_nii.to_filename(outputname_final)
else:
# If not, we test the net with the training set to look for misclassified negative with a high
# probability of being positives according to the net.
# These voxels will be the input of the second training iteration.
''' Here we get the seeds '''
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + '<Looking for seeds for the final iteration>' + c['nc'])
patients_names = zip(np.rollaxis(names_lou, 1), np.rollaxis(defo_names_lou, 1)) if defo\
else np.rollaxis(names_lou, 1)
for patient in patients_names:
if defo:
patient, d_patient = patient
else:
d_patient = None
patient_path = '/'.join(patient[0].rsplit('/')[:-1])
outputname = os.path.join(patient_path, 't' + case + sufix + '.nii.gz')
mask_nii = load_nii(os.path.join('/'.join(patient[0].rsplit('/')[:-3]), wm_name))
try:
load_nii(outputname)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + ' Patient ' + patient[0].rsplit('/')[-4] + ' already done' + c['nc'])
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + ' Testing with patient ' + c['b'] + patient[0].rsplit('/')[-4] + c['nc'])
image_nii = load_nii(patient[0])
image = test_net(
net,
patient,
mask_nii.get_data(),
batch_size,
patch_size,
defo_size,
image_nii.get_data().shape,
images,
d_patient
)
print(c['g'] + ' -- Saving image ' + c['b'] + outputname + c['nc'])
image_nii.get_data()[:] = image
image_nii.to_filename(outputname)
''' Here we perform the last iteration '''
# Finally we perform the final iteration. After refactoring the code, the code looks almost exactly
# the same as the training of the first iteration.
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Running iteration ' + c['b'] + '2' + c['nc'] + c['g'] + '>' + c['nc'])
f_s = '.f' if freeze else ''
ub_s = '.ub' if not balanced else ''
final_s = f_s + ub_s
outputname2 = os.path.join(path, 't' + case + final_s + sufix + '.iter2.nii.gz')
net_name = os.path.join(path, 'deep-longitudinal.final' + final_s + sufix + '.')
if multi:
net = create_cnn3d_det_string(
cnn_path=layers,
input_shape=(None, names.shape[0], patch_width, patch_width, patch_width),
convo_size=conv_size,
padding=padding,
pool_size=2,
dense_size=dense_size,
number_filters=n_filters,
patience=50,
multichannel=True,
name=net_name,
epochs=epochs
)
else:
if not freeze:
net = create_cnn3d_longitudinal(
convo_blocks=conv_blocks,
input_shape=(None, names.shape[0], patch_width, patch_width, patch_width),
images=images,
convo_size=conv_size,
pool_size=pool_size,
dense_size=dense_size,
number_filters=n_filters,
padding=padding,
drop=0.5,
register=register,
defo=defo,
patience=50,
name=net_name,
epochs=epochs
)
else:
net.max_epochs = epochs
net.on_epoch_finished[0].name = net_name + 'model_weights.pkl'
for layer in net.get_all_layers():
if not isinstance(layer, DenseLayer):
for param in layer.params:
layer.params[param].discard('trainable')
try:
net.load_params_from(net_name + 'model_weights.pkl')
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
c['g'] + 'Loading the data for ' + c['b'] + 'iteration 2' + c['nc'])
roi_paths = ['/'.join(name.rsplit('/')[:-1]) for name in names_lou[0, :]]
paths = [os.path.join(dir_name, p) for p in np.concatenate([patients[:i], patients[i + 1:]])]
ipr_names = [os.path.join(p_path, sub_folder, sub_name) for p_path in paths] if freeze else None
pr_names = [os.path.join(p_path, 't' + case + sufix + '.nii.gz') for p_path in roi_paths]
mask_names = [os.path.join(p_path, mask_name) for p_path in paths]
wm_names = [os.path.join(p_path, wm_name) for p_path in paths]
x_train, y_train = load_lesion_cnn_data(
names=names_lou,
mask_names=mask_names,
defo_names=defo_names_lou,
roi_names=wm_names,
init_pr_names=ipr_names,
pr_names=pr_names,
patch_size=patch_size,
defo_size=defo_size,
random_state=seed,
balanced=balanced
)
train_net(net, x_train, y_train, images)
with open(net_name + 'layers.pkl', 'wb') as fnet:
pickle.dump(net.layers, fnet, -1)
try:
image_nii = load_nii(outputname2)
image2 = image_nii.get_data()
except IOError:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + '2' + c['nc'] + c['g'] + '>' + c['nc'])
image_nii = load_nii(os.path.join(path, options['image_folder'], options['flair_f']))
mask_nii = load_nii(os.path.join(path, wm_name))
image2 = test_net(
net,
names_test,
mask_nii.get_data(),
batch_size,
patch_size,
defo_size,
image_nii.get_data().shape,
images,
defo_names_test
)
image_nii.get_data()[:] = image2
image_nii.to_filename(outputname2)
image = image1 * image2
image_nii.get_data()[:] = image
outputname_mult = os.path.join(path, 't' + case + final_s + sufix + '.iter1_x_2.nii.gz')
image_nii.to_filename(outputname_mult)
image = (image1 * image2) > 0.5
image_nii.get_data()[:] = image
outputname_final = os.path.join(path, 't' + case + final_s + sufix + '.final.nii.gz')
image_nii.to_filename(outputname_final)
# Finally we compute some metrics that are stored in the metrics file defined above.
# I plan on replicating Challenge's 2008 evaluation measures here.
gt = load_nii(os.path.join(path, mask_name)).get_data().astype(dtype=np.bool)
seg1 = image1 > 0.5
if not greenspan:
seg2 = image2 > 0.5
dsc1 = dsc_seg(gt, seg1)
if not greenspan:
dsc2 = dsc_seg(gt, seg2)
if not greenspan:
dsc_final = dsc_seg(gt, image)
else:
dsc_final = dsc1
tpf1 = tp_fraction_seg(gt, seg1)
if not greenspan:
tpf2 = tp_fraction_seg(gt, seg2)
if not greenspan:
tpf_final = tp_fraction_seg(gt, image)
fpf1 = fp_fraction_seg(gt, seg1)
if not greenspan:
fpf2 = fp_fraction_seg(gt, seg2)
if not greenspan:
fpf_final = fp_fraction_seg(gt, image)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<DSC ' + c['c'] + case + c['g'] + ' = ' + c['b'] + str(dsc_final) + c['nc'] + c['g'] + '>' + c['nc'])
f.write('%s;Test 1; %f;%f;%f\n' % (case, dsc1, tpf1, fpf1))
if not greenspan:
f.write('%s;Test 2; %f;%f;%f\n' % (case, dsc2, tpf2, fpf2))
if not greenspan:
f.write('%s;Final; %f;%f;%f\n' % (case, dsc_final, tpf_final, fpf_final))
def brats_main():
options = parse_inputs()
c = color_codes()
# Prepare the net architecture parameters
dfactor = options['dfactor']
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
# Data loading parameters
preload = options['preload']
# Prepare the sufix that will be added to the results for the net and images
path = options['dir_train']
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
ub_s = '.ub' if not balanced else ''
params_s = (ub_s, dfactor, patch_width, conv_s, filters_s, dense_size,
epochs)
sufix = '%s.D%d.p%d.c%s.n%s.d%d.e%d' % params_s
preload_s = ' (with ' + c['b'] + 'preloading' + c['nc'] + c[
'c'] + ')' if preload else ''
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting training' +
preload_s + c['nc'])
# N-fold cross validation main loop (we'll do 2 training iterations with testing for each patient)
train_data, train_labels = get_names_from_path(options)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + c['g'] +
'Number of training images (%d=%d)' %
(len(train_data), len(train_labels)) + c['nc'])
# Also, prepare the network
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Creating and compiling the model ' + c['nc'])
input_shape = (train_data.shape[1], ) + patch_size
net1, net2 = get_brats_nets(input_shape, filters_list, kernel_size_list,
dense_size, [2, 5])
net_name1 = os.path.join(path, 'brats2017-roi.tf' + sufix)
train_net(net1, net_name1, 2)
net_name2 = os.path.join(path, 'brats2017-full.tf' + sufix)
train_net(net2, net_name2, 5)
test_data, test_labels = get_names_from_path(options, False)
dsc_results = list()
for i, (p, gt_name) in enumerate(zip(test_data, test_labels)):
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + 'Case ' +
c['c'] + c['b'] + p_name + c['nc'] + c['c'] + ' (%d/%d):' %
(i + 1, len(test_data)) + c['nc'])
image_name = os.path.join(patient_path, p_name + '.tensorflow.test')
try:
image = load_nii(image_name + '.nii.gz').get_data()
except IOError:
image = test_net(net2, p, image_name)
results = check_dsc(gt_name, image)
dsc_string = c['g'] + '/'.join(['%f'] * len(results)) + c['nc']
print(''.join([' '] * 14) + c['c'] + c['b'] + p_name + c['nc'] +
' DSC: ' + dsc_string % tuple(results))
dsc_results.append(results)
f_dsc = tuple([
np.array([dsc[i] for dsc in dsc_results if len(dsc) > i]).mean()
for i in range(3)
])
print('Final results DSC: (%f/%f/%f)' % f_dsc)
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
# Data loading parameters
downsample = options['downsample']
preload = options['preload']
shuffle = options['shuffle']
# Prepare the sufix that will be added to the results for the net and images
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
unbalanced_s = '.ub' if not balanced else ''
shuffle_s = '.s' if shuffle else ''
params_s = (unbalanced_s, shuffle_s, patch_width, conv_s, filters_s, dense_size, downsample)
sufix = '%s%s.p%d.c%s.n%s.d%d.D%d' % params_s
preload_s = ' (with %spreloading%s%s)' % (c['b'], c['nc'], c['c']) if preload else ''
print('%s[%s] Starting training%s%s' % (c['c'], strftime("%H:%M:%S"), preload_s, c['nc']))
train_data, _ = get_names_from_path(options)
test_data, test_labels = get_names_from_path(options, False)
input_shape = (train_data.shape[1],) + patch_size
dsc_results = list()
dsc_results_pr = list()
train_data, train_labels = get_names_from_path(options)
centers_s = np.random.permutation(
get_cnn_centers(train_data[:, 0], train_labels, balanced=balanced)
)[::downsample]
x_seg, y_seg = load_patches_ganseg_by_batches(
image_names=train_data,
label_names=train_labels,
source_centers=centers_s,
size=patch_size,
nlabels=2,
preload=preload,
)
for i, (p, gt_name) in enumerate(zip(test_data, test_labels)):
p_name = p[0].rsplit('/')[-3]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
print('%s[%s] %sCase %s%s%s%s%s (%d/%d):%s' % (
c['c'], strftime("%H:%M:%S"), c['nc'],
c['c'], c['b'], p_name, c['nc'],
c['c'], i + 1, len(test_data), c['nc']
))
# NO DSC objective
image_cnn_name = os.path.join(patient_path, p_name + '.cnn.test%s.e%d' % (shuffle_s, epochs))
image_gan_name = os.path.join(patient_path, p_name + '.gan.test%s.e%d' % (shuffle_s, epochs))
# DSC objective
image_cnn_dsc_name = os.path.join(patient_path, p_name + '.dsc-cnn.test%s.e%d' % (shuffle_s, epochs))
image_gan_dsc_name = os.path.join(patient_path, p_name + '.dsc-gan.test%s.e%d' % (shuffle_s, epochs))
try:
# NO DSC objective
image_cnn = load_nii(image_cnn_name + '.nii.gz').get_data()
image_cnn_pr = load_nii(image_cnn_name + '.pr.nii.gz').get_data()
image_gan = load_nii(image_gan_name + '.nii.gz').get_data()
image_gan_pr = load_nii(image_gan_name + '.pr.nii.gz').get_data()
# DSC objective
image_cnn_dsc = load_nii(image_cnn_dsc_name + '.nii.gz').get_data()
image_cnn_dsc_pr = load_nii(image_cnn_dsc_name + '.pr.nii.gz').get_data()
image_gan_dsc = load_nii(image_gan_dsc_name + '.nii.gz').get_data()
image_gan_dsc_pr = load_nii(image_gan_dsc_name + '.pr.nii.gz').get_data()
except IOError:
# Lesion segmentation
adversarial_w = K.variable(0)
# NO DSC objective
cnn, gan, gan_test = get_wmh_nets(
input_shape=input_shape,
filters_list=filters_list,
kernel_size_list=kernel_size_list,
dense_size=dense_size,
lambda_var=adversarial_w
)
# DSC objective
cnn_dsc, gan_dsc, gan_dsc_test = get_wmh_nets(
input_shape=input_shape,
filters_list=filters_list,
kernel_size_list=kernel_size_list,
dense_size=dense_size,
lambda_var=adversarial_w,
dsc_obj=True
)
train_nets(
gan=gan,
gan_dsc=gan_dsc,
cnn=cnn,
cnn_dsc=cnn_dsc,
p=p,
x=x_seg,
y=y_seg,
name='wmh2017' + sufix,
adversarial_w=adversarial_w
)
# NO DSC objective
image_cnn = test_net(cnn, p, image_cnn_name)
image_cnn_pr = load_nii(image_cnn_name + '.pr.nii.gz').get_data()
image_gan = test_net(gan_test, p, image_gan_name)
image_gan_pr = load_nii(image_gan_name + '.pr.nii.gz').get_data()
# DSC objective
image_cnn_dsc = test_net(cnn_dsc, p, image_cnn_dsc_name)
image_cnn_dsc_pr = load_nii(image_cnn_dsc_name + '.pr.nii.gz').get_data()
image_gan_dsc = test_net(gan_dsc_test, p, image_gan_dsc_name)
image_gan_dsc_pr = load_nii(image_gan_dsc_name + '.pr.nii.gz').get_data()
# NO DSC objective
seg_cnn = image_cnn.astype(np.bool)
seg_gan = image_gan.astype(np.bool)
# DSC objective
seg_cnn_dsc = image_cnn_dsc.astype(np.bool)
seg_gan_dsc = image_gan_dsc.astype(np.bool)
seg_gt = load_nii(gt_name).get_data()
not_roi = np.logical_not(seg_gt == 2)
results_cnn_dsc = dsc_seg(seg_gt == 1, np.logical_and(seg_cnn_dsc, not_roi))
results_cnn_dsc_pr = probabilistic_dsc_seg(seg_gt == 1, image_cnn_dsc_pr * not_roi)
results_cnn = dsc_seg(seg_gt == 1, np.logical_and(seg_cnn, not_roi))
results_cnn_pr = probabilistic_dsc_seg(seg_gt == 1, image_cnn_pr * not_roi)
results_gan_dsc = dsc_seg(seg_gt == 1, np.logical_and(seg_gan_dsc, not_roi))
results_gan_dsc_pr = probabilistic_dsc_seg(seg_gt == 1, image_gan_dsc_pr * not_roi)
results_gan = dsc_seg(seg_gt == 1, np.logical_and(seg_gan, not_roi))
results_gan_pr = probabilistic_dsc_seg(seg_gt == 1, image_gan_pr * not_roi)
whites = ''.join([' '] * 14)
print('%sCase %s%s%s%s %sCNN%s vs %sGAN%s DSC: %s%f%s (%s%f%s) vs %s%f%s (%s%f%s)' % (
whites, c['c'], c['b'], p_name, c['nc'],
c['lgy'], c['nc'],
c['y'], c['nc'],
c['lgy'], results_cnn_dsc, c['nc'],
c['lgy'], results_cnn, c['nc'],
c['y'], results_gan_dsc, c['nc'],
c['y'], results_gan, c['nc']
))
print('%sCase %s%s%s%s %sCNN%s vs %sGAN%s DSC Pr: %s%f%s (%s%f%s) vs %s%f%s (%s%f%s)' % (
whites, c['c'], c['b'], p_name, c['nc'],
c['lgy'], c['nc'],
c['y'], c['nc'],
c['lgy'], results_cnn_dsc_pr, c['nc'],
c['lgy'], results_cnn_pr, c['nc'],
c['y'], results_gan_dsc_pr, c['nc'],
c['y'], results_gan_pr, c['nc']
))
dsc_results.append((results_cnn_dsc, results_cnn, results_gan_dsc, results_gan))
dsc_results_pr.append((results_cnn_dsc_pr, results_cnn_pr, results_gan_dsc_pr, results_gan_pr))
final_dsc = tuple(np.mean(dsc_results, axis=0))
final_dsc_pr = tuple(np.mean(dsc_results_pr, axis=0))
print('Final results DSC: %s%f%s (%s%f%s) vs %s%f%s (%s%f%s)' % (
c['lgy'], final_dsc[0], c['nc'],
c['lgy'], final_dsc[1], c['nc'],
c['y'], final_dsc[2], c['nc'],
c['y'], final_dsc[3], c['nc']
))
print('Final results DSC Pr: %s%f%s (%s%f%s) vs %s%f%s (%s%f%s)' % (
c['lgy'], final_dsc_pr[0], c['nc'],
c['lgy'], final_dsc_pr[1], c['nc'],
c['y'], final_dsc_pr[2], c['nc'],
c['y'], final_dsc_pr[3], c['nc']
))
def train_nets(gan, gan_dsc, cnn, cnn_dsc, p, x, y, name, adversarial_w):
options = parse_inputs()
c = color_codes()
# Data stuff
patient_path = '/'.join(p[0].rsplit('/')[:-1])
train_data, train_labels = get_names_from_path(options)
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
preload = options['preload']
batch_size = options['batch_size']
print('%s[%s] %sTraining the networks%s (%sCNN%s vs %sGAN%s: %s%s%s/%s%d%s parameters)' % (
c['c'], strftime("%H:%M:%S"),
c['g'], c['nc'],
c['lgy'], c['nc'],
c['y'], c['nc'],
c['b'], gan.count_params(), c['nc'],
c['b'], cnn.count_params(), c['nc']
))
net_name = os.path.join(patient_path, name)
checkpoint_name = os.path.join(patient_path, net_name + '.weights')
try:
gan.load_weights(checkpoint_name + '.gan.e%d' % epochs)
gan_dsc.load_weights(checkpoint_name + '.gan-dsc.e%d' % epochs)
cnn.load_weights(checkpoint_name + '.net.e%d' % epochs)
cnn_dsc.load_weights(checkpoint_name + '.net-dsc.e%d' % epochs)
except IOError:
x_disc, y_disc = load_patches_gandisc_by_batches(
source_names=train_data,
target_names=[p],
n_centers=len(x),
size=patch_size,
preload=preload,
)
print('%s[%s]%s %sStarting the training process%s' % (
c['c'], strftime("%H:%M:%S"), c['nc'],
c['g'], c['nc']
))
for e in range(epochs):
print(' '.join([''] * 16) + c['g'] + 'Epoch ' +
c['b'] + '%d' % (e + 1) + c['nc'] + c['g'] + '/%d' % epochs + c['nc'])
try:
cnn.load_weights(checkpoint_name + '.net.e%d' % (e + 1))
cnn_dsc.load_weights(checkpoint_name + '.net-dsc.e%d' % (e + 1))
gan.load_weights(checkpoint_name + '.gan.e%d' % (e + 1))
gan_dsc.load_weights(checkpoint_name + '.gan-dsc.e%d' % (e + 1))
except IOError:
print(c['lgy'], end='\r')
cnn.fit(x, y, batch_size=batch_size, epochs=1)
cnn_dsc.fit(x, y, batch_size=batch_size, epochs=1)
print(c['y'], end='\r')
gan.fit([x, x_disc], [y, y_disc], batch_size=batch_size, epochs=1)
gan_dsc.fit([x, x_disc], [y, y_disc], batch_size=batch_size, epochs=1)
print(c['nc'], end='\r')
cnn.save_weights(checkpoint_name + '.net.e%d' % (e + 1))
cnn_dsc.save_weights(checkpoint_name + '.net-dsc.e%d' % (e + 1))
gan.save_weights(checkpoint_name + '.gan.e%d' % (e + 1))
gan_dsc.save_weights(checkpoint_name + '.gan-dsc.e%d' % (e + 1))
adversarial_weight = min([K.eval(adversarial_w) + 0.1, 1.0])
K.set_value(adversarial_w, adversarial_weight)
def train_nets(gan, cnn, caps, x, y, p, name, adversarial_w):
options = parse_inputs()
c = color_codes()
# Data stuff
patient_path = '/'.join(p[0].rsplit('/')[:-1])
train_data, train_labels = get_names_from_path(options)
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
preload = options['preload']
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Training the networks ' + c['nc'] +
c['lgy'] + '(' + 'CNN' + c['nc'] + '/' +
c['r'] + 'CAPS' + c['nc'] + '/' +
c['y'] + 'GAN' + c['nc'] + ': ' +
c['b'] + '%d' % gan.count_params() + c['nc'] + '/' + c['b'] + '%d ' % cnn.count_params() + c['nc'] +
'parameters)')
net_name = os.path.join(patient_path, name)
checkpoint_name = os.path.join(patient_path, net_name + '.weights')
try:
cnn.load_weights(checkpoint_name + '.net.e%d' % epochs)
caps.load_weights(checkpoint_name + '.caps.e%d' % epochs)
gan.load_weights(checkpoint_name + '.gan.e%d' % epochs)
except IOError:
x_disc, y_disc = load_patches_gandisc_by_batches(
source_names=train_data,
target_names=[p],
n_centers=len(x),
size=patch_size,
preload=preload,
batch_size=51200
)
print(' '.join([''] * 15) + c['g'] + 'Starting the training process' + c['nc'])
for e in range(epochs):
print(' '.join([''] * 16) + c['g'] + 'Epoch ' +
c['b'] + '%d' % (e + 1) + c['nc'] + c['g'] + '/%d' % epochs + c['nc'])
try:
cnn.load_weights(checkpoint_name + '.net.e%d' % (e + 1))
except IOError:
print(c['lgy'], end='\r')
cnn.fit(x, y, batch_size=batch_size, epochs=1)
try:
caps.load_weights(checkpoint_name + '.caps.e%d' % (e + 1))
except IOError:
print(c['r'], end='\r')
caps.fit(x, y, batch_size=batch_size, epochs=1)
try:
gan.load_weights(checkpoint_name + '.gan.e%d' % (e + 1))
except IOError:
print(c['y'], end='\r')
gan.fit([x, x_disc], [y, y_disc], batch_size=batch_size, epochs=1)
print(c['nc'], end='\r')
cnn.save_weights(checkpoint_name + '.net.e%d' % (e + 1))
caps.save_weights(checkpoint_name + '.caps.e%d' % (e + 1))
gan.save_weights(checkpoint_name + '.gan.e%d' % (e + 1))
adversarial_weight = min([np.array(K.eval(adversarial_w)) + 0.1, 1.0])
K.set_value(adversarial_w, adversarial_weight)
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net architecture parameters
sequential = options['sequential']
dfactor = options['dfactor']
# Prepare the net hyperparameters
num_classes = 5
epochs = options['epochs']
padding = options['padding']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
recurrent = options['recurrent']
# Data loading parameters
preload = options['preload']
queue = options['queue']
# Prepare the sufix that will be added to the results for the net and images
path = options['dir_name']
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
s_s = '.s' if sequential else '.f'
ub_s = '.ub' if not balanced else ''
params_s = (ub_s, dfactor, s_s, patch_width, conv_s, filters_s, dense_size,
epochs, padding)
sufix = '%s.D%d%s.p%d.c%s.n%s.d%d.e%d.pad_%s.' % params_s
n_channels = np.count_nonzero([
options['use_flair'], options['use_t2'], options['use_t1'],
options['use_t1ce']
])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
'Starting cross-validation' + c['nc'])
# N-fold cross validation main loop (we'll do 2 training iterations with testing for each patient)
data_names, label_names = get_names_from_path(options)
folds = options['folds']
fold_generator = izip(
nfold_cross_validation(data_names, label_names, n=folds,
val_data=0.25), xrange(folds))
dsc_results = list()
for (train_data, train_labels, val_data, val_labels, test_data,
test_labels), i in fold_generator:
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] +
'Fold %d/%d: ' % (i + 1, folds) + c['g'] +
'Number of training/validation/testing images (%d=%d/%d=%d/%d)' %
(len(train_data), len(train_labels), len(val_data),
len(val_labels), len(test_data)) + c['nc'])
# Prepare the data relevant to the leave-one-out (subtract the patient from the dataset and set the path)
# Also, prepare the network
net_name = os.path.join(
path, 'baseline-brats2017.fold%d' % i + sufix + 'mdl')
# First we check that we did not train for that patient, in order to save time
try:
net = keras.models.load_model(net_name)
except IOError:
# NET definition using Keras
train_centers = get_cnn_centers(train_data[:, 0],
train_labels,
balanced=balanced)
val_centers = get_cnn_centers(val_data[:, 0],
val_labels,
balanced=balanced)
train_samples = len(train_centers) / dfactor
val_samples = len(val_centers) / dfactor
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Creating and compiling the model ' + c['b'] +
'(%d samples)' % train_samples + c['nc'])
train_steps_per_epoch = -(-train_samples / batch_size)
val_steps_per_epoch = -(-val_samples / batch_size)
input_shape = (n_channels, ) + patch_size
if sequential:
# Sequential model that merges all 4 images. This architecture is just a set of convolutional blocks
# that end in a dense layer. This is supposed to be an original baseline.
net = Sequential()
net.add(
Conv3D(filters_list[0],
kernel_size=kernel_size_list[0],
input_shape=input_shape,
activation='relu',
data_format='channels_first'))
for filters, kernel_size in zip(filters_list[1:],
kernel_size_list[1:]):
net.add(Dropout(0.5))
net.add(
Conv3D(filters,
kernel_size=kernel_size,
activation='relu',
data_format='channels_first'))
net.add(Dropout(0.5))
net.add(Flatten())
net.add(Dense(dense_size, activation='relu'))
net.add(Dropout(0.5))
net.add(Dense(num_classes, activation='softmax'))
else:
# This architecture is based on the functional Keras API to introduce 3 output paths:
# - Whole tumor segmentation
# - Core segmentation (including whole tumor)
# - Whole segmentation (tumor, core and enhancing parts)
# The idea is to let the network work on the three parts to improve the multiclass segmentation.
merged_inputs = Input(shape=(4, ) + patch_size,
name='merged_inputs')
flair = Reshape((1, ) + patch_size)(Lambda(
lambda l: l[:, 0, :, :, :],
output_shape=(1, ) + patch_size)(merged_inputs), )
t2 = Reshape((1, ) + patch_size)(Lambda(
lambda l: l[:, 1, :, :, :],
output_shape=(1, ) + patch_size)(merged_inputs))
t1 = Lambda(lambda l: l[:, 2:, :, :, :],
output_shape=(2, ) + patch_size)(merged_inputs)
for filters, kernel_size in zip(filters_list,
kernel_size_list):
flair = Conv3D(filters,
kernel_size=kernel_size,
activation='relu',
data_format='channels_first')(flair)
t2 = Conv3D(filters,
kernel_size=kernel_size,
activation='relu',
data_format='channels_first')(t2)
t1 = Conv3D(filters,
kernel_size=kernel_size,
activation='relu',
data_format='channels_first')(t1)
flair = Dropout(0.5)(flair)
t2 = Dropout(0.5)(t2)
t1 = Dropout(0.5)(t1)
# We only apply the RCNN to the multioutput approach (we keep the simple one, simple)
if recurrent:
flair = Conv3D(dense_size,
kernel_size=(1, 1, 1),
activation='relu',
data_format='channels_first',
name='fcn_flair')(flair)
flair = Dropout(0.5)(flair)
t2 = concatenate([flair, t2], axis=1)
t2 = Conv3D(dense_size,
kernel_size=(1, 1, 1),
activation='relu',
data_format='channels_first',
name='fcn_t2')(t2)
t2 = Dropout(0.5)(t2)
t1 = concatenate([t2, t1], axis=1)
t1 = Conv3D(dense_size,
kernel_size=(1, 1, 1),
activation='relu',
data_format='channels_first',
name='fcn_t1')(t1)
t1 = Dropout(0.5)(t1)
flair = Dropout(0.5)(flair)
t2 = Dropout(0.5)(t2)
t1 = Dropout(0.5)(t1)
lstm_instance = LSTM(dense_size,
implementation=1,
name='rf_layer')
flair = lstm_instance(
Permute((2, 1))(Reshape((dense_size, -1))(flair)))
t2 = lstm_instance(
Permute((2, 1))(Reshape((dense_size, -1))(t2)))
t1 = lstm_instance(
Permute((2, 1))(Reshape((dense_size, -1))(t1)))
else:
flair = Flatten()(flair)
t2 = Flatten()(t2)
t1 = Flatten()(t1)
flair = Dense(dense_size, activation='relu')(flair)
flair = Dropout(0.5)(flair)
t2 = concatenate([flair, t2])
t2 = Dense(dense_size, activation='relu')(t2)
t2 = Dropout(0.5)(t2)
t1 = concatenate([t2, t1])
t1 = Dense(dense_size, activation='relu')(t1)
t1 = Dropout(0.5)(t1)
tumor = Dense(2, activation='softmax', name='tumor')(flair)
core = Dense(3, activation='softmax', name='core')(t2)
enhancing = Dense(num_classes,
activation='softmax',
name='enhancing')(t1)
net = Model(inputs=merged_inputs,
outputs=[tumor, core, enhancing])
net.compile(optimizer='adadelta',
loss='categorical_crossentropy',
metrics=['accuracy'])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Training the model with a generator for ' + c['b'] +
'(%d parameters)' % net.count_params() + c['nc'])
print(net.summary())
net.fit_generator(
generator=load_patch_batch_train(image_names=train_data,
label_names=train_labels,
centers=train_centers,
batch_size=batch_size,
size=patch_size,
nlabels=num_classes,
dfactor=dfactor,
preload=preload,
split=not sequential,
datatype=np.float32),
validation_data=load_patch_batch_train(image_names=val_data,
label_names=val_labels,
centers=val_centers,
batch_size=batch_size,
size=patch_size,
nlabels=num_classes,
dfactor=dfactor,
preload=preload,
split=not sequential,
datatype=np.float32),
steps_per_epoch=train_steps_per_epoch,
validation_steps=val_steps_per_epoch,
max_q_size=queue,
epochs=epochs)
net.save(net_name)
# Then we test the net.
use_gt = options['use_gt']
for p, gt_name in zip(test_data, test_labels):
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname = os.path.join(patient_path,
'deep-brats17' + sufix + 'test.nii.gz')
try:
load_nii(outputname)
except IOError:
roi_nii = load_nii(p[0])
roi = roi_nii.get_data().astype(dtype=np.bool)
centers = get_mask_voxels(roi)
test_samples = np.count_nonzero(roi)
image = np.zeros_like(roi).astype(dtype=np.uint8)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + p_name +
c['nc'] + c['g'] + ' (%d samples)>' % test_samples +
c['nc'])
test_steps_per_epoch = -(-test_samples / batch_size)
y_pr_pred = net.predict_generator(
generator=load_patch_batch_generator_test(
image_names=p,
centers=centers,
batch_size=batch_size,
size=patch_size,
preload=preload,
),
steps=test_steps_per_epoch,
max_q_size=queue)
[x, y, z] = np.stack(centers, axis=1)
if not sequential:
tumor = np.argmax(y_pr_pred[0], axis=1)
y_pr_pred = y_pr_pred[-1]
roi = np.zeros_like(roi).astype(dtype=np.uint8)
roi[x, y, z] = tumor
roi_nii.get_data()[:] = roi
roiname = os.path.join(
patient_path,
'deep-brats17' + sufix + 'test.roi.nii.gz')
roi_nii.to_filename(roiname)
y_pred = np.argmax(y_pr_pred, axis=1)
image[x, y, z] = y_pred
# Post-processing (Basically keep the biggest connected region)
image = get_biggest_region(image)
if use_gt:
gt_nii = load_nii(gt_name)
gt = np.copy(gt_nii.get_data()).astype(dtype=np.uint8)
labels = np.unique(gt.flatten())
results = (p_name, ) + tuple(
[dsc_seg(gt == l, image == l) for l in labels[1:]])
text = 'Subject %s DSC: ' + '/'.join(
['%f' for _ in labels[1:]])
print(text % results)
dsc_results.append(results)
print(c['g'] + ' -- Saving image ' + c['b'] +
outputname + c['nc'])
roi_nii.get_data()[:] = image
roi_nii.to_filename(outputname)
def train_seg(
net, model_name, train_patients, val_patients,
dropout=0.99, lr=1e-2
# lr=1.
):
# Init
c = color_codes()
options = parse_inputs()
epochs = options['epochs']
patience = options['patience']
batch_size = options['batch_size']
patch_size = options['patch_size']
d_path = options['loo_dir']
try:
net.load_model(os.path.join(d_path, model_name))
except IOError:
n_params = sum(
p.numel() for p in net.parameters() if p.requires_grad
)
print(
'%sStarting refinement with a unet%s (%d parameters)' %
(c['c'], c['nc'], n_params)
)
num_workers = 8
# Training
targets = get_labels(train_patients)
rois, data = get_images(train_patients)
print('< Training dataset >')
if patch_size is None:
train_dataset = BBImageDataset(
data, targets, rois, flip=True
)
else:
# train_dataset = BoundarySegmentationCroppingDataset(
# data, targets, rois, patch_size
# )
train_dataset = BratsDataset(
data, targets, rois, patch_size
)
print('Dataloader creation <with validation>')
train_loader = DataLoader(
train_dataset, batch_size, True, num_workers=num_workers,
)
# Validation
targets = get_labels(val_patients)
rois, data = get_images(val_patients)
print('< Validation dataset >')
val_dataset = BBImageDataset(
data, targets, rois
)
print('Dataloader creation <val>')
val_loader = DataLoader(
val_dataset, 1
)
print(
'Training / validation samples = %d / %d' % (
len(train_dataset), len(val_dataset)
)
)
net.fit(
train_loader, val_loader, initial_dropout=dropout, initial_lr=lr,
epochs=epochs, patience=patience
)
net.save_model(os.path.join(d_path, model_name))
def train_survival_function(image_names,
survival,
features,
slices,
save_path,
thresholds,
sufix=''):
# Init
options = parse_inputs()
c = color_codes()
# Prepare the net hyperparameters
epochs = options['sepochs']
n_slices = options['n_slices']
''' Net preparation '''
net = get_brats_survival(thresholds=thresholds,
n_slices=n_slices,
n_features=features.shape[-1])
net_name = os.path.join(save_path, 'brats2018-survival%s.mdl' % sufix)
net.save(net_name)
# checkpoint = 'brats2018-survival%s.hdf5' % sufix
# net.summary()
''' Training '''
train_steps = options['train_steps']
try:
net.load_weights(
os.path.join(save_path,
'brats2018-survival%s.hdf5' % (train_steps - 1)))
print('%s[%s] %sSurvival network weights %sloaded%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], c['nc']))
except IOError:
trainable_params = int(
np.sum([K.count_params(w) for w in set(net.trainable_weights)]))
print('%s[%s] %sTraining the survival network %s(%s%d %sparameters)' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc'], c['b'],
trainable_params, c['nc']))
# Data preparation
x_vol = get_reshaped_data(image_names,
slices, (224, 224),
n_slices=n_slices,
verbose=True)
print('%s- Concatenating the data' % ' '.join([''] * 12))
x_vol = preprocess_input(np.stack(x_vol, axis=0))
print('%s-- X (volume) shape: (%s)' %
(' '.join([''] * 12), ', '.join(map(str, x_vol.shape))))
print('%s-- X (features) shape: (%s)' %
(' '.join([''] * 12), ', '.join(map(str, features.shape))))
print('%s-- Y shape: (%s)' %
(' '.join([''] * 12), ', '.join(map(str, survival.shape))))
checkpoint = 'brats2018-survival%s-step.hdf5' % sufix
try:
# net.load_weights(os.path.join(save_path, checkpoint))
net.load_weights(
os.path.join(
save_path, 'brats2018-survival%s-step%d.hdf5' %
(sufix, train_steps - 1)))
print('%s[%s] %sSurvival network weights %sloaded%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['b'], c['nc']))
except IOError:
''' Normal version '''
# callbacks = [
# EarlyStopping(
# monitor='val_loss',
# patience=options['spatience']
# ),
# ModelCheckpoint(
# os.path.join(save_path, checkpoint),
# monitor='val_loss',
# save_best_only=True
# )
# ]
# callbacks = [
# EarlyStopping(
# monitor='cat_survival_acc',
# patience=options['spatience']
# ),
# ModelCheckpoint(
# os.path.join(save_path, checkpoint),
# monitor='cat_survival_acc',
# save_best_only=True
# )
# ]
callbacks = [
ModelCheckpoint(os.path.join(save_path, checkpoint),
monitor='survival_loss',
save_best_only=True)
]
print('%s- Randomising the training data' % ' '.join([''] * 12))
idx = np.random.permutation(range(len(features)))
x_vol = x_vol[idx].astype(np.float32)
x_feat = features[idx].astype(np.float32)
x = [x_vol, x_feat]
y = survival[idx].astype(np.float32)
y_cat = np.squeeze(
np.stack(
[y < 300,
np.logical_and(y >= 300, y < 450), y >= 450],
axis=1))
# net.fit(x, y, batch_size=8, validation_split=options['sval_rate'], epochs=epochs, callbacks=callbacks)
net.fit(x, [y, y_cat],
batch_size=8,
epochs=epochs,
callbacks=callbacks)
net.load_weights(os.path.join(save_path, checkpoint))
''' Curriculum learning version '''
# sorted_idx = np.squeeze(np.argsort(survival, axis=0))
# for i in range(train_steps):
# checkpoint = 'brats2018-survival%s-step%d.hdf5' % (sufix, i)
# net = get_brats_survival(n_slices=n_slices, n_features=features.shape[-1])
# net_name = os.path.join(save_path, 'brats2018-survival%s.mdl' % sufix)
# net.save(net_name)
# try:
# net.load_weights(os.path.join(save_path, checkpoint))
# print(
# '%s[%s] %sSurvival network weights %sloaded%s' % (
# c['c'], strftime("%H:%M:%S"), c['g'],
# c['b'], c['nc']
# )
# )
# except IOError:
# print('%s- Selecting the next set of samples (step %d/%d)' %
# (' '.join([''] * 12), i + 1, train_steps))
# current_idx = np.concatenate([
# sorted_idx[:((i + 1) * len(sorted_idx)) / (2 * train_steps)],
# sorted_idx[-((i + 1) * len(sorted_idx)) / (2 * train_steps):],
# ])
#
# callbacks = [
# EarlyStopping(
# monitor='loss',
# patience=options['spatience']
# ),
# ModelCheckpoint(
# os.path.join(save_path, checkpoint),
# monitor='loss',
# save_best_only=True
# )
# ]
#
# idx = np.random.permutation(range(len(current_idx)))
#
# x = [x_vol[current_idx][idx].astype(np.float32), features[current_idx][idx].astype(np.float32)]
#
# y = survival[current_idx][idx].astype(np.float32)
#
# net.fit(x, y, batch_size=8, epochs=epochs, callbacks=callbacks)
# net.load_weights(os.path.join(save_path, checkpoint))
''' Average version '''
# for i in range(train_steps):
# checkpoint = 'brats2018-survival%s-step%d.hdf5' % (sufix, i)
# try:
# net.load_weights(os.path.join(save_path, checkpoint))
# print(
# '%s[%s] %sSurvival network weights %sloaded%s' % (
# c['c'], strftime("%H:%M:%S"), c['g'],
# c['b'], c['nc']
# )
# )
# except IOError:
# callbacks = [
# EarlyStopping(
# monitor='loss',
# patience=options['spatience']
# ),
# ModelCheckpoint(
# os.path.join(save_path, checkpoint),
# monitor='loss',
# save_best_only=True
# )
# ]
#
# print('%s- Randomising the training data' % ' '.join([''] * 12))
# idx = np.random.permutation(range(len(features)))
#
# x_vol = x_vol[idx].astype(np.float32)
# x_feat = features[idx].astype(np.float32)
# x = [x_vol, x_feat]
#
# y = survival[idx].astype(np.float32)
#
# net.fit(x, y, batch_size=8, epochs=epochs, callbacks=callbacks)
# net.load_weights(os.path.join(save_path, checkpoint))
return net
def train_test_seg(net_name, n_folds, val_split=0.1):
# Init
c = color_codes()
options = parse_inputs()
depth = options['blocks']
filters = options['filters']
d_path = options['loo_dir']
unc_path = os.path.join(d_path, 'uncertainty')
if not os.path.isdir(unc_path):
os.mkdir(unc_path)
seg_path = os.path.join(d_path, 'segmentation')
if not os.path.isdir(seg_path):
os.mkdir(seg_path)
patients = get_dirs(d_path)
cbica = filter(lambda p: 'CBICA' in p, patients)
tcia = filter(lambda p: 'TCIA' in p, patients)
tmc = filter(lambda p: 'TMC' in p, patients)
b2013 = filter(lambda p: '2013' in p, patients)
for i in range(n_folds):
print(
'%s[%s] %sFold %s(%s%d%s%s/%d)%s' % (
c['c'], strftime("%H:%M:%S"), c['g'],
c['c'], c['b'], i + 1, c['nc'], c['c'], n_folds, c['nc']
)
)
# Training itself
# Data split (using the patient names) for train and validation.
# We also compute the number of batches for both training and
# validation according to the batch size.
''' Training '''
ini_cbica = len(cbica) * i / n_folds
end_cbica = len(cbica) * (i + 1) / n_folds
fold_cbica = cbica[:ini_cbica] + cbica[end_cbica:]
n_fold_cbica = len(fold_cbica)
n_cbica = int(n_fold_cbica * (1 - val_split))
ini_tcia = len(tcia) * i / n_folds
end_tcia = len(tcia) * (i + 1) / n_folds
fold_tcia = tcia[:ini_tcia] + tcia[end_tcia:]
n_fold_tcia = len(fold_tcia)
n_tcia = int(n_fold_tcia * (1 - val_split))
ini_tmc = len(tmc) * i / n_folds
end_tmc = len(tmc) * (i + 1) / n_folds
fold_tmc = tmc[:ini_tmc] + tmc[end_tmc:]
n_fold_tmc = len(fold_tmc)
n_tmc = int(n_fold_tmc * (1 - val_split))
ini_b2013 = len(b2013) * i / n_folds
end_b2013 = len(b2013) * (i + 1) / n_folds
fold_b2013 = b2013[:ini_b2013] + b2013[end_b2013:]
n_fold_b2013 = len(fold_b2013)
n_b2013 = int(n_fold_b2013 * (1 - val_split))
training_n = n_fold_cbica + n_fold_tcia + n_fold_tmc + n_fold_b2013
testing_n = len(patients) - training_n
print(
'Training / testing samples = %d / %d' % (
training_n, testing_n
)
)
# Training
train_cbica = fold_cbica[:n_cbica]
train_tcia = fold_tcia[:n_tcia]
train_tmc = fold_tmc[:n_tmc]
train_b2013 = fold_b2013[:n_b2013]
train_patients = train_cbica + train_tcia + train_tmc + train_b2013
# Validation
val_cbica = fold_cbica[n_cbica:]
val_tcia = fold_tcia[n_tcia:]
val_tmc = fold_tmc[n_tmc:]
val_b2013 = fold_b2013[n_b2013:]
val_patients = val_cbica + val_tcia + val_tmc + val_b2013
model_name = '%s-f%d.mdl' % (net_name, i)
net = BratsSegmentationNet(depth=depth, filters=filters)
# train_seg(net, model_name, train_patients, val_patients)
train_seg(
net, model_name, train_patients, val_patients, dropout=0
)
# model_name = '%s-f%d-R.mdl' % (net_name, i)
# train_seg(
# net, model_name, train_patients, val_patients,
# refine=True, dropout=0.5, lr=1e-2
# )
# Testing data (with GT)
test_cbica = cbica[ini_cbica:end_cbica]
test_tcia = tcia[ini_tcia:end_tcia]
test_tmc = tmc[ini_tmc:end_tmc]
test_b2013 = b2013[ini_b2013:end_b2013]
test_patients = test_cbica + test_tcia + test_tmc + test_b2013
patient_paths = map(lambda p: os.path.join(d_path, p), test_patients)
_, test_x = get_images(test_patients)
print(
'Testing patients (with GT) = %d' % (
len(test_patients)
)
)
# The sub-regions considered for evaluation are:
# 1) the "enhancing tumor" (ET)
# 2) the "tumor core" (TC)
# 3) the "whole tumor" (WT)
#
# The provided segmentation labels have values of 1 for NCR & NET,
# 2 for ED, 4 for ET, and 0 for everything else.
# The participants are called to upload their segmentation labels
# as a single multi-label file in nifti (.nii.gz) format.
#
# The participants are called to upload 4 nifti (.nii.gz) volumes
# (3 uncertainty maps and 1 multi-class segmentation volume from
# Task 1) onto CBICA's Image Processing Portal format. For example,
# for each ID in the dataset, participants are expected to upload
# following 4 volumes:
# 1. {ID}.nii.gz (multi-class label map)
# 2. {ID}_unc_whole.nii.gz (Uncertainty map associated with whole tumor)
# 3. {ID}_unc_core.nii.gz (Uncertainty map associated with tumor core)
# 4. {ID}_unc_enhance.nii.gz (Uncertainty map associated with enhancing tumor)
for p, (path_i, p_i, test_i) in enumerate(zip(
patient_paths, test_patients, test_x
)):
pred_i = net.segment([test_i])[0]
# unc_i = net.uncertainty([test_i], steps=25)[0]
# whole_i = np.sum(unc_i[1:])
# core_i = unc_i[1] + unc_i[-1]
# enhance_i = unc_i[-1]
seg_i = np.argmax(pred_i, axis=0)
seg_i[seg_i == 3] = 4
# seg_unc_i = np.argmax(unc_i, axis=0)
# seg_unc_i[seg_unc_i == 3] = 4
# tumor_mask = remove_small_regions(
# seg_i.astype(np.bool), min_size=30
# )
#
# seg_i[log_not(tumor_mask)] = 0
# seg_unc_i[log_not(tumor_mask)] = 0
#
# whole_i *= tumor_mask.astype(np.float32)
# core_i *= tumor_mask.astype(np.float32)
# enhance_i *= tumor_mask.astype(np.float32)
niiname = os.path.join(path_i, p_i + '_seg.nii.gz')
nii = load_nii(niiname)
seg = nii.get_data()
dsc = map(
lambda label: dsc_seg(seg == label, seg_i == label),
[1, 2, 4]
)
# dsc_unc = map(
# lambda label: dsc_seg(seg == label, seg_unc_i == label),
# [1, 2, 4]
# )
nii.get_data()[:] = seg_i
save_nii(nii, os.path.join(seg_path, p_i + '.nii.gz'))
# nii.get_data()[:] = seg_unc_i
# save_nii(nii, os.path.join(unc_path, p_i + '.nii.gz'))
# niiname = os.path.join(d_path, p_i, p_i + '_flair.nii.gz')
# nii = load_nii(niiname)
# nii.get_data()[:] = whole_i
# save_nii(nii, os.path.join(unc_path, p_i + '_unc_whole.nii.gz'))
# nii.get_data()[:] = core_i
# save_nii(nii, os.path.join(unc_path, p_i + '_unc_core.nii.gz'))
# nii.get_data()[:] = enhance_i
# save_nii(nii, os.path.join(unc_path, p_i + '_unc_enhance.nii.gz'))
print(
'Segmentation - Patient %s (%d/%d): %s' % (
p_i, p, len(test_x), ' / '.join(map(str, dsc))
)
)
def train_seg_function(image_names, label_names, brain_centers, save_path):
# Init
options = parse_inputs()
c = color_codes()
# Prepare the net hyperparameters
epochs = options['epochs']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
# Prepare the sufix that will be added to the results for the net and images
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
flair_s = '.noflair' if not options['use_flair'] else ''
t1_s = '.not1' if not options['use_t1'] else ''
t1ce_s = '.not1ce' if not options['use_t1ce'] else ''
t2_s = '.not2' if not options['use_t2'] else ''
images_s = flair_s + t1_s + t1ce_s + t2_s
params_s = (options['netname'], images_s, options['nlabels'], patch_width,
conv_s, filters_s, epochs)
sufix = '.%s%s.l%d.p%d.c%s.n%s.e%d' % params_s
'''Tumor ROI stuff'''
# Training for the ROI
input_shape = (image_names.shape[-1], ) + patch_size
net = options['net'](input_shape=input_shape,
filters_list=filters_list,
kernel_size_list=kernel_size_list,
nlabels=options['nlabels'])
train_seg(image_names=image_names,
label_names=label_names,
train_centers=brain_centers,
net=net,
save_path=save_path,
sufix=sufix,
nlabels=options['nlabels'])
'''Tumor segmentation stuff'''
# Training for the tumor inside the ROI
# First we should retest each training image and get the test tumor mask
# and join it with the GT mask. This new mask after dilation will give us
# the training centers.
print('%s%s<Creating the tumor masks for the training data>%s' %
(''.join([' '] * 14), c['g'], c['nc']))
masks = map(lambda labels: load_nii(labels).get_data().astype(np.bool),
label_names)
# > Ensemble training
#
# I should probably try a sliding window and the random sampling version to see which one is better.
# Sadly, I have a feeling that the sliding window approach should be worse.
print('%s- Extracting centers from the tumor ROI' % ' '.join([''] * 15))
train_centers = get_mask_centers(masks)
train_centers = map(
lambda centers: map(
tuple,
np.random.permutation(centers)[::options['down_sampling']].tolist(
)), train_centers)
print('%s- %d centers will be used' %
(' '.join([''] * 15), sum(map(len, train_centers))))
dense_size = options['dense_size']
conv_blocks_seg = options['conv_blocks_seg']
nets, unet, cnn, fcnn, ucnn = get_brats_nets(
n_channels=image_names.shape[-1],
filters_list=n_filters * conv_blocks_seg,
kernel_size_list=[conv_width] * conv_blocks_seg,
nlabels=options['nlabels'],
dense_size=dense_size)
# First we train the nets inside the cluster net (I don't know what other name I could
# give to that architecture).
train_seg(image_names=image_names,
label_names=label_names,
train_centers=train_centers,
net=nets,
save_path=save_path,
sufix='-nets-%s.d%d' % (sufix, dense_size),
nlabels=options['nlabels'],
net_type='nets')
# Then we train the Dense/Fully Connected layer that defines the ensemble.
# The previous networks should be frozen here.
ensemble = get_brats_ensemble(n_channels=image_names.shape[-1],
n_blocks=conv_blocks_seg,
unet=unet,
cnn=cnn,
fcnn=fcnn,
ucnn=ucnn,
nlabels=options['nlabels'])
train_seg(image_names=image_names,
label_names=label_names,
train_centers=train_centers,
net=ensemble,
save_path=save_path,
sufix='-ensemble-%s.d%d' % (sufix, dense_size),
nlabels=options['nlabels'],
net_type='ensemble')
return net, ensemble
def train_test_survival(net_name, n_folds, val_split=0.1):
# Init
c = color_codes()
options = parse_inputs()
epochs = options['epochs']
batch_size = options['batch_size']
patch_size = options['patch_size']
patience = options['patience']
depth = options['blocks']
filters = options['filters']
d_path = options['loo_dir']
patients = get_dirs(d_path)
survival_dict = get_survival_data()
seg_patients = filter(lambda p: p not in survival_dict.keys(), patients)
survival_patients = survival_dict.keys()
low_patients = filter(
lambda p: int(survival_dict[p]['Survival']) < 300,
survival_patients
)
mid_patients = filter(
lambda p: (int(survival_dict[p]['Survival']) >= 300) &
(int(survival_dict[p]['Survival']) < 450),
survival_patients
)
hi_patients = filter(
lambda p: int(survival_dict[p]['Survival']) >= 450,
survival_patients
)
t_survival_dict = get_survival_data(True)
t_survival_patients = t_survival_dict.keys()
test_survivals = np.zeros(len(t_survival_patients))
t_survival_ages = map(lambda v: float(v['Age']), t_survival_dict.values())
''' Segmentation training'''
# The goal here is to pretrain a unique segmentation network for all
# the survival folds. We only do it once. Then we split the survival
# patients accordingly.
net = BratsSurvivalNet(depth_seg=depth, depth_pred=depth, filters=filters)
print(
'Training segmentation samples = %d' % (
len(seg_patients)
)
)
# Training itself
model_name = '%s-init.mdl' % net_name
num_workers = 8
try:
net.base_model.load_model(os.path.join(d_path, model_name))
except IOError:
n_params = sum(
p.numel() for p in net.base_model.parameters() if p.requires_grad
)
print(
'%sStarting segmentation training with a unet%s (%d parameters)' %
(c['c'], c['nc'], n_params)
)
targets = get_labels(seg_patients)
rois, data = get_images(seg_patients)
print('< Training dataset >')
if patch_size is None:
seg_dataset = BBImageDataset(
data, targets, rois, flip=True
)
else:
seg_dataset = BratsDataset(
data, targets, rois, patch_size
)
print('Dataloader creation <with validation>')
train_loader = DataLoader(
seg_dataset, batch_size, True, num_workers=num_workers,
)
targets = get_labels(survival_patients)
rois, data = get_images(survival_patients)
# Validation
print('< Validation dataset >')
val_dataset = BBImageDataset(
data, targets, rois
)
val_loader = DataLoader(
val_dataset, 1, num_workers=num_workers
)
print(
'%s[%s] %sInitial %s%spretraining%s' % (
c['c'], strftime("%H:%M:%S"), c['g'],
c['c'], c['b'], c['nc']
)
)
net.base_model.fit(train_loader, val_loader, epochs=epochs, patience=patience)
net.base_model.save_model(os.path.join(d_path, model_name))
tr_csv = os.path.join(options['loo_dir'], 'survival_results.csv')
val_csv = os.path.join(options['val_dir'], 'survival_results.csv')
with open(tr_csv, 'w') as csvfile, open(val_csv, 'w') as val_csvfile:
csvwriter = csv.writer(csvfile, delimiter=',')
val_csvwriter = csv.writer(val_csvfile, delimiter=',')
# First we'll define the maximum bounding box, for training and testing
masks, _ = get_images(survival_patients)
indices = map(lambda mask: np.where(mask > 0), masks)
min_bb = np.min(
map(
lambda idx: np.min(idx, axis=-1),
indices
),
axis=0
)
max_bb = np.max(
map(
lambda idx: np.max(idx, axis=-1),
indices
),
axis=0
)
bb = map(
lambda (min_i, max_i): slice(min_i, max_i),
zip(min_bb, max_bb)
)
# After that, we can finally train the model to predict the survival.
for i in range(n_folds):
model_name = '%s-init.mdl' % net_name
net = BratsSurvivalNet(depth_seg=depth, depth_pred=depth, filters=filters)
net.base_model.load_model(os.path.join(d_path, model_name))
''' Survival training'''
model_name = '%s_f%d.mdl' % (net_name, i)
print(
'%s[%s] %sFold %s(%s%d%s%s/%d)%s' % (
c['c'], strftime("%H:%M:%S"), c['g'],
c['c'], c['b'], i + 1, c['nc'], c['c'], n_folds, c['nc']
)
)
try:
net.load_model(os.path.join(d_path, model_name))
except IOError:
# Train/Test patient split according to survival classes
hi_ini = len(hi_patients) * i / n_folds
mid_ini = len(mid_patients) * i / n_folds
low_ini = len(low_patients) * i / n_folds
hi_end = len(hi_patients) * (i + 1) / n_folds
mid_end = len(mid_patients) * (i + 1) / n_folds
low_end = len(low_patients) * (i + 1) / n_folds
high_i = hi_patients[:hi_ini] + hi_patients[hi_end:]
mid_i = mid_patients[:mid_ini] + mid_patients[mid_end:]
low_i = low_patients[:low_ini] + low_patients[low_end:]
# Split of the target survivals
hi_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), high_i
)
mid_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), mid_i
)
low_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), low_i
)
# Split of the age feature
hi_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), high_i
)
mid_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), mid_i
)
low_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), low_i
)
# Data split (using numpy) for train and validation.
# We also compute the number of batches for both training and
# validation according to the batch size.
# Train/Validation split
n_hi = len(high_i)
n_mid = len(mid_i)
n_low = len(low_i)
n_hitrain = int(n_hi * (1 - val_split))
n_midtrain = int(n_mid * (1 - val_split))
n_lowtrain = int(n_low * (1 - val_split))
# Training
hi_train = high_i[:n_hitrain]
mid_train = mid_i[:n_midtrain]
low_train = low_i[:n_lowtrain]
train_i = hi_train + mid_train + low_train
hi_train_ages = hi_ages[:n_hitrain]
mid_train_ages = mid_ages[:n_midtrain]
low_train_ages = low_ages[:n_lowtrain]
train_ages = hi_train_ages + mid_train_ages + low_train_ages
hi_train_surv = hi_survival[:n_hitrain]
mid_train_surv = mid_survival[:n_midtrain]
low_train_surv = low_survival[:n_lowtrain]
train_surv = hi_train_surv + mid_train_surv + low_train_surv
print('< Training dataset >')
train_rois, train_data = get_images(train_i)
train_dataset = BBImageValueDataset(
train_data, train_ages, train_surv, train_rois, bb=bb
)
print('Dataloader creation <train>')
train_loader = DataLoader(
train_dataset, 8, True, num_workers=num_workers,
)
# Validation
hi_train = high_i[n_hitrain:]
mid_train = mid_i[n_midtrain:]
low_train = low_i[n_lowtrain:]
val_i = hi_train + mid_train + low_train
hi_train_ages = hi_ages[n_hitrain:]
mid_train_ages = mid_ages[n_midtrain:]
low_train_ages = low_ages[n_lowtrain:]
val_ages = hi_train_ages + mid_train_ages + low_train_ages
hi_train_surv = hi_survival[n_hitrain:]
mid_train_surv = mid_survival[n_midtrain:]
low_train_surv = low_survival[n_lowtrain:]
val_surv = hi_train_surv + mid_train_surv + low_train_surv
print('< Validation dataset >')
val_rois, val_data = get_images(val_i)
val_dataset = BBImageValueDataset(
val_data, val_ages, val_surv, val_rois, bb=bb
)
print('Dataloader creation <val>')
val_loader = DataLoader(
val_dataset, 1, num_workers=num_workers
)
n_train = len(train_dataset)
n_val = len(val_dataset)
n_test = len(survival_patients) - n_train - n_val
print(
'Train / validation / test samples = %d / %d / %d' % (
n_train, n_val, n_test
)
)
net.fit(
train_loader, val_loader, epochs=epochs, patience=patience
)
net.save_model(os.path.join(d_path, model_name))
''' Survival testing '''
# Testing data
high_i = hi_patients[hi_ini:hi_end]
mid_i = mid_patients[mid_ini:mid_end]
low_i = low_patients[low_ini:low_end]
test_patients = high_i + mid_i + low_i
# Split of the target survivals
hi_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), high_i
)
mid_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), mid_i
)
low_survival = map(
lambda h_i: float(survival_dict[h_i]['Survival']), low_i
)
test_survival = hi_survival + mid_survival + low_survival
# Split of the age feature
hi_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), high_i
)
mid_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), mid_i
)
low_ages = map(
lambda h_i: float(survival_dict[h_i]['Age']), low_i
)
test_ages = hi_ages + mid_ages + low_ages
_, test_data = get_images(test_patients)
print(
'Testing patients (with GT) = %d' % (
len(test_patients)
)
)
pred_y = net.predict(test_data, test_ages, bb)
for p, survival_out, survival in zip(
test_patients, pred_y, test_survival
):
print(
'Estimated survival = %f (%f)' % (survival_out, survival)
)
csvwriter.writerow([p, '%f' % float(survival_out)])
_, test_data = get_images(t_survival_patients, True)
print(
'Testing patients = %d' % (
len(t_survival_patients)
)
)
pred_y = net.predict(test_data, t_survival_ages, bb)
test_survivals += np.array(pred_y)
for p, survival_out, s in zip(
t_survival_patients, pred_y, test_survivals
):
print(
'Estimated survival = %f (%f)' % (
survival_out, s / (i + 1)
)
)
test_survivals = test_survivals / n_folds
for p, survival_out in zip(test_patients, test_survivals):
print('Final estimated survival = %f' % survival_out)
val_csvwriter.writerow([p, '%f' % float(survival_out)])
def test_seg(net, p, outputname, nlabels, mask=None, verbose=True):
c = color_codes()
options = parse_inputs()
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname_path = os.path.join(patient_path, outputname + '.nii.gz')
try:
roi_nii = load_nii(outputname_path)
if verbose:
print('%s%s<%s%s%s%s - probability map loaded>%s' % (''.join(
[' '] * 14), c['g'], c['b'], p_name, c['nc'], c['g'], c['nc']))
except IOError:
roi_nii = load_nii(p[0])
# Image loading
if mask is None:
# This is the unet path
x = np.expand_dims(np.stack(load_images(p), axis=0), axis=0)
# Network parameters
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(
n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
image_net = options['net'](x.shape[1:], filters_list,
kernel_size_list, nlabels)
# We should copy the weights here (if not using roinet)
for l_new, l_orig in zip(image_net.layers[1:], net.layers[1:]):
l_new.set_weights(l_orig.get_weights())
# Now we can test
if verbose:
print('%s[%s] %sTesting the network%s' %
(c['c'], strftime("%H:%M:%S"), c['g'], c['nc']))
# Load only the patient images
if verbose:
print(
'%s%s<Creating the probability map for %s%s%s%s - %s%s%s %s>%s'
% (''.join(
[' '] * 14), c['g'], c['b'], p_name, c['nc'], c['g'],
c['b'], outputname_path, c['nc'], c['g'], c['nc']))
pr_maps = image_net.predict(x, batch_size=options['test_size'])
image = np.argmax(pr_maps, axis=-1).reshape(x.shape[2:])
image = get_biggest_region(image)
else:
# This is the ensemble path
image = np.zeros_like(mask, dtype=np.int8)
options = parse_inputs()
conv_blocks = options['conv_blocks_seg']
test_centers = get_mask_blocks(mask)
x = get_data(
image_names=[p],
list_of_centers=[test_centers],
patch_size=(conv_blocks * 2 + 3, ) * 3,
verbose=verbose,
)
if verbose:
print('%s- Concatenating the data x' % ' '.join([''] * 12))
x = np.concatenate(x)
pr_maps = net.predict(x, batch_size=options['test_size'])
[x, y, z] = np.stack(test_centers, axis=1)
image[x, y, z] = np.argmax(pr_maps, axis=1).astype(dtype=np.int8)
roi_nii.get_data()[:] = image
roi_nii.to_filename(outputname_path)
return roi_nii
def test_seg_validation(net_name):
# Init
c = color_codes()
options = parse_inputs()
depth = options['blocks']
filters = options['filters']
d_path = options['loo_dir']
v_path = options['val_dir']
seg_path = os.path.join(v_path, 'segmentation')
if not os.path.isdir(seg_path):
os.mkdir(seg_path)
unc_path = os.path.join(v_path, 'uncertainty')
if not os.path.isdir(unc_path):
os.mkdir(unc_path)
p = get_dirs(d_path)[0]
test_patients = filter(lambda p: 'BraTS19' in p, get_dirs(v_path))
_, test_x = get_images(test_patients, True)
print(
'Testing patients = %d' % (
len(test_patients)
)
)
# The sub-regions considered for evaluation are:
# 1) the "enhancing tumor" (ET)
# 2) the "tumor core" (TC)
# 3) the "whole tumor" (WT)
#
# The provided segmentation labels have values of 1 for NCR & NET,
# 2 for ED, 4 for ET, and 0 for everything else.
# The participants are called to upload their segmentation labels
# as a single multi-label file in nifti (.nii.gz) format.
#
# The participants are called to upload 4 nifti (.nii.gz) volumes
# (3 uncertainty maps and 1 multi-class segmentation volume from
# Task 1) onto CBICA's Image Processing Portal format. For example,
# for each ID in the dataset, participants are expected to upload
# following 4 volumes:
# 1. {ID}.nii.gz (multi-class label map)
# 2. {ID}_unc_whole.nii.gz (Uncertainty map associated with whole tumor)
# 3. {ID}_unc_core.nii.gz (Uncertainty map associated with tumor core)
# 4. {ID}_unc_enhance.nii.gz (Uncertainty map associated with enhancing tumor)
for i, (p_i, test_i) in enumerate(zip(test_patients, test_x)):
t_in = time.time()
# unc_i = np.zeros((4,) + test_i.shape[1:])
pred_i = np.zeros((4,) + test_i.shape[1:])
for f in range(5):
model_name = '%s-f%d.mdl' % (net_name, f)
net = BratsSegmentationNet(depth=depth, filters=filters)
net.load_model(os.path.join(d_path, model_name))
#
# unc_i += net.uncertainty([test_i], steps=10)[0] * 0.2
pred_i += net.segment([test_i])[0]
seg_i = np.argmax(pred_i, axis=0)
seg_i[seg_i == 3] = 4
# seg_unc_i = np.argmax(unc_i, axis=0)
# seg_unc_i[seg_unc_i == 3] = 4
tumor_mask = remove_small_regions(
seg_i.astype(np.bool), min_size=30
)
seg_i[log_not(tumor_mask)] = 0
# seg_unc_i[log_not(tumor_mask)] = 0
#
# whole_i = np.sum(unc_i[1:]) * tumor_mask.astype(np.float32)
# core_i = unc_i[1] + unc_i[-1] * tumor_mask.astype(np.float32)
# enhance_i = unc_i[-1] * tumor_mask.astype(np.float32)
#
# seg_unc_i = np.argmax(unc_i, axis=0)
# seg_unc_i[seg_unc_i == 3] = 4
niiname = os.path.join(d_path, p, p + '_seg.nii.gz')
nii = load_nii(niiname)
nii.get_data()[:] = seg_i
save_nii(nii, os.path.join(seg_path, p_i + '.nii.gz'))
# nii.get_data()[:] = seg_unc_i
# save_nii(nii, os.path.join(unc_path, p_i + '.nii.gz'))
# niiname = os.path.join(v_path, p_i, p_i + '_flair.nii.gz')
# nii = load_nii(niiname)
# nii.get_data()[:] = whole_i
# save_nii(
# nii, os.path.join(unc_path, p_i + '_unc_whole.nii.gz')
# )
# nii.get_data()[:] = core_i
# save_nii(
# nii, os.path.join(unc_path, p_i + '_unc_core.nii.gz')
# )
# nii.get_data()[:] = enhance_i
# save_nii(
# nii, os.path.join(unc_path, p_i + '_unc_enhance.nii.gz')
# )
t_s = time_to_string(time.time() - t_in)
print(
'Finished patient %s (%d/%d) %s' % (p_i, i + 1, len(test_x), t_s)
)
def main():
options = parse_inputs()
c = color_codes()
# Prepare the net architecture parameters
sequential = options['sequential']
dfactor = options['dfactor']
# Prepare the net hyperparameters
num_classes = 5
epochs = options['epochs']
padding = options['padding']
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
dense_size = options['dense_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters * conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(
conv_width, list) else [conv_width] * conv_blocks
balanced = options['balanced']
# Data loading parameters
preload = options['preload']
queue = options['queue']
# Prepare the sufix that will be added to the results for the net and images
path = options['dir_name']
filters_s = 'n'.join(['%d' % nf for nf in filters_list])
conv_s = 'c'.join(['%d' % cs for cs in kernel_size_list])
s_s = '.s' if sequential else '.f'
ub_s = '.ub' if not balanced else ''
params_s = (ub_s, dfactor, s_s, patch_width, conv_s, filters_s, dense_size,
epochs, padding)
sufix = '%s.D%d%s.p%d.c%s.n%s.d%d.e%d.pad_%s.' % params_s
n_channels = np.count_nonzero([
options['use_flair'], options['use_t2'], options['use_t1'],
options['use_t1ce']
])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' +
'Starting cross-validation' + c['nc'])
# N-fold cross validation main loop (we'll do 2 training iterations with testing for each patient)
data_names, label_names = get_names_from_path(options)
folds = options['folds']
fold_generator = izip(
nfold_cross_validation(data_names, label_names, n=folds,
val_data=0.25), xrange(folds))
dsc_results = list()
for (train_data, train_labels, val_data, val_labels, test_data,
test_labels), i in fold_generator:
print(
c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] +
'Fold %d/%d: ' % (i + 1, folds) + c['g'] +
'Number of training/validation/testing images (%d=%d/%d=%d/%d)' %
(len(train_data), len(train_labels), len(val_data),
len(val_labels), len(test_data)) + c['nc'])
# Prepare the data relevant to the leave-one-out (subtract the patient from the dataset and set the path)
# Also, prepare the network
net_name = os.path.join(
path, 'baseline-brats2017.fold%d' % i + sufix + 'mdl')
# First we check that we did not train for that patient, in order to save time
try:
# net_name_before = os.path.join(path,'baseline-brats2017.fold0.D500.f.p13.c3c3c3c3c3.n32n32n32n32n32.d256.e1.pad_valid.mdl')
net = keras.models.load_model(net_name)
except IOError:
print '==============================================================='
# NET definition using Keras
train_centers = get_cnn_centers(train_data[:, 0],
train_labels,
balanced=balanced)
val_centers = get_cnn_centers(val_data[:, 0],
val_labels,
balanced=balanced)
train_samples = len(train_centers) / dfactor
val_samples = len(val_centers) / dfactor
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Creating and compiling the model ' + c['b'] +
'(%d samples)' % train_samples + c['nc'])
train_steps_per_epoch = -(-train_samples / batch_size)
val_steps_per_epoch = -(-val_samples / batch_size)
input_shape = (n_channels, ) + patch_size
# This architecture is based on the functional Keras API to introduce 3 output paths:
# - Whole tumor segmentation
# - Core segmentation (including whole tumor)
# - Whole segmentation (tumor, core and enhancing parts)
# The idea is to let the network work on the three parts to improve the multiclass segmentation.
# merged_inputs = Input(shape=(4,) + patch_size, name='merged_inputs')
# flair = merged_inputs
model = Sequential()
model.add(
Conv3D(64, (3, 3, 3),
strides=1,
padding='same',
activation='relu',
data_format='channels_first',
input_shape=(4, options['patch_width'],
options['patch_width'],
options['patch_width'])))
model.add(
Conv3D(64, (3, 3, 3),
strides=1,
padding='same',
activation='relu',
data_format='channels_first'))
model.add(
MaxPooling3D(pool_size=(3, 3, 3),
strides=2,
data_format='channels_first'))
model.add(
Conv3D(128, (3, 3, 3),
strides=1,
padding='same',
activation='relu',
data_format='channels_first'))
model.add(
Conv3D(128, (3, 3, 3),
strides=1,
padding='same',
activation='relu',
data_format='channels_first'))
model.add(
MaxPooling3D(pool_size=(3, 3, 3),
strides=2,
data_format='channels_first'))
model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes, activation='softmax'))
net = model
# net_name_before = os.path.join(path,'baseline-brats2017.fold0.D500.f.p13.c3c3c3c3c3.n32n32n32n32n32.d256.e1.pad_valid.mdl')
# net = keras.models.load_model(net_name_before)
net.compile(optimizer='sgd',
loss='categorical_crossentropy',
metrics=['accuracy'])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'Training the model with a generator for ' + c['b'] +
'(%d parameters)' % net.count_params() + c['nc'])
print(net.summary())
net.fit_generator(
generator=load_patch_batch_train(
image_names=train_data,
label_names=train_labels,
centers=train_centers,
batch_size=batch_size,
size=patch_size,
# fc_shape = patch_size,
nlabels=num_classes,
dfactor=dfactor,
preload=preload,
split=not sequential,
datatype=np.float32),
validation_data=load_patch_batch_train(
image_names=val_data,
label_names=val_labels,
centers=val_centers,
batch_size=batch_size,
size=patch_size,
# fc_shape = patch_size,
nlabels=num_classes,
dfactor=dfactor,
preload=preload,
split=not sequential,
datatype=np.float32),
# workers=queue,
steps_per_epoch=train_steps_per_epoch,
validation_steps=val_steps_per_epoch,
max_q_size=queue,
epochs=epochs)
net.save(net_name)
# Then we test the net.
for p, gt_name in zip(test_data, test_labels):
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
outputname = os.path.join(patient_path,
'deep-brats17' + sufix + 'test.nii.gz')
gt_nii = load_nii(gt_name)
gt = np.copy(gt_nii.get_data()).astype(dtype=np.uint8)
try:
load_nii(outputname)
except IOError:
roi_nii = load_nii(p[0])
roi = roi_nii.get_data().astype(dtype=np.bool)
centers = get_mask_voxels(roi)
test_samples = np.count_nonzero(roi)
image = np.zeros_like(roi).astype(dtype=np.uint8)
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] +
'<Creating the probability map ' + c['b'] + p_name +
c['nc'] + c['g'] + ' (%d samples)>' % test_samples +
c['nc'])
test_steps_per_epoch = -(-test_samples / batch_size)
y_pr_pred = net.predict_generator(
generator=load_patch_batch_generator_test(
image_names=p,
centers=centers,
batch_size=batch_size,
size=patch_size,
preload=preload,
),
steps=test_steps_per_epoch,
max_q_size=queue)
[x, y, z] = np.stack(centers, axis=1)
if not sequential:
tumor = np.argmax(y_pr_pred[0], axis=1)
y_pr_pred = y_pr_pred[-1]
roi = np.zeros_like(roi).astype(dtype=np.uint8)
roi[x, y, z] = tumor
roi_nii.get_data()[:] = roi
roiname = os.path.join(
patient_path,
'deep-brats17' + sufix + 'test.roi.nii.gz')
roi_nii.to_filename(roiname)
y_pred = np.argmax(y_pr_pred, axis=1)
image[x, y, z] = y_pred
# Post-processing (Basically keep the biggest connected region)
image = get_biggest_region(image)
labels = np.unique(gt.flatten())
results = (p_name, ) + tuple(
[dsc_seg(gt == l, image == l) for l in labels[1:]])
text = 'Subject %s DSC: ' + '/'.join(
['%f' for _ in labels[1:]])
print(text % results)
dsc_results.append(results)
print(c['g'] + ' -- Saving image ' + c['b'] +
outputname + c['nc'])
roi_nii.get_data()[:] = image
roi_nii.to_filename(outputname)
def main():
options = parse_inputs()
c = color_codes()
path = options['dir_name']
test_data, test_labels = get_names_from_path(path, options)
train_data, train_labels = get_names_from_path(os.path.join(path, '../Brats17Test-Training'), options)
net_name = os.path.join(path, 'baseline-brats2017.D50.f.p13.c3c3c3c3c3.n32n32n32n32n32.d256.e50.mdl')
# Prepare the net hyperparameters
patch_width = options['patch_width']
patch_size = (patch_width, patch_width, patch_width)
batch_size = options['batch_size']
conv_blocks = options['conv_blocks']
n_filters = options['n_filters']
filters_list = n_filters if len(n_filters) > 1 else n_filters*conv_blocks
conv_width = options['conv_width']
kernel_size_list = conv_width if isinstance(conv_width, list) else [conv_width]*conv_blocks
options_s = 'e%d.E%d.D%d.' % (options['epochs'], options['net_epochs'], options['down_factor'])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + 'Starting testing' + c['nc'])
# Testing. We retrain the convolutionals and then apply testing. We also check the results without doing it.
dsc_results_o = list()
dsc_results_d = list()
net_roi_name = os.path.join(path, 'CBICA-brats2017.D25.p13.c3c3c3c3c3.n32n32n32n32n32.d256.e50.mdl')
net_roi = keras.models.load_model(net_roi_name)
for i, (p, gt_name) in enumerate(zip(test_data, test_labels)):
p_name = p[0].rsplit('/')[-2]
patient_path = '/'.join(p[0].rsplit('/')[:-1])
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['nc'] + 'Case ' + c['c'] + c['b'] + p_name + c['nc'] +
c['c'] + ' (%d/%d):' % (i + 1, len(test_data)) + c['nc'])
try:
image_o = load_nii(os.path.join(patient_path, p_name + '.nii.gz')).get_data()
except IOError:
# First let's test the original network
net_orig = keras.models.load_model(net_name)
net_orig_conv_layers = sorted(
[l for l in net_orig.layers if 'conv' in l.name],
cmp=lambda x, y: int(x.name[7:]) - int(y.name[7:])
)
image_o = test_network(net_orig, p, batch_size, patch_size, sufix='original', filename=p_name)
try:
outputname = 'deep-brats17.test.' + options_s + 'domain'
image_d = load_nii(os.path.join(patient_path, outputname + '.nii.gz')).get_data()
except IOError:
net_orig = keras.models.load_model(net_name)
net_orig_conv_layers = sorted(
[l for l in net_orig.layers if 'conv' in l.name],
cmp=lambda x, y: int(x.name[7:]) - int(y.name[7:])
)
# Now let's create the domain network and train it
net_new_name = os.path.join(path, 'domain-exp-brats2017.' + options_s + p_name + '.mdl')
try:
net_new = keras.models.load_model(net_new_name)
net_new_conv_layers = [l for l in net_new.layers if 'conv' in l.name]
except IOError:
# First we get the tumor ROI
image_r = test_network(net_roi, p, batch_size, patch_size, sufix='tumor', filename=outputname)
roi = np.logical_and(image_r.astype(dtype=np.bool), image_o.astype(dtype=np.bool))
p_images = np.stack(load_norm_list(p)).astype(dtype=np.float32)
data, clip = clip_to_roi(p_images, roi) if np.count_nonzero(roi) > 0 else clip_to_roi(p_images, image_r)
data_s = c['g'] + c['b'] + 'x'.join(['%d' % i_len for i_len in data.shape[1:]]) + c['nc']
print(c['c'] + '[' + strftime("%H:%M:%S") + '] ' + c['g'] + 'Preparing ' + c['b'] + 'domain' + c['nc'] +
c['g'] + ' data' + c['nc'] + c['g'] + '(shape = ' + data_s + c['g'] + ')' + c['nc'])
# We prepare the zoomed tumors for training
roi_name = os.path.join(path, p_name + '.roi.pkl')
mask_name = os.path.join(path, p_name + '.mask.pkl')
image_name = os.path.join(path, p_name + '.image.pkl')
rate_name = os.path.join(path, p_name + '.rate.pkl')
try:
train_roi = pickle.load(open(roi_name, 'rb'))
train_mask = pickle.load(open(mask_name, 'rb'))
train_image = pickle.load(open(image_name, 'rb'))
train_rate = pickle.load(open(rate_name, 'rb'))
except IOError:
train_num, train_roi, train_rate = get_best_roi(data, train_data, train_labels)
train_image = np.stack(load_norm_list(train_data[train_num])).astype(dtype=np.float32)
train_mask = load_nii(train_labels[train_num]).get_data().astype(dtype=np.uint8)
pickle.dump(train_roi, open(roi_name, 'wb'))
pickle.dump(train_mask, open(mask_name, 'wb'))
pickle.dump(train_image, open(image_name, 'wb'))
pickle.dump(train_rate, open(rate_name, 'wb'))
_, train_clip = clip_to_roi(train_image, train_mask)
train_x = zoom(train_image, train_rate)
train_y = zoom(train_mask, train_rate[1:], order=0)
# We create the domain network
net_new = create_new_network(data.shape[1:], filters_list, kernel_size_list)
net_new_conv_layers = [l for l in net_new.layers if 'conv' in l.name]
for l_new, l_orig in zip(net_new_conv_layers, net_orig_conv_layers):
l_new.set_weights(l_orig.get_weights())
# Transfer learning
train_centers_r = [range(int(cl[0] * tr), int(cl[1] * tr)) for cl, tr in zip(train_clip, train_rate[1:])]
train_centers = list(product(*train_centers_r))
transfer_learning(net_new, net_orig, data, train_x, train_y, train_roi, train_centers, options)
net_new.save(net_new_name)
# Now we transfer the new weights an re-test
for l_new, l_orig in zip(net_new_conv_layers, net_orig_conv_layers):
l_orig.set_weights(l_new.get_weights())
image_d = test_network(net_orig, p, batch_size, patch_size, sufix=options_s + 'domain')
if options['use_dsc']:
results_o = check_dsc(gt_name, image_o)
dsc_results_o.append(results_o)
results_d = check_dsc(gt_name, image_d)
dsc_results_d.append(results_d)
subject_name = c['c'] + c['b'] + '%s' + c['nc']
dsc_string = c['g'] + '/'.join(['%f']*len(results_o)) + c['nc']
text = subject_name + ' DSC: ' + dsc_string
results = (p_name,) + tuple(results_o)
print(''.join([' ']*14) + 'Original ' + text % results)
results = (p_name,) + tuple(results_d)
print(''.join([' ']*14) + 'Domain ' + text % results)
if options['use_dsc']:
f_dsc_o = tuple([np.array([dsc[i] for dsc in dsc_results_o if len(dsc) > i]).mean() for i in range(3)])
f_dsc_d = tuple([np.array([dsc[i] for dsc in dsc_results_d if len(dsc) > i]).mean() for i in range(3)])
f_dsc = f_dsc_o + f_dsc_d
print('Final results DSC: (%f/%f/%f) vs (%f/%f/%f)' % f_dsc)
