Scripts for training on the COMET CT Dataset

COMET/Build_test_set.py

@@ -0,0 +1,135 @@
+import os, sys
+
+import shutil
+
+import matplotlib.pyplot as plt
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import tensorflow as tf
+# tf.enable_eager_execution(config=config)
+
+import numpy as np
+import h5py
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.layers import AugmentationLayer
+from DeepDeformationMapRegistration.utils.visualization import save_disp_map_img, plot_predictions
+from DeepDeformationMapRegistration.utils.misc import get_segmentations_centroids
+from tqdm import tqdm
+
+from Brain_study.data_generator import BatchGenerator
+
+from skimage.measure import regionprops
+from scipy.interpolate import griddata
+
+import argparse
+
+
+DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/test'
+
+POINTS = None
+MISSING_CENTROID = np.asarray([[np.nan]*3])
+
+
+def get_mov_centroids(fix_seg, disp_map, nb_labels=28, brain_study=True):
+    fix_centroids, _ = get_segmentations_centroids(fix_seg[0, ...], ohe=True, expected_lbls=range(0, nb_labels), brain_study=brain_study)
+    disp = griddata(POINTS, disp_map.reshape([-1, 3]), fix_centroids, method='linear')
+    return fix_centroids, fix_centroids + disp, disp
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-d', '--dir', type=str, help='Directory where to store the files', default='')
+    parser.add_argument('--reldir', type=str, help='Relative path to dataset, in where to store the files', default='')
+    parser.add_argument('--gpu', type=int, help='GPU', default=0)
+    parser.add_argument('--dataset', type=str, help='Dataset to build the test set', default='')
+    parser.add_argument('--erase', type=bool, help='Erase the content of the output folder', default=False)
+    args = parser.parse_args()
+
+    assert args.dataset != '', "Missing original dataset dataset"
+    if args.dir == '' and args.reldir != '':
+        OUTPUT_FOLDER_DIR = os.path.join(args.dataset, 'test_dataset')
+    elif args.dir != '' and args.reldir == '':
+        OUTPUT_FOLDER_DIR = args.dir
+    else:
+        raise ValueError("Either provide 'dir' or 'reldir'")
+
+    if args.erase:
+        shutil.rmtree(OUTPUT_FOLDER_DIR, ignore_errors=True)
+    os.makedirs(OUTPUT_FOLDER_DIR, exist_ok=True)
+    print('DESTINATION FOLDER: ', OUTPUT_FOLDER_DIR)
+
+    DATASET = args.dataset
+
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)  # Check availability before running using 'nvidia-smi'
+
+    data_generator = BatchGenerator(DATASET, 1, False, 1.0, False, ['all'])
+
+    img_generator = data_generator.get_train_generator()
+    nb_labels = len(img_generator.get_segmentation_labels())
+    image_input_shape = img_generator.get_data_shape()[-1][:-1]
+    image_output_shape = [64] * 3
+    # Build model
+
+    xx = np.linspace(0, image_output_shape[0], image_output_shape[0], endpoint=False)
+    yy = np.linspace(0, image_output_shape[1], image_output_shape[2], endpoint=False)
+    zz = np.linspace(0, image_output_shape[2], image_output_shape[1], endpoint=False)
+
+    xx, yy, zz = np.meshgrid(xx, yy, zz)
+
+    POINTS = np.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=0).T
+
+    input_augm = tf.keras.Input(shape=img_generator.get_data_shape()[0], name='input_augm')
+    augm_layer = AugmentationLayer(max_displacement=C.MAX_AUG_DISP,  # Max 30 mm in isotropic space
+                                   max_deformation=C.MAX_AUG_DEF,  # Max 6 mm in isotropic space
+                                   max_rotation=C.MAX_AUG_ANGLE,  # Max 10 deg in isotropic space
+                                   num_control_points=C.NUM_CONTROL_PTS_AUG,
+                                   num_augmentations=C.NUM_AUGMENTATIONS,
+                                   gamma_augmentation=C.GAMMA_AUGMENTATION,
+                                   brightness_augmentation=C.BRIGHTNESS_AUGMENTATION,
+                                   in_img_shape=image_input_shape,
+                                   out_img_shape=image_output_shape,
+                                   only_image=False,
+                                   only_resize=False,
+                                   trainable=False,
+                                   return_displacement_map=True)
+    augm_model = tf.keras.Model(inputs=input_augm, outputs=augm_layer(input_augm))
+
+    config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
+    config.gpu_options.allow_growth = True
+    config.log_device_placement = False  ## to log device placement (on which device the operation ran)
+
+    sess = tf.Session(config=config)
+    tf.keras.backend.set_session(sess)
+    with sess.as_default():
+        sess.run(tf.global_variables_initializer())
+        progress_bar = tqdm(enumerate(img_generator, 1), desc='Generating samples', total=len(img_generator))
+        for step, (in_batch, _) in progress_bar:
+            fix_img, mov_img, fix_seg, mov_seg, disp_map = augm_model.predict(in_batch)
+
+            fix_centroids, mov_centroids, disp_centroids = get_mov_centroids(fix_seg, disp_map, nb_labels, False)
+
+            out_file = os.path.join(OUTPUT_FOLDER_DIR, 'test_sample_{:04d}.h5'.format(step))
+            out_file_dm = os.path.join(OUTPUT_FOLDER_DIR, 'test_sample_dm_{:04d}.h5'.format(step))
+            img_shape = fix_img.shape
+            segm_shape = fix_seg.shape
+            disp_shape = disp_map.shape
+            centroids_shape = fix_centroids.shape
+            with h5py.File(out_file, 'w') as f:
+                f.create_dataset('fix_image', shape=img_shape[1:], dtype=np.float32, data=fix_img[0, ...])
+                f.create_dataset('mov_image', shape=img_shape[1:], dtype=np.float32, data=mov_img[0, ...])
+                f.create_dataset('fix_segmentations', shape=segm_shape[1:], dtype=np.uint8, data=fix_seg[0, ...])
+                f.create_dataset('mov_segmentations', shape=segm_shape[1:], dtype=np.uint8, data=mov_seg[0, ...])
+                f.create_dataset('fix_centroids', shape=centroids_shape, dtype=np.float32, data=fix_centroids)
+                f.create_dataset('mov_centroids', shape=centroids_shape, dtype=np.float32, data=mov_centroids)
+
+            with h5py.File(out_file_dm, 'w') as f:
+                f.create_dataset('disp_map', shape=disp_shape[1:], dtype=np.float32, data=disp_map)
+                f.create_dataset('disp_centroids', shape=centroids_shape, dtype=np.float32, data=disp_centroids)
+
+    print('Done')

COMET/COMET_train.py

