Scripts for training on the IXI T1 MRI Dataset

Brain_study/Build_test_set.py

@@ -0,0 +1,137 @@
+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/IXI_dataset/T1/training'
+MODEL_FILE = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/checkpoints/best_model.h5'
+DATA_ROOT_DIR = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/'
+
+POINTS = None
+MISSING_CENTROID = np.asarray([[np.nan]*3])
+
+
+def get_mov_centroids(fix_seg, disp_map):
+    fix_centroids, _ = get_segmentations_centroids(fix_seg[0, ...], ohe=True, expected_lbls=range(0, 28))
+    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)
+
+            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')

Brain_study/Evaluate_network.py

@@ -0,0 +1,187 @@
+import os, sys
+
+import shutil
+
+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.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.layers import AugmentationLayer
+from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC, GeneralizedDICEScore, HausdorffDistanceErosion
+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 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
+
+from Brain_study.data_generator import BatchGenerator
+
+import argparse
+
+from skimage.transform import warp
+import neurite as ne
+
+DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+MODEL_FILE = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/checkpoints/best_model.h5'
+DATA_ROOT_DIR = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/'
+
+OUTPUT_FOLDER_NAME = 'Evaluate'
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-m', '--model', type=str, help='.h5 of the model', default='')
+    parser.add_argument('-d', '--dir', type=str, help='Directory where ./checkpoints/best_model.h5 is located', default='')
+    parser.add_argument('--gpu', type=int, help='GPU', default=0)
+    parser.add_argument('--dataset', type=str, help='Dataset to run predictions on',
+                        default='/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training')
+    parser.add_argument('--erase', type=bool, help='Erase the content of the output folder', default=False)
+    args = parser.parse_args()
+    if args.model != '':
+        assert '.h5' in args.model, 'No checkpoint file provided, use -d/--dir instead'
+        MODEL_FILE = args.model
+        DATA_ROOT_DIR = os.path.split(args.model)[0]
+    elif args.dir != '':
+        assert '.h5' not in args.model, 'Provided checkpoint file, user -m/--model instead'
+        MODEL_FILE = os.path.join(args.dir, 'checkpoints', 'best_model.h5')
+        DATA_ROOT_DIR = args.dir
+    else:
+        raise ValueError("Provide either the model file or the directory ./containing checkpoints/best_model.h5")
+
+    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
+
+    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, OUTPUT_FOLDER_NAME)  # '/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)
+
+    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
+
+    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)
+    augm_model = tf.keras.Model(inputs=input_augm, outputs=augm_layer(input_augm))
+
+    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 + [img_generator.get_data_shape()[2][-1]], num_labels=nb_labels).loss,
+                   HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [img_generator.get_data_shape()[2][-1]]).loss]
+
+    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)
+
+    # Needed for VxmDense type of network
+    warp_segmentation = vxm.networks.Transform(image_output_shape, interp_method='nearest', nb_feats=nb_labels)
+
+    # Record metrics
+    metrics = pd.DataFrame(columns=['File', 'SSIM', 'MS-SSIM', 'MSE', 'DICE', 'HD'])
+    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())
+        network.load_weights(MODEL_FILE, by_name=True)
+        progress_bar = tqdm(enumerate(img_generator, 1), desc='Evaluation', total=len(img_generator))
+        for step, (in_batch, _) in progress_bar:
+            fix_img, mov_img, fix_seg, mov_seg = augm_model.predict(in_batch)
+
+            if network.name == 'vxm_dense_semi_supervised_seg':
+                pred_img, disp_map, pred_seg = network.predict([mov_img, fix_img, mov_seg, fix_seg])    # predict([source, target])
+            else:
+                pred_img, disp_map = network.predict([mov_img, fix_img])
+                pred_seg = warp_segmentation.predict([mov_seg, disp_map])
+
+            # I need the labels to be OHE to compute the segmentation metrics.
+            dice = GeneralizedDICEScore(image_output_shape + [img_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric(fix_seg, pred_seg).eval()
+            hd = HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [img_generator.get_data_shape()[2][-1]]).metric(fix_seg, pred_seg).eval()
+
+            pred_seg = np.argmax(pred_seg, axis=-1)[..., np.newaxis].astype(np.float32)
+            mov_seg = np.argmax(mov_seg, axis=-1)[..., np.newaxis].astype(np.float32)
+            fix_seg = np.argmax(fix_seg, axis=-1)[..., np.newaxis].astype(np.float32)
+
+            mov_coords = np.stack(np.meshgrid(*[np.arange(0, 64)]*3), axis=-1)
+            dest_coords = mov_coords + disp_map[0, ...]
+
+            ssim = StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric(fix_img, pred_img).eval()
+            ms_ssim = MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric(fix_img, pred_img).eval()[0]
+            mse = vxm.losses.MSE().loss(fix_img, pred_img).eval()
+
+            metrics.append({'File': step,
+                            'SSIM': ssim,
+                            'MS-SSIM': ms_ssim,
+                            'MSE': mse,
+                            'DICE': dice,
+                            'HD': hd}, ignore_index=True)
+            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[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[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[0, ...], os.path.join(output_folder, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+
+            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.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}\tDICE: {:.04f}'.format(ssim, dice))
+
+    metrics.to_csv(os.path.join(output_folder, 'metrics.csv'))
+    print('Done')

Brain_study/Evaluate_network__test_fixed.py

@@ -0,0 +1,249 @@
+import os, sys
+
+import shutil
+import time
+
+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.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
+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
+
+DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+MODEL_FILE = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/checkpoints/best_model.h5'
+DATA_ROOT_DIR = '/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE/MS_SSIM/BASELINE_L_ssim__MET_mse_ncc_ssim_162756-29062021/'
+
+
+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='/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training')
+    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:
+        raise ValueError("Provide either the model file or the directory ./containing checkpoints/best_model.h5")
+
+    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()
+
+                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(0, 28))
+                # 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[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[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[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.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}\tDICE: {:.04f}'.format(ssim, dice))
+
+        print('Summary\n=======\n')
+        print(pd.read_csv(metrics_file, sep=';', header=0).mean(axis=0))
+        print('\n=======\n')
+        tf.keras.backend.clear_session()
+        # sess.close()
+        del network
+    print('Done')

Brain_study/MultiTrain_Baseline.py

@@ -0,0 +1,32 @@
+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 Brain_study.Train_Baseline import launch_train
+import argparse
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    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', type=str, help='Similarity metric: mse, ncc, ssim')
+    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)
+
+    args = parser.parse_args()
+
+    print('TRAIN ' + args.dataset)
+    launch_train(dataset_folder=args.dataset,
+                 validation_folder=args.validation,
+                 output_folder=os.path.join(args.output, 'BASELINE_L_{}__MET_mse_ncc_ssim'.format(args.similarity)),
+                 gpu_num=args.gpu,
+                 lr=args.lr,
+                 rw=args.rw,
+                 simil=args.similarity)

Brain_study/MultiTrain_SegGuided.py

@@ -0,0 +1,34 @@
+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 Brain_study.Train_SegmentationGuided import launch_train
+import argparse
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    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', type=str, help='Similarity loss function: mse, ncc, ssim')
+    parser.add_argument('--segmentation', type=str, help='Segmentation loss function: hd, 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=2e-2)
+
+    args = parser.parse_args()
+
+    print('TRAIN ' + args.dataset)
+    launch_train(dataset_folder=args.dataset,
+                 validation_folder=args.validation,
+                 output_folder=os.path.join(args.output, 'SEGGUIDED_Lsim_{}__Lseg_{}__MET_mse_ncc_ssim'.format(args.similarity, args.segmentation)),
+                 gpu_num=args.gpu,
+                 lr=args.lr,
+                 rw=args.rw,
+                 simil=args.similarity,
+                 segm=args.segmentation)

Brain_study/MultiTrain_UW.py

@@ -0,0 +1,38 @@
+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 Brain_study.Train_UncertaintyWeighted import launch_train
+import argparse
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    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='+', type=str, help='Similarity loss function: mse, ncc, ssim', default=[])
+    parser.add_argument('--segmentation', nargs='+', type=str, help='Segmentation loss function: hd, dice', default=[])
+    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)
+
+    args = parser.parse_args()
+
+    output_folder = os.path.join(args.output, 'UW_Lsim_{}__Lseg_{}__MET_mse_ncc_ssim'.format('__'.join(args.similarity),
+                                                                                             '__'.join(args.segmentation)))
+    print('TRAIN ' + args.dataset)
+    launch_train(dataset_folder=args.dataset,
+                 validation_folder=args.validation,
+                 output_folder=output_folder,
+                 gpu_num=args.gpu,
+                 prior_reg_w=args.rw,
+                 lr=args.lr,
+                 simil=args.similarity,
+                 segm=args.segmentation)
+
+

