Esempio n. 1
0
def train_Segment_GBM(data_directory):

    # Define input modalities to load.
    training_modality_dict = {'input_modalities': 
    ['*FLAIR_pp.*', '*T2_pp.*', '*T1_pp.*', '*T1post_pp.*', '*full_edemamask_pp.*'],
    'ground_truth': ['*full_edemamask_pp.*']}

    # Create a Data Collection
    training_data_collection = DataCollection(data_directory, training_modality_dict, 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'])
    training_data_collection.append_augmentation(flip_augmentation, multiplier=2)

    # 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

    # Add patch augmentation
    patch_augmentation = ExtractPatches(patch_shape=(32,32,32), patch_region_conditions=[[brain_region, .3], [roi_region, .7]], data_groups=['input_modalities', 'ground_truth'])
    training_data_collection.append_augmentation(patch_augmentation, multiplier=70)

    # Write the data to hdf5
    training_data_collection.write_data_to_file('./test.h5')

    # Define model parameters
    model_parameters = {input_shape: (32, 32, 32, 4),
                    downsize_filters_factor: 1,
                    pool_size: (2, 2, 2), 
                    filter_shape: (3, 3, 3), 
                    dropout: .1, 
                    batch_norm: False, 
                    initial_learning_rate: 0.00001, 
                    output_type: 'binary_label',
                    num_outputs: 1, 
                    activation: 'relu',
                    padding: 'same', 
                    implementation: 'keras',
                    depth: 4,
                    max_filter=512}

    # Create U-Net
    if True:
        unet_model = UNet(**model_parameters)

    # Or load an old one
    else:
        unet_model = load_old_model('model.h5')

    # Define training parameters
    training_parameters = {}

    # Define training generators
    training_generator = None
Esempio n. 2
0
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)
Esempio n. 3
0
def train_Segment_GBM(data_directory, val_data_directory):

    # Define input modalities to load.
    training_modality_dict = {
        'input_modalities': [
            '*FLAIR*', ['*T2SPACE*', '*T2_pp*'], ['*T1_pp.*', '*MPRAGE_Pre*'],
            ['*T1post_pp.*', '*MPRAGE_POST*'], ['enhancing*'],
            ['wholetumor*', 'full_edemamask*']
        ],
        'ground_truth': [['enhancing*'], ['wholetumor*', 'full_edemamask*']]
    }

    load_data = False
    train_model = False
    load_test_data = True
    predict = True

    training_data = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/enhancing_label_upsampling_323232.h5'
    model_file = 'label_upsampling_323232_model_correct.h5'
    testing_data = './FLAIR_upsampling_323232_test.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

        # Add patch augmentation
        patch_augmentation = ExtractPatches(
            patch_shape=(32, 32, 32),
            patch_region_conditions=[[roi_region, 1]],
            data_groups=['input_modalities', 'ground_truth'],
            patch_dimensions={
                'ground_truth': [0, 1, 2],
                'input_modalities': [0, 1, 2]
            })
        training_data_collection.append_augmentation(patch_augmentation,
                                                     multiplier=70)

        # 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()

        # Choose a modality
        choice_augmentation = ChooseData(
            axis={
                'input_modalities': -1,
                'ground_truth': -1
            },
            choices=[-1, -2],
            data_groups=['input_modalities', 'ground_truth'],
            random_sample=False)
        training_data_collection.append_augmentation(choice_augmentation,
                                                     multiplier=2)

        # Add down-sampling
        mask_augmentation = Downsample(channel=4,
                                       axes={'input_modalities': [-4, -3, -2]},
                                       factor=3,
                                       data_groups=['input_modalities'])
        training_data_collection.append_augmentation(mask_augmentation,
                                                     multiplier=4)

        # Add left-right flips
        flip_augmentation = Flip_Rotate_2D(
            flip=True,
            rotate=False,
            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, 5),
            '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': 'binary_label',
            '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)

    # Load testing data..
    if not os.path.exists(testing_data) or load_test_data:
        # Create a Data Collection
        testing_data_collection = DataCollection(
            val_data_directory,
            modality_dict=training_modality_dict,
            verbose=True)
        testing_data_collection.fill_data_groups()
        # Write data to hdf5
        testing_data_collection.write_data_to_file(testing_data)

    if predict:
        testing_data_collection = DataCollection(data_storage=testing_data,
                                                 verbose=True)
        testing_data_collection.fill_data_groups()

        # Choose a modality
        choice_augmentation = ChooseData(
            axis={
                'input_modalities': -1,
                'ground_truth': -1
            },
            choices=[-1, -2],
            data_groups=['input_modalities', 'ground_truth'],
            random_sample=False)
        testing_data_collection.append_augmentation(choice_augmentation,
                                                    multiplier=2)

        # Add down-sampling
        mask_augmentation = Downsample(channel=4,
                                       axes={'input_modalities': [-4, -3, -2]},
                                       factor=3,
                                       data_groups=['input_modalities'],
                                       random_sample=False)
        testing_data_collection.append_augmentation(mask_augmentation,
                                                    multiplier=3)

        testing_parameters = {
            'inputs': ['input_modalities'],
            'output_filename': 'deepneuro-label.nii.gz',
            'batch_size': 250,
            'patch_overlaps': 6,
            'output_patch_shape': (26, 26, 26, 4)
        }

        prediction = ModelPatchesInference(testing_data_collection,
                                           **testing_parameters)

        unet_model.append_output([prediction])
        unet_model.generate_outputs()