@@ -0,0 +1,426 @@
+import os, sys
+
+import keras
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+from datetime import datetime
+
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.python.keras.utils import Progbar
+from tensorflow.keras import Input
+from tensorflow.keras.models import Model
+from tensorflow.python.framework.errors import InvalidArgumentError
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC, GeneralizedDICEScore, HausdorffDistanceErosion
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
+from DeepDeformationMapRegistration.ms_ssim_tf import _MSSSIM_WEIGHTS
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+from DeepDeformationMapRegistration.utils.misc import function_decorator
+from DeepDeformationMapRegistration.layers import AugmentationLayer
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+
+from Brain_study.data_generator import BatchGenerator
+from Brain_study.utils import SummaryDictionary, named_logs
+
+import COMET.augmentation_constants as COMET_C
+
+import numpy as np
+import tensorflow as tf
+import voxelmorph as vxm
+import h5py
+import re
+import itertools
+
+
+def launch_train(dataset_folder, validation_folder, output_folder, model_file, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim',
+                 segm='dice', max_epochs=C.EPOCHS, freeze_layers=None,
+                 acc_gradients=1, batch_size=16):
+    # 0. Input checks
+    assert dataset_folder is not None and output_folder is not None and model_file is not None
+    assert '.h5' in model_file, 'The model must be an H5 file'
+
+    USE_SEGMENTATIONS = bool(re.search('SEGGUIDED', model_file))
+
+    # 1. Load variables
+    os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_num) # Check availability before running using 'nvidia-smi'
+    C.GPU_NUM = str(gpu_num)
+
+    if freeze_layers is not None:
+        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in freeze_layers), \
+            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
+        freeze_layers = [list(COMET_C.LAYER_RANGES[l]) for l in list(set(freeze_layers))]
+        if len(freeze_layers) > 1:
+            freeze_layers = list(itertools.chain.from_iterable(freeze_layers))
+
+    os.makedirs(output_folder, exist_ok=True)
+    # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
+    log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
+    C.TRAINING_DATASET = dataset_folder #dataset_copy.copy_dataset()
+    C.VALIDATION_DATASET = validation_folder
+    C.ACCUM_GRADIENT_STEP = acc_gradients
+    C.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
+    C.EARLY_STOP_PATIENCE = 5 * (C.ACCUM_GRADIENT_STEP / 2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.LEARNING_RATE = lr
+    C.LIMIT_NUM_SAMPLES = None
+    C.EPOCHS = max_epochs
+
+    aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
+          "GPU: {}\n" \
+          "BATCH SIZE: {}\n" \
+          "LR: {}\n" \
+          "SIMILARITY: {}\n" \
+          "SEGMENTATION: {}\n"\
+          "REG. WEIGHT: {}\n" \
+          "EPOCHS: {:d}\n" \
+          "ACCUM. GRAD: {}\n" \
+          "EARLY STOP PATIENCE: {}\n" \
+          "FROZEN LAYERS: {}".format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'),
+                                           C.TRAINING_DATASET,
+                                           C.VALIDATION_DATASET,
+                                           C.GPU_NUM,
+                                           C.BATCH_SIZE,
+                                           C.LEARNING_RATE,
+                                           simil,
+                                           segm,
+                                           rw,
+                                           C.EPOCHS,
+                                           C.ACCUM_GRADIENT_STEP,
+                                           C.EARLY_STOP_PATIENCE,
+                                     freeze_layers)
+
+    log_file.write(aux)
+    print(aux)
+
+    # 2. Data generator
+    used_labels = 'all' if USE_SEGMENTATIONS else 'none'
+    data_generator = BatchGenerator(C.TRAINING_DATASET, C.BATCH_SIZE if C.ACCUM_GRADIENT_STEP == 1 else 1, True,
+                                    C.TRAINING_PERC, labels=[used_labels], combine_segmentations=not USE_SEGMENTATIONS,
+                                    directory_val=C.VALIDATION_DATASET)
+
+    train_generator = data_generator.get_train_generator()
+    validation_generator = data_generator.get_validation_generator()
+
+    image_input_shape = train_generator.get_data_shape()[-1][:-1]
+    image_output_shape = [64] * 3
+    nb_labels = len(train_generator.get_segmentation_labels())
+
+    # 3. Load model
+    # IMPORTANT: the mode MUST be loaded AFTER setting up the session configuration
+    config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
+    config.gpu_options.allow_growth = True
+    config.log_device_placement = False  ## to log device placement (on which device the operation ran)
+    sess = tf.Session(config=config)
+    tf.keras.backend.set_session(sess)
+
+    loss_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).loss,
+                 NCC(image_input_shape).loss,
+                 vxm.losses.MSE().loss,
+                 MultiScaleStructuralSimilarity(max_val=1., filter_size=3).loss,
+                 HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).loss,
+                 GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss,
+                 GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss_macro
+                 ]
+
+    metric_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric,
+                   NCC(image_input_shape).metric,
+                   vxm.losses.MSE().loss,
+                   MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric,
+                   GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric,
+                   HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).metric,
+                   GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric_macro,]
+    print('MODEL LOCATION: ', model_file)
+
+    try:
+        network = tf.keras.models.load_model(model_file, {#'VxmDenseSemiSupervisedSeg': vxm.networks.VxmDenseSemiSupervisedSeg,
+                                                          'VxmDense': vxm.networks.VxmDense,
+                                                          'AdamAccumulated': AdamAccumulated,
+                                                          'loss': loss_fncs,
+                                                          'metric': metric_fncs},
+                                             compile=False)
+    except ValueError as e:
+        enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+        dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+        nb_features = [enc_features, dec_features]
+        if USE_SEGMENTATIONS:
+            network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
+                                                             nb_labels=nb_labels,
+                                                             nb_unet_features=nb_features,
+                                                             int_steps=0,
+                                                             int_downsize=1,
+                                                             seg_downsize=1)
+        else:
+            network = vxm.networks.VxmDense(inshape=image_output_shape,
+                                            nb_unet_features=nb_features,
+                                            int_steps=0)
+        network.load_weights(model_file, by_name=True)
+    # 4. Freeze/unfreeze model layers
+    # freeze_layers = range(0, len(network.layers) - 8)  # Do not freeze the last layers after the UNet (8 last layers)
+    # for l in freeze_layers:
+    #     network.layers[l].trainable = False
+    # msg = "[INF]: Frozen layers {} to {}".format(0, len(network.layers) - 8)
+    # print(msg)
+    # log_file.write("INF: Frozen layers {} to {}".format(0, len(network.layers) - 8))
+    if freeze_layers is not None:
+        aux = list()
+        for r in freeze_layers:
+            for l in range(*r):
+                network.layers[l].trainable = False
+                aux.append(l)
+        aux.sort()
+        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in aux]))
+    else:
+        msg = "[INF] None frozen layers"
+    print(msg)
+    log_file.write(msg)
+    # network.trainable = False  # Freeze the base model
+    # # Create a new model on top
+    # input_new_model = keras.Input(network.input_shape)
+    # x = base_model(input_new_model, training=False)
+    # x =
+    # network = keras.Model(input_new_model, x)
+
+    network.summary()
+    network.summary(print_fn=log_file.writelines)
+    #   Complete the model with the augmentation layer
+    augm_train_input_shape = train_generator.get_data_shape()[0] if USE_SEGMENTATIONS else train_generator.get_data_shape()[-1]
+    input_layer_train = Input(shape=augm_train_input_shape, name='input_train')
+    augm_layer_train = AugmentationLayer(max_displacement=COMET_C.MAX_AUG_DISP,   # Max 30 mm in isotropic space
+                                         max_deformation=COMET_C.MAX_AUG_DEF,  # Max 6 mm in isotropic space
+                                         max_rotation=COMET_C.MAX_AUG_ANGLE,   # Max 10 deg in isotropic space
+                                         num_control_points=COMET_C.NUM_CONTROL_PTS_AUG,
+                                         num_augmentations=COMET_C.NUM_AUGMENTATIONS,
+                                         gamma_augmentation=COMET_C.GAMMA_AUGMENTATION,
+                                         brightness_augmentation=COMET_C.BRIGHTNESS_AUGMENTATION,
+                                         in_img_shape=image_input_shape,
+                                         out_img_shape=image_output_shape,
+                                         only_image=not USE_SEGMENTATIONS,  # If baseline then True
+                                         only_resize=False,
+                                         trainable=False)
+    augm_model_train = Model(inputs=input_layer_train, outputs=augm_layer_train(input_layer_train))
+
+    input_layer_valid = Input(shape=validation_generator.get_data_shape()[0], name='input_valid')
+    augm_layer_valid = AugmentationLayer(max_displacement=COMET_C.MAX_AUG_DISP,   # Max 30 mm in isotropic space
+                                         max_deformation=COMET_C.MAX_AUG_DEF,  # Max 6 mm in isotropic space
+                                         max_rotation=COMET_C.MAX_AUG_ANGLE,   # Max 10 deg in isotropic space
+                                         num_control_points=COMET_C.NUM_CONTROL_PTS_AUG,
+                                         num_augmentations=COMET_C.NUM_AUGMENTATIONS,
+                                         gamma_augmentation=COMET_C.GAMMA_AUGMENTATION,
+                                         brightness_augmentation=COMET_C.BRIGHTNESS_AUGMENTATION,
+                                         in_img_shape=image_input_shape,
+                                         out_img_shape=image_output_shape,
+                                         only_image=False,
+                                         only_resize=False,
+                                         trainable=False)
+    augm_model_valid = Model(inputs=input_layer_valid, outputs=augm_layer_valid(input_layer_valid))
+
+    # 5. Setup training environment: loss, optimizer, callbacks, evaluation
+
+    # Losses and loss weights
+    SSIM_KER_SIZE = 5
+    MS_SSIM_WEIGHTS = _MSSSIM_WEIGHTS[:3]
+    MS_SSIM_WEIGHTS /= np.sum(MS_SSIM_WEIGHTS)
+    if simil.lower() == 'mse':
+        loss_fnc = vxm.losses.MSE().loss
+    elif simil.lower() == 'ncc':
+        loss_fnc = NCC(image_input_shape).loss
+    elif simil.lower() == 'ssim':
+        loss_fnc = StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss
+    elif simil.lower() == 'ms_ssim':
+        loss_fnc = MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss
+    elif simil.lower() == 'mse__ms_ssim' or simil.lower() == 'ms_ssim__mse':
+        @function_decorator('MSSSIM_MSE__loss')
+        def loss_fnc(y_true, y_pred):
+            return vxm.losses.MSE().loss(y_true, y_pred) + \
+                   MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss(y_true, y_pred)
+    elif simil.lower() == 'ncc__ms_ssim' or simil.lower() == 'ms_ssim__ncc':
+        @function_decorator('MSSSIM_NCC__loss')
+        def loss_fnc(y_true, y_pred):
+            return NCC(image_input_shape).loss(y_true, y_pred) + \
+                   MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss(y_true, y_pred)
+    elif simil.lower() == 'mse__ssim' or simil.lower() == 'ssim__mse':
+        @function_decorator('SSIM_MSE__loss')
+        def loss_fnc(y_true, y_pred):
+            return vxm.losses.MSE().loss(y_true, y_pred) + \
+                   StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss(y_true, y_pred)
+    elif simil.lower() == 'ncc__ssim' or simil.lower() == 'ssim__ncc':
+        @function_decorator('SSIM_NCC__loss')
+        def loss_fnc(y_true, y_pred):
+            return NCC(image_input_shape).loss(y_true, y_pred) + \
+                   StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss(y_true, y_pred)
+    else:
+        raise ValueError('Unknown similarity metric: ' + simil)
+
+    if USE_SEGMENTATIONS:
+        if segm == 'hd':
+            loss_segm = HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).loss
+        elif segm == 'dice':
+            loss_segm = GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss
+        elif segm == 'dice_macro':
+            loss_segm = GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss_macro
+        else:
+            raise ValueError('No valid value for segm')
+
+    os.makedirs(os.path.join(output_folder, 'checkpoints'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'tensorboard'), exist_ok=True)
+    callback_tensorboard = TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                                       batch_size=C.BATCH_SIZE, write_images=False, histogram_freq=0,
+                                       update_freq='epoch',     # or 'batch' or integer
+                                       write_graph=True, write_grads=True
+                                       )
+    callback_early_stop = EarlyStopping(monitor='val_loss', verbose=1, patience=C.EARLY_STOP_PATIENCE, min_delta=0.00001)
+
+    callback_best_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                                          save_best_only=True, monitor='val_loss', verbose=1, mode='min')
+    callback_save_checkpoint = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.h5'),
+                                               save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    if USE_SEGMENTATIONS:
+        losses = {'transformer': loss_fnc,
+                  'seg_transformer': loss_segm,
+                  'flow': vxm.losses.Grad('l2').loss}
+        metrics = {'transformer': [StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).metric,
+                                   MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).metric,
+                                   tf.keras.losses.MSE,
+                                   NCC(image_input_shape).metric],
+                   'seg_transformer': [GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric,
+                                       HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [train_generator.get_data_shape()[2][-1]]).metric,
+                                       GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric_macro,
+                                       ],
+                   #'flow': vxm.losses.Grad('l2').loss
+                   }
+        loss_weights = {'transformer': 1.,
+                        'seg_transformer': 1.,
+                        'flow': rw}
+    else:
+        losses = {'transformer': loss_fnc,
+                  'flow': vxm.losses.Grad('l2').loss}
+        metrics = {'transformer': [StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).metric,
+                                   MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).metric,
+                                   tf.keras.losses.MSE,
+                                   NCC(image_input_shape).metric],
+                   #'flow': vxm.losses.Grad('l2').loss
+                   }
+        loss_weights = {'transformer': 1.,
+                        'flow': rw}
+
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    network.compile(optimizer=optimizer,
+                    loss=losses,
+                    loss_weights=loss_weights,
+                    metrics=metrics)
+
+    # 6. Training loop
+    callback_tensorboard.set_model(network)
+    callback_early_stop.set_model(network)
+    callback_best_model.set_model(network)
+    callback_save_checkpoint.set_model(network)
+
+    summary = SummaryDictionary(network, C.BATCH_SIZE)
+    names = network.metrics_names
+    log_file.write('\n\n[{}]\tINFO:\tStart training\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y')))
+
+    with sess.as_default():
+        # tf.global_variables_initializer()
+        callback_tensorboard.on_train_begin()
+        callback_early_stop.on_train_begin()
+        callback_best_model.on_train_begin()
+        callback_save_checkpoint.on_train_begin()
+
+        for epoch in range(C.EPOCHS):
+            callback_tensorboard.on_epoch_begin(epoch)
+            callback_early_stop.on_epoch_begin(epoch)
+            callback_best_model.on_epoch_begin(epoch)
+            callback_save_checkpoint.on_epoch_begin(epoch)
+            print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
+            print("TRAIN")
+
+            log_file.write('\n\n[{}]\tINFO:\tTraining epoch {}\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), epoch))
+            progress_bar = Progbar(len(train_generator), width=30, verbose=1)
+            for step, (in_batch, _) in enumerate(train_generator, 1):
+                callback_best_model.on_train_batch_begin(step)
+                callback_save_checkpoint.on_train_batch_begin(step)
+                callback_early_stop.on_train_batch_begin(step)
+
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augm_model_train.predict(in_batch)
+                    np.nan_to_num(fix_img, copy=False)
+                    np.nan_to_num(mov_img, copy=False)
+                    if np.isnan(np.sum(mov_img)) or np.isnan(np.sum(fix_img)) or np.isinf(np.sum(mov_img)) or np.isinf(np.sum(fix_img)):
+                        msg = 'CORRUPTED DATA!! Unique: Fix: {}\tMoving: {}'.format(np.unique(fix_img),
+                                                                                    np.unique(mov_img))
+                        print(msg)
+                        log_file.write('\n\n[{}]\tWAR: {}'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), msg))
+
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+
+                if USE_SEGMENTATIONS:
+                    in_data = (mov_img, fix_img, mov_seg)
+                    out_data = (fix_img, fix_img, fix_seg)
+                else:
+                    in_data = (mov_img, fix_img)
+                    out_data = (fix_img, fix_img)
+                ret = network.train_on_batch(x=in_data, y=out_data)  # The second element doesn't matter
+                if np.isnan(ret).any():
+                    os.makedirs(os.path.join(output_folder, 'corrupted'), exist_ok=True)
+                    save_nifti(mov_img, os.path.join(output_folder, 'corrupted', 'mov_img_nan.nii.gz'))
+                    save_nifti(fix_img, os.path.join(output_folder, 'corrupted', 'fix_img_nan.nii.gz'))
+                    pred_img, dm = network((mov_img, fix_img))
+                    save_nifti(pred_img, os.path.join(output_folder, 'corrupted', 'pred_img_nan.nii.gz'))
+                    save_nifti(dm, os.path.join(output_folder, 'corrupted', 'dm_nan.nii.gz'))
+                    log_file.write('\n\n[{}]\tERR: Corruption error'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y')))
+                    raise ValueError('CORRUPTION ERROR: Halting training')
+
+                summary.on_train_batch_end(ret)
+                callback_best_model.on_train_batch_end(step, named_logs(network, ret))
+                callback_save_checkpoint.on_train_batch_end(step, named_logs(network, ret))
+                callback_early_stop.on_train_batch_end(step, named_logs(network, ret))
+                progress_bar.update(step, zip(names, ret))
+                log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
+            val_values = progress_bar._values.copy()
+            ret = [val_values[x][0]/val_values[x][1] for x in names]
+
+            print('\nVALIDATION')
+            log_file.write('\n\n[{}]\tINFO:\tValidation epoch {}\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), epoch))
+            progress_bar = Progbar(len(validation_generator), width=30, verbose=1)
+            for step, (in_batch, _) in enumerate(validation_generator, 1):
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augm_model_valid.predict(in_batch)
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+
+                if USE_SEGMENTATIONS:
+                    in_data = (mov_img, fix_img, mov_seg)
+                    out_data = (fix_img, fix_img, fix_seg)
+                else:
+                    in_data = (mov_img, fix_img)
+                    out_data = (fix_img, fix_img)
+                ret = network.test_on_batch(x=in_data,
+                                            y=out_data)
+
+                summary.on_validation_batch_end(ret)
+                progress_bar.update(step, zip(names, ret))
+                log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
+            val_values = progress_bar._values.copy()
+            ret = [val_values[x][0]/val_values[x][1] for x in names]
+
+            train_generator.on_epoch_end()
+            validation_generator.on_epoch_end()
+            epoch_summary = summary.on_epoch_end()  # summary resets after on_epoch_end() call
+            callback_tensorboard.on_epoch_end(epoch, epoch_summary)
+            callback_best_model.on_epoch_end(epoch, epoch_summary)
+            callback_save_checkpoint.on_epoch_end(epoch, epoch_summary)
+            callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            print('End of epoch {}: '.format(epoch), ret, '\n')
+
+        callback_tensorboard.on_train_end()
+        callback_best_model.on_train_end()
+        callback_save_checkpoint.on_train_end()
+        callback_early_stop.on_train_end()
+# 7. Wrap up

COMET/COMET_train_UW.py

@@ -0,0 +1,342 @@
+import os, sys
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+from datetime import datetime
+
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.python.keras.utils import Progbar
+from tensorflow.keras import Input
+from tensorflow.keras.models import Model
+from tensorflow.python.framework.errors import InvalidArgumentError
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC, GeneralizedDICEScore, \
+    HausdorffDistanceErosion
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
+from DeepDeformationMapRegistration.ms_ssim_tf import _MSSSIM_WEIGHTS
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+from DeepDeformationMapRegistration.utils.misc import function_decorator
+from DeepDeformationMapRegistration.layers import AugmentationLayer, UncertaintyWeighting
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+
+from Brain_study.data_generator import BatchGenerator
+from Brain_study.utils import SummaryDictionary, named_logs
+
+import COMET.augmentation_constants as COMET_C
+
+import numpy as np
+import tensorflow as tf
+import voxelmorph as vxm
+import h5py
+import re
+import itertools
+
+
+def launch_train(dataset_folder, validation_folder, output_folder, model_file, gpu_num=0, lr=1e-4, rw=5e-3,
+                 simil=['ssim'], segm=['dice'], max_epochs=C.EPOCHS, prior_reg_w=5e-3, freeze_layers=None,
+                 acc_gradients=1, batch_size=16):
+    # 0. Input checks
+    assert dataset_folder is not None and output_folder is not None and model_file is not None
+    assert '.h5' in model_file, 'The model must be an H5 file'
+
+    # 1. Load variables
+    os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_num)  # Check availability before running using 'nvidia-smi'
+    C.GPU_NUM = str(gpu_num)
+
+    if freeze_layers is not None:
+        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in freeze_layers), \
+            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
+        multiple_ranges = 'TOP' in freeze_layers
+        freeze_layers = [list(COMET_C.LAYER_RANGES[l]) for l in list(set(freeze_layers))]
+        freeze_layers = freeze_layers[0] if multiple_ranges else freeze_layers
+
+        # if len(freeze_layers) > 1:
+        #     freeze_layers = list(itertools.chain.from_iterable(freeze_layers))
+
+    os.makedirs(output_folder, exist_ok=True)
+    # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
+    log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
+    C.TRAINING_DATASET = dataset_folder  # dataset_copy.copy_dataset()
+    C.VALIDATION_DATASET = validation_folder
+    C.ACCUM_GRADIENT_STEP = acc_gradients
+    C.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
+    C.EARLY_STOP_PATIENCE = 5 * (C.ACCUM_GRADIENT_STEP / 2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.LEARNING_RATE = lr
+    C.LIMIT_NUM_SAMPLES = None
+    C.EPOCHS = max_epochs
+
+    aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
+          "GPU: {}\n" \
+          "BATCH SIZE: {}\n" \
+          "LR: {}\n" \
+          "SIMILARITY: {}\n" \
+          "SEGMENTATION: {}\n" \
+          "REG. WEIGHT: {}\n" \
+          "EPOCHS: {:d}\n" \
+          "ACCUM. GRAD: {}\n" \
+          "EARLY STOP PATIENCE: {}\n" \
+          "FROZEN LAYERS: {}".format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'),
+                                     C.TRAINING_DATASET,
+                                     C.VALIDATION_DATASET,
+                                     C.GPU_NUM,
+                                     C.BATCH_SIZE,
+                                     C.LEARNING_RATE,
+                                     simil,
+                                     segm,
+                                     rw,
+                                     C.EPOCHS,
+                                     C.ACCUM_GRADIENT_STEP,
+                                     C.EARLY_STOP_PATIENCE,
+                                     freeze_layers)
+
+    log_file.write(aux)
+    print(aux)
+
+    # 2. Data generator
+    data_generator = BatchGenerator(C.TRAINING_DATASET, C.BATCH_SIZE if C.ACCUM_GRADIENT_STEP == 1 else 1, True,
+                                    C.TRAINING_PERC, labels=['all'], combine_segmentations=False,
+                                    directory_val=C.VALIDATION_DATASET)
+
+    train_generator = data_generator.get_train_generator()
+    validation_generator = data_generator.get_validation_generator()
+
+    image_input_shape = train_generator.get_data_shape()[-1][:-1]
+    image_output_shape = [64] * 3
+    nb_labels = len(train_generator.get_segmentation_labels())
+
+    # 3. Load model
+    # IMPORTANT: the mode MUST be loaded AFTER setting up the session configuration
+    config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
+    config.gpu_options.allow_growth = True
+    config.log_device_placement = False  ## to log device placement (on which device the operation ran)
+    sess = tf.Session(config=config)
+    tf.keras.backend.set_session(sess)
+
+    print('MODEL LOCATION: ', model_file)
+
+    enc_features = [16, 32, 32, 32]  # const.ENCODER_FILTERS
+    dec_features = [32, 32, 32, 32, 32, 16, 16]  # const.ENCODER_FILTERS[::-1]
+    nb_features = [enc_features, dec_features]
+    network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
+                                                     nb_labels=nb_labels,
+                                                     nb_unet_features=nb_features,
+                                                     int_steps=0,
+                                                     int_downsize=1,
+                                                     seg_downsize=1)
+    network.load_weights(model_file, by_name=True)
+
+    # 4. Freeze/unfreeze model layers
+    if freeze_layers is not None:
+        aux = list()
+        for r in freeze_layers:
+            for l in range(*r):
+                network.layers[l].trainable = False
+                aux.append(l)
+        aux.sort()
+        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in aux]))
+    else:
+        msg = "[INF] None frozen layers"
+    print(msg)
+    log_file.write(msg)
+
+    network.summary()
+    network.summary(print_fn=log_file.write)
+    #   Complete the model with the augmentation layer
+    input_layer_train = Input(shape=train_generator.get_data_shape()[0], name='input_train')
+    augm_layer = AugmentationLayer(max_displacement=COMET_C.MAX_AUG_DISP,  # Max 30 mm in isotropic space
+                                   max_deformation=COMET_C.MAX_AUG_DEF,  # Max 6 mm in isotropic space
+                                   max_rotation=COMET_C.MAX_AUG_ANGLE,  # Max 10 deg in isotropic space
+                                   num_control_points=COMET_C.NUM_CONTROL_PTS_AUG,
+                                   num_augmentations=COMET_C.NUM_AUGMENTATIONS,
+                                   gamma_augmentation=COMET_C.GAMMA_AUGMENTATION,
+                                   brightness_augmentation=COMET_C.BRIGHTNESS_AUGMENTATION,
+                                   in_img_shape=image_input_shape,
+                                   out_img_shape=image_output_shape,
+                                   only_image=False,  # If baseline then True
+                                   only_resize=False,
+                                   trainable=False)
+    augm_model = Model(inputs=input_layer_train, outputs=augm_layer(input_layer_train))
+
+    # 5. Setup training environment: loss, optimizer, callbacks, evaluation
+
+    # Losses and loss weights
+    SSIM_KER_SIZE = 5
+    MS_SSIM_WEIGHTS = _MSSSIM_WEIGHTS[:3]
+    MS_SSIM_WEIGHTS /= np.sum(MS_SSIM_WEIGHTS)
+    loss_simil = []
+    prior_loss_w = []
+    for s in simil:
+        if s == 'ssim':
+            loss_simil.append(StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss)
+            prior_loss_w.append(1.)
+        elif s == 'ms_ssim':
+            loss_simil.append(MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE,
+                                                             power_factors=MS_SSIM_WEIGHTS).loss)
+            prior_loss_w.append(1.)
+        elif s == 'ncc':
+            loss_simil.append(NCC(image_input_shape).loss)
+            prior_loss_w.append(1.)
+        elif s == 'mse':
+            loss_simil.append(vxm.losses.MSE().loss)
+            prior_loss_w.append(1.)
+        else:
+            raise ValueError('Unknown similarity function: ', s)
+
+    loss_segm = []
+    for s in segm:
+        if s == 'dice':
+            loss_segm.append(GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]],
+                                                  num_labels=nb_labels).loss)
+            prior_loss_w.append(1.)
+        elif s == 'dice_macro':
+            loss_segm.append(GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]],
+                                                  num_labels=nb_labels).loss_macro)
+            prior_loss_w.append(1.)
+        elif s == 'hd':
+            loss_segm.append(HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [
+                train_generator.get_data_shape()[2][-1]]).loss)
+            prior_loss_w.append(1.)
+        else:
+            raise ValueError('Unknown similarity function: ', s)
+
+    # Uncertainty weigthing module
+    grad = tf.keras.Input(shape=(*image_output_shape, 3), name='multiLoss_grad_input', dtype=tf.float32)
+    fix_seg = tf.keras.Input(shape=(*image_output_shape, len(train_generator.get_segmentation_labels())),
+                             name='multiLoss_fix_seg_input', dtype=tf.float32)
+
+    multiLoss = UncertaintyWeighting(num_loss_fns=len(loss_simil) + len(loss_segm),
+                                     num_reg_fns=1,
+                                     loss_fns=[*loss_simil,
+                                               *loss_segm],
+                                     reg_fns=[vxm.losses.Grad('l2').loss],
+                                     prior_loss_w=prior_loss_w,
+                                     # prior_loss_w=[1., 0.1, 1., 1.],
+                                     prior_reg_w=[prior_reg_w],
+                                     name='MultiLossLayer')
+    loss = multiLoss([*[network.inputs[1]] * len(loss_simil), *[fix_seg] * len(loss_segm),
+                      *[network.outputs[0]] * len(loss_simil), *[network.outputs[2]] * len(loss_simil),
+                      grad,
+                      network.outputs[1]])
+
+    # inputs = [mov_img, fix_img, mov_segm, fix_segm, zero_grads]
+    # outputs = [pred_img, flow, pred_segm, loss]
+    full_model = tf.keras.Model(inputs=network.inputs + [fix_seg, grad],
+                                outputs=network.outputs + [loss])
+
+    os.makedirs(os.path.join(output_folder, 'checkpoints'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'tensorboard'), exist_ok=True)
+    callback_tensorboard = TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                                       batch_size=C.BATCH_SIZE, write_images=False, histogram_freq=0,
+                                       update_freq='epoch',  # or 'batch' or integer
+                                       write_graph=True, write_grads=True
+                                       )
+    callback_early_stop = EarlyStopping(monitor='val_loss', verbose=1, patience=C.EARLY_STOP_PATIENCE,
+                                        min_delta=0.00001)
+
+    callback_best_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                                          save_best_only=True, monitor='val_loss', verbose=1, mode='min')
+    callback_save_checkpoint = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.h5'),
+                                               save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    full_model.compile(optimizer=optimizer,
+                       loss=None, )
+
+    # 6. Training loop
+    callback_tensorboard.set_model(full_model)
+    callback_early_stop.set_model(full_model)
+    callback_best_model.set_model(network)  # ONLY SAVE THE NETWORK!!!
+    callback_save_checkpoint.set_model(network)  # ONLY SAVE THE NETWORK!!!
+
+    summary = SummaryDictionary(full_model, C.BATCH_SIZE)
+    names = full_model.metrics_names
+    zero_grads = tf.zeros_like(network.references.pos_flow, name='dummy_zero_grads')  # Dummy zeros-tensor
+    log_file.write('\n\n[{}]\tINFO:\tStart training\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y')))
+
+    with sess.as_default():
+        # tf.global_variables_initializer()
+        callback_tensorboard.on_train_begin()
+        callback_early_stop.on_train_begin()
+        callback_best_model.on_train_begin()
+        callback_save_checkpoint.on_train_begin()
+
+        for epoch in range(C.EPOCHS):
+            callback_tensorboard.on_epoch_begin(epoch)
+            callback_early_stop.on_epoch_begin(epoch)
+            callback_best_model.on_epoch_begin(epoch)
+            callback_save_checkpoint.on_epoch_begin(epoch)
+            print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
+            print("TRAIN")
+
+            log_file.write(
+                '\n\n[{}]\tINFO:\tTraining epoch {}\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), epoch))
+            progress_bar = Progbar(len(train_generator), width=30, verbose=1)
+            for step, (in_batch, _) in enumerate(train_generator, 1):
+                callback_best_model.on_train_batch_begin(step)
+                callback_save_checkpoint.on_train_batch_begin(step)
+                callback_early_stop.on_train_batch_begin(step)
+
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augm_model.predict(in_batch)
+                    np.nan_to_num(fix_img, copy=False)
+                    np.nan_to_num(mov_img, copy=False)
+                    if np.isnan(np.sum(mov_img)) or np.isnan(np.sum(fix_img)) or np.isinf(np.sum(mov_img)) or np.isinf(
+                            np.sum(fix_img)):
+                        msg = 'CORRUPTED DATA!! Unique: Fix: {}\tMoving: {}'.format(np.unique(fix_img),
+                                                                                    np.unique(mov_img))
+                        print(msg)
+                        log_file.write('\n\n[{}]\tWAR: {}'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), msg))
+
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+
+                ret = full_model.train_on_batch(
+                    x=(mov_img, fix_img, mov_seg, fix_seg, zero_grads))  # The second element doesn't matter
+
+                summary.on_train_batch_end(ret)
+                callback_best_model.on_train_batch_end(step, named_logs(full_model, ret))
+                callback_save_checkpoint.on_train_batch_end(step, named_logs(full_model, ret))
+                callback_early_stop.on_train_batch_end(step, named_logs(full_model, ret))
+                progress_bar.update(step, zip(names, ret))
+                log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
+            val_values = progress_bar._values.copy()
+            ret = [val_values[x][0] / val_values[x][1] for x in names]
+
+            print('\nVALIDATION')
+            log_file.write(
+                '\n\n[{}]\tINFO:\tValidation epoch {}\n\n'.format(datetime.now().strftime('%H:%M:%S\t%d/%m/%Y'), epoch))
+            progress_bar = Progbar(len(validation_generator), width=30, verbose=1)
+            for step, (in_batch, _) in enumerate(validation_generator, 1):
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augm_model.predict(in_batch)
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+
+                ret = full_model.test_on_batch(x=(mov_img, fix_img, mov_seg, fix_seg, zero_grads))
+
+                summary.on_validation_batch_end(ret)
+                progress_bar.update(step, zip(names, ret))
+                log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
+            val_values = progress_bar._values.copy()
+            ret = [val_values[x][0] / val_values[x][1] for x in names]
+
+            train_generator.on_epoch_end()
+            validation_generator.on_epoch_end()
+            epoch_summary = summary.on_epoch_end()  # summary resets after on_epoch_end() call
+            callback_tensorboard.on_epoch_end(epoch, epoch_summary)
+            callback_best_model.on_epoch_end(epoch, epoch_summary)
+            callback_save_checkpoint.on_epoch_end(epoch, epoch_summary)
+            callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            print('End of epoch {}: '.format(epoch), ret, '\n')
+
+        callback_tensorboard.on_train_end()
+        callback_best_model.on_train_end()
+        callback_save_checkpoint.on_train_end()
+        callback_early_stop.on_train_end()
+# 7. Wrap up