Brain_study/Train_Baseline.py

@@ -0,0 +1,322 @@
+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 numpy as np
+import tensorflow as tf
+
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras import Input
+from tensorflow.keras.models import Model
+from tensorflow.python.keras.utils import Progbar
+from tensorflow.python.framework.errors import InvalidArgumentError
+
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+import pickle
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.losses import NCC, StructuralSimilarity, StructuralSimilarity_simplified
+from DeepDeformationMapRegistration.utils.misc import try_mkdir, DatasetCopy, function_decorator
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+from DeepDeformationMapRegistration.layers import AugmentationLayer
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity, _MSSSIM_WEIGHTS
+
+from Brain_study.data_generator import BatchGenerator
+from Brain_study.utils import SummaryDictionary, named_logs
+
+from tqdm import tqdm
+from datetime import datetime
+
+
+def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim'):
+    assert dataset_folder is not None and output_folder is not None
+
+    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)
+
+    output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+    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 = 16
+    C.BATCH_SIZE = 16 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 = 10000
+
+    aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
+          "GPU: {}\n" \
+          "BATCH SIZE: {}\n" \
+          "LR: {}\n" \
+          "SIMILARITY: {}\n" \
+          "REG. WEIGHT: {}\n" \
+          "EPOCHS: {:d}\n" \
+          "ACCUM. GRAD: {}\n" \
+          "EARLY STOP PATIENCE: {}".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,
+                                           rw,
+                                           C.EPOCHS,
+                                           C.ACCUM_GRADIENT_STEP,
+                                           C.EARLY_STOP_PATIENCE)
+
+    log_file.write(aux)
+    print(aux)
+
+    # Load data
+    # Build data generator
+    data_generator = BatchGenerator(C.TRAINING_DATASET, C.BATCH_SIZE if C.ACCUM_GRADIENT_STEP == 1 else 1, True,
+                                    C.TRAINING_PERC, True, ['none'], 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
+
+    # Config the training sessions
+    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)
+
+    # Build model
+    input_layer_train = Input(shape=train_generator.get_data_shape()[-1], name='input_train')
+    augm_layer_train = 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=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=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)
+    augm_model_valid = Model(inputs=input_layer_valid, outputs=augm_layer_valid(input_layer_valid))
+
+    # Build model
+    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.VxmDense(inshape=image_output_shape,
+                                    nb_unet_features=nb_features,
+                                    int_steps=0)
+
+    # 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)
+
+    # Train
+    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'checkpoints'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'tensorboard'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'history'), exist_ok=True)
+
+    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_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+    #                save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    # CSVLogger(train_log_name, ';'),
+    # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+    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)
+
+    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}
+
+    # Compile the model
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    network.compile(optimizer=optimizer,
+                    loss=losses,
+                    loss_weights=loss_weights,
+                    metrics=metrics)
+
+    callback_tensorboard.set_model(network)
+    callback_best_model.set_model(network)
+    # callback_save_model.set_model(network)
+    callback_early_stop.set_model(network)
+    # TODO: https://towardsdatascience.com/writing-tensorflow-2-custom-loops-438b1ab6eb6c
+
+    summary = SummaryDictionary(network, C.BATCH_SIZE)
+    names = network.metrics_names  # It give both the loss and metric 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():
+        callback_tensorboard.on_train_begin()
+        callback_early_stop.on_train_begin()
+        callback_best_model.on_train_begin()
+        # callback_save_model.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_model.on_epoch_begin(epoch)
+
+            print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
+            print('TRAINING')
+
+            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_tensorboard.on_train_batch_begin(step)
+                callback_best_model.on_train_batch_begin(step)
+                # callback_save_model.on_train_batch_begin(step)
+                callback_early_stop.on_train_batch_begin(step)
+
+                try:
+                    fix_img, mov_img, *_ = 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
+
+                ret = network.train_on_batch(x=(mov_img, fix_img),
+                                             y=(fix_img, fix_img))  # 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_tensorboard.on_train_batch_end(step, named_logs(network, ret))
+                callback_best_model.on_train_batch_end(step, named_logs(network, ret))
+                # callback_save_model.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):
+                # callback_tensorboard.on_test_batch_begin(step)    # This is cursed, don't do it again
+                # callback_early_stop.on_test_batch_begin(step)
+                try:
+                    fix_img, mov_img, *_ = augm_model_valid.predict(in_batch)
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+
+                ret = network.test_on_batch(x=(mov_img, fix_img),
+                                            y=(fix_img, fix_img))
+                # pred_segm = network.register(mov_segm, fix_segm)
+                summary.on_validation_batch_end(ret)
+                # callback_early_stop.on_test_batch_end(step, named_logs(network, ret))
+                # callback_tensorboard.on_test_batch_end(step, named_logs(network, ret))    # This is cursed, don't do it again
+                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_early_stop.on_epoch_end(epoch, epoch_summary)
+            callback_best_model.on_epoch_end(epoch, epoch_summary)
+            # callback_save_model.on_epoch_end(epoch, epoch_summary)
+            print('End of epoch {}: '.format(epoch), ret, '\n')
+
+        callback_tensorboard.on_train_end()
+        # callback_save_model.on_train_end()
+        callback_best_model.on_train_end()
+        callback_early_stop.on_train_end()
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # Check availability before running using 'nvidia-smi'
+
+    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)
+    tf.keras.backend.set_session(tf.Session(config=config))
+
+    launch_train('/mnt/EncryptedData1/Users/javier/vessel_registration/LiTS/None',
+                 'TrainOutput/THESIS/UW_None_mse_ssim_haus', 0, mse=True)

Brain_study/Train_SegmentationGuided.py