Esempio n. 4
0
def train_Segment_GBM(data_directory, val_data_directory):

    # Define input modalities to load.
    training_modality_dict = {'input_modalities': [['T1_pp*']]}

    load_data = False
    train_model = True
    load_test_data = False
    predict = False

    training_data = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/GAN_data_corrected.h5'
    model_file = 'GAN_model.h5'
    testing_data = './GAN_test.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)
            return data['input_modalities'] != 0

        def roi_region(data):
            return data['ground_truth'] == 1

        # Add patch augmentation
        patch_augmentation = ExtractPatches(
            patch_shape=(64, 64, 8),
            patch_region_conditions=[[brain_region, 1]],
            data_groups=['input_modalities'],
            patch_dimensions={'input_modalities': [0, 1, 2]})
        training_data_collection.append_augmentation(patch_augmentation,
                                                     multiplier=200)

        # 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'])
        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.5,
            'batch_norm': True,
            'initial_learning_rate': 0.0001,
            'output_type': 'regression',
            'num_outputs': 1,
            'activation': 'relu',
            'padding': 'same',
            'implementation': 'keras',
            'depth': 4,
            'max_filter': 256
        }

        # Create U-Net
        GAN_model = GAN(**model_parameters)
        # plot_model(GAN_model.model, to_file='model_image_dn.png', show_shapes=True)
        training_parameters = {
            'input_groups': ['input_modalities'],
            'output_model_filepath': model_file,
            'training_batch_size': 32,
            'num_epochs': 10000,
            'training_steps_per_epoch': 20
        }
        GAN_model.train(training_data_collection, **training_parameters)
    else:
        GAN_model = load_old_model('DCGAN_150.model.meta', backend='tf')

    if predict:
        print GAN_model
        for i in GAN_model.graph.get_operations():
            print i