COMET/Evaluate_network.py

@@ -0,0 +1,264 @@
+import os, sys
+
+import shutil
+import time
+import tkinter
+
+import h5py
+import matplotlib.pyplot as plt
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import tensorflow as tf
+# tf.enable_eager_execution(config=config)
+
+import numpy as np
+import pandas as pd
+import voxelmorph as vxm
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.layers import AugmentationLayer, UncertaintyWeighting
+from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC, GeneralizedDICEScore, HausdorffDistanceErosion, target_registration_error
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+from DeepDeformationMapRegistration.utils.visualization import save_disp_map_img, plot_predictions
+from DeepDeformationMapRegistration.utils.misc import DisplacementMapInterpolator, get_segmentations_centroids, segmentation_ohe_to_cardinal, segmentation_cardinal_to_ohe
+from EvaluationScripts.Evaluate_class import EvaluationFigures, resize_pts_to_original_space, resize_img_to_original_space, resize_transformation
+from scipy.interpolate import RegularGridInterpolator
+from tqdm import tqdm
+
+import h5py
+import re
+from Brain_study.data_generator import BatchGenerator
+
+import argparse
+
+from skimage.transform import warp
+import neurite as ne
+
+# from tkinter import filedialog as fd
+
+
+DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/test_fixed'
+MODEL_FILE = '/mnt/EncryptedData1/Users/javier/train_output/COMET/ERASE/COMET_L_ssim__MET_mse_ncc_ssim_141343-01122021/checkpoints/best_model.h5'
+DATA_ROOT_DIR = '/mnt/EncryptedData1/Users/javier/train_output/COMET/ERASE/COMET_L_ssim__MET_mse_ncc_ssim_141343-01122021/'
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-m', '--model', nargs='+', type=str, help='.h5 of the model', default=None)
+    parser.add_argument('-d', '--dir', nargs='+', type=str, help='Directory where ./checkpoints/best_model.h5 is located', default=None)
+    parser.add_argument('--gpu', type=int, help='GPU', default=0)
+    parser.add_argument('--dataset', type=str, help='Dataset to run predictions on', default=DATASET)
+    parser.add_argument('--erase', type=bool, help='Erase the content of the output folder', default=False)
+    parser.add_argument('--outdirname', type=str, default='Evaluate')
+    args = parser.parse_args()
+    if args.model is not None:
+        assert '.h5' in args.model[0], 'No checkpoint file provided, use -d/--dir instead'
+        MODEL_FILE_LIST = args.model
+        DATA_ROOT_DIR_LIST = [os.path.split(model_path)[0] for model_path in args.model]
+    elif args.dir is not None:
+        assert '.h5' not in args.dir[0], 'Provided checkpoint file, user -m/--model instead'
+        MODEL_FILE_LIST = [os.path.join(dir_path, 'checkpoints', 'best_model.h5') for dir_path in args.dir]
+        DATA_ROOT_DIR_LIST = args.dir
+    else:
+        # try:
+        #     MODEL_FILE_LIST = fd.askopenfilenames(title='Select .h model file',
+        #                                           initialdir='/mnt/EncryptedData1/Users/javier/train_output/COMET',
+        #                                           filetypes=(('Model', '*.h5'),))
+        #     if len(MODEL_FILE_LIST):
+        #         DATA_ROOT_DIR_LIST = [os.path.split(model_path)[0] for model_path in MODEL_FILE_LIST]
+        #     else:
+        #         raise ValueError('No model selected')
+        # except tkinter.TclError as e:
+        #     raise ValueError('Cannot launch TkInter file explorer. User the -m or -d arguments')
+        pass
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)  # Check availability before running using 'nvidia-smi'
+    DATASET = args.dataset
+    list_test_files = [os.path.join(DATASET, f) for f in os.listdir(DATASET) if f.endswith('h5') and 'dm' not in f]
+    list_test_files.sort()
+
+    with h5py.File(list_test_files[0], 'r') as f:
+        image_input_shape = image_output_shape = list(f['fix_image'][:].shape[:-1])
+        nb_labels = f['fix_segmentations'][:].shape[-1]
+
+    # Header of the metrics csv file
+    csv_header = ['File', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'DICE_MACRO', 'HD', 'Time', 'TRE', 'No_missing_lbls', 'Missing_lbls']
+
+    # TF stuff
+    config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
+    config.gpu_options.allow_growth = True
+    config.log_device_placement = False  ## to log device placement (on which device the operation ran)
+    config.allow_soft_placement = True
+
+    sess = tf.Session(config=config)
+    tf.keras.backend.set_session(sess)
+
+    # Loss and metric functions. Common to all models
+    loss_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).loss,
+                 NCC(image_input_shape).loss,
+                 vxm.losses.MSE().loss,
+                 MultiScaleStructuralSimilarity(max_val=1., filter_size=3).loss]
+
+    metric_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric,
+                   NCC(image_input_shape).metric,
+                   vxm.losses.MSE().loss,
+                   MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric,
+                   GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric,
+                   HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).metric,
+                   GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric_macro]
+
+    ### METRICS GRAPH ###
+    fix_img_ph = tf.placeholder(tf.float32, (1, *image_output_shape, 1), name='fix_img')
+    pred_img_ph = tf.placeholder(tf.float32, (1, *image_output_shape, 1), name='pred_img')
+    fix_seg_ph = tf.placeholder(tf.float32, (1, *image_output_shape, nb_labels), name='fix_seg')
+    pred_seg_ph = tf.placeholder(tf.float32, (1, *image_output_shape, nb_labels), name='pred_seg')
+
+    ssim_tf = metric_fncs[0](fix_img_ph, pred_img_ph)
+    ncc_tf = metric_fncs[1](fix_img_ph, pred_img_ph)
+    mse_tf = metric_fncs[2](fix_img_ph, pred_img_ph)
+    ms_ssim_tf = metric_fncs[3](fix_img_ph, pred_img_ph)
+    dice_tf = metric_fncs[4](fix_seg_ph, pred_seg_ph)
+    hd_tf = metric_fncs[5](fix_seg_ph, pred_seg_ph)
+    dice_macro_tf = metric_fncs[6](fix_seg_ph, pred_seg_ph)
+    # hd_exact_tf = HausdorffDistance_exact(fix_seg_ph, pred_seg_ph, ohe=True)
+
+    # Needed for VxmDense type of network
+    warp_segmentation = vxm.networks.Transform(image_output_shape, interp_method='nearest', nb_feats=nb_labels)
+
+    dm_interp = DisplacementMapInterpolator(image_output_shape, 'griddata')
+
+    for MODEL_FILE, DATA_ROOT_DIR in zip(MODEL_FILE_LIST, DATA_ROOT_DIR_LIST):
+        print('MODEL LOCATION: ', MODEL_FILE)
+
+        # data_folder = '/mnt/EncryptedData1/Users/javier/train_output/DDMR/THESIS/BASELINE_Affine_ncc___mse_ncc_160606-25022021'
+        output_folder = os.path.join(DATA_ROOT_DIR, args.outdirname)  # '/mnt/EncryptedData1/Users/javier/train_output/DDMR/THESIS/eval/BASELINE_TRAIN_Affine_ncc_EVAL_Affine'
+        # os.makedirs(os.path.join(output_folder, 'images'), exist_ok=True)
+        if args.erase:
+            shutil.rmtree(output_folder, ignore_errors=True)
+        os.makedirs(output_folder, exist_ok=True)
+        print('DESTINATION FOLDER: ', output_folder)
+
+        try:
+            network = tf.keras.models.load_model(MODEL_FILE, {'VxmDenseSemiSupervisedSeg': vxm.networks.VxmDenseSemiSupervisedSeg,
+                                                              'VxmDense': vxm.networks.VxmDense,
+                                                              'AdamAccumulated': AdamAccumulated,
+                                                              'loss': loss_fncs,
+                                                              'metric': metric_fncs},
+                                                 compile=False)
+        except ValueError as e:
+            enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+            dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+            nb_features = [enc_features, dec_features]
+            if re.search('^UW|SEGGUIDED_', MODEL_FILE):
+                network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
+                                                                 nb_labels=nb_labels,
+                                                                 nb_unet_features=nb_features,
+                                                                 int_steps=0,
+                                                                 int_downsize=1,
+                                                                 seg_downsize=1)
+            else:
+                network = vxm.networks.VxmDense(inshape=image_output_shape,
+                                                nb_unet_features=nb_features,
+                                                int_steps=0)
+            network.load_weights(MODEL_FILE, by_name=True)
+        # Record metrics
+        metrics_file = os.path.join(output_folder, 'metrics.csv')
+        with open(metrics_file, 'w') as f:
+            f.write(';'.join(csv_header)+'\n')
+
+        ssim = ncc = mse = ms_ssim = dice = hd = 0
+        with sess.as_default():
+            sess.run(tf.global_variables_initializer())
+            network.load_weights(MODEL_FILE, by_name=True)
+            progress_bar = tqdm(enumerate(list_test_files, 1), desc='Evaluation', total=len(list_test_files))
+            for step, in_batch in progress_bar:
+                with h5py.File(in_batch, 'r') as f:
+                    fix_img = f['fix_image'][:][np.newaxis, ...]    # Add batch axis
+                    mov_img = f['mov_image'][:][np.newaxis, ...]
+                    fix_seg = f['fix_segmentations'][:][np.newaxis, ...].astype(np.float32)
+                    mov_seg = f['mov_segmentations'][:][np.newaxis, ...].astype(np.float32)
+                    fix_centroids = f['fix_centroids'][:]
+
+                if network.name == 'vxm_dense_semi_supervised_seg':
+                    t0 = time.time()
+                    pred_img, disp_map, pred_seg = network.predict([mov_img, fix_img, mov_seg, fix_seg])    # predict([source, target])
+                    t1 = time.time()
+                else:
+                    t0 = time.time()
+                    pred_img, disp_map = network.predict([mov_img, fix_img])
+                    pred_seg = warp_segmentation.predict([mov_seg, disp_map])
+                    t1 = time.time()
+
+                pred_img = min_max_norm(pred_img)
+                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(0, nb_labels), brain_study=False)
+                # pred_centroids = dm_interp(disp_map, mov_centroids, backwards=True)  # with tps, it returns the pred_centroids directly
+                pred_centroids = dm_interp(disp_map, mov_centroids, backwards=True) + mov_centroids
+
+                # I need the labels to be OHE to compute the segmentation metrics.
+                dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf], {'fix_seg:0': fix_seg, 'pred_seg:0': pred_seg})
+
+                pred_seg_card = segmentation_ohe_to_cardinal(pred_seg).astype(np.float32)
+                mov_seg_card = segmentation_ohe_to_cardinal(mov_seg).astype(np.float32)
+                fix_seg_card = segmentation_ohe_to_cardinal(fix_seg).astype(np.float32)
+
+                ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf], {'fix_img:0': fix_img, 'pred_img:0': pred_img})
+                ms_ssim = ms_ssim[0]
+
+                # Rescale the points back to isotropic space, where we have a correspondence voxel <-> mm
+                upsample_scale = 128 / 64
+                fix_centroids_isotropic = fix_centroids * upsample_scale
+                # mov_centroids_isotropic = mov_centroids * upsample_scale
+                pred_centroids_isotropic = pred_centroids * upsample_scale
+
+                fix_centroids_isotropic = np.divide(fix_centroids_isotropic, C.IXI_DATASET_iso_to_cubic_scales)
+                # mov_centroids_isotropic = np.divide(mov_centroids_isotropic, C.IXI_DATASET_iso_to_cubic_scales)
+                pred_centroids_isotropic = np.divide(pred_centroids_isotropic, C.IXI_DATASET_iso_to_cubic_scales)
+                # Now we can measure the TRE in mm
+                tre_array = target_registration_error(fix_centroids_isotropic, pred_centroids_isotropic, False).eval()
+                tre = np.mean([v for v in tre_array if not np.isnan(v)])
+                # ['File', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'HD', 'Time', 'TRE', 'No_missing_lbls', 'Missing_lbls']
+
+                new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd, t1-t0, tre, len(missing_lbls), missing_lbls]
+                with open(metrics_file, 'a') as f:
+                    f.write(';'.join(map(str, new_line))+'\n')
+
+                save_nifti(fix_img[0, ...], os.path.join(output_folder, '{:03d}_fix_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(mov_img[0, ...], os.path.join(output_folder, '{:03d}_mov_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(pred_img[0, ...], os.path.join(output_folder, '{:03d}_pred_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(fix_seg_card[0, ...], os.path.join(output_folder, '{:03d}_fix_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(mov_seg_card[0, ...], os.path.join(output_folder, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(pred_seg_card[0, ...], os.path.join(output_folder, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+
+                # with h5py.File(os.path.join(output_folder, '{:03d}_centroids.h5'.format(step)), 'w') as f:
+                    # f.create_dataset('fix_centroids', dtype=np.float32, data=fix_centroids)
+                    # f.create_dataset('mov_centroids', dtype=np.float32, data=mov_centroids)
+                    # f.create_dataset('pred_centroids', dtype=np.float32, data=pred_centroids)
+                    # f.create_dataset('fix_centroids_isotropic', dtype=np.float32, data=fix_centroids_isotropic)
+                    # f.create_dataset('mov_centroids_isotropic', dtype=np.float32, data=mov_centroids_isotropic)
+
+                # magnitude = np.sqrt(np.sum(disp_map[0, ...] ** 2, axis=-1))
+                # _ = plt.hist(magnitude.flatten())
+                # plt.title('Histogram of disp. magnitudes')
+                # plt.show(block=False)
+                # plt.savefig(os.path.join(output_folder, '{:03d}_hist_mag_ssim_{:.03f}_dice_{:.03f}.png'.format(step, ssim, dice)))
+                # plt.close()
+
+                plot_predictions(fix_img, mov_img, disp_map, pred_img, os.path.join(output_folder, '{:03d}_figures_img.png'.format(step)), show=False)
+                plot_predictions(fix_seg, mov_seg, disp_map, pred_seg, os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False)
+                save_disp_map_img(disp_map, 'Displacement map', os.path.join(output_folder, '{:03d}_disp_map_fig.png'.format(step)), show=False)
+
+                progress_bar.set_description('SSIM {:.04f}\tG_DICE: {:.04f}\tM_DICE: {:.04f}'.format(ssim, dice, dice_macro))
+
+        print('Summary\n=======\n')
+        print('\nAVG:\n' + str(pd.read_csv(metrics_file, sep=';', header=0).mean(axis=0)) + '\nSTD:\n' + str(pd.read_csv(metrics_file, sep=';', header=0).std(axis=0)))
+        print('\n=======\n')
+        tf.keras.backend.clear_session()
+        # sess.close()
+        del network
+    print('Done')

COMET/MultiTrain_cli.py

@@ -0,0 +1,61 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import argparse
+from datetime import datetime
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model', type=str, help='Path to the model file', required=True)
+    parser.add_argument('--dataset', type=str, help='Location of the training data', default=TRAIN_DATASET)
+    parser.add_argument('--validation', type=str, help='Location of the validation data', default=None)
+    parser.add_argument('--similarity', nargs='+', help='Similarity metric: mse, ncc, ssim', default=['ncc'])
+    parser.add_argument('--segmentation', nargs='+', help='Segmentation loss function: hd, dice', default=['dice'])
+    parser.add_argument('--output', type=str, help='Output directory', default=TRAIN_DATASET)
+    parser.add_argument('--gpu', type=str, help='GPU number', default='0')
+    parser.add_argument('--lr', type=float, help='Learning rate', default=1e-4)
+    parser.add_argument('--rw', type=float, help='Regularization weigh', default=5e-3)
+    parser.add_argument('--epochs', type=int, help='Max number of epochs', default=1000)
+    parser.add_argument('--name', type=str, default='COMET')
+    parser.add_argument('--uw', type=bool, default=False)
+    parser.add_argument('--freeze', nargs='+', help='What layers to freeze: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM', default=None)
+    parser.add_argument('--epochs', default=1000)
+    parser.add_argument('--batch', default=16)
+    args = parser.parse_args()
+
+    print('TRAIN ' + args.dataset)
+
+    if args.uw:
+        from COMET.COMET_train_UW import launch_train
+        simil = args.similarity
+        segm = args.segmentation
+        output_folder = os.path.join(args.output, '{}_Lsim_{}__Lseg_{}'.format(args.name, '_'.join(simil), '_'.join(segm)))
+    else:
+        from COMET.COMET_train import launch_train
+        simil = args.similarity[0]
+        segm = args.segmentation[0]
+        output_folder = os.path.join(args.output, '{}_Lsim_{}__Lseg_{}'.format(args.name, simil, segm))
+    output_folder = output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y")
+
+    if args.freeze is not None:
+        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in args.freeze),\
+            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
+        args.freeze = list(set(args.freeze))
+
+    launch_train(dataset_folder=args.dataset,
+                 validation_folder=args.validation,
+                 output_folder=output_folder,
+                 gpu_num=args.gpu,
+                 lr=args.lr,
+                 rw=args.rw,
+                 simil=simil,
+                 segm=segm,
+                 max_epochs=args.epochs,
+                 model_file=args.model,
+                 freeze_layers=args.freeze)

COMET/MultiTrain_config.py

@@ -0,0 +1,71 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import argparse
+from configparser import ConfigParser
+from shutil import copy2
+import os
+from datetime import datetime
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ini', help='Configuration file')
+    args = parser.parse_args()
+
+    configFile = ConfigParser()
+    configFile.read(args.ini)
+    print('Loaded configuration file: ' + args.ini)
+    print({section: dict(configFile[section]) for section in configFile.sections()})
+    print('\n\n')
+
+    trainConfig = configFile['TRAIN']
+    lossesConfig = configFile['LOSSES']
+    datasetConfig = configFile['DATASETS']
+    othersConfig = configFile['OTHERS']
+
+    print('TRAIN MODEL IN' + trainConfig['model'])
+
+    simil = lossesConfig['similarity'].split(',')
+    segm = lossesConfig['segmentation'].split(',')
+    if trainConfig['name'].lower() == 'uw':
+        from COMET.COMET_train_UW import launch_train
+        output_folder = os.path.join(othersConfig['outputFolder'], '{}_Lsim_{}__Lseg_{}'.format(trainConfig['name'], '_'.join(simil), '_'.join(segm)))
+    else:
+        from COMET.COMET_train import launch_train
+        simil = simil[0]
+        segm = segm[0]
+        output_folder = os.path.join(othersConfig['outputFolder'], '{}_Lsim_{}__Lseg_{}'.format(trainConfig['name'], simil, segm))
+    output_folder = output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y")
+
+    try:
+        froozen_layers = eval(trainConfig['freeze'])
+    except NameError as err:
+        froozen_layers = [trainConfig['freeze'].upper()]
+    if froozen_layers is not None:
+        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in froozen_layers),\
+            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
+        froozen_layers = list(set(froozen_layers))  # Unique elements
+
+    # copy the configuration file to the destionation folder
+    os.makedirs(output_folder, exist_ok=True)
+    copy2(args.ini, os.path.join(output_folder, os.path.split(args.ini)[-1]))
+
+    launch_train(dataset_folder=datasetConfig['train'],
+                 validation_folder=datasetConfig['validation'],
+                 output_folder=output_folder,
+                 gpu_num=eval(trainConfig['gpu']),
+                 lr=eval(trainConfig['learningRate']),
+                 rw=eval(trainConfig['regularizationWeight']),
+                 simil=simil,
+                 segm=segm,
+                 max_epochs=eval(trainConfig['epochs']),
+                 model_file=trainConfig['model'],
+                 freeze_layers=froozen_layers,
+                 acc_gradients=eval(trainConfig['accumulativeGradients']),
+                 batch_size=eval(trainConfig['batchSize']))

COMET/augmentation_constants.py

@@ -0,0 +1,34 @@
+import numpy as np
+
+# Constants for augmentation layer
+# .../T1/training/zoom_factors.csv contain the scale factors of all the training samples from isotropic to 128x128x128
+#   The augmentation values will be scaled using the average+std
+ZOOM_FACTORS = np.asarray([0.5032864535069749, 0.5363100665659675, 0.6292598243796296])
+MAX_AUG_DISP_ISOT = 30
+MAX_AUG_DEF_ISOT = 6
+MAX_AUG_DISP = np.max(MAX_AUG_DISP_ISOT * ZOOM_FACTORS)  # Scaled displacements
+MAX_AUG_DEF = np.max(MAX_AUG_DEF_ISOT * ZOOM_FACTORS)  # Scaled deformations
+MAX_AUG_ANGLE = np.max([np.arctan(np.tan(10*np.pi/180) * ZOOM_FACTORS[1] / ZOOM_FACTORS[0]) * 180 / np.pi,
+                        np.arctan(np.tan(10*np.pi/180) * ZOOM_FACTORS[2] / ZOOM_FACTORS[1]) * 180 / np.pi,
+                        np.arctan(np.tan(10*np.pi/180) * ZOOM_FACTORS[2] / ZOOM_FACTORS[0]) * 180 / np.pi])  # Scaled angles
+GAMMA_AUGMENTATION = False
+BRIGHTNESS_AUGMENTATION = False
+NUM_CONTROL_PTS_AUG = 10
+NUM_AUGMENTATIONS = 5
+
+IN_LAYERS = (0, 3)
+OUT_LAYERS = (33, 39)
+
+ENCONDER_LAYERS = (3, 17)
+DECODER_LAYERS = (17, 33)
+
+TOP_LAYERS_ENC = (3, 9)
+TOP_LAYERS_DEC = (22, 29)
+BOTTOM_LAYERS = (9, 22)
+
+LAYER_RANGES = {'INPUT': (IN_LAYERS),
+                'OUTPUT': (OUT_LAYERS),
+                'ENCODER': (ENCONDER_LAYERS),
+                'DECODER': (DECODER_LAYERS),
+                'TOP': (TOP_LAYERS_ENC, TOP_LAYERS_DEC),
+                'BOTTOM': (BOTTOM_LAYERS)}

COMET/format_dataset.py

@@ -0,0 +1,114 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import h5py
+import nibabel as nib
+from nilearn.image import resample_img
+import re
+import numpy as np
+from scipy.ndimage import zoom
+from skimage.measure import regionprops
+from tqdm import tqdm
+from argparse import ArgumentParser
+from scipy.ndimage.morphology import binary_dilation, generate_binary_structure
+
+import pandas as pd
+
+from DeepDeformationMapRegistration.utils import constants as C
+from DeepDeformationMapRegistration.utils.misc import segmentation_cardinal_to_ohe, segmentation_ohe_to_cardinal
+
+SEGMENTATION_NR2LBL_LUT = {0: 'background',
+                           1: 'parenchyma',
+                           2: 'vessel'}
+SEGMENTATION_LBL2NR_LUT = {v: k for k, v in SEGMENTATION_NR2LBL_LUT.items()}
+
+IMG_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Volumes'
+SEG_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Segmentations' # '/home/jpdefrutos/workspace/LiverSegmentation_UNet3D/data/prediction'
+
+IMG_NAME_PATTERN = '(.*).nii.gz'
+SEG_NAME_PATTERN = '(.*).nii.gz'
+
+OUT_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128'
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('--crop', action='store_true')  # If present, args.crop = True, else args.crop = False
+    parser.add_argument('--offset', type=int, default=C.MAX_AUG_DISP_ISOT + 10, help='Crop offset in mm')
+    parser.add_argument('--dilate-segmentations', type=bool, default=False)
+    args = parser.parse_args()
+
+    img_list = [os.path.join(IMG_DIRECTORY, f) for f in os.listdir(IMG_DIRECTORY) if f.endswith('.nii.gz')]
+    img_list.sort()
+
+    seg_list = [os.path.join(SEG_DIRECTORY, f) for f in os.listdir(SEG_DIRECTORY) if f.endswith('.nii.gz')]
+    seg_list.sort()
+
+    zoom_file = pd.DataFrame(columns=['scale_i', 'scale_j', 'scale_k'])
+    os.makedirs(OUT_DIRECTORY, exist_ok=True)
+    binary_ball = generate_binary_structure(3, 1)
+    for seg_file in tqdm(seg_list):
+        img_name = re.match(SEG_NAME_PATTERN, os.path.split(seg_file)[-1])[1]
+        img_file = os.path.join(IMG_DIRECTORY, img_name + '.nii.gz')
+
+        img = resample_img(nib.load(img_file), np.eye(3))
+        seg = resample_img(nib.load(seg_file), np.eye(3), interpolation='nearest')
+
+        img = np.asarray(img.dataobj)
+        seg = np.asarray(seg.dataobj)
+
+        segs_are_ohe = bool(len(seg.shape) > 3 and seg.shape[3] > 1)
+        if args.crop:
+            parenchyma = regionprops(seg[..., 0])[0]
+            bbox = np.asarray(parenchyma.bbox) + [*[-args.offset]*3, *[args.offset]*3]
+            # check that the new bbox is within the image limits!
+            bbox[:3] = np.maximum(bbox[:3], [0, 0, 0])
+            bbox[3:] = np.minimum(bbox[3:], img.shape)
+            img = img[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]]
+            seg = seg[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5], ...]
+        # Resize to 128x128x128
+        isot_shape = img.shape
+
+        zoom_factors = (np.asarray([128]*3) / np.asarray(img.shape)).tolist()
+
+        img = zoom(img, zoom_factors, order=3)
+        if args.dilate_segmentations:
+            seg = binary_dilation(seg, binary_ball, iterations=1)
+        seg = zoom(seg, zoom_factors + [1]*(len(seg.shape) - len(img.shape)), order=0)
+        zoom_file = zoom_file.append({'scale_i': zoom_factors[0],
+                                      'scale_j': zoom_factors[1],
+                                      'scale_k': zoom_factors[2]}, ignore_index=True)
+
+        # seg -> cardinal
+        # seg_expanded -> OHE
+        if segs_are_ohe:
+            seg_expanded = seg.copy()
+            seg = segmentation_ohe_to_cardinal(seg)  # Ordinal encoded. argmax returns the first ocurrence of the maximum. Hence the previoous multiplication operation
+        else:
+            seg_expanded = segmentation_cardinal_to_ohe(seg)
+
+        h5_file = h5py.File(os.path.join(OUT_DIRECTORY, img_name + '.h5'), 'w')
+
+        h5_file.create_dataset('image', data=img[..., np.newaxis], dtype=np.float32)
+        h5_file.create_dataset('segmentation', data=seg.astype(np.uint8), dtype=np.uint8)
+        h5_file.create_dataset('segmentation_expanded', data=seg_expanded.astype(np.uint8), dtype=np.uint8)
+        h5_file.create_dataset('segmentation_labels', data=np.unique(seg)[1:])  # Remove the 0 (background label)
+        h5_file.create_dataset('isotropic_shape', data=isot_shape)
+
+        print('{}: Segmentation labels {}'.format(img_name, np.unique(seg)[1:]))
+        h5_file.close()
+
+    zoom_file.to_csv(os.path.join(OUT_DIRECTORY, 'zoom_factors.csv'))
+    print("Average")
+    print(zoom_file.mean().to_list())
+
+    print("Standard deviation")
+    print(zoom_file.std().to_list())
+
+    print("Average + STD")
+    print((zoom_file.mean() + zoom_file.std()).to_list())
+
+
+

COMET/spit_dataset.py

@@ -0,0 +1,40 @@
+from shutil import move, copy2
+import os
+
+OR_DIR = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128'
+val_split = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/For_validation.txt'
+test_split = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/For_testing.txt'
+
+# Create out dirs
+os.makedirs(os.path.join(OR_DIR, 'train'), exist_ok=True)
+os.makedirs(os.path.join(OR_DIR, 'validation'), exist_ok=True)
+os.makedirs(os.path.join(OR_DIR, 'test'), exist_ok=True)
+
+# Copy all to train and then split into validation and test
+list_of_files = [os.path.join(OR_DIR, f) for f in os.listdir(OR_DIR) if f.endswith('.h5')]
+list_of_files.sort()
+for f in list_of_files:
+    copy2(f, os.path.join(OR_DIR, 'train'))
+
+#   Get the indices for the validation and test subsets
+with open(val_split, 'r') as f:
+    val_idcs = f.readlines()[0]
+    val_idcs = [int(e) for e in val_idcs.split(',')]
+
+with open(test_split, 'r') as f:
+    test_indcs = f.readlines()[0]
+    test_indcs = [int(e) for e in test_indcs.split(',')]
+
+#   move the files from train to validation and test
+for i in val_idcs:
+    move(os.path.join(OR_DIR, 'train', '{:05d}_CT.h5'.format(i)), os.path.join(OR_DIR, 'validation'))
+print('Done moving the validation subset.')
+
+for i in test_indcs:
+    move(os.path.join(OR_DIR, 'train', '{:05d}_CT.h5'.format(i)), os.path.join(OR_DIR, 'test'))
+print('Done moving the validation subset.')
+
+print('Done splitting the data')
+print('Training samples: '+str(len(os.listdir(os.path.join(OR_DIR, 'train')))))
+print('Validation samples: '+str(len(val_idcs)))
+print('Test samples: '+str(len(test_indcs)))

COMET/train_config_files/Config_BASELINE_None_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+similarity = ncc
+segmentation =
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/BASELINE_L_ncc__MET_mse_ncc_ssim_232329-01092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 1
+regularizationWeight = 1e-5
+epochs = 10000
+name = BASELINE
+freeze = None
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/NONE_FROZEN

COMET/train_config_files/Config_BASELINE_bottom_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+similarity = ncc
+segmentation =
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/BASELINE_L_ncc__MET_mse_ncc_ssim_232329-01092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 0
+regularizationWeight = 1e-5
+epochs = 10000
+name = BASELINE
+freeze = BOTTOM
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/BOTTOM_FROZEN

COMET/train_config_files/Config_BASELINE_top_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+similarity = ncc
+segmentation =
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/BASELINE_L_ncc__MET_mse_ncc_ssim_232329-01092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 1
+regularizationWeight = 1e-5
+epochs = 10000
+name = BASELINE
+freeze = TOP
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/TOP_FROZEN

COMET/train_config_files/Config_SEGGUIDED_None_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice
+similarity = ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/SEGGUIDED_Lsim_ncc__Lseg_dice__MET_mse_ncc_ssim_013319-11092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 1
+regularizationWeight = 1e-5
+epochs = 10000
+name = SEGGUIDED
+freeze = None
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/NONE_FROZEN

COMET/train_config_files/Config_SEGGUIDED_bottom_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice
+similarity = ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/SEGGUIDED_Lsim_ncc__Lseg_dice__MET_mse_ncc_ssim_013319-11092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 2
+regularizationWeight = 1e-5
+epochs = 10000
+name = SEGGUIDED
+freeze = BOTTOM
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/BOTTOM_FROZEN

COMET/train_config_files/Config_SEGGUIDED_top_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice
+similarity = ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/SEGGUIDED_Lsim_ncc__Lseg_dice__MET_mse_ncc_ssim_013319-11092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 1
+regularizationWeight = 1e-5
+epochs = 10000
+name = SEGGUIDED
+freeze = TOP
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/TOP_FROZEN

COMET/train_config_files/Config_UW_None_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice,hd
+similarity = ncc,ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/UW_Lsim_ncc__ssim__Lseg_dice__MET_mse_ncc_ssim_204557-02092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 0
+regularizationWeight = 1e-5
+epochs = 10000
+name = UW
+freeze = None
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/NONE_FROZEN

COMET/train_config_files/Config_UW_bottom_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice,hd
+similarity = ncc,ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/UW_Lsim_ncc__ssim__Lseg_dice__MET_mse_ncc_ssim_204557-02092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 2
+regularizationWeight = 1e-5
+epochs = 10000
+name = UW
+freeze = BOTTOM
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/BOTTOM_FROZEN

COMET/train_config_files/Config_UW_top_froozen.ini

@@ -0,0 +1,21 @@
+[LOSSES]
+segmentation = dice,hd
+similarity = ncc,ssim
+
+[TRAIN]
+model = /mnt/EncryptedData1/Users/javier/train_output/Brain_study/No_gamma/UW_Lsim_ncc__ssim__Lseg_dice__MET_mse_ncc_ssim_204557-02092021/checkpoints/best_model.h5
+batchSize = 8
+learningRate = 1e-5
+accumulativeGradients = 1
+gpu = 0
+regularizationWeight = 1e-5
+epochs = 10000
+name = UW
+freeze = TOP
+
+[DATASETS]
+train = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train
+validation = /mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/validation
+
+[OTHERS]
+outputFolder = /mnt/EncryptedData1/Users/javier/train_output/COMET/TOP_FROZEN