@@ -0,0 +1,323 @@
+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 numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras import Input
+from tensorflow.keras.models import Model
+from tensorflow.python.keras.utils import Progbar
+from tensorflow.python.framework.errors import InvalidArgumentError
+
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+import pickle
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.misc import try_mkdir, function_decorator
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.losses import NCC, HausdorffDistanceErosion, GeneralizedDICEScore, StructuralSimilarity_simplified
+from DeepDeformationMapRegistration.layers import AugmentationLayer
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity, _MSSSIM_WEIGHTS
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+
+from Brain_study.data_generator import BatchGenerator
+from Brain_study.utils import SummaryDictionary, named_logs
+
+import time
+
+def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim', segm='hd'):
+    assert dataset_folder is not None and output_folder is not None
+
+    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)
+
+    output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+    os.makedirs(output_folder, exist_ok=True)
+    log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
+    C.TRAINING_DATASET = dataset_folder
+    C.VALIDATION_DATASET = validation_folder
+    C.ACCUM_GRADIENT_STEP = 16
+    C.BATCH_SIZE = 2 if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
+    C.EARLY_STOP_PATIENCE = 10 * (C.ACCUM_GRADIENT_STEP / 2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.LEARNING_RATE = lr
+    C.LIMIT_NUM_SAMPLES = None
+    C.EPOCHS = 10000
+
+    aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
+          "GPU: {}\n" \
+          "BATCH SIZE: {}\n" \
+          "LR: {}\n" \
+          "SIMILARITY: {}\n" \
+          "REG. WEIGHT: {}\n" \
+          "EPOCHS: {:d}\n" \
+          "ACCUM. GRAD: {}\n" \
+          "EARLY STOP PATIENCE: {}".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,
+                                           rw,
+                                           C.EPOCHS,
+                                           C.ACCUM_GRADIENT_STEP,
+                                           C.EARLY_STOP_PATIENCE)
+    log_file.write(aux)
+    print(aux)
+
+    # Load data
+    # Build 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())
+
+    # Config the training sessions
+    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  # https://github.com/tensorflow/tensorflow/issues/30782
+    sess = tf.Session(config=config)
+    tf.keras.backend.set_session(sess)
+
+    # Build model
+    input_layer_augm = Input(shape=train_generator.get_data_shape()[0], name='input_augmentation')
+    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)
+    augm_model = Model(inputs=input_layer_augm, outputs=augm_layer(input_layer_augm))
+
+    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)
+
+    # Compile the model
+    SSIM_KER_SIZE = 5
+    MS_SSIM_WEIGHTS = _MSSSIM_WEIGHTS[:3]
+    MS_SSIM_WEIGHTS /= np.sum(MS_SSIM_WEIGHTS)
+    if simil=='ssim':
+        loss_simil = StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss
+    elif simil=='ms_ssim':
+        loss_simil = MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss
+    elif simil=='ncc':
+        loss_simil = NCC(image_input_shape).loss
+    elif simil=='ms_ssim__ncc' or simil=='ncc__ms_ssim':
+        @function_decorator('MS_SSIM_NCC__loss')
+        def loss_simil(y_true, y_pred):
+            return MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss(y_true, y_pred) + NCC(image_input_shape).loss(y_true, y_pred)
+    elif simil=='ms_ssim__mse' or simil=='mse__ms_ssim':
+        @function_decorator('MS_SSIM_MSE__loss')
+        def loss_simil(y_true, y_pred):
+            return MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss(y_true, y_pred) + vxm.losses.MSE().loss(y_true, y_pred)
+    elif simil=='ssim__ncc' or simil=='ncc__ssim' :
+        @function_decorator('SSIM_NCC__loss')
+        def loss_simil(y_true, y_pred):
+            return StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss(y_true, y_pred) + NCC(image_input_shape).loss(y_true, y_pred)
+    elif simil=='ssim__mse' or simil=='mse__ssim':
+        @function_decorator('SSIM_MSE__loss')
+        def loss_simil(y_true, y_pred):
+            return StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss(y_true, y_pred) + vxm.losses.MSE().loss(y_true, y_pred)
+    else:
+        loss_simil = vxm.losses.MSE().loss
+
+    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
+    else:
+        raise ValueError('No valid value for segm')
+
+    losses = {'transformer': loss_simil,
+            'seg_transformer': loss_segm,
+            'flow': vxm.losses.Grad('l2').loss}
+    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],
+               '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
+                                     ]}
+    metrics_weights = {'transformer': 1,
+                       'seg_transformer': 1,
+                       'flow': rw}
+
+    # Train
+    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'checkpoints'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'tensorboard'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'history'), exist_ok=True)
+
+    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_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+    #                save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    # CSVLogger(train_log_name, ';'),
+    # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+    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)
+
+    # Compile the model
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, lr=C.LEARNING_RATE)
+    network.compile(optimizer=optimizer,
+                    loss=losses,
+                    loss_weights=loss_weights,
+                    metrics=metrics)
+
+    callback_tensorboard.set_model(network)
+    callback_best_model.set_model(network)
+    # callback_save_model.set_model(network)
+    callback_early_stop.set_model(network)
+
+    summary = SummaryDictionary(network, C.BATCH_SIZE, C.ACCUM_GRADIENT_STEP)
+    names = network.metrics_names  # It give both the loss and metric 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():
+        #sess.run(tf.global_variables_initializer())
+        callback_tensorboard.on_train_begin()
+        callback_early_stop.on_train_begin()
+        callback_best_model.on_train_begin()
+        # callback_save_model.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_model.on_epoch_begin(epoch)
+
+            print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
+            print('TRAINING')
+
+            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)
+            t0 = time.time()
+            for step, (in_batch, _) in enumerate(train_generator, 1):
+                #print('Loaded in {} s'.format(time.time() - t0))
+                # callback_tensorboard.on_train_batch_begin(step)
+                callback_best_model.on_train_batch_begin(step)
+                # callback_save_model.on_train_batch_begin(step)
+                callback_early_stop.on_train_batch_begin(step)
+
+                try:
+                    t0 = time.time()
+                    fix_img, mov_img, fix_seg, mov_seg = augm_model.predict(in_batch)
+                    #print('Augmented in {} s'.format(time.time() - t0))
+                    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
+
+                t0 = time.time()
+                ret = network.train_on_batch(x=(mov_img, fix_img, mov_seg),
+                                             y=(fix_img, fix_img, fix_seg))
+                # print("Trained on batch in {} s".format(time.time() - t0))
+
+                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_tensorboard.on_train_batch_end(step, named_logs(network, ret))
+                callback_best_model.on_train_batch_end(step, named_logs(network, ret))
+                # callback_save_model.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))
+                t0 = time.time()
+            print('End of epoch{}: '.format(step), ret, '\n')
+            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):
+                # callback_tensorboard.on_test_batch_begin(step)    # This is cursed, don't do it again
+                # callback_early_stop.on_test_batch_begin(step)
+                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 = network.test_on_batch(x=(mov_img, fix_img, mov_seg),
+                                            y=(fix_img, fix_img, fix_seg))
+                # pred_segm = network.register(mov_segm, fix_segm)
+                summary.on_validation_batch_end(ret)
+                # callback_early_stop.on_test_batch_end(step, named_logs(network, ret))
+                # callback_tensorboard.on_test_batch_end(step, named_logs(network, ret))    # This is cursed, don't do it again
+                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()
+            callback_tensorboard.on_epoch_end(epoch, epoch_summary)
+            callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            callback_best_model.on_epoch_end(epoch, epoch_summary)
+            # callback_save_model.on_epoch_end(epoch, named_logs(network, ret, True))
+
+        callback_tensorboard.on_train_end()
+        # callback_save_model.on_train_end()
+        callback_best_model.on_train_end()
+        callback_early_stop.on_train_end()
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # Check availability before running using 'nvidia-smi'
+
+    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)
+    tf.keras.backend.set_session(tf.Session(config=config))
+
+    launch_train('/mnt/EncryptedData1/Users/javier/Brain_study/ERASE',
+                     'TrainOutput/THESIS/UW_None_mse_ssim_haus',
+                     0)

Brain_study/Train_UncertaintyWeighted.py

