def train_Segment_GBM(data_directory, val_data_directory):
# Define input modalities to load.
training_modality_dict = {
'input_modalities':
['FLAIR_pp.*', 'T2_pp.*', 'T1_pp.*', 'T1post_pp.*'],
'ground_truth': ['enhancingmask_pp.nii.gz']
}
load_data = False
train_model = False
load_test_data = False
predict = True
training_data = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/BRATS_enhancing_prediction_only_data.h5'
model_file = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/BRATS_enhancing_prediction_only_model.h5'
testing_data = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/BRATS_enhancing_prediction_only_data.h5'
# Write the data to hdf5
if (not os.path.exists(training_data) and train_model) or load_data:
# Create a Data Collection
training_data_collection = DataCollection(
data_directory, modality_dict=training_modality_dict, verbose=True)
training_data_collection.fill_data_groups()
# Define patch sampling regions
def brain_region(data):
return (data['ground_truth'] != 1) & (data['input_modalities'] !=
0)
def roi_region(data):
return data['ground_truth'] == 1
def empty_region(data):
return data['input_modalities'] == 0
# Add patch augmentation
patch_augmentation = ExtractPatches(
patch_shape=(32, 32, 32),
patch_region_conditions=[[empty_region, .05], [brain_region, .25],
[roi_region, .7]],
data_groups=['input_modalities', 'ground_truth'],
patch_dimensions={
'ground_truth': [1, 2, 3],
'input_modalities': [1, 2, 3]
})
training_data_collection.append_augmentation(patch_augmentation,
multiplier=2000)
# Write data to hdf5
training_data_collection.write_data_to_file(training_data)
if train_model:
# Or load pre-loaded data.
training_data_collection = DataCollection(data_storage=training_data,
verbose=True)
training_data_collection.fill_data_groups()
# Add left-right flips
flip_augmentation = Flip_Rotate_2D(
flip=True,
rotate=False,
data_groups=['input_modalities', 'ground_truth'])
# flip_augmentation = Flip_Rotate_3D(data_groups=['input_modalities', 'ground_truth'])
training_data_collection.append_augmentation(flip_augmentation,
multiplier=2)
# Define model parameters
model_parameters = {
'input_shape': (32, 32, 32, 4),
'downsize_filters_factor': 1,
'pool_size': (2, 2, 2),
'filter_shape': (5, 5, 5),
'dropout': 0,
'batch_norm': True,
'initial_learning_rate': 0.000001,
'output_type': 'regression',
'num_outputs': 1,
'activation': 'relu',
'padding': 'same',
'implementation': 'keras',
'depth': 4,
'max_filter': 512
}
# Create U-Net
unet_model = UNet(**model_parameters)
plot_model(unet_model.model,
to_file='model_image_dn.png',
show_shapes=True)
training_parameters = {
'input_groups': ['input_modalities', 'ground_truth'],
'output_model_filepath': model_file,
'training_batch_size': 64,
'num_epochs': 1000,
'training_steps_per_epoch': 20
}
unet_model.train(training_data_collection, **training_parameters)
else:
unet_model = load_old_model(model_file)
# Define input modalities to load.
testing_modality_dict = {
'input_modalities':
['FLAIR_pp.*', 'T2_pp.*', 'T1_pp.*', 'T1post_pp.*']
}
if predict:
testing_data_collection = DataCollection(
val_data_directory,
modality_dict=testing_modality_dict,
verbose=True)
testing_data_collection.fill_data_groups()
if load_test_data:
# Write data to hdf5
testing_data_collection.write_data_to_file(testing_data)
testing_parameters = {
'inputs': ['input_modalities'],
'output_filename': 'brats_enhancing_only_prediction.nii.gz',
'batch_size': 250,
'patch_overlaps': 1,
'output_patch_shape': (26, 26, 26, 4),
'save_all_steps': True
}
prediction = ModelPatchesInference(**testing_parameters)
label_binarization = BinarizeLabel(postprocessor_string='_label')
prediction.append_postprocessor([label_binarization, largest_island])
unet_model.append_output([prediction])
unet_model.generate_outputs(testing_data_collection)
def skull_strip(output_folder, T1POST=None, FLAIR=None, ground_truth=None, input_directory=None, bias_corrected=True, resampled=False, registered=False, normalized=False, preprocessed=False, save_preprocess=False, save_all_steps=False, mask_output='skullstrip_mask.nii.gz', verbose=True):
#--------------------------------------------------------------------#
# Step 1, Load Data
#--------------------------------------------------------------------#
input_data = {'input_modalities': [FLAIR, T1POST]}
if ground_truth is not None:
input_data['ground_truth'] = [ground_truth]
if input_directory is None:
if any(data is None for data in input_data):
raise ValueError("Cannot segment GBM. Please specify all four modalities.")
data_collection = DataCollection(verbose=verbose)
data_collection.add_case(input_data, case_name=output_folder)
else:
data_collection = DataCollection(input_directory, modality_dict=input_data, verbose=verbose)
data_collection.fill_data_groups()
#--------------------------------------------------------------------#
# Step 2, Preprocess Data
#--------------------------------------------------------------------#
if not preprocessed:
print 'ABOUT TO PREPROCESS....'
# Random hack to save DICOMs to niftis for further processing.
preprocessing_steps = [Preprocessor(data_groups=['input_modalities'], save_output=save_all_steps, verbose=verbose, output_folder=output_folder)]
if not bias_corrected:
preprocessing_steps += [N4BiasCorrection(data_groups=['input_modalities'], save_output=save_all_steps, verbose=verbose, output_folder=output_folder)]
if not resampled:
preprocessing_steps += [Resample(data_groups=['input_modalities'], save_output=save_all_steps, verbose=verbose, output_folder=output_folder)]
if not registered:
preprocessing_steps += [Coregister(data_groups=['input_modalities'], save_output=(save_preprocess or save_all_steps), verbose=verbose, output_folder=output_folder, reference_channel=0)]
if not normalized:
preprocessing_steps += [ZeroMeanNormalization(data_groups=['input_modalities'], save_output=save_all_steps, verbose=verbose, output_folder=output_folder, preprocessor_string='_preprocessed')]
data_collection.append_preprocessor(preprocessing_steps)
#--------------------------------------------------------------------#
# Step 3, Skullstripping
#--------------------------------------------------------------------#
skullstrip_prediction_parameters = {'inputs': ['input_modalities'],
'output_filename': os.path.join(output_folder, mask_output),
'batch_size': 25,
'patch_overlaps': 8,
'channels_first': True,
'patch_dimensions': [-3, -2, -1],
'output_patch_shape': (1, 64, 64, 32),
# 'input_channels': [0, 3],
}
skull_stripping_model = load_old_model(load('Skull_Strip_T1Post_FLAIR'))
skull_stripping_prediction = ModelPatchesInference(**skullstrip_prediction_parameters)
label_binarization = BinarizeLabel()
largest_component = LargestComponents()
hole_filler = FillHoles(postprocessor_string='_label')
skull_stripping_prediction.append_postprocessor([label_binarization, largest_component, hole_filler])
skull_stripping_model.append_output([skull_stripping_prediction])
for case in data_collection.cases:
print '\nStarting New Case...\n'
skull_stripping_prediction.case = case
skull_stripping_mask = skull_stripping_model.generate_outputs(data_collection)[0]['filenames'][-1]
if not save_preprocess:
for index, file in enumerate(data_collection.data_groups['input_modalities'].preprocessed_case):
os.remove(file)
def predict_GBM(output_folder,
T2=None,
T1=None,
T1POST=None,
FLAIR=None,
ground_truth=None,
input_directory=None,
bias_corrected=True,
resampled=False,
registered=False,
skullstripped=False,
normalized=False,
preprocessed=False,
save_preprocess=False,
save_all_steps=False,
output_wholetumor_filename='wholetumor_segmentation.nii.gz',
output_enhancing_filename='enhancing_segmentation.nii.gz',
verbose=True):
#--------------------------------------------------------------------#
# Step 1, Load Data
#--------------------------------------------------------------------#
input_data = {'input_modalities': [FLAIR, T2, T1, T1POST]}
if ground_truth is not None:
input_data['ground_truth'] = [ground_truth]
if input_directory is None:
if any(data is None for data in input_data):
raise ValueError(
"Cannot segment GBM. Please specify all four modalities.")
data_collection = DataCollection(verbose=verbose)
data_collection.add_case(input_data, case_name=output_folder)
else:
data_collection = DataCollection(input_directory,
modality_dict=input_data,
verbose=verbose)
data_collection.fill_data_groups()
#--------------------------------------------------------------------#
# Step 2, Preprocess Data
#--------------------------------------------------------------------#
if not preprocessed or True:
# Random hack to save DICOMs to niftis for further processing.
preprocessing_steps = [
DICOMConverter(data_groups=['input_modalities'],
save_output=save_all_steps,
verbose=verbose,
output_folder=output_folder)
]
if not bias_corrected:
preprocessing_steps += [
N4BiasCorrection(data_groups=['input_modalities'],
save_output=save_all_steps,
verbose=verbose,
output_folder=output_folder)
]
if not resampled:
preprocessing_steps += [
Resample(data_groups=['input_modalities'],
save_output=save_all_steps,
verbose=verbose,
output_folder=output_folder)
]
if not registered:
preprocessing_steps += [
Coregister(data_groups=['input_modalities'],
save_output=(save_preprocess or save_all_steps),
verbose=verbose,
output_folder=output_folder,
reference_channel=1)
]
if not skullstripped:
preprocessing_steps += [
SkullStrip(data_groups=['input_modalities'],
save_output=save_all_steps,
verbose=verbose,
output_folder=output_folder,
reference_channel=1)
]
if not normalized:
preprocessing_steps += [
ZeroMeanNormalization(
data_groups=['input_modalities'],
save_output=save_all_steps,
verbose=verbose,
mask_preprocessor=preprocessing_steps[-1],
preprocessor_string='_preprocessed')
]
data_collection.append_preprocessor(preprocessing_steps)
#--------------------------------------------------------------------#
# Step 3, Segmentation
#--------------------------------------------------------------------#
wholetumor_prediction_parameters = {
'inputs': ['input_modalities'],
'output_filename': os.path.join(output_folder,
output_wholetumor_filename),
'batch_size': 75,
'patch_overlaps': 1,
'channels_first': True,
'patch_dimensions': [-3, -2, -1],
'output_patch_shape': (1, 26, 26, 26),
# 'input_channels': [0, 3],
}
enhancing_prediction_parameters = {
'inputs': ['input_modalities'],
'output_filename': os.path.join(output_folder,
output_enhancing_filename),
'batch_size': 75,
'patch_overlaps': 1,
'channels_first': True,
'output_patch_shape': (1, 26, 26, 26),
'patch_dimensions': [-3, -2, -1]
}
wholetumor_model = load_old_model(load('Segment_GBM_wholetumor'))
enhancing_model = load_old_model(load('Segment_GBM_enhancing'))
wholetumor_prediction = ModelPatchesInference(
**wholetumor_prediction_parameters)
wholetumor_model.append_output([wholetumor_prediction])
enhancing_prediction = ModelPatchesInference(
**enhancing_prediction_parameters)
enhancing_model.append_output([enhancing_prediction])
label_binarization = BinarizeLabel(postprocessor_string='_label')
wholetumor_prediction.append_postprocessor([label_binarization])
enhancing_prediction.append_postprocessor([label_binarization])
for case in data_collection.cases:
print '\nStarting New Case...\n'
wholetumor_file = wholetumor_model.generate_outputs(
data_collection, case)[0]['filenames'][-1]
data_collection.add_channel(case, wholetumor_file)
enhancing_file = enhancing_model.generate_outputs(
data_collection, case)[0]['filenames'][-1]
data_collection.clear_outputs()