Update

Brain_study/Train_SegmentationGuided.py

@@ -173,7 +173,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     loss_weights = {'transformer': 1,
                     'seg_transformer': 1.,
                     'flow': 5e-3}
-    metrics = {'transformer': [vxm.losses.MSE().loss, NCC(image_input_shape).metric, MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).metric],
+    metrics = {'transformer': [vxm.losses.MSE().loss, NCC(image_input_shape).metric, 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],
                'seg_transformer': [GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric_macro,
                                      #HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [train_generator.get_data_shape()[2][-1]]).metric
                                      ]}

COMET/COMET_train.py

@@ -34,22 +34,26 @@ import voxelmorph as vxm
 import h5py
 import re
 import itertools
+import warnings
 
 
 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):
+                 segm='dice', max_epochs=C.EPOCHS, early_stop_patience=1000, freeze_layers=None,
+                 acc_gradients=1, batch_size=16, image_size=64,
+                 unet=[16, 32, 64, 128, 256], head=[16, 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))
+    assert dataset_folder is not None and output_folder is not None
+    if model_file != '':
+        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 batch_size != 1 and acc_gradients != 1:
+        warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
+
     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'
@@ -63,8 +67,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     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.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else 1
+    C.EARLY_STOP_PATIENCE = early_stop_patience
     C.LEARNING_RATE = lr
     C.LIMIT_NUM_SAMPLES = None
     C.EPOCHS = max_epochs
@@ -97,16 +101,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     print(aux)
 
     # 2. Data generator
-    used_labels = 'all' if USE_SEGMENTATIONS else 'none'
+    used_labels = '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,
+                                    C.TRAINING_PERC, labels=[used_labels], combine_segmentations=True,
                                     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
+    image_output_shape = [image_size] * 3
     nb_labels = len(train_generator.get_segmentation_labels())
 
     # 3. Load model
@@ -133,7 +137,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                    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,
@@ -143,21 +147,18 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                                           '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]
+        # enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+        # dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+        enc_features = unet  # const.ENCODER_FILTERS
+        dec_features = enc_features[::-1] + head  # 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)
+
+        network = vxm.networks.VxmDense(inshape=image_output_shape,
+                                        nb_unet_features=nb_features,
+                                        int_steps=0)
+        if model_file != '':
+            network.load_weights(model_file, by_name=True)
+            print('MODEL LOCATION: ', model_file)
     # 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:
@@ -187,7 +188,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     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]
+    augm_train_input_shape = 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
@@ -198,7 +199,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                          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_image=True,  # If baseline then True
                                          only_resize=False,
                                          trainable=False)
     augm_model_train = Model(inputs=input_layer_train, outputs=augm_layer_train(input_layer_train))
@@ -255,16 +256,6 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     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'),
@@ -278,34 +269,17 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                           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}
+
+    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,
@@ -359,12 +333,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                     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)
+                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)
@@ -395,12 +366,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                     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)
+                in_data = (mov_img, fix_img)
+                out_data = (fix_img, fix_img)
                 ret = network.test_on_batch(x=in_data,
                                             y=out_data)
 

COMET/COMET_train_UW.py

@@ -33,20 +33,26 @@ import voxelmorph as vxm
 import h5py
 import re
 import itertools
+import warnings
 
 
 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):
+                 simil=['ssim'], segm=['dice'], max_epochs=C.EPOCHS, early_stop_patience=1000, prior_reg_w=5e-3,
+                 freeze_layers=None, acc_gradients=1, batch_size=16, image_size=64,
+                 unet=[16, 32, 64, 128, 256], head=[16, 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'
+    assert dataset_folder is not None and output_folder is not None
+    if model_file != '':
+        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 batch_size != 1 and acc_gradients != 1:
+        warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
+
     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'
@@ -63,8 +69,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     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.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else 1
+    C.EARLY_STOP_PATIENCE = early_stop_patience
     C.LEARNING_RATE = lr
     C.LIMIT_NUM_SAMPLES = None
     C.EPOCHS = max_epochs
@@ -105,7 +111,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     validation_generator = data_generator.get_validation_generator()
 
     image_input_shape = train_generator.get_data_shape()[-1][:-1]
-    image_output_shape = [64] * 3
+    image_output_shape = [image_size] * 3
     nb_labels = len(train_generator.get_segmentation_labels())
 
     # 3. Load model
@@ -116,10 +122,10 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     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]
+    # enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+    # dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    enc_features = unet  # const.ENCODER_FILTERS
+    dec_features = enc_features[::-1] + head  # const.ENCODER_FILTERS[::-1]
     nb_features = [enc_features, dec_features]
     network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
                                                      nb_labels=nb_labels,