@@ -0,0 +1,304 @@
+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 numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras import Input
+from tensorflow.keras.models import Model
+from tensorflow.python.keras.utils import Progbar
+from tensorflow.python.framework.errors import InvalidArgumentError
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+import pickle
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.misc import try_mkdir, DatasetCopy, function_decorator
+from DeepDeformationMapRegistration.networks import WeaklySupervised
+from DeepDeformationMapRegistration.losses import HausdorffDistanceErosion, NCC, StructuralSimilarity_simplified, GeneralizedDICEScore
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity, _MSSSIM_WEIGHTS
+from DeepDeformationMapRegistration.layers import UncertaintyWeighting, AugmentationLayer
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+
+from Brain_study.data_generator import BatchGenerator
+from Brain_study.utils import SummaryDictionary, named_logs
+
+
+def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5e-3, lr=1e-4,
+                 gpu_num=0, simil=['mse'], segm=['dice']):
+    assert dataset_folder is not None and output_folder is not None
+
+    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)
+
+    output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+    # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
+    os.makedirs(output_folder, exist_ok=True)
+    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 = 16
+    C.BATCH_SIZE = 2 if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
+    C.EARLY_STOP_PATIENCE = 10 * (C.ACCUM_GRADIENT_STEP/2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.LEARNING_RATE = lr
+    C.LIMIT_NUM_SAMPLES = None
+    C.EPOCHS = 10000
+
+    aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
+          "GPU: {}\n" \
+          "BATCH SIZE: {}\n" \
+          "LR: {}\n" \
+          "SIMILARITY {:d}: {}\n" \
+          "SEGMENTATION {:d}: {}\n" \
+          "EPOCHS: {:d}" \
+          "ACCUM. GRAD: {}" \
+          "EARLY STOP PATIENCE: {}".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,
+                                           len(simil), ', '.join(simil),
+                                           len(segm), ', '.join(segm),
+                                           C.EPOCHS,
+                                           C.ACCUM_GRADIENT_STEP,
+                                           C.EARLY_STOP_PATIENCE)
+    log_file.write(aux)
+    print(aux)
+
+    # Load data
+    # Build 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())
+
+    # Config the training sessions
+    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)
+
+    # 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=='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)
+
+    # Build augmentation layer model
+    input_layer_augm = Input(shape=train_generator.get_data_shape()[0], name='input_augmentation')
+    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,
+                                   in_img_shape=image_input_shape,
+                                   out_img_shape=image_output_shape,
+                                   only_image=False,
+                                   only_resize=False,
+                                   trainable=False)
+    augmentation_model = Model(inputs=input_layer_augm, outputs=augm_layer(input_layer_augm))
+
+    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 inputs: mov_img, fix_img, mov_seg
+    # Network outputs: pred_img, disp_map, pred_seg
+    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])
+
+    # Train
+    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'checkpoints'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'tensorboard'), exist_ok=True)
+    os.makedirs(os.path.join(output_folder, 'history'), exist_ok=True)
+
+    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_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+    #                save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    # CSVLogger(train_log_name, ';'),
+    # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+    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)
+
+    # Compile the model
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, lr=C.LEARNING_RATE)
+    full_model.compile(optimizer=optimizer, loss=None)
+
+    callback_tensorboard.set_model(full_model)
+    callback_best_model.set_model(network)  # ONLY SAVE THE NETWORK!!!
+    # callback_save_model.set_model(network)
+    callback_early_stop.set_model(full_model)
+    # TODO: https://towardsdatascience.com/writing-tensorflow-2-custom-loops-438b1ab6eb6c
+
+    summary = SummaryDictionary(full_model, C.BATCH_SIZE)
+    names = full_model.metrics_names  # It give both the loss and metric 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():
+        callback_tensorboard.on_train_begin()
+        callback_early_stop.on_train_begin()
+        callback_best_model.on_train_begin()
+        # callback_save_model.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_model.on_epoch_begin(epoch)
+
+            print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
+            print('TRAINING')
+
+            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_tensorboard.on_train_batch_begin(step)
+                callback_best_model.on_train_batch_begin(step)
+                # callback_save_model.on_train_batch_begin(step)
+                callback_early_stop.on_train_batch_begin(step)
+
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augmentation_model.predict(in_batch)
+                    np.nan_to_num(fix_img, copy=False)
+                    np.nan_to_num(mov_img, copy=False)
+                except InvalidArgumentError as err:
+                    print('TF Error : {}'.format(str(err)))
+                    continue
+                # inputs = [mov_img, fix_img, mov_segm, fix_segm, zero_grads]
+                # outputs = [pred_img, flow, pred_segm, loss]
+                ret = full_model.train_on_batch(x=(mov_img, fix_img, mov_seg, fix_seg, zero_grads))
+
+                summary.on_train_batch_end(ret)
+                # callback_tensorboard.on_train_batch_end(step, named_logs(full_model, ret))
+                callback_best_model.on_train_batch_end(step, named_logs(full_model, ret))
+                # callback_save_model.on_train_batch_end(step, named_logs(network, ret))
+                callback_early_stop.on_train_batch_end(step, named_logs(full_model, ret))
+                log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
+                # print(ret, '\n')
+                progress_bar.update(step, zip(names, ret))
+            print('End of epoch{}: '.format(step), ret, '\n')
+            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):
+                # callback_tensorboard.on_test_batch_begin(step)    # This is cursed, don't do it again
+                # callback_early_stop.on_test_batch_begin(step)
+                try:
+                    fix_img, mov_img, fix_seg, mov_seg = augmentation_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))
+                # pred_segm = network.register(mov_segm, fix_segm)
+                summary.on_validation_batch_end(ret)
+                # callback_early_stop.on_test_batch_end(step, named_logs(full_model, ret))
+                # callback_tensorboard.on_test_batch_end(step, named_logs(network, ret))    # This is cursed, don't do it again
+                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_early_stop.on_epoch_end(epoch, epoch_summary)
+            # callback_save_model.on_train_end(epoch, epoch_summary)
+
+        callback_tensorboard.on_train_end()
+        # callback_save_model.on_train_end()
+        callback_best_model.on_train_end()
+        callback_early_stop.on_train_end()
+
+
+if __name__ == '__main__':
+    os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # Check availability before running using 'nvidia-smi'
+
+    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)
+    tf.keras.backend.set_session(tf.Session(config=config))
+
+    launch_train('/mnt/EncryptedData1/Users/javier/Brain_study/ERASE',
+                 'TrainOutput/THESIS/UW_None_mse_ssim_haus',
+                 0)

Brain_study/data_generator.py