Esempio n. 5
0
def train_Segment_GBM(data_directory, val_data_directory):

    # Define input modalities to load.
    if True:
        training_modality_dict = {
            'input_modalities':
            ['*FLAIR_pp.*', '*T2_pp.*', '*T1_pp.*', '*T1post_pp.*'],
            'ground_truth': ['*full_edemamask_pp.*']
        }
    else:
        training_modality_dict = {
            'input_modalities': [['*FLAIR_pp.*', 'FLAIR_norm2*'],
                                 ['*T1post_pp.*', 'T1post_norm2*']],
            'ground_truth': ['*full_edemamask_pp.*', 'FLAIRmask-label.nii.gz']
        }

    load_data = True
    train_model = True
    load_test_data = True
    predict = True

    training_data = './wholetumor_predict_patches_test3.h5'
    model_file = 'wholetumor_segnet-58-0.38.h5'
    testing_data = './brats_test_case.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

        # Add patch augmentation
        patch_augmentation = ExtractPatches(
            patch_shape=(32, 32, 32),
            patch_region_conditions=[[brain_region, 1]],
            data_groups=['input_modalities', 'ground_truth'])
        training_data_collection.append_augmentation(patch_augmentation,
                                                     multiplier=200)

        # Add left-right flips
        flip_augmentation = Flip_Rotate_2D(
            flip=True,
            rotate=False,
            data_groups=['input_modalities', 'ground_truth'])
        training_data_collection.append_augmentation(flip_augmentation,
                                                     multiplier=2)

        # Write data to hdf5
        training_data_collection.write_data_to_file(training_data)

    # Or load pre-loaded data.
    training_data_collection = DataCollection(data_storage=training_data,
                                              verbose=True)
    training_data_collection.fill_data_groups()

    # Define model parameters
    model_parameters = {
        'input_shape': (32, 32, 32, 4),
        'downsize_filters_factor': 1,
        'pool_size': (2, 2, 2),
        'filter_shape': (3, 3, 3),
        'dropout': 0,
        'batch_norm': True,
        'initial_learning_rate': 0.000001,
        'output_type': 'binary_label',
        'num_outputs': 1,
        'activation': 'relu',
        'padding': 'same',
        'implementation': 'keras',
        'depth': 4,
        'max_filter': 512
    }

    # Create U-Net
    if train_model:
        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':
            'wholetumor_segnet-{epoch:02d}-{loss:.2f}.h5',
            'training_batch_size': 2,
            'num_epochs': 100,
            'training_steps_per_epoch': 200,
            'save_best_only': False
        }
        unet_model.train(training_data_collection, **training_parameters)
    else:
        unet_model = load_old_model(model_file)

    # Load testing data..
    if not os.path.exists(testing_data) or load_test_data:
        # Create a Data Collection
        testing_data_collection = DataCollection(
            val_data_directory,
            modality_dict=training_modality_dict,
            verbose=True)
        testing_data_collection.fill_data_groups()
        # Write data to hdf5
        testing_data_collection.write_data_to_file(testing_data)

    if predict:
        testing_data_collection = DataCollection(data_storage=testing_data,
                                                 verbose=True)
        testing_data_collection.fill_data_groups()

        testing_parameters = {
            'inputs': ['input_modalities'],
            'output_filename': 'deepneuro.nii.gz',
            'batch_size': 200,
            'patch_overlaps': 1
        }

        prediction = ModelPatchesInference(testing_data_collection,
                                           **testing_parameters)

        unet_model.append_output([prediction])
        unet_model.generate_outputs()
Esempio n. 6
0
def train_Segment_GBM(data_directory, val_data_directory):

    # Define input modalities to load.
    training_modality_dict = {'input_modalities': 
    ['*FLAIR*nii.gz', ['*T2SPACE*nii.gz'], ['*MPRAGE_POST*nii.gz'], ['*MPRAGE_Pre*nii.gz']],
    'ground_truth': ['*SUV_r_T2_raw.nii.gz*']}

    load_data = False
    train_model = False
    load_test_data = True
    predict = True

    training_data = '/mnt/jk489/QTIM_Databank/DeepNeuro_Datasets/TMZ_4_323232.h5'
    model_file = 'TMZ_4_323232_model.h5'
    testing_data = './TMZ_4_323232_test.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.5

        # Add patch augmentation
        patch_augmentation = ExtractPatches(patch_shape=(32, 32, 32), patch_region_conditions=[[brain_region, .5], [roi_region, .5]], data_groups=['input_modalities', 'ground_truth'], patch_dimensions={'ground_truth': [0,1,2], 'input_modalities': [0,1,2]})
        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'])
        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)

    # Load testing data..
    if not os.path.exists(testing_data) or load_test_data:
        # Create a Data Collection
        testing_data_collection = DataCollection(val_data_directory, modality_dict=training_modality_dict, verbose=True)
        testing_data_collection.fill_data_groups()
        # Write data to hdf5
        testing_data_collection.write_data_to_file(testing_data)

    if predict:
        testing_data_collection = DataCollection(data_storage=testing_data, verbose=True)
        testing_data_collection.fill_data_groups()

        flip_augmentation = Copy(data_groups=['input_modalities', 'ground_truth'])
        testing_data_collection.append_augmentation(flip_augmentation, multiplier=1)

        testing_parameters = {'inputs': ['input_modalities'], 
                        'output_filename': 'deepneuro_suv_4.nii.gz',
                        'batch_size': 50,
                        'patch_overlaps': 6,
                        'output_patch_shape': (26,26,26,4)}

        prediction = ModelPatchesInference(testing_data_collection, **testing_parameters)

        unet_model.append_output([prediction])
        unet_model.generate_outputs()
Esempio n. 7
0
from deepneuro.models.weighted_cat_cross_entropy import WeightedCategoricalCrossEntropy

TrainingDataCollection = DataCollection(data_sources={'csv': 'Metastases_Data_Train.csv'})
TestingDataCollection = DataCollection(data_sources={'csv': 'Metastases_Data_Test.csv'})

Normalization = ZeroMeanNormalization(data_groups=['input_data'])
TrainingDataCollection.append_preprocessor(Normalization)