@@ -127,7 +133,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                                      int_steps=0,
                                                      int_downsize=1,
                                                      seg_downsize=1)
-    network.load_weights(model_file, by_name=True)
+    if model_file != '':
+        print('MODEL LOCATION: ', model_file)
+        network.load_weights(model_file, by_name=True)
 
     # 4. Freeze/unfreeze model layers
     if freeze_layers is not None:

DeepDeformationMapRegistration/layers/augmentation.py

@@ -182,14 +182,15 @@ class AugmentationLayer(kl.Layer):
         return tf.squeeze(disp_map, axis=0)
 
     def gamma_augmentation(self, in_img: tf.Tensor):
-        in_img += 1e-5  # To prvent NaNs
-        gamma = tf.random.uniform((), 0.5, 2, tf.float32)
+        in_img += 1e-5  # To prevent NaNs
+        f = tf.random.uniform((), -1, 1, tf.float32)  # gamma [0.5, 2]
+        gamma = tf.pow(2.0, f)
 
         return tf.clip_by_value(tf.pow(in_img, gamma), 0, 1)
 
     def brightness_augmentation(self, in_img: tf.Tensor):
-        c = tf.random.uniform((), 0.5, 2, tf.float32)
-        return tf.clip_by_value(c*in_img, 0, 1)
+        c = tf.random.uniform((), -0.2, 0.2, tf.float32)  # 20% shift
+        return tf.clip_by_value(c + in_img, 0, 1)
 
     def min_max_normalization(self, in_img: tf.Tensor):
         return tf.div(tf.subtract(in_img, tf.reduce_min(in_img)),
@@ -228,6 +229,7 @@ class AugmentationLayer(kl.Layer):
             except InvalidArgumentError as err:
                 # If the transformation raises a non-invertible error,
                 # try again until we get a valid transformation
+                tf.print('TPS non invertible matrix', output_stream=sys.stdout)
                 continue
             else:
                 valid_trf = True

DeepDeformationMapRegistration/utils/constants.py

@@ -518,7 +518,7 @@ MAX_AUG_DEF = np.max(MAX_AUG_DEF_ISOT * IXI_DATASET_iso_to_cubic_scales)  # Scal
 MAX_AUG_ANGLE = np.max([np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[1] / IXI_DATASET_iso_to_cubic_scales[0]) * 180 / np.pi,
                         np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[2] / IXI_DATASET_iso_to_cubic_scales[1]) * 180 / np.pi,
                         np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[2] / IXI_DATASET_iso_to_cubic_scales[0]) * 180 / np.pi])  # Scaled angles
-GAMMA_AUGMENTATION = False
+GAMMA_AUGMENTATION = True
 BRIGHTNESS_AUGMENTATION = False
 NUM_CONTROL_PTS_AUG = 10
 NUM_AUGMENTATIONS = 1

DeepDeformationMapRegistration/utils/misc.py

@@ -7,6 +7,7 @@ from skimage.measure import regionprops
 from DeepDeformationMapRegistration.layers.b_splines import interpolate_spline
 from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
 from tensorflow import squeeze
+from scipy.ndimage import zoom
 
 
 def try_mkdir(dir, verbose=True):
@@ -148,3 +149,34 @@ def segmentation_cardinal_to_ohe(segmentation):
     for ch, lbl in enumerate(np.unique(segmentation)[1:]):
         cpy[segmentation == lbl, ch] = 1
     return cpy