@@ -0,0 +1,453 @@
+import warnings
+
+import numpy as np
+from tensorflow import keras
+import os
+import h5py
+import random
+from PIL import Image
+import nibabel as nib
+from nilearn.image import resample_img
+from skimage.exposure import equalize_adapthist
+from scipy.ndimage import zoom
+import tensorflow as tf
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
+from voxelmorph.tf.layers import SpatialTransformer
+from Brain_study.format_dataset import SEGMENTATION_NR2LBL_LUT, SEGMENTATION_LBL2NR_LUT
+
+from tensorflow.python.keras.preprocessing.image import Iterator
+from tensorflow.python.keras.utils import Sequence
+import sys
+
+#import concurrent.futures
+#import multiprocessing as mp
+import time
+
+class BatchGenerator:
+    def __init__(self,
+                 directory,
+                 batch_size,
+                 shuffle=True,
+                 split=0.7,
+                 combine_segmentations=True,
+                 labels=['all'],
+                 directory_val=None):
+        self.file_directory = directory
+        self.batch_size = batch_size
+        self.combine_segmentations = combine_segmentations
+        self.labels = labels
+        self.shuffle = shuffle
+        self.split = split
+
+        if directory_val is None:
+            self.file_list = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(('h5', 'hd5'))]
+            random.shuffle(self.file_list) if self.shuffle else self.file_list.sort()
+            self.num_samples = len(self.file_list)
+            training_samples = self.file_list[:int(self.num_samples * self.split)]
+
+            self.train_iter = BatchIterator(training_samples, batch_size, shuffle, combine_segmentations, labels)
+            if self.split < 1.:
+                validation_samples = list(set(self.file_list) - set(training_samples))
+                self.validation_iter = BatchIterator(validation_samples, batch_size, shuffle, combine_segmentations, ['all'],
+                                                     validation=True)
+            else:
+                self.validation_iter = None
+        else:
+            training_samples = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(('h5', 'hd5'))]
+            random.shuffle(training_samples) if self.shuffle else training_samples.sort()
+
+            validation_samples = [os.path.join(directory_val, f) for f in os.listdir(directory_val) if f.endswith(('h5', 'hd5'))]
+            random.shuffle(validation_samples) if self.shuffle else validation_samples.sort()
+
+            self.num_samples = len(training_samples) + len(validation_samples)
+            self.file_list = training_samples + validation_samples
+
+            self.train_iter = BatchIterator(training_samples, batch_size, shuffle, combine_segmentations, labels)
+            self.validation_iter = BatchIterator(validation_samples, batch_size, shuffle, combine_segmentations, ['all'],
+                                                 validation=True)
+
+    def get_train_generator(self):
+        return self.train_iter
+
+    def get_validation_generator(self):
+        if self.validation_iter is not None:
+            return self.validation_iter
+        else:
+            raise ValueError('No validation iterator. Split must be < 1.0')
+
+    def get_file_list(self):
+        return self.file_list
+
+    def get_data_shape(self):
+        return self.train_iter.get_data_shape()
+
+
+ALL_LABELS = {2., 3., 4., 6., 8., 9., 11., 12., 14., 16., 20., 23., 29., 33., 39., 53., 67., 76., 102., 203., 210.,
+              211., 218., 219., 232., 233., 254., 255.}
+ALL_LABELS_LOC = {label: loc for label, loc in zip(ALL_LABELS, range(0, len(ALL_LABELS)))}
+
+
+class BatchIterator(Sequence):
+    def __init__(self, file_list, batch_size, shuffle, combine_segmentations=True, labels=['all'],
+                                            zero_grads=[64, 64, 64, 3], validation=False, **kwargs):
+        # super(BatchIterator, self).__init__(n=len(file_list),
+        #                                     batch_size=batch_size,
+        #                                     shuffle=shuffle,
+        #                                     seed=None,
+        #                                     **kwargs)
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.file_list = file_list
+        self.combine_segmentations = combine_segmentations
+        self.labels = labels
+        self.zero_grads = zero_grads
+        self.idx_list = np.arange(0, len(self.file_list))
+        self.validation = validation
+        self._initialize()
+        self.shuffle_samples()
+
+    def _initialize(self):
+        with h5py.File(self.file_list[0], 'r') as f:
+            self.image_shape = list(f['image'][:].shape)
+            self.segm_shape = list(f['segmentation'][:].shape)
+            if not self.combine_segmentations:
+                self.segm_shape[-1] = len(f['segmentation_labels'][:]) if self.labels[0].lower() == 'all' else len(self.labels)
+
+        self.batch_shape = self.image_shape.copy()
+        if self.labels[0].lower() != 'none':
+            self.batch_shape[-1] = 2 if self.combine_segmentations else 1 + self.segm_shape[-1]  # +1 because we have the fix and the moving images
+
+            if self.labels[0] != 'all':
+                if isinstance(self.labels[0], str):
+                    self.labels = [SEGMENTATION_LBL2NR_LUT[lbl] for lbl in self.labels]
+
+        self.num_steps = len(self.file_list) // self.batch_size + (1 if len(self.file_list) % self.batch_size else 0)
+        #self.executor = concurrent.futures.ProcessPoolExecutor(max_workers=self.batch_size)
+        #self.mp_pool = mp.Pool(self.batch_size)
+
+    def shuffle_samples(self):
+        np.random.shuffle(self.idx_list)
+
+    def __len__(self):
+        return self.num_steps
+
+    def _filter_segmentations(self, segm, segm_labels):
+        if self.combine_segmentations:
+            # TODO
+            warnings.warn('Cannot select labels when combinine_segmentations options is active')
+        if self.labels[0] != 'all':
+            if set(self.labels).issubset(set(segm_labels)):
+                # If labels in self.labels are in segm
+                idx = [ALL_LABELS_LOC[l] for l in self.labels]
+                segm = segm[..., idx]
+            else:
+                # Else we have to collect those labels that are contained and complete with zeros
+                idx = [ALL_LABELS_LOC[l] for l in list(set(self.labels).intersection(set(segm_labels)))]
+                aux = segm.copy()
+                segm = np.zeros(self.segm_shape)
+                segm[..., :len(idx)] = aux[..., idx]
+                # TODO: leave the zero-ed segmentations before or after the selected labels based on the order
+        return segm
+
+    def _load_sample(self, file_path):
+        with h5py.File(file_path, 'r') as f:
+            img = f['image'][:]
+            segm_labels = f['segmentation_labels'][:]
+            if self.combine_segmentations:
+                segm = f['segmentation'][:]
+            else:
+                segm = f['segmentation_expanded'][:]
+                if segm.shape[-1] != self.segm_shape[-1]:
+                    aux = np.zeros(self.segm_shape)
+                    aux[..., :segm.shape[-1]] = segm     # Ensure the same shape in case there are missing labels in aux
+                    segm = aux
+                # TODO: selection label segm = aux[..., self.labels]  but:
+                #       what if aux does not have a label in self.labels??
+
+        if self.labels[0].lower() != 'none' or self.validation:  # I expect to ask for the segmentations during val
+            segm = self._filter_segmentations(segm, segm_labels)
+
+            if self.validation:
+                ret_val = np.concatenate([img, segm], axis=-1), (img, segm, np.zeros(self.zero_grads))
+            else:
+                ret_val = np.concatenate([img, segm], axis=-1), (img, np.zeros(self.zero_grads))
+        else:
+            ret_val = img, (img, np.zeros(self.zero_grads))
+        return ret_val
+
+    def __getitem__(self, idx):
+        in_batch = list()
+        # out_batch = list()
+
+        batch_idxs = self.idx_list[idx * self.batch_size:(idx + 1) * self.batch_size]
+        file_list = [self.file_list[i] for i in batch_idxs]
+        # if self.batch_size > 1:
+        #     # Multiprocessing to speed up laoding
+        #
+        #     for ret in self.executor.map(self._load_sample, file_list):
+        #         b, i = ret
+        #         in_batch.append(b)
+        #         # out_batch.append(i)
+        # else:
+            # No need for multithreading, we are loading a single file
+        for f in file_list:
+            b, i = self._load_sample(f)
+            in_batch.append(b)
+            # out_batch.append(i)
+
+        in_batch = np.asarray(in_batch)
+        # out_batch = np.asarray(out_batch)
+        return in_batch, in_batch
+
+    def __iter__(self):
+        """Create a generator that iterate over the Sequence."""
+        for item in (self[i] for i in range(len(self))):
+            yield item
+
+    def get_data_shape(self):
+        return self.batch_shape, self.image_shape, self.segm_shape
+
+    def on_epoch_end(self):
+        self.shuffle_samples()
+
+    def get_segmentation_labels(self):
+        if self.combine_segmentations:
+            labels = [1]
+        else:
+            with h5py.File(self.file_list[0], 'r') as f:
+                labels = np.unique(f['segmentation'][:])
+            labels = np.sort(labels)[1:]  # Ignore the background
+        return labels
+
+
+
+
+
+
+'''
+def get_size(obj, seen=None):
+    """Recursively finds size of objects"""
+    size = sys.getsizeof(obj)
+    if seen is None:
+        seen = set()
+    obj_id = id(obj)
+    if obj_id in seen:
+        return 0
+    # Important mark as seen *before* entering recursion to gracefully handle
+    # self-referential objects
+    seen.add(obj_id)
+    if isinstance(obj, dict):
+        size += sum([get_size(v, seen) for v in obj.values()])
+        size += sum([get_size(k, seen) for k in obj.keys()])
+    elif hasattr(obj, '__dict__'):
+        size += get_size(obj.__dict__, seen)
+    elif hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes, bytearray)):
+        size += sum([get_size(i, seen) for i in obj])
+    return size
+
+
+class BatchIterator(Iterator):
+    def __init__(self, generator, file_list, input_shape, output_shape, batch_size, shuffle, all_files_in_batch):
+        self.file_list = file_list
+        self.generator = generator
+        self.input_shape = input_shape
+        self.nr_of_inputs = len(input_shape)
+        self.output_shape = output_shape
+        self.nr_of_outputs = len(output_shape)
+        self.all_files_in_batch = all_files_in_batch
+        self.preload_to_memory = False
+        self.file_cache = {}
+        self.max_cache_size = 10*1024
+        self.verbose = False
+        if self.preload_to_memory:
+            for filename, file_index in self.file_list:
+                file = h5py.File(filename, 'r')
+                inputs = {}
+                for name, data in file['input'].items():
+                    inputs[name] = np.copy(data)
+                self.file_cache[filename] = {'input': inputs, 'output': np.copy(file['output'])}
+                file.close()
+                if get_size(self.file_cache) / (1024*1024) >= self.max_cache_size:
+                    print('File cache has reached limit of', self.max_cache_size, 'MBs')
+                    break
+        epoch_size = len(file_list)
+        if all_files_in_batch:
+            epoch_size = len(file_list) * 10
+        super(BatchIterator, self).__init__(epoch_size, batch_size, shuffle, None)
+
+    def _get_sample(self, index):
+        filename, file_index = self.file_list[index]
+        if filename in self.file_cache:
+            file = self.file_cache[filename]
+        else:
+            file = h5py.File(filename, 'r')
+        inputs = []
+        outputs = []
+        for name, data in file['input'].items():
+            inputs.append(data[file_index, :])
+        for name, data in file['output'].items():
+            if len(data.shape) > 1:
+                outputs.append(data[file_index, :])
+            else:
+                outputs.append(data[file_index])
+        #outputs.append(file['output'][file_index, :]) # TODO fix
+        if filename not in self.file_cache:
+            file.close()
+        return inputs, outputs
+
+    def _get_random_sample_in_file(self, file_index):
+        filename = self.file_list[file_index]
+        file = h5py.File(filename, 'r')
+        x = file['output/0']
+        sample = np.random.randint(0, x.shape[0])
+        #print('Sampling image', sample, 'from file', filename)
+        inputs = []
+        outputs = []
+        for name, data in file['input'].items():
+            inputs.append(data[sample, :])
+        for name, data in file['output'].items():
+            outputs.append(data[file_index, :])
+        #outputs.append(file['output'][sample, :]) # TODO FIX output
+        file.close()
+        return inputs, outputs
+
+    def next(self):
+
+        with self.lock:
+            index_array = next(self.index_generator)
+
+        #print(len(index_array))
+        return self._get_batches_of_transformed_samples(index_array)
+
+    def _get_batches_of_transformed_samples(self, index_array):
+        start_batch = time.time()
+        batches_x = []
+        batches_y = []
+        for input_index in range(self.nr_of_inputs):
+            batches_x.append(np.zeros(tuple([len(index_array)] + list(self.input_shape[input_index]))))
+        for output_index in range(self.nr_of_outputs):
+            batches_y.append(np.zeros(tuple([len(index_array)] + list(self.output_shape[output_index]))))
+
+        timings_sampling = np.zeros((len(index_array,)))
+        timings_transform = np.zeros((len(index_array,)))
+        for batch_index, sample_index in enumerate(index_array):
+            # Have to copy here in order to not modify original data
+            start = time.time()
+            if self.all_files_in_batch:
+                input, output = self._get_random_sample_in_file(batch_index)
+            else:
+                input, output = self._get_sample(sample_index)
+            timings_sampling[batch_index] = time.time() - start
+            start = time.time()
+            input, output = self.generator.transform(input, output)
+            timings_transform[batch_index] = time.time() - start
+
+            #print('inputs', self.nr_of_inputs, len(input))
+            for input_index in range(self.nr_of_inputs):
+                batches_x[input_index][batch_index] = input[input_index]
+            for output_index in range(self.nr_of_outputs):
+                batches_y[output_index][batch_index] = output[output_index]
+
+        elapsed = time.time() - start_batch
+        if self.verbose:
+            print('Time to prepare batch:', round(elapsed,3), 'seconds')
+            print('Sampling mean:', round(timings_sampling.mean(), 3), 'seconds')
+            print('Transform mean:', round(timings_transform.mean(), 3), 'seconds')
+
+        return batches_x, batches_y
+
+
+CLASSIFICATION = 'classification'
+SEGMENTATION = 'segmentation'
+
+
+class BatchGenerator():
+    def __init__(self, filelist, all_files_in_batch=False):
+        self.methods = []
+        self.args = []
+        self.crop_width_to = None
+        self.image_list = []
+        self.input_shape = []
+        self.output_shape = []
+        self.all_files_in_batch = all_files_in_batch
+        self.transforms = []
+
+        if all_files_in_batch:
+            file = h5py.File(filelist[0], 'r')
+            for name, data in file['input'].items():
+                self.input_shape.append(data.shape[1:])
+            for name, data in file['output'].items():
+                self.output_shape.append(data.shape[1:])
+            # TODO fix
+            #self.output_shape.append(file['output'].shape[1:])
+            file.close()
+            self.image_list = filelist
+            return
+
+        # Go through filelist
+        first = True
+        for filename in filelist:
+            samples = None
+            # Open file to see how many samples it has
+            file = h5py.File(filename, 'r')
+            for name, data in file['input'].items():
+                if first:
+                    self.input_shape.append(data.shape[1:])
+                samples = data.shape[0]
+            # TODO fix
+            for name, data in file['output'].items():
+                if first:
+                    self.output_shape.append(data.shape[1:])
+                if samples != data.shape[0]:
+                    raise ValueError()
+            #self.output_shape.append(file['output'].shape[1:])
+            if len(self.output_shape) == 1:
+                self.problem_type = CLASSIFICATION
+            else:
+                self.problem_type = SEGMENTATION
+
+            file.close()
+            if samples is None:
+                raise ValueError()
+            # Append a tuple to image_list for each image consisting of filename and index
+            print(filename, samples)
+            for i in range(samples):
+                self.image_list.append((filename, i))
+            first = False
+
+        print('Image generator with', len(self.image_list), ' image samples created')
+
+    def flow(self, batch_size, shuffle=True):
+
+        return BatchIterator(self, self.image_list, self.input_shape, self.output_shape, batch_size, shuffle, self.all_files_in_batch)
+
+    def transform(self, inputs, outputs):
+        #input = input.astype(np.float32) # TODO
+        #output = output.astype(np.float32)
+        for input_indices, output_indices, transform in self.transforms:
+            transform.randomize()
+            inputs, outputs = transform.transform_all(inputs, outputs, input_indices, output_indices)
+        return inputs, outputs
+
+    def add_transform(self, input_indices: Union[int, List[int], None], output_indices: Union[int, List[int], None], transform: Transform):
+        if type(input_indices) is int:
+            input_indices = [input_indices]
+        if type(output_indices) is int:
+            output_indices = [output_indices]
+
+        self.transforms.append((
+            input_indices,
+            output_indices,
+            transform
+        ))
+
+    def get_size(self):
+        if self.all_files_in_batch:
+            return 10*len(self.image_list)
+        else:
+            return len(self.image_list)
+
+'''