def BrainRegion(data):
    return data['input_data'] != 0
def TumorRegion(data):
    return data['ground_truth'] == 1

PatchAugmentation = ExtractPatches(patch_shape=(32, 32, 32), 
    patch_region_conditions=[[BrainRegion, 0.70], [TumorRegion, 0.30]])
TrainingDataCollection.append_augmentation(PatchAugmentation, multiplier=20)
TrainingDataCollection.write_data_to_file('training_data.hdf5')

ModelParameters = {'input_shape': (32, 32, 32, 1),
                'cost_function': 'weighted_categorical_label',
                ''}
UNETModel = UNet(**ModelParameters)

TrainingParameters = {'output_model_filepath': 'unet_metastases.py',
                'training_batch_size': 16,
                'num_epochs': 50}
UNETModel.train(TrainingDataCollection, **TrainingParameters)

TestingParameters = {'inputs': ['input_data'],
               'output_filename': '_segmentation.nii.gz',
               'batch_size': 20,
Esempio n. 8
0
def train_Segment_GBM(data_directory, val_data_directory):

    # Define input modalities to load.
    training_modality_dict = {
        'input_modalities': ['*phantom*'],
        'ground_truth': ['*ktrans*']
    }

    load_data = False
    train_model = False
    load_test_data = False
    predict = True

    training_data = './dce_mri_ktrans_training_884_1.h5'
    model_file = 'ktrans_net_884_1_3layer_conv_separated_sym.h5'
    testing_data = './dce_mri_ktrans_testing_884_1.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)

        # Add patch augmentation
        patch_augmentation = ExtractPatches(
            patch_shape=(8, 8, 4),
            patch_region_conditions=[[brain_region, 1]],
            data_groups=['input_modalities', 'ground_truth'],
            patch_dimensions={
                'ground_truth': [0, 1, 2],
                'input_modalities': [1, 2, 3]
            })
        training_data_collection.append_augmentation(patch_augmentation,
                                                     multiplier=5000)

        # Add left-right flips
        flip_augmentation = Flip_Rotate_2D(
            flip=True,
            rotate=False,
            data_groups=['input_modalities', 'ground_truth'])
        training_data_collection.append_augmentation(flip_augmentation,
                                                     multiplier=2)

        # Write data to hdf5
        training_data_collection.write_data_to_file(training_data)

    # Or load pre-loaded data.
    training_data_collection = DataCollection(data_storage=training_data,
                                              verbose=True)
    training_data_collection.fill_data_groups()

    # Define model parameters
    model_parameters = {
        'input_shape': (65, 8, 8, 4, 1),
        'downsize_filters_factor': 4,
        'pool_size': (2, 2, 2),
        'filter_shape': (3, 3, 3),
        'dropout': .1,
        'batch_norm': True,
        'initial_learning_rate': 0.000001,
        'output_type': 'regression',
        'num_outputs': 1,
        'activation': 'relu',
        'padding': 'same',
        'implementation': 'keras',
        'depth': 1,
        'max_filter': 32
    }

    # Create U-Net
    if train_model:
        timenet_model = TimeNet(**model_parameters)
        plot_model(timenet_model.model,
                   to_file='timenet_model.png',
                   show_shapes=True)
        training_parameters = {
            'input_groups': ['input_modalities', 'ground_truth'],
            'output_model_filepath': model_file,
            'training_batch_size': 32,
            'num_epochs': 100,
            'training_steps_per_epoch': 200,
            'save_best_only': True
        }
        timenet_model.train(training_data_collection, **training_parameters)
    else:
        timenet_model = load_old_model(model_file)

    # Load testing data..
    if not os.path.exists(testing_data) or load_test_data:
        # Create a Data Collection
        testing_data_collection = DataCollection(
            val_data_directory,
            modality_dict=training_modality_dict,
            verbose=True)
        testing_data_collection.fill_data_groups()
        # Write data to hdf5
        testing_data_collection.write_data_to_file(testing_data)

    if predict:
        testing_data_collection = DataCollection(data_storage=testing_data,
                                                 verbose=True)
        testing_data_collection.fill_data_groups()

        testing_parameters = {
            'inputs': ['input_modalities'],
            'output_filename': 'deepneuro.nii.gz',
            'batch_size': 200,
            'patch_overlaps': 8,
            'output_patch_shape': (6, 6, 2, 1)
        }

        prediction = ModelPatchesInference(testing_data_collection,
                                           **testing_parameters)

        timenet_model.append_output([prediction])
        timenet_model.generate_outputs()