+
+
+def resize_displacement_map(displacement_map: np.ndarray, dest_shape: [list, np.ndarray, tuple], scale_trf: np.ndarray=None):
+    if scale_trf is None:
+        scale_trf = scale_transformation(displacement_map.shape, dest_shape)
+    else:
+        assert isinstance(scale_trf, np.ndarray) and scale_trf.shape == (4, 4), 'Invalid transformation: {}'.format(scale_trf)
+    zoom_factors = scale_trf.diagonal()
+    # First scale the values, so we cut down the number of multiplications
+    dm_resized = np.copy(displacement_map)
+    dm_resized[..., 0] *= zoom_factors[0]
+    dm_resized[..., 1] *= zoom_factors[1]
+    dm_resized[..., 2] *= zoom_factors[2]
+    # Then rescale using zoom
+    dm_resized = zoom(dm_resized, zoom_factors)
+    return dm_resized
+
+
+def scale_transformation(original_shape: [list, tuple, np.ndarray], dest_shape: [list, tuple, np.ndarray]) -> np.ndarray:
+    if isinstance(original_shape, (list, tuple)):
+        original_shape = np.asarray(original_shape, dtype=int)
+    if isinstance(dest_shape, (list, tuple)):
+        dest_shape = np.asarray(dest_shape, dtype=int)
+    original_shape = original_shape.astype(int)
+    dest_shape = dest_shape.astype(int)
+
+    trf = np.eye(4)
+    np.fill_diagonal(trf, [*np.divide(dest_shape, original_shape), 1])
+
+    return trf
+

SoA_methods/eval_ants.py

@@ -35,6 +35,9 @@ WARPED_FIX = 'warpedfixout'
 FWD_TRFS = 'fwdtransforms'
 INV_TRFS = 'invtransforms'
 
+os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+os.environ['CUDA_VISIBLE_DEVICES'] = '2'
+
 if __name__ == '__main__':
     parser = ArgumentParser()
     parser.add_argument('--dataset', type=str, help='Directory with the images')
@@ -42,11 +45,13 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     os.makedirs(args.outdir, exist_ok=True)
+    os.makedirs(os.path.join(args.outdir, 'SyN'), exist_ok=True)
+    os.makedirs(os.path.join(args.outdir, 'SyNCC'), exist_ok=True)
     dataset_files = os.listdir(args.dataset)
     dataset_files.sort()
     dataset_files = [os.path.join(args.dataset, f) for f in dataset_files if re.match(DATASET_NAMES, f)]
 
-    dataset_iterator = tqdm(dataset_files)
+    dataset_iterator = tqdm(enumerate(dataset_files), desc="Running ANTs")
 
     f = h5py.File(dataset_files[0], 'r')
     image_shape = list(f['fix_image'][:].shape[:-1])
@@ -87,17 +92,19 @@ if __name__ == '__main__':
     print("Running ANTs using {} threads".format(os.environ.get("ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS")))
     dm_interp = DisplacementMapInterpolator(image_shape, 'griddata')
     # Header of the metrics csv file
-    csv_header = ['File', 'Method', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'DICE_MACRO', 'HD', 'Time_SyN', 'Time_SyNCC', 'TRE']
+    csv_header = ['File', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'DICE_MACRO', 'HD', 'Time', 'TRE']
 
-    metrics_file = os.path.join(args.outdir, 'metrics.csv')
-    with open(metrics_file, 'w') as f:
-        f.write(';'.join(csv_header)+'\n')
+    metrics_file = {'SyN': os.path.join(args.outdir, 'SyN', 'metrics.csv'),
+                    'SyNCC': os.path.join(args.outdir, 'SyNCC', 'metrics.csv')}
+    for k in metrics_file.keys():
+        with open(metrics_file[k], 'w') as f:
+            f.write(';'.join(csv_header)+'\n')
 
     print('Starting the loop')
-    for step, file_path in tqdm(enumerate(dataset_iterator), desc="Running ANTs"):
+    for step, file_path in dataset_iterator:
         file_num = int(re.findall('(\d+)', os.path.split(file_path)[-1])[0])
 
-        dataset_iterator.set_description('{} ({}): loading data'.format(file_num, args.dataset))
+        dataset_iterator.set_description('{} ({}): loading data'.format(file_num, file_path))
         with h5py.File(file_path, 'r') as vol_file:
             fix_img = vol_file['fix_image'][:]
             mov_img = vol_file['mov_image'][:]
@@ -112,10 +119,12 @@ if __name__ == '__main__':
         fix_img_ants = ants.make_image(fix_img.shape[:-1], np.squeeze(fix_img))  # SoA doesn't work fine with 1-ch images
         mov_img_ants = ants.make_image(mov_img.shape[:-1], np.squeeze(mov_img))  # SoA doesn't work fine with 1-ch images
 