Brain_study/format_dataset.py

@@ -0,0 +1,91 @@
+import h5py
+import nibabel as nib
+from nilearn.image import resample_img
+import os
+import re
+import numpy as np
+from scipy.ndimage import zoom
+from tqdm import tqdm
+
+SEGMENTATION_NR2LBL_LUT = {0: 'background',
+                           2: 'parietal-right-gm',
+                           3: 'lateral-ventricle-left',
+                           4: 'occipital-right-gm',
+                           6: 'parietal-left-gm',
+                           8: 'occipital-left-gm',
+                           9: 'lateral-ventricle-right',
+                           11: 'globus-pallidus-right',
+                           12: 'globus-pallidus-left',
+                           14: 'putamen-left',
+                           16: 'putamen-right',
+                           20: 'brain-stem',
+                           23: 'subthalamic-nucleus-right',
+                           29: 'fornix-left',
+                           33: 'subthalamic-nucleus-left',
+                           39: 'caudate-left',
+                           53: 'caudate-right',
+                           67: 'cerebellum-left',
+                           76: 'cerebellum-right',
+                           102: 'thalamus-left',
+                           203: 'thalamus-right',
+                           210: 'frontal-left-gm',
+                           211: 'frontal-right-gm',
+                           218: 'temporal-left-gm',
+                           219: 'temporal-right-gm',
+                           232: '3rd-ventricle',
+                           233: '4th-ventricle',
+                           254: 'fornix-right',
+                           255: 'csf'}
+SEGMENTATION_LBL2NR_LUT = {v: k for k, v in SEGMENTATION_NR2LBL_LUT.items()}
+
+IMG_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1'
+SEG_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/anatomical_masks'
+
+IMG_NAME_PATTERN = '(.*).nii.gz'
+SEG_NAME_PATTERN = '(.*)_lobes.nii.gz'
+
+OUT_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/test'
+
+if __name__ == '__main__':
+    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()
+
+    os.makedirs(OUT_DIRECTORY, exist_ok=True)
+    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')
+
+        isot_shape = img.shape
+
+        # Resize to 128x128x128
+        img = np.asarray(img.dataobj)
+        img = zoom(img, np.asarray([128]*3) / np.asarray(isot_shape), order=3)
+
+        seg = np.asarray(seg.dataobj)
+        seg = zoom(seg, np.asarray([128]*3) / np.asarray(isot_shape), order=0)
+
+        unique_lbls = np.unique(seg)[1:]     # Omit background
+        seg_expanded = np.tile(np.zeros_like(seg)[..., np.newaxis], (1, 1, 1, len(unique_lbls)))
+        for ch, lbl in enumerate(unique_lbls):
+            seg_expanded[seg == lbl, ch] = 1
+
+        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[..., np.newaxis].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=unique_lbls)
+        h5_file.create_dataset('isotropic_shape', data=isot_shape)
+
+        h5_file.close()
+
+
+
+
+

Brain_study/split_dataset.py

@@ -0,0 +1,68 @@
+import os
+import argparse
+import random
+import warnings
+
+import math
+from shutil import copyfile
+from tqdm import tqdm
+import concurrent.futures
+import numpy as np
+
+
+def copy_file(s_d):
+    s, d = s_d
+    file_name = os.path.split(s)[-1]
+    copyfile(s, os.path.join(d, file_name))
+    return int(os.path.exists(d))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--train', '-t', type=float, default=.70, help='Train percentage. Default: 0.70')
+    parser.add_argument('--validation', '-v', type=float, default=0.15, help='Validation percentage. Default: 0.15')
+    parser.add_argument('--test', '-s', type=float, default=0.15, help='Test percentage. Default: 0.15')
+    parser.add_argument('-d', '--dir', type=str, help='Directory where the data is')
+    parser.add_argument('-f', '--format', type=str, help='Format of the data files. Default: h5', default='h5')
+    parser.add_argument('-r', '--random', type=bool, help='Randomly split the dataset or not. Default: True', default=True)
+
+    args = parser.parse_args()
+
+    assert args.train + args.validation + args.test == 1.0, 'Train+Validation+Test != 1 (100%)'
+
+    file_set = [os.path.join(args.dir, f) for f in os.listdir(args.dir) if f.endswith(args.format)]
+    random.shuffle(file_set) if args.random else file_set.sort()
+
+    num_files = len(file_set)
+    num_validation = math.floor(num_files * args.validation)
+    num_test = math.floor(num_files * args.test)
+    num_train = num_files - num_test - num_validation
+
+    dataset_root, dataset_name = os.path.split(args.dir)
+    dst_train = os.path.join(dataset_root, 'SPLIT_'+dataset_name, 'train_set')
+    dst_validation = os.path.join(dataset_root, 'SPLIT_'+dataset_name, 'validation_set')
+    dst_test = os.path.join(dataset_root, 'SPLIT_'+dataset_name, 'test_set')
+
+    print('OUTPUT INFORMATION\n=============')
+    print('Train:\t\t{}'.format(num_train))
+    print('Validation:\t{}'.format(num_validation))
+    print('Test:\t\t{}'.format(num_test))
+    print('Num. samples\t{}'.format(num_files))
+    print('Path:\t\t', os.path.join(dataset_root, 'SPLIT_'+dataset_name))
+
+    dest = [dst_train] * num_train + [dst_validation] * num_validation + [dst_test] * num_test
+
+    os.makedirs(dst_train, exist_ok=True)
+    os.makedirs(dst_validation, exist_ok=True)
+    os.makedirs(dst_test, exist_ok=True)
+
+    progress_bar = tqdm(zip(file_set, dest), desc='Copying files', total=num_files)
+    with concurrent.futures.ProcessPoolExecutor(max_workers=10) as ex:
+        results = list(tqdm(ex.map(copy_file, zip(file_set, dest)), desc='Copying files', total=num_files))
+
+    num_copies = np.sum(results)
+    if num_copies == num_files:
+        print('Done successfully')
+    else:
+        warnings.warn('Missing files: {}'.format(num_files - num_copies))
+

Brain_study/test_datagenerator.py

@@ -0,0 +1,91 @@
+from Brain_study.data_generator import BatchGenerator
+
+import DeepDeformationMapRegistration.utils.constants as C
+from tqdm import tqdm
+
+from tensorflow import keras
+import tensorflow as tf
+from tensorflow.keras.callbacks import TensorBoard
+import os
+import voxelmorph as vxm
+from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
+from DeepDeformationMapRegistration.losses import NCC, StructuralSimilarity, StructuralSimilarity_simplified
+
+
+def named_logs(model, logs, validation=False):
+    result = {'size': C.BATCH_SIZE} # https://gist.github.com/erenon/91f526302cd8e9d21b73f24c0f9c4bb8#gistcomment-3041181
+    for l in zip(model.metrics_names, logs):
+        k = ('val_' if validation else '') + l[0]
+        result[k] = l[1]
+    return result
+
+if __name__ == '__main__':
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(1)
+
+    C.BATCH_SIZE = 12
+    C.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+    output_folder = "/mnt/EncryptedData1/Users/javier/train_output/Brain_study/ERASE"
+
+    data_generator = BatchGenerator(C.TRAINING_DATASET, C.BATCH_SIZE, True, C.TRAINING_PERC, True, ['none'])
+    train_generator = data_generator.get_train_generator()
+    val_generator = data_generator.get_validation_generator()
+
+    e_iter = tqdm(range(100))
+    t_iter = tqdm(train_generator)
+    v_iter = tqdm(val_generator)
+
+    e_iter.set_description('Epoch')
+    t_iter.set_description('Train')
+    v_iter.set_description('Val')
+    #
+    # for s in e_iter:
+    #     for b in t_iter:
+    #         continue
+    #
+    #     for b in v_iter:
+    #         continue
+
+    # Build model
+    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.VxmDense(inshape=(64, 64, 64),
+                                    nb_unet_features=nb_features,
+                                    int_steps=0)
+
+    d = os.path.join(os.getcwd(), 'tensorboard_test')
+    os.makedirs(d, exist_ok=True)
+    callback_tensorboard = TensorBoard(log_dir=d,
+                                       batch_size=C.BATCH_SIZE, write_images=False, histogram_freq=0, update_freq='epoch',
+                                       write_graph=True,
+                                       write_grads=True)
+
+    losses = {'transformer': StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).loss,
+              'flow': vxm.losses.Grad('l2').loss}
+    metrics = {'transformer': [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric,
+                               tf.keras.losses.MSE],
+               # 'flow': vxm.losses.Grad('l2').loss
+               }
+    loss_weights = {'transformer': 1.,
+                    'flow': 5e-3}
+
+    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+
+    network.compile(optimizer=optimizer,
+                    loss=losses,
+                    loss_weights=loss_weights,
+                    metrics=metrics)
+    callback_tensorboard.set_model(network)
+    dummy = lambda x: named_logs(network, [x, 0, x, 0, 0])
+
+    callback_tensorboard.on_train_begin()
+    for s in e_iter:
+        callback_tensorboard.on_epoch_begin(s)
+
+        for n in range(100):
+            callback_tensorboard.on_train_batch_begin(n)
+            input('Press enter')
+            callback_tensorboard.on_train_batch_end(n, dummy(n))
+
+        callback_tensorboard.on_epoch_end(s, dummy)
+    callback_tensorboard.on_train_end()

Brain_study/utils.py

@@ -0,0 +1,41 @@
+import numpy as np
+import DeepDeformationMapRegistration.utils.constants as C
+
+class SummaryDictionary:
+    def __init__(self, model, batch_size, accumulative_gradients_step=None):
+        self.train_names = model.metrics_names
+        self.val_names = ['val_'+n for n in self.train_names]
+        self.batch_size = batch_size
+        self.acc_grad_step = accumulative_gradients_step
+        self._reset()
+
+    def _reset(self):
+        self.summary_dict = {'size': self.batch_size}
+        if self.acc_grad_step is not None:
+            self.summary_dict = {'accumulative_grad_step': self.acc_grad_step}
+        for k in self.train_names + self.val_names:
+            self.summary_dict[k] = list()
+
+    def on_train_batch_end(self, values):
+        for k, v in zip(self.train_names, values):
+            self.summary_dict[k].append(v)
+
+    def on_validation_batch_end(self, values):
+        for k, v in zip(self.val_names, values):
+            self.summary_dict[k].append(v)
+
+    def on_epoch_end(self):
+        for k, v in self.summary_dict.items():
+            self.summary_dict[k] = np.asarray(v).mean()
+
+        ret_val = self.summary_dict.copy()
+        self._reset()
+        return ret_val
+
+
+def named_logs(model, logs, validation=False):
+    result = {'size': C.BATCH_SIZE} # https://gist.github.com/erenon/91f526302cd8e9d21b73f24c0f9c4bb8#gistcomment-3041181
+    for l in zip(model.metrics_names, logs):
+        k = ('val_' if validation else '') + l[0]
+        result[k] = l[1]
+    return result

EvaluationScripts/Evaluate_class.py