+        dataset_iterator.set_description('{} ({}): running ANTs SyN'.format(file_num, file_path))
         t0_syn = time.time()
         reg_output_syn = ants.registration(fix_img_ants, mov_img_ants, 'SyN')
         t1_syn = time.time()
 
+        dataset_iterator.set_description('{} ({}): running ANTs SyN'.format(file_num, file_path))
         t0_syncc = time.time()
         reg_output_syncc = ants.registration(fix_img_ants, mov_img_ants, 'SyNCC')
         t1_syncc = time.time()
@@ -135,6 +144,8 @@ if __name__ == '__main__':
                 pred_seg = ants.apply_transforms(fixed=fix_seg_ants, moving=mov_seg_ants,
                                                  transformlist=mov_to_fix_trf_list).numpy()
                 pred_seg = np.squeeze(pred_seg)  # SoA adds an extra axis which shouldn't be there
+
+                dataset_iterator.set_description('{} ({}): Getting metrics {}'.format(file_num, file_path, reg_method))
                 with sess.as_default():
                     dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf],
                                                     {'fix_seg:0': fix_seg[np.newaxis, ...],  # Batch axis
@@ -163,22 +174,26 @@ if __name__ == '__main__':
                     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)])
 
-                new_line = [step, reg_method, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd, t1_syn-t0_syn, t1_syncc-t0_syncc, tre]
-                with open(metrics_file, 'a') as f:
+                dataset_iterator.set_description('{} ({}): Saving data {}'.format(file_num, file_path, reg_method))
+                new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd,
+                            t1_syn-t0_syn if reg_method == 'SyN' else t1_syncc-t0_syncc,
+                            tre]
+                with open(metrics_file[reg_method], 'a') as f:
                     f.write(';'.join(map(str, new_line))+'\n')
 
-                save_nifti(fix_img[0, ...], os.path.join(args.outdir, '{:03d}_fix_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(mov_img[0, ...], os.path.join(args.outdir, '{:03d}_mov_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(pred_img[0, ...], os.path.join(args.outdir, '{:03d}_pred_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(fix_seg_card[0, ...], os.path.join(args.outdir, '{:03d}_fix_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(mov_seg_card[0, ...], os.path.join(args.outdir, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(pred_seg_card[0, ...], os.path.join(args.outdir, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-
-                plot_predictions(fix_img[np.newaxis, ...], mov_img[np.newaxis, ...], disp_map[np.newaxis, ...], pred_img[np.newaxis, ...], os.path.join(args.outdir, '{:03d}_figures_img.png'.format(step)), show=False)
-                plot_predictions(fix_seg[np.newaxis, ...], mov_seg[np.newaxis, ...], disp_map[np.newaxis, ...], pred_seg[np.newaxis, ...], os.path.join(args.outdir, '{:03d}_figures_seg.png'.format(step)), show=False)
-                save_disp_map_img(disp_map[np.newaxis, ...], 'Displacement map', os.path.join(args.outdir, '{:03d}_disp_map_fig.png'.format(step)), show=False)
-
-    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')
+                save_nifti(fix_img[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_fix_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(mov_img[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_mov_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(pred_img[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_pred_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(fix_seg_card[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_fix_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(mov_seg_card[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                save_nifti(pred_seg_card[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+
+                plot_predictions(fix_img[np.newaxis, ...], mov_img[np.newaxis, ...], disp_map[np.newaxis, ...], pred_img[np.newaxis, ...], os.path.join(args.outdir, reg_method, '{:03d}_figures_img.png'.format(step)), show=False)
+                plot_predictions(fix_seg[np.newaxis, ...], mov_seg[np.newaxis, ...], disp_map[np.newaxis, ...], pred_seg[np.newaxis, ...], os.path.join(args.outdir, reg_method, '{:03d}_figures_seg.png'.format(step)), show=False)
+                save_disp_map_img(disp_map[np.newaxis, ...], 'Displacement map', os.path.join(args.outdir, reg_method, '{:03d}_disp_map_fig.png'.format(step)), show=False)
+
+    for k in metrics_file.keys():
+        print('Summary {}\n=======\n'.format(k))
+        print('\nAVG:\n' + str(pd.read_csv(metrics_file[k], sep=';', header=0).mean(axis=0)) + '\nSTD:\n' + str(
+            pd.read_csv(metrics_file[k], sep=';', header=0).std(axis=0)))
+        print('\n=======\n')