@@ -0,0 +1,184 @@
+from scipy.signal import correlate as cc
+from scipy.spatial.distance import euclidean
+from skimage.metrics import mean_squared_error as mse
+from skimage.metrics import structural_similarity as ssim
+from medpy.metric.binary import dc, hd95
+import numpy as np
+import pandas as pd
+import os
+from DeepDeformationMapRegistration.utils.constants import EPS
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from skimage.transform import resize
+from skimage.measure import regionprops, label
+
+
+def ncc(y_true, y_pred, eps=EPS):
+    f_yt = np.reshape(y_true, [-1])
+    f_yp = np.reshape(y_pred, [-1])
+    mean_yt = np.mean(f_yt)
+    mean_yp = np.mean(f_yp)
+
+    n_f_yt = f_yt - mean_yt
+    n_f_yp = f_yp - mean_yp
+    norm_yt = np.linalg.norm(f_yt, ord=2)
+    norm_yp = np.linalg.norm(f_yp, ord=2)
+    numerator = np.sum(np.multiply(n_f_yt, n_f_yp))
+    denominator = norm_yt * norm_yp + eps
+    return np.divide(numerator, denominator)
+
+
+class EvaluationFigures:
+    def __init__(self, output_folder):
+        pd.set_option('display.max_columns', None)
+        self.__metrics_df = pd.DataFrame(columns=['Name', 'Train_ds', 'Eval_ds', 'MSE', 'NCC', 'SSIM',
+                                                  'DICE_PAR', 'DICE_TUM', 'DICE_VES',
+                                                  'HD95_PAR', 'HD95_TUM', 'HD95_VES', 'TRE'])
+        self.__output_folder = output_folder
+
+    def add_sample(self, name, train_ds, eval_ds, fix_t, par_t, tum_t, ves_t, centroid_t, fix_p, par_p, tum_p, ves_p, centroid_p, scale_transform=None):
+
+        n_fix_t = self.__mean_centred_img(fix_t)
+        n_fix_p = self.__mean_centred_img(fix_p)
+
+        if scale_transform is not None:
+            s_centroid_t = self.__scale_point(centroid_t, scale_transform)
+            s_centroid_p = self.__scale_point(centroid_p, scale_transform)
+        else:
+            s_centroid_t = centroid_t
+            s_centroid_p = centroid_p
+
+        new_row = {'Name': name,
+                   'Train_ds': train_ds,
+                   'Eval_ds': eval_ds,
+                   'MSE': mse(fix_t, fix_p),
+                   'NCC': ncc(n_fix_t, n_fix_p),
+                   'SSIM': ssim(fix_t, fix_p, multichannel=True),
+                   'DICE_PAR': dc(par_p, par_t),
+                   'DICE_TUM': dc(tum_p, tum_t),
+                   'DICE_VES': dc(ves_p, ves_t),
+                   'HD95_PAR': hd95(par_p, par_t) if np.sum(par_p) else 64,
+                   'HD95_TUM': hd95(tum_p, tum_t) if np.sum(tum_p) else 64,
+                   'HD95_VES': hd95(ves_p, ves_t) if np.sum(ves_p) else 64,
+                   'TRE': euclidean(s_centroid_t, s_centroid_p)}
+
+        self.__metrics_df = self.__metrics_df.append(new_row, ignore_index=True)
+
+    @staticmethod
+    def __mean_centred_img(img):
+        return img - np.mean(img)
+
+    @staticmethod
+    def __scale_point(point, scale_matrix):
+        assert scale_matrix.shape == (4, 4), 'Transformation matrix is expected to have shape (4, 4)'
+        aux_aug = np.ones((4,))
+        aux_aug[:3] = point
+        return np.matmul(scale_matrix, aux_aug)[:1]
+
+    def save_metrics(self, dest_folder=None):
+        if dest_folder is None:
+            dest_folder = self.__output_folder
+        self.__metrics_df.to_csv(os.path.join(dest_folder, 'metrics.csv'))
+        self.__metrics_df.to_latex(os.path.join(dest_folder, 'table.txt'), sparsify=True)
+        print('Metrics saved in: ' + os.path.join(dest_folder))
+
+    def print_summary(self):
+        print(self.__metrics_df[['MSE', 'NCC', 'SSIM',
+                                 'DICE_PAR', 'DICE_TUM', 'DICE_VES',
+                                 'HD95_PAR', 'HD95_TUM', 'HD95_VES', 'TRE']].describe())
+
+
+def resize_img_to_original_space(img, bb, first_reshape, original_shape, clip_img=False, flow=False):
+    first_reshape = first_reshape.astype(int)
+    bb = bb.astype(int)
+    original_shape = original_shape.astype(int)
+
+    if flow:
+        # Multiply before resizing to reduce the number of multiplications
+        img = _rescale_flow_values(img, bb, img.shape, first_reshape, original_shape)
+
+    min_i, min_j, min_k, bb_i, bb_j, bb_k = bb
+    max_i = min_i + bb_i
+    max_j = min_j + bb_j
+    max_k = min_k + bb_k
+
+    img_bb = resize(img, (bb_i, bb_j, bb_k))   # Get the original bounding box shape
+
+    # Place the bounding box again in the cubic volume
+    img_copy = np.zeros((*first_reshape, img.shape[-1]) if len(img.shape) > 3 else first_reshape)  # Get channels if any
+    img_copy[min_i:max_i, min_j:max_j, min_k:max_k, ...] = img_bb
+
+    # Now resize to the original shape
+    resized_img = resize(img_copy, original_shape, preserve_range=True, anti_aliasing=False)
+    if clip_img or flow:
+        # clip_mask = np.zeros(img_copy.shape[:3], np.int)
+        # clip_mask[min_i:max_i, min_j:max_j, min_k:max_k] = 1
+        # clip_mask = resize(clip_mask, original_shape, preserve_range=True, anti_aliasing=False)
+        # clip_mask[clip_mask > 0.5] = 1
+        # clip_mask[clip_mask < 1] = 0
+        #
+        # [min_i, min_j, min_k, max_i, max_j, max_k] = regionprops(label(clip_mask))[0].bbox
+        #
+        # resized_img = resized_img[min_i:max_i, min_j:max_j, min_k:max_k, ...]
+
+        # Compute the coordinates of the boundix box in the upsampled volume, instead of resizing a mask image
+        S = resize_transformation(img.shape, bb=None, first_reshape=first_reshape, original_shape=original_shape, translate=True)
+        bb_coords = np.asarray([[min_i, min_j, min_k], [max_i, max_j, max_k]])
+        bb_coords = np.hstack([bb_coords, np.ones((2, 1))])
+
+        upsamp_bbox_coords = np.around(np.matmul(S, bb_coords.T)[:-1, :].T).astype(np.int)
+        min_i = upsamp_bbox_coords[0][0]
+        min_j = upsamp_bbox_coords[0][1]
+        min_k = upsamp_bbox_coords[0][2]
+        max_i = upsamp_bbox_coords[1][0]
+        max_j = upsamp_bbox_coords[1][1]
+        max_k = upsamp_bbox_coords[1][2]
+        resized_img = resized_img[min_i:max_i, min_j:max_j, min_k:max_k, ...]
+
+        if flow:
+            # Return also the origin of the bb in the resized volume for the following interpolation
+            return resized_img, np.asarray([min_i, min_j, min_k])
+    return resized_img
+    # This is supposed to be an isotropic image with voxel size 1 mm
+
+
+def _rescale_flow_values(flow, bb, current_img_shape, first_reshape, original_shape):
+    S = resize_transformation(current_img_shape, bb, first_reshape, original_shape, translate=False)
+
+    [si, sj, sk] = np.diag(S[:3, :3])
+    flow[..., 0] *= si
+    flow[..., 1] *= sj
+    flow[..., 2] *= sk
+
+    return flow
+
+
+def resize_pts_to_original_space(pt, bb, current_img_shape, first_reshape, original_shape):
+    T = resize_transformation(current_img_shape, bb, first_reshape, original_shape)
+    if len(pt.shape) > 1:
+        pt_aug = np.ones((4, pt.shape[0]))
+        pt_aug[0:3, :] = pt.T
+    else:
+        pt_aug = np.ones((4,))
+        pt_aug[0:3] = pt
+    trf_pt = np.matmul(T, pt_aug)[:-1, ...].T
+
+    return trf_pt
+
+
+def resize_transformation(current_img_shape, bb=None, first_reshape=None, original_shape=None, translate=True):
+    first_reshape = first_reshape.astype(int)
+    original_shape = original_shape.astype(int)
+
+    first_resize_trf = np.eye(4)
+    if bb is not None:
+        bb = bb.astype(int)
+        min_i, min_j, min_k, bb_i, bb_j, bb_k = bb
+        np.fill_diagonal(first_resize_trf, [bb_i / current_img_shape[0], bb_j / current_img_shape[1], bb_k / current_img_shape[2], 1])
+        if translate:
+            first_resize_trf[:3, -1] = np.asarray([min_i, min_j, min_k])
+
+    original_resize_trf = np.eye(4)
+    np.fill_diagonal(original_resize_trf, [original_shape[0] / first_reshape[0], original_shape[1] / first_reshape[1], original_shape[2] / first_reshape[2], 1])
+
+    return np.matmul(original_resize_trf, first_resize_trf)
+