Updating latest changes

Brain_study/ABSTRACT/figures.py

@@ -0,0 +1,158 @@
+import os
+import argparse
+import re
+import warnings
+
+import nibabel as nib
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from matplotlib.colors import ListedColormap, LinearSegmentedColormap
+
+segm_cm = cm.get_cmap('Dark2', 256)
+segm_cm = segm_cm(np.linspace(0, 1, 28))
+segm_cm[0, :] = np.asarray([0, 0, 0, 0])
+segm_cm = ListedColormap(segm_cm)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-d', '--dir', type=str, help='Directories where the models are stored', default=None)
+    parser.add_argument('-o', '--output', type=str, help='Output directory', default=os.getcwd())
+    parser.add_argument('--overwrite', type=bool, default=True)
+    args = parser.parse_args()
+    assert args.dir is not None, "No directories provided. Stopping"
+
+    list_fix_img = list()
+    list_mov_img = list()
+    list_fix_seg = list()
+    list_mov_seg = list()
+    list_pred_img = list()
+    list_pred_seg = list()
+    print('Fetching data...')
+    for r, d, f in os.walk(args.dir):
+        if os.path.split(r)[1] == 'Evaluation_paper':
+            for name in f:
+                if re.search('^050', name) and name.endswith('nii.gz'):
+                    if re.search('fix_img', name) and name.endswith('nii.gz'):
+                        list_fix_img.append(os.path.join(r, name))
+                    elif re.search('mov_img', name):
+                        list_mov_img.append(os.path.join(r, name))
+                    elif re.search('fix_seg', name):
+                        list_fix_seg.append(os.path.join(r, name))
+                    elif re.search('mov_seg', name):
+                        list_mov_seg.append(os.path.join(r, name))
+                    elif re.search('pred_img', name):
+                        list_pred_img.append(os.path.join(r, name))
+                    elif re.search('pred_seg', name):
+                        list_pred_seg.append(os.path.join(r, name))
+
+    # Figure: all coronal views
+    # Fix img | Mov img
+    # BASELINE 1 | BASELINE 2 | SEGGUIDED
+    # UW 1 | UW 2 | UW 3
+    list_fix_img.sort()
+    list_fix_seg.sort()
+    list_mov_img.sort()
+    list_mov_seg.sort()
+    list_pred_img.sort()
+    list_pred_seg.sort()
+    print('Making Test_data.png...')
+    selected_slice = 30
+    fix_img = np.asarray(nib.load(list_fix_img[0]).dataobj)[..., selected_slice, 0]
+    mov_img = np.asarray(nib.load(list_mov_img[0]).dataobj)[..., selected_slice, 0]
+    fix_seg = np.asarray(nib.load(list_fix_seg[0]).dataobj)[..., selected_slice, 0]
+    mov_seg = np.asarray(nib.load(list_mov_seg[0]).dataobj)[..., selected_slice, 0]
+
+    fig, ax = plt.subplots(nrows=1, ncols=4, figsize=(9, 3), dpi=200)
+
+    for i, (img, title) in enumerate(zip([(fix_img, fix_seg), (mov_img, mov_seg)],
+                                         [('Fixed image', 'Fixed Segms.'), ('Moving image', 'Moving Segms.')])):
+
+        ax[i].imshow(img[0], origin='lower', cmap='Greys_r')
+        ax[i+2].imshow(img[0], origin='lower', cmap='Greys_r')
+        ax[i+2].imshow(img[1], origin='lower', cmap=segm_cm, alpha=0.6)
+
+        ax[i].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+        ax[i+2].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+
+        ax[i].set_xlabel(title[0], fontsize=16)
+        ax[i+2].set_xlabel(title[1], fontsize=16)
+
+    plt.tight_layout()
+    if not args.overwrite and os.path.exists(os.path.join(args.output, 'Test_data.png')):
+        warnings.warn('File Test_data.png already exists. Skipping')
+    else:
+        plt.savefig(os.path.join(args.output, 'Test_data.png'), format='png')
+    plt.close()
+
+    print('Making Pred_data.png...')
+    fig, ax = plt.subplots(nrows=2, ncols=6, figsize=(9, 3), dpi=200)
+
+    for i, (pred_img_path, pred_seg_path) in enumerate(zip(list_pred_img, list_pred_seg)):
+        img = np.asarray(nib.load(pred_img_path).dataobj)[..., selected_slice, 0]
+        seg = np.asarray(nib.load(pred_seg_path).dataobj)[..., selected_slice, 0]
+
+        ax[0, i].imshow(img, origin='lower', cmap='Greys_r')
+        ax[1, i].imshow(img, origin='lower', cmap='Greys_r')
+        ax[1, i].imshow(seg, origin='lower', cmap=segm_cm, alpha=0.6)
+
+        ax[0, i].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+        ax[1, i].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+
+        model = re.search('((UW|SEGGUIDED|BASELINE).*)_{2,}MET', pred_img_path).group(1).rstrip('_')
+        model = model.replace('_Lsim', ' ')
+        model = model.replace('_Lseg', ' ')
+        model = model.replace('_L', ' ')
+        model = model.replace('_', ' ')
+        model = model.upper()
+        model = ' '.join(model.split())
+
+        ax[1, i].set_xlabel(model, fontsize=9)
+    plt.tight_layout()
+    if not args.overwrite and os.path.exists(os.path.join(args.output, 'Pred_data.png')):
+        warnings.warn('File Pred_data.png already exists. Skipping')
+    else:
+        plt.savefig(os.path.join(args.output, 'Pred_data.png'), format='png')
+    plt.close()
+
+    print('Making Pred_data_large.png...')
+    fig, ax = plt.subplots(nrows=2, ncols=8, figsize=(9, 3), dpi=200)
+    list_pred_img = [list_mov_img[0]] + list_pred_img
+    list_pred_img = [list_fix_img[0]] + list_pred_img
+    list_pred_seg = [list_mov_seg[0]] + list_pred_seg
+    list_pred_seg = [list_fix_seg[0]] + list_pred_seg
+
+    for i, (pred_img_path, pred_seg_path) in enumerate(zip(list_pred_img, list_pred_seg)):
+        img = np.asarray(nib.load(pred_img_path).dataobj)[..., selected_slice, 0]
+        seg = np.asarray(nib.load(pred_seg_path).dataobj)[..., selected_slice, 0]
+
+        ax[0, i].imshow(img, origin='lower', cmap='Greys_r')
+        ax[1, i].imshow(img, origin='lower', cmap='Greys_r')
+        ax[1, i].imshow(seg, origin='lower', cmap=segm_cm, alpha=0.6)
+
+        ax[0, i].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+        ax[1, i].tick_params(axis='both', which='both', bottom=False, left=False, labelleft=False, labelbottom=False)
+
+        if i > 1:
+            model = re.search('((UW|SEGGUIDED|BASELINE).*)_{2,}MET', pred_img_path).group(1).rstrip('_')
+            model = model.replace('_Lsim', ' ')
+            model = model.replace('_Lseg', ' ')
+            model = model.replace('_L', ' ')
+            model = model.replace('_', ' ')
+            model = model.upper()
+            model = ' '.join(model.split())
+        elif i == 0:
+            model = 'Moving image'
+        else:
+            model = 'Fixed image'
+
+        ax[1, i].set_xlabel(model, fontsize=7)
+    plt.tight_layout()
+    if not args.overwrite and os.path.exists(os.path.join(args.output, 'Pred_data_large.png')):
+        warnings.warn('File Pred_data.png already exists. Skipping')
+    else:
+        plt.savefig(os.path.join(args.output, 'Pred_data_large.png'), format='png')
+    plt.close()
+
+    print('...done!')

Brain_study/ABSTRACT/format_tables_abstract.py

@@ -0,0 +1,111 @@
+import warnings
+
+import pandas as pd
+import os
+import argparse
+import re
+import shutil
+
+DICT_MODEL_NAMES = {'BASELINE': 'BL',
+                    'SEGGUIDED': 'SG',
+                    'UW': 'UW'}
+
+DICT_METRICS_NAMES = {'NCC': 'N',
+                      'SSIM': 'S',
+                      'DICE': 'D',
+                      'DICE MACRO': 'D',
+                      'HD': 'H', }
+
+
+def row_name(in_path: str):
+    model = re.search('((UW|SEGGUIDED|BASELINE).*)_\d', in_path).group(1).rstrip('_')
+    model = model.replace('_Lsim', '')
+    model = model.replace('_Lseg', '')
+    model = model.replace('_L', '')
+    model = model.replace('_', ' ')
+    model = model.upper()
+    elements = model.split()
+    model = elements[0]
+    metrics = list()
+    model = DICT_MODEL_NAMES[model]
+    for m in elements[1:]:
+        if m != 'MACRO':
+            metrics.append(DICT_METRICS_NAMES[m])
+
+    return '{}-{}'.format(model, ''.join(metrics))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('-d', '--dir', nargs='+', type=str, help='List of directories where metrics.csv file is',
+                        default=None)
+    parser.add_argument('-o', '--output', type=str, help='Output directory', default=os.getcwd())
+    parser.add_argument('--overwrite', type=bool, default=True)
+    parser.add_argument('--filename', type=str, help='Output file name', default='metrics')
+    parser.add_argument('--removemetrics', nargs='+', type=str, default=None)
+    args = parser.parse_args()
+    assert args.dir is not None, "No directories provided. Stopping"
+
+    if len(args.dir) == 1:
+        list_files = list()
+        for r, d, f in os.walk(args.dir[0]):
+            for name in f:
+                if 'metrics.csv' == name:  # and os.path.split(r)[1] == 'Evaluation_paper':
+                    list_files.append(os.path.join(r, name))
+    else:
+        list_files = [os.path.join(d, 'metrics.csv') for d in args.dir]
+
+        for d in list_files:
+            assert os.path.exists(d), "Missing metrics.csv file in: " + os.path.split(d)[0]
+
+    print('Metric files found: {}'.format(list_files))
+
+    dataframes = list()
+    if len(list_files):
+        for d in list_files:
+            df = pd.read_csv(d, sep=';', header=0)
+            model = row_name(d)
+
+            df.insert(0, "Model", model)
+            df.drop(columns=list(df.filter(regex='Unnamed')), inplace=True)
+            df.drop(columns=['File', 'MSE', 'No_missing_lbls'], inplace=True)
+            dataframes.append(df)
+
+        full_table = pd.concat(dataframes)
+        if args.removemetrics is not None:
+            full_table = full_table.drop(columns=args.removemetrics)
+        mean_table = full_table.copy()
+        # mean_table.insert(column='Type', value='Avg.', loc=1)
+        # mean_table = mean_table.groupby(['Type', 'Model']).mean().round(3)
+        mean_table = mean_table.groupby(['Model'])
+        hd95 = mean_table.HD.quantile(0.95).map('{:.2f}'.format)
+        mean_table = mean_table.mean().round(3)
+
+        std_table = full_table.copy()
+        # std_table.insert(column='Type', value='STD', loc=1)
+        # std_table = std_table.groupby(['Type', 'Model']).std().round(3)
+        std_table = std_table.groupby(['Model']).std().round(3)
+
+        # metrics_table = pd.concat([mean_table, std_table]).swaplevel(axis='rows')
+        metrics_table = mean_table.applymap('{:.2f}'.format) + u"\u00B1" + std_table.applymap('{:.2f}'.format)
+        time_col = metrics_table.pop('Time')
+        metrics_table.insert(len(metrics_table.columns), 'Time', time_col)
+        metrics_table.insert(5, 'HD 95%ile', hd95)
+
+        metrics_file = os.path.join(args.output, args.filename + '.tex')
+        if os.path.exists(metrics_file) and args.overwrite:
+            shutil.rmtree(metrics_file, ignore_errors=True)
+            metrics_table.to_latex(metrics_file,
+                                   column_format='l' + 'c' * len(metrics_table.columns),
+                                   caption='Average and standard deviation of the metrics: MSE, NCC, SSIM, DICE and HD. As well as the number of missing labels in the predicted images.')
+        elif os.path.exists(metrics_file):
+            warnings.warn('File {} already exists. Skipping'.format(metrics_file))
+        else:
+            metrics_table.to_latex(metrics_file,
+                                   column_format='l' + 'c' * len(metrics_table.columns),
+                                   caption='Average and standard deviation of the metrics: MSE, NCC, SSIM, DICE and HD. As well as the number of missing labels in the predicted images.')
+
+        print('Done')
+    else:
+        print('No files found in {}!'.format(args.dir))

Brain_study/Build_test_set.py

@@ -18,7 +18,8 @@ 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 DeepDeformationMapRegistration.utils.misc import get_segmentations_centroids, DisplacementMapInterpolator
+
 from tqdm import tqdm
 
 from Brain_study.data_generator import BatchGenerator
@@ -37,9 +38,15 @@ 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')
+def get_mov_centroids(fix_seg, disp_map, nb_labels=28, exclude_background_lbl=False, brain_study=True, dm_interp=None):
+    if exclude_background_lbl:
+        fix_centroids, _ = get_segmentations_centroids(fix_seg[0, ..., 1:], ohe=True, expected_lbls=range(1, nb_labels), brain_study=brain_study)
+    else:
+        fix_centroids, _ = get_segmentations_centroids(fix_seg[0, ...], ohe=True, expected_lbls=range(1, nb_labels), brain_study=brain_study)
+    if dm_interp is None:
+        disp = griddata(POINTS, disp_map.reshape([-1, 3]), fix_centroids, method='linear')
+    else:
+        disp = dm_interp(disp_map, fix_centroids, backwards=False)
     return fix_centroids, fix_centroids + disp, disp
 
 
@@ -50,6 +57,7 @@ if __name__ == '__main__':
     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)
+    parser.add_argument('--output_shape', help='If an int, a cubic shape is presumed. Otherwise provide it as a space separated sequence', nargs='+', default=128)
     args = parser.parse_args()
 
     assert args.dataset != '', "Missing original dataset dataset"
@@ -70,12 +78,20 @@ if __name__ == '__main__':
     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'])
+    data_generator = BatchGenerator(DATASET, 1, False, 1.0, False, ['all'], return_isotropic_shape=True)
 
     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
+
+    if isinstance(args.output_shape, int):
+        image_output_shape = [args.output_shape] * 3
+    elif isinstance(args.output_shape, list):
+        assert len(args.output_shape) == 3, 'Invalid output shape, expected three values and got {}'.format(len(args.output_shape))
+        image_output_shape = [int(s) for s in args.output_shape]
+    else:
+        raise ValueError('Invalid output_shape. Must be an int or a space-separated sequence of ints')
+    print('Scaling to: ', image_output_shape)
     # Build model
 
     xx = np.linspace(0, image_output_shape[0], image_output_shape[0], endpoint=False)
@@ -102,19 +118,27 @@ if __name__ == '__main__':
                                    return_displacement_map=True)
     augm_model = tf.keras.Model(inputs=input_augm, outputs=augm_layer(input_augm))
 
+    fix_img_ph = tf.placeholder(dtype=tf.float32, shape=[1,] + image_input_shape + [1+nb_labels,], name='fix_image')
+
+    augmentation_pipeline = augm_model(fix_img_ph)
+
     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)
 
+    dm_interp = DisplacementMapInterpolator(image_output_shape, 'griddata', step=8)
+
     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)
+        for step, (in_batch, _, isotropic_shape) in progress_bar:
+            # fix_img, mov_img, fix_seg, mov_seg, disp_map = augm_model.predict(in_batch)
+            fix_img, mov_img, fix_seg, mov_seg, disp_map = sess.run(augmentation_pipeline,
+                                                                    feed_dict={'fix_image:0': in_batch})
 
-            fix_centroids, mov_centroids, disp_centroids = get_mov_centroids(fix_seg, disp_map)
+            fix_centroids, mov_centroids, disp_centroids = get_mov_centroids(fix_seg, disp_map, nb_labels, dm_interp=dm_interp)
 
             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))
@@ -129,7 +153,7 @@ if __name__ == '__main__':
                 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)
-
+                f.create_dataset('isotropic_shape', data=np.squeeze(isotropic_shape))
             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)

Brain_study/Evaluate_network.py

@@ -23,6 +23,7 @@ from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplifie
 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 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
@@ -147,9 +148,9 @@ if __name__ == '__main__':
             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)
+            pred_seg = segmentation_ohe_to_cardinal(pred_seg).astype(np.float32)
+            mov_seg = segmentation_ohe_to_cardinal(mov_seg).astype(np.float32)
+            fix_seg = segmentation_ohe_to_cardinal(fix_seg).astype(np.float32)
 
             mov_coords = np.stack(np.meshgrid(*[np.arange(0, 64)]*3), axis=-1)
             dest_coords = mov_coords + disp_map[0, ...]
@@ -178,8 +179,8 @@ if __name__ == '__main__':
             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)
-            plot_predictions(fix_seg, mov_seg, disp_map, pred_seg, os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False)
+            plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, seg_batches=[fix_seg, mov_seg, pred_seg], filename=os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False)
+            plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, filename=os.path.join(output_folder, '{:03d}_figures_img.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))

Brain_study/Evaluate_network__test_fixed.py

@@ -18,20 +18,23 @@ import pandas as pd
 import voxelmorph as vxm
 
 import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.operators import min_max_norm, safe_medpy_metric
 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 resize_displacement_map, scale_transformation
 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 scipy.ndimage import zoom
 from tqdm import tqdm
-
+import medpy.metric as medpy_metrics
 import h5py
 import re
 from Brain_study.data_generator import BatchGenerator
+from voxelmorph.tf.layers import SpatialTransformer
 
 import argparse
 
@@ -49,9 +52,10 @@ if __name__ == '__main__':
     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')
+                        default='/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/test')
     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'
@@ -83,8 +87,8 @@ if __name__ == '__main__':
     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)
+    sess = tf.compat.v1.Session(config=config)
+    tf.compat.v1.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,
@@ -101,10 +105,10 @@ if __name__ == '__main__':
                    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')
+    fix_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='fix_img')
+    pred_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='pred_img')
+    fix_seg_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, nb_labels), name='fix_seg')
+    pred_seg_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 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)
@@ -118,7 +122,7 @@ if __name__ == '__main__':
     # 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')
+    dm_interp = DisplacementMapInterpolator(image_output_shape, 'griddata', step=4)
 
     for MODEL_FILE, DATA_ROOT_DIR in zip(MODEL_FILE_LIST, DATA_ROOT_DIR_LIST):
         print('MODEL LOCATION: ', MODEL_FILE)
@@ -171,6 +175,8 @@ if __name__ == '__main__':
                     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'][:]
+                    isotropic_shape = f['isotropic_shape'][:]
+                    voxel_size = np.divide(fix_img.shape[1:-1], isotropic_shape)
 
                 if network.name == 'vxm_dense_semi_supervised_seg':
                     t0 = time.time()
@@ -182,29 +188,42 @@ if __name__ == '__main__':
                     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_img = min_max_norm(pred_img)
+                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(1, nb_labels + 1))
                 # 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
 
+                # Up sample the segmentation masks to isotropic resolution
+                zoom_factors = np.diag(scale_transformation(image_output_shape, isotropic_shape))
+                pred_seg_isot = zoom(pred_seg[0, ...], zoom_factors, order=0)[np.newaxis, ...]
+                fix_seg_isot = zoom(fix_seg[0, ...], zoom_factors, order=0)[np.newaxis, ...]
+
+                pred_img_isot = zoom(pred_img[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+                fix_img_isot = zoom(fix_img[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+
                 # 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})
+                # dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf], {'fix_seg:0': fix_seg, 'pred_seg:0': pred_seg})
+                dice = np.mean([medpy_metrics.dc(pred_seg_isot[0, ..., l], fix_seg_isot[0, ..., l]) / np.sum(fix_seg_isot[0, ..., l]) for l in range(nb_labels)])
+                hd = np.mean(safe_medpy_metric(pred_seg_isot[0, ...], fix_seg_isot[0, ...], nb_labels, medpy_metrics.hd, {'voxelspacing': voxel_size}))
+                dice_macro = np.mean([medpy_metrics.dc(pred_seg_isot[0, ..., l], fix_seg_isot[0, ..., l]) for l in range(nb_labels)])
 
                 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]
+                ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf], {'fix_img:0': fix_img_isot, 'pred_img:0': pred_img_isot})
+                ssim = np.mean(ssim)    # returns a list of values, which correspond to the ssim of each patch
+                ms_ssim = ms_ssim[0]    # returns an array of shape (1,)
 
                 # 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
+                # upsample_scale = 128 / 64
+                fix_centroids_isotropic = fix_centroids * voxel_size
                 # mov_centroids_isotropic = mov_centroids * upsample_scale
-                pred_centroids_isotropic = pred_centroids * upsample_scale
+                pred_centroids_isotropic = pred_centroids * voxel_size
 
-                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)
+                # 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)])
@@ -235,14 +254,20 @@ if __name__ == '__main__':
                 # 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)
-                plot_predictions(fix_seg, mov_seg, disp_map, pred_seg, os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False)
-                save_disp_map_img(disp_map, 'Displacement map', os.path.join(output_folder, '{:03d}_disp_map_fig.png'.format(step)), show=False)
+                plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, seg_batches=[fix_seg_card, mov_seg_card, pred_seg_card], filename=os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False, step=16)
+                plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, filename=os.path.join(output_folder, '{:03d}_figures_img.png'.format(step)), show=False, step=16)
+                save_disp_map_img(disp_map, 'Displacement map', os.path.join(output_folder, '{:03d}_disp_map_fig.png'.format(step)), show=False, step=16)
 
                 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))
+        metrics_df = pd.read_csv(metrics_file, sep=';', header=0)
+        print('\nAVG:\n')
+        print(metrics_df.mean(axis=0))
+        print('\nSTD:\n')
+        print(metrics_df.std(axis=0))
+        print('\nHD95perc:\n')
+        print(metrics_df['HD'].describe(percentiles=[.95]))
         print('\n=======\n')
         tf.keras.backend.clear_session()
         # sess.close()

Brain_study/MultiTrain_config.py

@@ -62,6 +62,11 @@ if __name__ == '__main__':
     except KeyError as e:
         head = [16, 16]
 
+    try:
+        resume_checkpoint = trainConfig['resumeCheckpoint']
+    except KeyError as e:
+        resume_checkpoint = None
+
     launch_train(dataset_folder=datasetConfig['train'],
                  validation_folder=datasetConfig['validation'],
                  output_folder=output_folder,
@@ -75,4 +80,5 @@ if __name__ == '__main__':
                  early_stop_patience=eval(trainConfig['earlyStopPatience']),
                  unet=unet,
                  head=head,
+                 resume=resume_checkpoint,
                  **loss_config)

Brain_study/Train_Baseline.py

@@ -8,7 +8,7 @@ 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.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
 from tensorflow.python.keras.utils import Progbar
@@ -32,11 +32,12 @@ from Brain_study.utils import SummaryDictionary, named_logs
 
 from tqdm import tqdm
 from datetime import datetime
+import re
 
 
 def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim',
                  acc_gradients=16, batch_size=1, max_epochs=10000, early_stop_patience=1000, image_size=64,
-                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
@@ -46,7 +47,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
     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')
@@ -141,6 +151,26 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                     nb_unet_features=nb_features,
                                     int_steps=0)
     network.summary(line_length=150)
+
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
+
     # Losses and loss weights
     SSIM_KER_SIZE = 5
     MS_SSIM_WEIGHTS = _MSSSIM_WEIGHTS[:3]
@@ -178,14 +208,14 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
 
     # 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, 'checkpoints'), exist_ok=True)  # exist_ok=True leaves directory unaltered.
     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')
+    callback_save_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.{epoch:05d}-{val_loss:.2f}.h5'),
+                                            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'),
@@ -194,6 +224,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                        write_graph=True, write_grads=True
                                        )
     callback_early_stop = EarlyStopping(monitor='val_loss', verbose=1, patience=C.EARLY_STOP_PATIENCE, min_delta=0.00001)
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
 
     losses = {'transformer': loss_fnc,
               'flow': vxm.losses.Grad('l2').loss}
@@ -215,8 +246,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
 
     callback_tensorboard.set_model(network)
     callback_best_model.set_model(network)
-    # callback_save_model.set_model(network)
+    callback_save_model.set_model(network)
     callback_early_stop.set_model(network)
+    callback_lr.set_model(network)
     # TODO: https://towardsdatascience.com/writing-tensorflow-2-custom-loops-438b1ab6eb6c
 
     summary = SummaryDictionary(network, C.BATCH_SIZE)
@@ -226,13 +258,13 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
         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_save_model.on_train_begin()
+        for epoch in range(resume_epoch, 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)
-
+            callback_save_model.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print('TRAINING')
 
@@ -241,9 +273,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
             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_save_model.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
-
+                callback_lr.on_train_batch_begin(step)
                 try:
                     fix_img, mov_img, *_ = augm_model_train.predict(in_batch)
                     np.nan_to_num(fix_img, copy=False)
@@ -273,8 +305,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                 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_save_model.on_train_batch_end(step, named_logs(network, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(network, ret))
+                callback_lr.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()
@@ -309,13 +342,15 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
             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)
+            callback_save_model.on_epoch_end(epoch, epoch_summary)
+            callback_lr.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_save_model.on_train_end()
         callback_best_model.on_train_end()
         callback_early_stop.on_train_end()
+        callback_lr.on_train_end()
 
 
 if __name__ == '__main__':

Brain_study/Train_SegmentationGuided.py

@@ -1,11 +1,12 @@
 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.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
 from tensorflow.python.keras.utils import Progbar
@@ -30,11 +31,14 @@ from Brain_study.utils import SummaryDictionary, named_logs
 
 import time
 import warnings
+import re
+import tqdm
 
 
 def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim', segm='hd',
                  acc_gradients=16, batch_size=1, max_epochs=10000, early_stop_patience=1000, image_size=64,
-                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
+
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
@@ -44,7 +48,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
     os.makedirs(output_folder, exist_ok=True)
     log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
     C.TRAINING_DATASET = dataset_folder
@@ -85,6 +98,10 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                     directory_val=C.VALIDATION_DATASET)
 
     train_generator = data_generator.get_train_generator()
+
+    # for l in tqdm.tqdm(train_generator, smoothing=0):
+    #     pass
+    # exit()
     validation_generator = data_generator.get_validation_generator()
 
     image_input_shape = train_generator.get_data_shape()[1][:-1]
@@ -96,7 +113,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     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
+    # config.allow_soft_placement = False  # https://github.com/tensorflow/tensorflow/issues/30782
     sess = tf.Session(config=config)
     tf.keras.backend.set_session(sess)
 
@@ -128,6 +145,27 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                                      int_steps=0,
                                                      int_downsize=1,
                                                      seg_downsize=1)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH, print_fn=log_file.writelines)
+
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
 
     # Compile the model
     SSIM_KER_SIZE = 5
@@ -190,8 +228,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
 
     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')
+    callback_save_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.{epoch:05d}-{val_loss:.2f}.h5'),
+                   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'),
@@ -200,6 +238,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                        write_graph=True, write_grads=True
                                        )
     callback_early_stop = EarlyStopping(monitor='val_loss', verbose=1, patience=C.EARLY_STOP_PATIENCE, min_delta=0.00001)
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
 
     # Compile the model
     optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, lr=C.LEARNING_RATE)
@@ -210,9 +249,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
 
     callback_tensorboard.set_model(network)
     callback_best_model.set_model(network)
-    # callback_save_model.set_model(network)
+    callback_save_model.set_model(network)
     callback_early_stop.set_model(network)
-
+    callback_lr.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')))
@@ -221,13 +260,14 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
         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_save_model.on_train_begin()
+        callback_lr.on_train_begin()
+        for epoch in range(resume_epoch, 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)
-
+            callback_save_model.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print('TRAINING')
 
@@ -238,12 +278,13 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                 #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_save_model.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
-
+                callback_lr.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)
@@ -275,8 +316,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                 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_save_model.on_train_batch_end(step, named_logs(network, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(network, ret))
+                callback_lr.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()
@@ -313,13 +355,14 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
             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_save_model.on_epoch_end(epoch, epoch_summary)
+            callback_lr.on_epoch_end(epoch, epoch_summary)
 
         callback_tensorboard.on_train_end()
-        # callback_save_model.on_train_end()
+        callback_save_model.on_train_end()
         callback_best_model.on_train_end()
         callback_early_stop.on_train_end()
-
+        callback_lr.on_train_end()
 
 if __name__ == '__main__':
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER

Brain_study/Train_UncertaintyWeighted.py

@@ -5,7 +5,7 @@ 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.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
 from tensorflow.python.keras.utils import Progbar
@@ -27,11 +27,12 @@ from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccum
 from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
 import warnings
+import re
 
 
 def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5e-3, lr=1e-4, rw=5e-3,
                  gpu_num=0, simil=['mse'], segm=['dice'], acc_gradients=16, batch_size=1, max_epochs=10000,
-                 early_stop_patience=1000, image_size=64, unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 early_stop_patience=1000, image_size=64, unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
@@ -41,8 +42,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    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'))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
     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()
@@ -94,7 +103,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
     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
+    # config.allow_soft_placement = True
     sess = tf.Session(config=config)
     tf.keras.backend.set_session(sess)
 
@@ -161,6 +170,26 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
                                                      int_steps=0,
                                                      int_downsize=1,
                                                      seg_downsize=1)
+
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
+
     # 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)
@@ -194,8 +223,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
 
     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')
+    callback_save_model = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.{epoch:05d}-{val_loss:.2f}.h5'),
+                   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'),
@@ -204,6 +233,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
                                        write_graph=True, write_grads=True
                                        )
     callback_early_stop = EarlyStopping(monitor='val_loss', verbose=1, patience=C.EARLY_STOP_PATIENCE, min_delta=0.00001)
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
 
     # Compile the model
     optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, lr=C.LEARNING_RATE)
@@ -211,8 +241,10 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
 
     callback_tensorboard.set_model(full_model)
     callback_best_model.set_model(network)  # ONLY SAVE THE NETWORK!!!
-    # callback_save_model.set_model(network)
+    callback_save_model.set_model(network)
     callback_early_stop.set_model(full_model)
+    callback_lr.set_model(full_model)
+
     # TODO: https://towardsdatascience.com/writing-tensorflow-2-custom-loops-438b1ab6eb6c
 
     summary = SummaryDictionary(full_model, C.BATCH_SIZE)
@@ -223,12 +255,15 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
         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_save_model.on_train_begin()
+        callback_lr.on_train_begin()
+
+        for epoch in range(resume_epoch, 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)
+            callback_save_model.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
 
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print('TRAINING')
@@ -238,8 +273,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
             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_save_model.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
+                callback_lr.on_train_batch_begin(step)
 
                 try:
                     fix_img, mov_img, fix_seg, mov_seg = augmentation_model.predict(in_batch)
@@ -255,8 +291,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
                 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_save_model.on_train_batch_end(step, named_logs(full_model, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(full_model, ret))
+                callback_lr.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))
@@ -292,12 +329,14 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
             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_save_model.on_epoch_end(epoch, epoch_summary)
+            callback_lr.on_epoch_end(epoch, epoch_summary)
 
         callback_tensorboard.on_train_end()
-        # callback_save_model.on_train_end()
+        callback_save_model.on_train_end()
         callback_best_model.on_train_end()
         callback_early_stop.on_train_end()
+        callback_lr.on_train_end()
 
 
 if __name__ == '__main__':

Brain_study/data_generator.py

@@ -1,5 +1,5 @@
 import warnings
-
+import time
 import numpy as np
 from tensorflow import keras
 import os
@@ -14,6 +14,7 @@ import tensorflow as tf
 
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.misc import segmentation_cardinal_to_ohe
 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
@@ -22,6 +23,10 @@ from tensorflow.python.keras.preprocessing.image import Iterator
 from tensorflow.python.keras.utils import Sequence
 import sys
 
+from collections import defaultdict
+
+from Brain_study.format_dataset import SEGMENTATION_LOC
+
 #import concurrent.futures
 #import multiprocessing as mp
 import time
@@ -34,13 +39,15 @@ class BatchGenerator:
                  split=0.7,
                  combine_segmentations=True,
                  labels=['all'],
-                 directory_val=None):
+                 directory_val=None,
+                 return_isotropic_shape=False):
         self.file_directory = directory
         self.batch_size = batch_size
         self.combine_segmentations = combine_segmentations
         self.labels = labels
         self.shuffle = shuffle
         self.split = split
+        self.return_isotropic_shape=return_isotropic_shape
 
         if directory_val is None:
             self.file_list = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(('h5', 'hd5'))]
@@ -48,11 +55,11 @@ class BatchGenerator:
             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)
+            self.train_iter = BatchIterator(training_samples, batch_size, shuffle, combine_segmentations, labels, return_isotropic_shape=return_isotropic_shape)
             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)
+                                                     validation=True, return_isotropic_shape=return_isotropic_shape)
             else:
                 self.validation_iter = None
         else:
@@ -66,7 +73,7 @@ class BatchGenerator:
             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'],
+            self.validation_iter = BatchIterator(validation_samples, batch_size, shuffle, combine_segmentations, labels,
                                                  validation=True)
 
     def get_train_generator(self):
@@ -92,7 +99,8 @@ ALL_LABELS_LOC = {label: loc for label, loc in zip(ALL_LABELS, range(0, len(ALL_
 
 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):
+                                            zero_grads=[64, 64, 64, 3], validation=False, sequential_labels=True,
+                 return_isotropic_shape=False, **kwargs):
         # super(BatchIterator, self).__init__(n=len(file_list),
         #                                     batch_size=batch_size,
         #                                     shuffle=shuffle,
@@ -103,31 +111,51 @@ class BatchIterator(Sequence):
         self.file_list = file_list
         self.combine_segmentations = combine_segmentations
         self.labels = labels
-        self.zero_grads = zero_grads
+        self.zero_grads = np.zeros(zero_grads)
         self.idx_list = np.arange(0, len(self.file_list))
         self.validation = validation
+        self.sequential_labels = sequential_labels
+        self.return_isotropic_shape = return_isotropic_shape
         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 (isinstance(self.labels[0], str) and self.labels[0].lower() != 'none'):
             if self.labels[0] != 'all':
-                if isinstance(self.labels[0], str):
-                    self.labels = [SEGMENTATION_LBL2NR_LUT[lbl] for lbl in self.labels]
+                # Labels are tag names. Convert to numeric and check if the expected labels are in sequence or not
+                self.labels = [SEGMENTATION_LBL2NR_LUT[lbl] for lbl in self.labels]
+                if not self.sequential_labels:
+                    self.labels = [SEGMENTATION_LOC[lbl] for lbl in self.labels]
+                    self.labels_dict = lambda x: SEGMENTATION_LOC[x] if x in self.labels else 0
+                else:
+                    self.labels_dict = lambda x: ALL_LABELS_LOC[x] if x in self.labels else 0
+            else:
+                # Use all labels
+                if self.sequential_labels:
+                    self.labels = list(set(SEGMENTATION_LOC.values()))
+                    self.labels_dict = lambda x: SEGMENTATION_LOC[x] if x else 0
+                else:
+                    self.labels = list(ALL_LABELS)
+                    self.labels_dict = lambda x: ALL_LABELS_LOC[x] if x in self.labels else 0
+        elif hasattr(self.labels[0], 'lower') and self.labels[0].lower() == 'none':
+            # self.labels = list()
+            self.labels_dict = dict()
+        else:
+            assert np.all([isinstance(lbl, (int, float)) for lbl in self.labels]), "Labels must be a str, int or float"
+            # Nothing to do, the self.labels contains a list of numbers
 
         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)
 
+        with h5py.File(self.file_list[0], 'r') as f:
+            self.image_shape = list(f['image'][:].shape)
+            self.segm_shape = self.image_shape.copy()
+            self.segm_shape[-1] = len(self.labels) if not self.combine_segmentations else 1
+
+        self.batch_shape = self.image_shape.copy()
+        self.batch_shape[-1] = self.image_shape[-1] + self.segm_shape[-1]
+
     def shuffle_samples(self):
         np.random.shuffle(self.idx_list)
 
@@ -137,7 +165,7 @@ class BatchIterator(Sequence):
     def _filter_segmentations(self, segm, segm_labels):
         if self.combine_segmentations:
             # TODO
-            warnings.warn('Cannot select labels when combinine_segmentations options is active')
+            warnings.warn('Cannot select labels when combine_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
@@ -155,31 +183,38 @@ class BatchIterator(Sequence):
     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??
+            segm = f['segmentation'][:]
+            isot_shape = f['isotropic_shape'][:]
 
-        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 not self.combine_segmentations:
+            if self.sequential_labels:
+                # TODO: I am assuming I want all the labels
+                segm = np.squeeze(np.eye(len(self.labels))[segm])
+            else:
+                lbls_list = list(ALL_LABELS) if self.labels[0] == 'all' else self.labels
+                segm = segmentation_cardinal_to_ohe(segm, lbls_list)  # Filtering is done here
+            #     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??
+
+        img = np.asarray(img, dtype=np.float32)
+        segm = np.asarray(segm, dtype=np.float32)
+        if not isinstance(self.labels[0], str) or 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))
+                ret_val = np.concatenate([img, segm], axis=-1), (img, segm, self.zero_grads), isot_shape
             else:
-                ret_val = np.concatenate([img, segm], axis=-1), (img, np.zeros(self.zero_grads))
+                ret_val = np.concatenate([img, segm], axis=-1), (img, self.zero_grads), isot_shape
         else:
-            ret_val = img, (img, np.zeros(self.zero_grads))
+            ret_val = img, (img, self.zero_grads), isot_shape
         return ret_val
 
     def __getitem__(self, idx):
         in_batch = list()
+        isotropic_shape = list()
         # out_batch = list()
 
         batch_idxs = self.idx_list[idx * self.batch_size:(idx + 1) * self.batch_size]
@@ -193,14 +228,22 @@ class BatchIterator(Sequence):
         #         # 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 = np.zeros([self.batch_size] + self.batch_shape, dtype=np.float32)
+        for batch_idx, f in enumerate(file_list):
+            b, i, isot_shape = self._load_sample(f)
+            # in_batch[batch_idx, :, :, :, :] = b
+            if self.return_isotropic_shape:
+                isotropic_shape.append(isot_shape)
             in_batch.append(b)
             # out_batch.append(i)
 
-        in_batch = np.asarray(in_batch)
+        in_batch = np.asarray(in_batch, dtype=np.float32)
+        ret_val = (in_batch, in_batch)
+        if self.return_isotropic_shape:
+            isotropic_shape = np.asarray(isotropic_shape, dtype=np.int)
+            ret_val += (isotropic_shape,)
         # out_batch = np.asarray(out_batch)
-        return in_batch, in_batch
+        return ret_val
 
     def __iter__(self):
         """Create a generator that iterate over the Sequence."""
@@ -217,9 +260,7 @@ class BatchIterator(Sequence):
         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
+            labels = self.labels
         return labels
 
 

Brain_study/format_dataset.py

@@ -1,12 +1,18 @@
 import h5py
 import nibabel as nib
 from nilearn.image import resample_img
-import os
+import os, sys
 import re
 import numpy as np
 from scipy.ndimage import zoom
 from tqdm import tqdm
 
+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.split_dataset import split
+
 SEGMENTATION_NR2LBL_LUT = {0: 'background',
                            2: 'parietal-right-gm',
                            3: 'lateral-ventricle-left',
@@ -36,15 +42,26 @@ SEGMENTATION_NR2LBL_LUT = {0: 'background',
                            233: '4th-ventricle',
                            254: 'fornix-right',
                            255: 'csf'}
+
 SEGMENTATION_LBL2NR_LUT = {v: k for k, v in SEGMENTATION_NR2LBL_LUT.items()}
 
+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.}
+LABELS_COMBINED = {0, (2, 6), (3, 9), (4, 8), (11, 12), (14, 16), 20, (23, 33), (29, 254), (39, 53), (67, 76), (102, 203), (210, 211), (218, 219), 232, 233, 255}
+SEGMENTATION_LOC = {}
+for loc, label in enumerate(LABELS_COMBINED):
+    if isinstance(label, tuple):
+        SEGMENTATION_LOC.update(dict.fromkeys(label, loc))
+    else:
+        SEGMENTATION_LOC[label] = loc
+
 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'
+OUT_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/ERASEME_sequential'
 
 if __name__ == '__main__':
     img_list = [os.path.join(IMG_DIRECTORY, f) for f in os.listdir(IMG_DIRECTORY) if f.endswith('.nii.gz')]
@@ -54,6 +71,9 @@ if __name__ == '__main__':
     seg_list.sort()
 
     os.makedirs(OUT_DIRECTORY, exist_ok=True)
+
+    vectorize_fnc = np.vectorize(lambda x: SEGMENTATION_LOC[x] if x in SEGMENTATION_LOC.keys() else 0)
+    change_labels = lambda x: np.reshape(vectorize_fnc(x.ravel()), x.shape)
     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')
@@ -70,6 +90,8 @@ if __name__ == '__main__':
         seg = np.asarray(seg.dataobj)
         seg = zoom(seg, np.asarray([128]*3) / np.asarray(isot_shape), order=0)
 
+        seg = change_labels(seg)    # This way the segmentation numbering is continuous
+
         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):
@@ -79,12 +101,13 @@ if __name__ == '__main__':
 
         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_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()
-
+    # We should only have train and test. The val split is done by the batch generator
+    split(train_perc=0.70, validation_perc=0.15, test_perc=0.15, data_dir=OUT_DIRECTORY, move_files=True)
 
 
 

Brain_study/split_dataset.py

@@ -4,51 +4,55 @@ import random
 import warnings
 
 import math
-from shutil import copyfile
+from shutil import copyfile, move
 from tqdm import tqdm
 import concurrent.futures
 import numpy as np
 
 
-def copy_file(s_d):
+def copy_file_fnc(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)
+def move_file_fnc(s_d):
+    s, d = s_d
+    file_name = os.path.split(s)[-1]
+    move(s, os.path.join(d, file_name))
+    return int(os.path.exists(d))
 
-    args = parser.parse_args()
 
-    assert args.train + args.validation + args.test == 1.0, 'Train+Validation+Test != 1 (100%)'
+def split(train_perc: float=0.7,
+          validation_perc: float=0.15,
+          test_perc: float=0.15,
+          data_dir: str='',
+          file_format: str='h5',
+          random_split: bool=True,
+          move_files: bool=False):
+    assert train_perc + validation_perc + test_perc == 1.0, 'Train+Validation+Test != 1 (100%)'
+    assert train_perc > 0 and test_perc > 0, 'Train and test percentages must be greater than zero'
 
-    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()
+    file_set = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith(file_format)]
+    random.shuffle(file_set) if random_split 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_validation = math.floor(num_files * validation_perc)
+    num_test = math.floor(num_files * test_perc)
     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')
+    dataset_root, dataset_name = os.path.split(data_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))
+    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
 
@@ -57,8 +61,10 @@ if __name__ == '__main__':
     os.makedirs(dst_test, exist_ok=True)
 
     progress_bar = tqdm(zip(file_set, dest), desc='Copying files', total=num_files)
+    operation = move_file_fnc if move_files else copy_file_fnc
+    desc = 'Moving files' if move_files else 'Copying 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))
+        results = list(tqdm(ex.map(operation, zip(file_set, dest)), desc=desc, total=num_files))
 
     num_copies = np.sum(results)
     if num_copies == num_files:
@@ -66,3 +72,18 @@ if __name__ == '__main__':
     else:
         warnings.warn('Missing files: {}'.format(num_files - num_copies))
 
+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', help='Randomly split the dataset or not. Default: True', action='store_true', default=True)
+    parser.add_argument('-m', '--movefiles', help='Move files. Otherwise copy. Default: False', action='store_true', default=False)
+
+    args = parser.parse_args()
+
+    split(args.train, args.validation, args.test, args.dir, args.format, args.random, args.movefiles)
+
+

COMET/Build_test_set.py

@@ -18,10 +18,11 @@ 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 DeepDeformationMapRegistration.utils.misc import DisplacementMapInterpolator
 from tqdm import tqdm
 
 from Brain_study.data_generator import BatchGenerator
+from Brain_study.Build_test_set import get_mov_centroids
 
 from skimage.measure import regionprops
 from scipy.interpolate import griddata
@@ -35,12 +36,6 @@ POINTS = None
 MISSING_CENTROID = np.asarray([[np.nan]*3])
 
 
-def get_mov_centroids(fix_seg, disp_map, nb_labels=28, brain_study=True):
-    fix_centroids, _ = get_segmentations_centroids(fix_seg[0, ...], ohe=True, expected_lbls=range(0, nb_labels), brain_study=brain_study)
-    disp = griddata(POINTS, disp_map.reshape([-1, 3]), fix_centroids, method='linear')
-    return fix_centroids, fix_centroids + disp, disp
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('-d', '--dir', type=str, help='Directory where to store the files', default='')
@@ -48,6 +43,7 @@ if __name__ == '__main__':
     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)
+    parser.add_argument('--output_shape', help='If an int, a cubic shape is presumed. Otherwise provide it as a space separated sequence', nargs='+', default=128)
     args = parser.parse_args()
 
     assert args.dataset != '', "Missing original dataset dataset"
@@ -68,12 +64,20 @@ if __name__ == '__main__':
     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'])
+    data_generator = BatchGenerator(DATASET, 1, False, 1.0, False, [0, 1, 2], return_isotropic_shape=True)
 
     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
+
+    if isinstance(args.output_shape, int):
+        image_output_shape = [args.output_shape] * 3
+    elif isinstance(args.output_shape, list):
+        assert len(args.output_shape) == 3, 'Invalid output shape, expected three values and got {}'.format(len(args.output_shape))
+        image_output_shape = [int(s) for s in args.output_shape]
+    else:
+        raise ValueError('Invalid output_shape. Must be an int or a space-separated sequence of ints')
+    print('Scaling to: ', image_output_shape)
     # Build model
 
     xx = np.linspace(0, image_output_shape[0], image_output_shape[0], endpoint=False)
@@ -104,15 +108,17 @@ if __name__ == '__main__':
     config.gpu_options.allow_growth = True
     config.log_device_placement = False  ## to log device placement (on which device the operation ran)
 
+    dm_interp = DisplacementMapInterpolator(image_output_shape, 'griddata', step=8)
+
     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:
+        for step, (in_batch, _, isotropic_shape) in progress_bar:
             fix_img, mov_img, fix_seg, mov_seg, disp_map = augm_model.predict(in_batch)
 
-            fix_centroids, mov_centroids, disp_centroids = get_mov_centroids(fix_seg, disp_map, nb_labels, False)
+            fix_centroids, mov_centroids, disp_centroids = get_mov_centroids(fix_seg, disp_map, nb_labels, False, dm_interp=dm_interp)
 
             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))
@@ -127,7 +133,7 @@ if __name__ == '__main__':
                 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)
-
+                f.create_dataset('isotropic_shape', data=np.squeeze(isotropic_shape))
             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)

COMET/COMET_train.py

@@ -8,7 +8,7 @@ sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
 
 from datetime import datetime
 
-from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.python.keras.utils import Progbar
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
@@ -27,6 +27,7 @@ from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
 
 import COMET.augmentation_constants as COMET_C
+from COMET.utils import freeze_layers_by_group
 
 import numpy as np
 import tensorflow as tf
@@ -40,7 +41,7 @@ import warnings
 def launch_train(dataset_folder, validation_folder, output_folder, model_file, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim',
                  segm='dice', max_epochs=C.EPOCHS, early_stop_patience=1000, freeze_layers=None,
                  acc_gradients=1, batch_size=16, image_size=64,
-                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
     # 0. Input checks
     assert dataset_folder is not None and output_folder is not None
     if model_file != '':
@@ -54,12 +55,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    if freeze_layers is not None:
-        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in freeze_layers), \
-            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
-        freeze_layers = [list(COMET_C.LAYER_RANGES[l]) for l in list(set(freeze_layers))]
-        if len(freeze_layers) > 1:
-            freeze_layers = list(itertools.chain.from_iterable(freeze_layers))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
 
     os.makedirs(output_folder, exist_ok=True)
     # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
@@ -159,34 +164,37 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         if model_file != '':
             network.load_weights(model_file, by_name=True)
             print('MODEL LOCATION: ', model_file)
+
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
+
     # 4. Freeze/unfreeze model layers
-    # freeze_layers = range(0, len(network.layers) - 8)  # Do not freeze the last layers after the UNet (8 last layers)
-    # for l in freeze_layers:
-    #     network.layers[l].trainable = False
-    # msg = "[INF]: Frozen layers {} to {}".format(0, len(network.layers) - 8)
-    # print(msg)
-    # log_file.write("INF: Frozen layers {} to {}".format(0, len(network.layers) - 8))
-    if freeze_layers is not None:
-        aux = list()
-        for r in freeze_layers:
-            for l in range(*r):
-                network.layers[l].trainable = False
-                aux.append(l)
-        aux.sort()
-        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in aux]))
+    _, frozen_layers = freeze_layers_by_group(network, freeze_layers)
+    if frozen_layers is not None:
+        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in frozen_layers]))
     else:
         msg = "[INF] None frozen layers"
     print(msg)
     log_file.write(msg)
-    # network.trainable = False  # Freeze the base model
-    # # Create a new model on top
-    # input_new_model = keras.Input(network.input_shape)
-    # x = base_model(input_new_model, training=False)
-    # x =
-    # network = keras.Model(input_new_model, x)
-
-    network.summary()
-    network.summary(print_fn=log_file.writelines)
+
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH, print_fn=log_file.writelines)
     #   Complete the model with the augmentation layer
     augm_train_input_shape = train_generator.get_data_shape()[-1]
     input_layer_train = Input(shape=augm_train_input_shape, name='input_train')
@@ -269,6 +277,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                           save_best_only=True, monitor='val_loss', verbose=1, mode='min')
     callback_save_checkpoint = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.h5'),
                                                save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
 
     losses = {'transformer': loss_fnc,
               'flow': vxm.losses.Grad('l2').loss}
@@ -281,7 +290,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     loss_weights = {'transformer': 1.,
                     'flow': rw}
 
-    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    optimizer = AdamAccumulated(accumulation_steps=C.ACCUM_GRADIENT_STEP, learning_rate=C.LEARNING_RATE)
     network.compile(optimizer=optimizer,
                     loss=losses,
                     loss_weights=loss_weights,
@@ -292,6 +301,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     callback_early_stop.set_model(network)
     callback_best_model.set_model(network)
     callback_save_checkpoint.set_model(network)
+    callback_lr.set_model(network)
 
     summary = SummaryDictionary(network, C.BATCH_SIZE)
     names = network.metrics_names
@@ -303,12 +313,14 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         callback_early_stop.on_train_begin()
         callback_best_model.on_train_begin()
         callback_save_checkpoint.on_train_begin()
+        callback_lr.on_train_begin()
 
-        for epoch in range(C.EPOCHS):
+        for epoch in range(resume_epoch, C.EPOCHS):
             callback_tensorboard.on_epoch_begin(epoch)
             callback_early_stop.on_epoch_begin(epoch)
             callback_best_model.on_epoch_begin(epoch)
             callback_save_checkpoint.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print("TRAIN")
 
@@ -318,6 +330,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_begin(step)
                 callback_save_checkpoint.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
+                callback_lr.on_train_batch_begin(step)
 
                 try:
                     fix_img, mov_img, fix_seg, mov_seg = augm_model_train.predict(in_batch)
@@ -351,6 +364,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_end(step, named_logs(network, ret))
                 callback_save_checkpoint.on_train_batch_end(step, named_logs(network, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(network, ret))
+                callback_lr.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()
@@ -384,10 +398,12 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
             callback_best_model.on_epoch_end(epoch, epoch_summary)
             callback_save_checkpoint.on_epoch_end(epoch, epoch_summary)
             callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            callback_lr.on_epoch_end(epoch, epoch_summary)
             print('End of epoch {}: '.format(epoch), ret, '\n')
 
         callback_tensorboard.on_train_end()
         callback_best_model.on_train_end()
         callback_save_checkpoint.on_train_end()
         callback_early_stop.on_train_end()
+        callback_lr.on_train_end()
 # 7. Wrap up

COMET/COMET_train_UW.py

@@ -6,7 +6,7 @@ sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
 
 from datetime import datetime
 
-from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.python.keras.utils import Progbar
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
@@ -26,6 +26,7 @@ from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
 
 import COMET.augmentation_constants as COMET_C
+from COMET.utils import freeze_layers_by_group
 
 import numpy as np
 import tensorflow as tf
@@ -39,7 +40,7 @@ import warnings
 def launch_train(dataset_folder, validation_folder, output_folder, model_file, gpu_num=0, lr=1e-4, rw=5e-3,
                  simil=['ssim'], segm=['dice'], max_epochs=C.EPOCHS, early_stop_patience=1000, prior_reg_w=5e-3,
                  freeze_layers=None, acc_gradients=1, batch_size=16, image_size=64,
-                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
     # 0. Input checks
     assert dataset_folder is not None and output_folder is not None
     if model_file != '':
@@ -53,15 +54,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    if freeze_layers is not None:
-        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in freeze_layers), \
-            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
-        multiple_ranges = 'TOP' in freeze_layers
-        freeze_layers = [list(COMET_C.LAYER_RANGES[l]) for l in list(set(freeze_layers))]
-        freeze_layers = freeze_layers[0] if multiple_ranges else freeze_layers
-
-        # if len(freeze_layers) > 1:
-        #     freeze_layers = list(itertools.chain.from_iterable(freeze_layers))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
 
     os.makedirs(output_folder, exist_ok=True)
     # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
@@ -103,8 +105,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     print(aux)
 
     # 2. Data generator
+    used_labels = [0, 1, 2]
     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,
+                                    C.TRAINING_PERC, labels=used_labels, combine_segmentations=False,
                                     directory_val=C.VALIDATION_DATASET)
 
     train_generator = data_generator.get_train_generator()
@@ -137,22 +140,36 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         print('MODEL LOCATION: ', model_file)
         network.load_weights(model_file, by_name=True)
 
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
+
     # 4. Freeze/unfreeze model layers
-    if freeze_layers is not None:
-        aux = list()
-        for r in freeze_layers:
-            for l in range(*r):
-                network.layers[l].trainable = False
-                aux.append(l)
-        aux.sort()
-        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in aux]))
+    _, frozen_layers = freeze_layers_by_group(network, freeze_layers)
+    if frozen_layers is not None:
+        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in frozen_layers]))
     else:
         msg = "[INF] None frozen layers"
     print(msg)
     log_file.write(msg)
 
-    network.summary()
-    network.summary(print_fn=log_file.write)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH, print_fn=log_file.write)
     #   Complete the model with the augmentation layer
     input_layer_train = Input(shape=train_generator.get_data_shape()[0], name='input_train')
     augm_layer = AugmentationLayer(max_displacement=COMET_C.MAX_AUG_DISP,  # Max 30 mm in isotropic space
@@ -250,7 +267,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     callback_save_checkpoint = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.h5'),
                                                save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
 
-    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
+
+    optimizer = AdamAccumulated(accumulation_steps=C.ACCUM_GRADIENT_STEP, learning_rate=C.LEARNING_RATE)
     full_model.compile(optimizer=optimizer,
                        loss=None, )
 
@@ -259,6 +278,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     callback_early_stop.set_model(full_model)
     callback_best_model.set_model(network)  # ONLY SAVE THE NETWORK!!!
     callback_save_checkpoint.set_model(network)  # ONLY SAVE THE NETWORK!!!
+    callback_lr.set_model(full_model)
 
     summary = SummaryDictionary(full_model, C.BATCH_SIZE)
     names = full_model.metrics_names
@@ -271,12 +291,15 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         callback_early_stop.on_train_begin()
         callback_best_model.on_train_begin()
         callback_save_checkpoint.on_train_begin()
+        callback_lr.on_train_begin()
 
-        for epoch in range(C.EPOCHS):
+        for epoch in range(resume_epoch, C.EPOCHS):
             callback_tensorboard.on_epoch_begin(epoch)
             callback_early_stop.on_epoch_begin(epoch)
             callback_best_model.on_epoch_begin(epoch)
             callback_save_checkpoint.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
+
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print("TRAIN")
 
@@ -287,6 +310,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_begin(step)
                 callback_save_checkpoint.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
+                callback_lr.on_train_batch_begin(step)
 
                 try:
                     fix_img, mov_img, fix_seg, mov_seg = augm_model.predict(in_batch)
@@ -310,6 +334,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_end(step, named_logs(full_model, ret))
                 callback_save_checkpoint.on_train_batch_end(step, named_logs(full_model, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(full_model, ret))
+                callback_lr.on_train_batch_end(step, named_logs(full_model, ret))
                 progress_bar.update(step, zip(names, ret))
                 log_file.write('\t\tStep {:03d}: {}'.format(step, ret))
             val_values = progress_bar._values.copy()
@@ -341,10 +366,12 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
             callback_best_model.on_epoch_end(epoch, epoch_summary)
             callback_save_checkpoint.on_epoch_end(epoch, epoch_summary)
             callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            callback_lr.on_epoch_end(epoch, epoch_summary)
             print('End of epoch {}: '.format(epoch), ret, '\n')
 
         callback_tensorboard.on_train_end()
         callback_best_model.on_train_end()
         callback_save_checkpoint.on_train_end()
         callback_early_stop.on_train_end()
+        callback_lr.on_train_end()
 # 7. Wrap up

COMET/COMET_train_seggguided.py

@@ -8,7 +8,7 @@ sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
 
 from datetime import datetime
 
-from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping, ReduceLROnPlateau
 from tensorflow.python.keras.utils import Progbar
 from tensorflow.keras import Input
 from tensorflow.keras.models import Model
@@ -27,6 +27,7 @@ from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
 
 import COMET.augmentation_constants as COMET_C
+from COMET.utils import freeze_layers_by_group
 
 import numpy as np
 import tensorflow as tf
@@ -40,7 +41,7 @@ import warnings
 def launch_train(dataset_folder, validation_folder, output_folder, model_file, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim',
                  segm='dice', max_epochs=C.EPOCHS, early_stop_patience=1000, freeze_layers=None,
                  acc_gradients=1, batch_size=16, image_size=64,
-                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
+                 unet=[16, 32, 64, 128, 256], head=[16, 16], resume=None):
     # 0. Input checks
     assert dataset_folder is not None and output_folder is not None
     if model_file != '':
@@ -54,12 +55,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     if batch_size != 1 and acc_gradients != 1:
         warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
 
-    if freeze_layers is not None:
-        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in freeze_layers), \
-            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
-        freeze_layers = [list(COMET_C.LAYER_RANGES[l]) for l in list(set(freeze_layers))]
-        if len(freeze_layers) > 1:
-            freeze_layers = list(itertools.chain.from_iterable(freeze_layers))
+    if resume is not None:
+        try:
+            assert os.path.exists(resume) and len(os.listdir(os.path.join(resume, 'checkpoints'))), 'Invalid directory: ' + resume
+            output_folder = resume
+            resume = True
+        except AssertionError:
+            output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+            resume = False
+    else:
+        resume = False
 
     os.makedirs(output_folder, exist_ok=True)
     # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
@@ -101,9 +106,9 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     print(aux)
 
     # 2. Data generator
-    used_labels = 'all'
+    used_labels = [0, 1, 2]
     data_generator = BatchGenerator(C.TRAINING_DATASET, C.BATCH_SIZE if C.ACCUM_GRADIENT_STEP == 1 else 1, True,
-                                    C.TRAINING_PERC, labels=[used_labels], combine_segmentations=False,
+                                    C.TRAINING_PERC, labels=used_labels, combine_segmentations=False,
                                     directory_val=C.VALIDATION_DATASET)
 
     train_generator = data_generator.get_train_generator()
@@ -163,34 +168,37 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         if model_file != '':
             network.load_weights(model_file, by_name=True)
             print('MODEL LOCATION: ', model_file)
+
+    resume_epoch = 0
+    if resume:
+        cp_dir = os.path.join(output_folder, 'checkpoints')
+        cp_file_list = [os.path.join(cp_dir, f) for f in os.listdir(cp_dir) if (f.startswith('checkpoint') and f.endswith('.h5'))]
+        if len(cp_file_list):
+            cp_file_list.sort()
+            checkpoint_file = cp_file_list[-1]
+            if os.path.exists(checkpoint_file):
+                network.load_weights(checkpoint_file, by_name=True)
+                print('Loaded checkpoint file: ' + checkpoint_file)
+                try:
+                    resume_epoch = int(re.match('checkpoint\.(\d+)-*.h5', os.path.split(checkpoint_file)[-1])[1])
+                except TypeError:
+                    # Checkpoint file has no epoch number in the name
+                    resume_epoch = 0
+                print('Resuming from epoch: {:d}'.format(resume_epoch))
+            else:
+                warnings.warn('Checkpoint file NOT found. Training from scratch')
+
     # 4. Freeze/unfreeze model layers
-    # freeze_layers = range(0, len(network.layers) - 8)  # Do not freeze the last layers after the UNet (8 last layers)
-    # for l in freeze_layers:
-    #     network.layers[l].trainable = False
-    # msg = "[INF]: Frozen layers {} to {}".format(0, len(network.layers) - 8)
-    # print(msg)
-    # log_file.write("INF: Frozen layers {} to {}".format(0, len(network.layers) - 8))
-    if freeze_layers is not None:
-        aux = list()
-        for r in freeze_layers:
-            for l in range(*r):
-                network.layers[l].trainable = False
-                aux.append(l)
-        aux.sort()
-        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in aux]))
+    _, frozen_layers = freeze_layers_by_group(network, freeze_layers)
+    if frozen_layers is not None:
+        msg = "[INF]: Frozen layers {}".format(', '.join([str(a) for a in frozen_layers]))
     else:
         msg = "[INF] None frozen layers"
     print(msg)
     log_file.write(msg)
-    # network.trainable = False  # Freeze the base model
-    # # Create a new model on top
-    # input_new_model = keras.Input(network.input_shape)
-    # x = base_model(input_new_model, training=False)
-    # x =
-    # network = keras.Model(input_new_model, x)
-
-    network.summary()
-    network.summary(print_fn=log_file.writelines)
+
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH)
+    network.summary(line_length=C.SUMMARY_LINE_LENGTH, print_fn=log_file.writelines)
     #   Complete the model with the augmentation layer
     augm_train_input_shape = train_generator.get_data_shape()[0]
     input_layer_train = Input(shape=augm_train_input_shape, name='input_train')
@@ -282,6 +290,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                                           save_best_only=True, monitor='val_loss', verbose=1, mode='min')
     callback_save_checkpoint = ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'checkpoint.h5'),
                                                save_weights_only=True, monitor='val_loss', verbose=0, mode='min')
+    callback_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10)
 
     losses = {'transformer': loss_fnc,
               'seg_transformer': loss_segm,
@@ -300,8 +309,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                     'seg_transformer': 1.,
                     'flow': rw}
 
-
-    optimizer = AdamAccumulated(C.ACCUM_GRADIENT_STEP, C.LEARNING_RATE)
+    optimizer = AdamAccumulated(accumulation_steps=C.ACCUM_GRADIENT_STEP, learning_rate=C.LEARNING_RATE)
     network.compile(optimizer=optimizer,
                     loss=losses,
                     loss_weights=loss_weights,
@@ -312,6 +320,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
     callback_early_stop.set_model(network)
     callback_best_model.set_model(network)
     callback_save_checkpoint.set_model(network)
+    callback_lr.set_model(network)
 
     summary = SummaryDictionary(network, C.BATCH_SIZE)
     names = network.metrics_names
@@ -323,12 +332,15 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
         callback_early_stop.on_train_begin()
         callback_best_model.on_train_begin()
         callback_save_checkpoint.on_train_begin()
+        callback_lr.on_train_begin()
 
-        for epoch in range(C.EPOCHS):
+        for epoch in range(resume_epoch, C.EPOCHS):
             callback_tensorboard.on_epoch_begin(epoch)
             callback_early_stop.on_epoch_begin(epoch)
             callback_best_model.on_epoch_begin(epoch)
             callback_save_checkpoint.on_epoch_begin(epoch)
+            callback_lr.on_epoch_begin(epoch)
+
             print("\nEpoch {}/{}".format(epoch, C.EPOCHS))
             print("TRAIN")
 
@@ -338,6 +350,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_begin(step)
                 callback_save_checkpoint.on_train_batch_begin(step)
                 callback_early_stop.on_train_batch_begin(step)
+                callback_lr.on_train_batch_begin(step)
 
                 try:
                     fix_img, mov_img, fix_seg, mov_seg = augm_model_train.predict(in_batch)
@@ -371,6 +384,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
                 callback_best_model.on_train_batch_end(step, named_logs(network, ret))
                 callback_save_checkpoint.on_train_batch_end(step, named_logs(network, ret))
                 callback_early_stop.on_train_batch_end(step, named_logs(network, ret))
+                callback_lr.on_predict_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()
@@ -405,10 +420,13 @@ def launch_train(dataset_folder, validation_folder, output_folder, model_file, g
             callback_best_model.on_epoch_end(epoch, epoch_summary)
             callback_save_checkpoint.on_epoch_end(epoch, epoch_summary)
             callback_early_stop.on_epoch_end(epoch, epoch_summary)
+            callback_lr.on_epoch_end(epoch,epoch_summary)
+
             print('End of epoch {}: '.format(epoch), ret, '\n')
 
         callback_tensorboard.on_train_end()
         callback_best_model.on_train_end()
         callback_save_checkpoint.on_train_end()
         callback_early_stop.on_train_end()
+        callback_lr.on_train_end()
 # 7. Wrap up

COMET/Evaluate_network.py

@@ -17,9 +17,10 @@ import tensorflow as tf
 import numpy as np
 import pandas as pd
 import voxelmorph as vxm
+from voxelmorph.tf.layers import SpatialTransformer
 
 import DeepDeformationMapRegistration.utils.constants as C
-from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.operators import min_max_norm, safe_medpy_metric
 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
@@ -27,10 +28,14 @@ from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimila
 from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
 from DeepDeformationMapRegistration.utils.visualization import save_disp_map_img, plot_predictions
 from DeepDeformationMapRegistration.utils.misc import DisplacementMapInterpolator, get_segmentations_centroids, segmentation_ohe_to_cardinal, segmentation_cardinal_to_ohe
+from DeepDeformationMapRegistration.utils.misc import resize_displacement_map, scale_transformation, GaussianFilter
+import medpy.metric as medpy_metrics
 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
 
+from scipy.ndimage import gaussian_filter, zoom
+
 import h5py
 import re
 from Brain_study.data_generator import BatchGenerator
@@ -44,6 +49,7 @@ import neurite as ne
 
 
 DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/test_fixed'
+DATASET_FR = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/test_fixed/full_res'
 MODEL_FILE = '/mnt/EncryptedData1/Users/javier/train_output/COMET/ERASE/COMET_L_ssim__MET_mse_ncc_ssim_141343-01122021/checkpoints/best_model.h5'
 DATA_ROOT_DIR = '/mnt/EncryptedData1/Users/javier/train_output/COMET/ERASE/COMET_L_ssim__MET_mse_ncc_ssim_141343-01122021/'
 
@@ -56,6 +62,7 @@ if __name__ == '__main__':
     parser.add_argument('--dataset', type=str, help='Dataset to run predictions on', default=DATASET)
     parser.add_argument('--erase', type=bool, help='Erase the content of the output folder', default=False)
     parser.add_argument('--outdirname', type=str, default='Evaluate')
+    parser.add_argument('--fullres', action='store_true', default=False)
     args = parser.parse_args()
     if args.model is not None:
         assert '.h5' in args.model[0], 'No checkpoint file provided, use -d/--dir instead'
@@ -82,6 +89,9 @@ if __name__ == '__main__':
     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()
+    if args.fullres:
+        list_test_fr_files = [os.path.join(DATASET_FR, f) for f in os.listdir(DATASET_FR) if f.endswith('h5') and 'dm' not in f]
+        list_test_fr_files.sort()
 
     with h5py.File(list_test_files[0], 'r') as f:
         image_input_shape = image_output_shape = list(f['fix_image'][:].shape[:-1])
@@ -96,8 +106,8 @@ if __name__ == '__main__':
     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)
+    sess = tf.compat.v1.Session(config=config)
+    tf.compat.v1.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,
@@ -114,24 +124,24 @@ if __name__ == '__main__':
                    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')
+    fix_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='fix_img')
+    pred_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='pred_img')
+    fix_seg_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, nb_labels), name='fix_seg')
+    pred_seg_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 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)
+    # 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')
+    dm_interp = DisplacementMapInterpolator(image_output_shape, 'griddata', step=4)
 
     for MODEL_FILE, DATA_ROOT_DIR in zip(MODEL_FILE_LIST, DATA_ROOT_DIR_LIST):
         print('MODEL LOCATION: ', MODEL_FILE)
@@ -144,6 +154,14 @@ if __name__ == '__main__':
         os.makedirs(output_folder, exist_ok=True)
         print('DESTINATION FOLDER: ', output_folder)
 
+        if args.fullres:
+            output_folder_fr = os.path.join(DATA_ROOT_DIR, args.outdirname, 'full_resolution')  # '/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_fr, ignore_errors=True)
+            os.makedirs(output_folder_fr, exist_ok=True)
+            print('DESTINATION FOLDER FULL RESOLUTION: ', output_folder_fr)
+
         try:
             network = tf.keras.models.load_model(MODEL_FILE, {'VxmDenseSemiSupervisedSeg': vxm.networks.VxmDenseSemiSupervisedSeg,
                                                               'VxmDense': vxm.networks.VxmDense,
@@ -172,10 +190,16 @@ if __name__ == '__main__':
         with open(metrics_file, 'w') as f:
             f.write(';'.join(csv_header)+'\n')
 
+        if args.fullres:
+            metrics_file_fr = os.path.join(output_folder_fr, 'metrics.csv')
+            with open(metrics_file_fr, '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)
+            network.summary(line_length=C.SUMMARY_LINE_LENGTH)
             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:
@@ -184,6 +208,8 @@ if __name__ == '__main__':
                     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'][:]
+                    isotropic_shape = f['isotropic_shape'][:]
+                    voxel_size = np.divide(fix_img.shape[1:-1], isotropic_shape)
 
                 if network.name == 'vxm_dense_semi_supervised_seg':
                     t0 = time.time()
@@ -196,29 +222,40 @@ if __name__ == '__main__':
                     t1 = time.time()
 
                 pred_img = min_max_norm(pred_img)
-                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(0, nb_labels), brain_study=False)
+                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(1, nb_labels+1), brain_study=False)
                 # pred_centroids = dm_interp(disp_map, mov_centroids, backwards=True)  # with tps, it returns the pred_centroids directly
                 pred_centroids = dm_interp(disp_map, mov_centroids, backwards=True) + mov_centroids
 
+                # Up sample the segmentation masks to isotropic resolution
+                zoom_factors = np.diag(scale_transformation(image_output_shape, isotropic_shape))
+                pred_seg_isot = zoom(pred_seg[0, ...], zoom_factors, order=0)[np.newaxis, ...]
+                fix_seg_isot = zoom(fix_seg[0, ...], zoom_factors, order=0)[np.newaxis, ...]
+
+                pred_img_isot = zoom(pred_img[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+                fix_img_isot = zoom(fix_img[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+
                 # 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})
+                # dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf], {'fix_seg:0': fix_seg, 'pred_seg:0': pred_seg})
+                dice = np.mean([medpy_metrics.dc(pred_seg_isot[0, ..., l], fix_seg_isot[0, ..., l]) / np.sum(fix_seg_isot[0, ..., l]) for l in range(nb_labels)])
+                hd = np.mean(safe_medpy_metric(pred_seg_isot[0, ...], fix_seg_isot[0, ...], nb_labels, medpy_metrics.hd, {'voxelspacing': voxel_size}))
+                dice_macro = np.mean([medpy_metrics.dc(pred_seg_isot[0, ..., l], fix_seg_isot[0, ..., l]) for l in range(nb_labels)])
 
                 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})
+                ssim = np.mean(ssim)
                 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.COMET_DATASET_iso_to_cubic_scales)
-                # mov_centroids_isotropic = np.divide(mov_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
-                pred_centroids_isotropic = np.divide(pred_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
+                fix_centroids_isotropic = fix_centroids * voxel_size
+                # mov_centroids_isotropic = mov_centroids * voxel_size
+                pred_centroids_isotropic = pred_centroids * voxel_size
+
+                # fix_centroids_isotropic = np.divide(fix_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
+                # # mov_centroids_isotropic = np.divide(mov_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
+                # pred_centroids_isotropic = np.divide(pred_centroids_isotropic, C.COMET_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)])
@@ -249,15 +286,103 @@ if __name__ == '__main__':
                 # plt.savefig(os.path.join(output_folder, '{:03d}_hist_mag_ssim_{:.03f}_dice_{:.03f}.png'.format(step, ssim, dice)))
                 # plt.close()
 
-                plot_predictions(fix_img, mov_img, disp_map, pred_img, os.path.join(output_folder, '{:03d}_figures_img.png'.format(step)), show=False)
-                plot_predictions(fix_seg, mov_seg, disp_map, pred_seg, os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False)
-                save_disp_map_img(disp_map, 'Displacement map', os.path.join(output_folder, '{:03d}_disp_map_fig.png'.format(step)), show=False)
+                plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, seg_batches=[fix_seg_card, mov_seg_card, pred_seg_card], filename=os.path.join(output_folder, '{:03d}_figures_seg.png'.format(step)), show=False, step=16)
+                plot_predictions(img_batches=[fix_img, mov_img, pred_img], disp_map_batch=disp_map, filename=os.path.join(output_folder, '{:03d}_figures_img.png'.format(step)), show=False, step=16)
+                save_disp_map_img(disp_map, 'Displacement map', os.path.join(output_folder, '{:03d}_disp_map_fig.png'.format(step)), show=False, step=16)
+
+                progress_bar.set_description('SSIM {:.04f}\tM_DICE: {:.04f}'.format(ssim, dice_macro))
+
+                if args.fullres:
+                    with h5py.File(list_test_fr_files[step - 1], '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'][:]
+
+                    # Up sample the displacement map to the full res
+                    trf = scale_transformation(image_output_shape, fix_img.shape[1:-1])
+                    disp_map_fr = resize_displacement_map(np.squeeze(disp_map), None, trf)[np.newaxis, ...]
+                    disp_map_fr = gaussian_filter(disp_map_fr, 5)
+                    # disp_mad_fr = sess.run(smooth_filter, feed_dict={'dm:0': disp_map_fr})
+
+                    # Predicted image
+                    pred_img_fr = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([mov_img, disp_map_fr]).eval()
+                    pred_seg_fr = SpatialTransformer(interp_method='nearest', indexing='ij', single_transform=False)([mov_seg, disp_map_fr]).eval()
+
+                    # Predicted centroids
+                    dm_interp_fr = DisplacementMapInterpolator(fix_img.shape[1:-1], 'griddata', step=2)
+                    pred_centroids = dm_interp_fr(disp_map_fr, mov_centroids, backwards=True) + mov_centroids
+
+                    # Metrics - segmentation
+                    dice = np.mean([medpy_metrics.dc(pred_seg_fr[..., l], fix_seg[..., l]) / np.sum(fix_seg[..., l]) for l in range(nb_labels)])
+                    hd = np.mean(safe_medpy_metric(pred_seg[0, ...], fix_seg[0, ...], nb_labels, medpy_metrics.hd, {'voxelspacing': voxel_size}))
+                    dice_macro = np.mean([medpy_metrics.dc(pred_seg_fr[..., l], fix_seg[..., l]) for l in range(nb_labels)])
+
+                    pred_seg_card_fr = segmentation_ohe_to_cardinal(pred_seg_fr).astype(np.float32)
+                    mov_seg_card_fr = segmentation_ohe_to_cardinal(mov_seg).astype(np.float32)
+                    fix_seg_card_fr = segmentation_ohe_to_cardinal(fix_seg).astype(np.float32)
+
+                    # Metrics - image
+                    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_fr})
+                    ssim = np.mean(ssim)
+                    ms_ssim = ms_ssim[0]
+
+                    # Metrics - registration
+                    tre_array = target_registration_error(fix_centroids, pred_centroids, False).eval()
+
+                    new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd, t1 - t0, tre, len(missing_lbls),
+                                missing_lbls]
+                    with open(metrics_file_fr, 'a') as f:
+                        f.write(';'.join(map(str, new_line)) + '\n')
+
+                    save_nifti(fix_img[0, ...], os.path.join(output_folder_fr, '{: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_fr, '{: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_fr, '{:03d}_pred_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                    save_nifti(fix_seg_card[0, ...], os.path.join(output_folder_fr, '{:03d}_fix_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                    save_nifti(mov_seg_card[0, ...], os.path.join(output_folder_fr, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
+                    save_nifti(pred_seg_card[0, ...], os.path.join(output_folder_fr, '{: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(img_batches=[fix_img, mov_img, pred_img_fr], disp_map_batch=disp_map_fr, seg_batches=[fix_seg_card_fr, mov_seg_card_fr, pred_seg_card_fr], filename=os.path.join(output_folder_fr, '{:03d}_figures_seg.png'.format(step)), show=False, step=10)
+                    plot_predictions(img_batches=[fix_img, mov_img, pred_img_fr], disp_map_batch=disp_map_fr, filename=os.path.join(output_folder_fr, '{:03d}_figures_img.png'.format(step)), show=False, step=10)
+                    # save_disp_map_img(disp_map_fr, 'Displacement map', os.path.join(output_folder_fr, '{:03d}_disp_map_fig.png'.format(step)), show=False, step=10)
 
-                progress_bar.set_description('SSIM {:.04f}\tG_DICE: {:.04f}\tM_DICE: {:.04f}'.format(ssim, dice, dice_macro))
+                    progress_bar.set_description('[FR] SSIM {:.04f}\tM_DICE: {:.04f}'.format(ssim, dice_macro))
 
         print('Summary\n=======\n')
-        print('\nAVG:\n' + str(pd.read_csv(metrics_file, sep=';', header=0).mean(axis=0)) + '\nSTD:\n' + str(pd.read_csv(metrics_file, sep=';', header=0).std(axis=0)))
+        metrics_df = pd.read_csv(metrics_file, sep=';', header=0)
+        print('\nAVG:\n')
+        print(metrics_df.mean(axis=0))
+        print('\nSTD:\n')
+        print(metrics_df.std(axis=0))
+        print('\nHD95perc:\n')
+        print(metrics_df['HD'].describe(percentiles=[.95]))
         print('\n=======\n')
+        if args.fullres:
+            print('Summary full resolution\n=======\n')
+            metrics_df = pd.read_csv(metrics_file_fr, sep=';', header=0)
+            print('\nAVG:\n')
+            print(metrics_df.mean(axis=0))
+            print('\nSTD:\n')
+            print(metrics_df.std(axis=0))
+            print('\nHD95perc:\n')
+            print(metrics_df['HD'].describe(percentiles=[.95]))
+            print('\n=======\n')
         tf.keras.backend.clear_session()
         # sess.close()
         del network

COMET/MultiTrain_config.py

@@ -9,6 +9,8 @@ from shutil import copy2
 import os
 from datetime import datetime
 import DeepDeformationMapRegistration.utils.constants as C
+import re
+from COMET.augmentation_constants import LAYER_SELECTION
 TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/train'
 
 err = list()
@@ -58,10 +60,11 @@ if __name__ == '__main__':
     except KeyError as err:
         froozen_layers = None
     except NameError as err:
-        froozen_layers = [trainConfig['freeze'].upper()]
+        froozen_layers = list(filter(lambda x: x != '', re.split(';|\s|,|,\s|;\s', trainConfig['freeze'].upper())))
+
     if froozen_layers is not None:
-        assert all(s in ['INPUT', 'OUTPUT', 'ENCODER', 'DECODER', 'TOP', 'BOTTOM'] for s in froozen_layers),\
-            'Invalid option for "freeze". Expected one or several of: INPUT, OUTPUT, ENCODER, DECODER, TOP, BOTTOM'
+        assert all(s in LAYER_SELECTION.keys() for s in froozen_layers), \
+            'Invalid option for "freeze". Expected one or several of: ' + ', '.join(LAYER_SELECTION.keys())
         froozen_layers = list(set(froozen_layers))  # Unique elements
 
     if augmentationConfig:
@@ -70,11 +73,20 @@ if __name__ == '__main__':
 
 
     # copy the configuration file to the destionation folder
-    os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(output_folder, exist_ok=True) # TODO: move this within the "resume" if case, and bring here the creation of the resume-output folder!
     copy2(args.ini, os.path.join(output_folder, os.path.split(args.ini)[-1]))
 
-    unet = [int(x) for x in trainConfig['unet'].split(',')] if trainConfig['unet'] else [16, 32, 64, 128, 256]
-    head = [int(x) for x in trainConfig['head'].split(',')] if trainConfig['head'] else [16, 16]
+    try:
+        unet = [int(x) for x in trainConfig['unet'].split(',')] if trainConfig['unet'] else [16, 32, 64, 128, 256]
+        head = [int(x) for x in trainConfig['head'].split(',')] if trainConfig['head'] else [16, 16]
+    except KeyError as err:
+        unet = [16, 32, 64, 128, 256]
+        head = [16, 16]
+
+    try:
+        resume_checkpoint = trainConfig['resumeCheckpoint']
+    except KeyError as e:
+        resume_checkpoint = None
 
     launch_train(dataset_folder=datasetConfig['train'],
                  validation_folder=datasetConfig['validation'],
@@ -85,10 +97,12 @@ if __name__ == '__main__':
                  simil=simil,
                  segm=segm,
                  max_epochs=eval(trainConfig['epochs']),
+                 image_size=eval(trainConfig['imageSize']),
                  early_stop_patience=eval(trainConfig['earlyStopPatience']),
                  model_file=trainConfig['model'],
                  freeze_layers=froozen_layers,
                  acc_gradients=eval(trainConfig['accumulativeGradients']),
                  batch_size=eval(trainConfig['batchSize']),
                  unet=unet,
-                 head=head)
+                 head=head,
+                 resume=resume_checkpoint)

COMET/augmentation_constants.py

@@ -1,4 +1,5 @@
 import numpy as np
+import re
 
 # Constants for augmentation layer
 # .../T1/training/zoom_factors.csv contain the scale factors of all the training samples from isotropic to 128x128x128
@@ -31,4 +32,26 @@ LAYER_RANGES = {'INPUT': (IN_LAYERS),
                 'ENCODER': (ENCONDER_LAYERS),
                 'DECODER': (DECODER_LAYERS),
                 'TOP': (TOP_LAYERS_ENC, TOP_LAYERS_DEC),
-                'BOTTOM': (BOTTOM_LAYERS)}
+                'BOTTOM': (BOTTOM_LAYERS)}
+
+# LAYER names:
+IN_LAYER_REGEXP = '.*input'
+FC_LAYER_REGEXP = '.*final.*'
+OUT_LAYER_REGEXP = '(?:flow|transformer)'
+ENC_LAYER_REGEXP = '.*enc_(?:conv|pooling)_(\d).*'
+DEC_LAYER_REGEXP = '.*dec_(?:conv|upsample)_(\d).*'
+LEVEL_NUMBER = lambda x: re.match('.*(?:enc|dec)_(?:conv|upsample|pooling)_(\d).*', x)
+IS_TOP_LEVEL = lambda x: int(LEVEL_NUMBER(x)[1]) < 3 if LEVEL_NUMBER(x) is not None else False or bool(re.match(FC_LAYER_REGEXP, x))
+IS_BOTTOM_LEVEL = lambda x: int(LEVEL_NUMBER(x)[1]) >= 3 if LEVEL_NUMBER(x) is not None else False
+
+LAYER_SELECTION = {'INPUT': lambda x: bool(re.match(IN_LAYER_REGEXP, x)),
+                   'FULLYCONNECTED': lambda x: bool(re.match(FC_LAYER_REGEXP, x)),
+                   'ENCODER': lambda x: bool(re.match(ENC_LAYER_REGEXP, x)),
+                   'DECODER': lambda x: bool(re.match(DEC_LAYER_REGEXP, x)),
+                   'TOP': lambda x: IS_TOP_LEVEL(x),
+                   'BOTTOM': lambda x: IS_BOTTOM_LEVEL(x)
+                   }
+
+# STUPID IDEA THAT COMPLICATES THINGS. The points was to allow combinations of the layer groups
+# OR_GROUPS = ['ENCODER', 'DECODER', 'INPUT', 'OUTPUT', 'FULLYCONNECTED']     # These groups can be OR'ed with the AND_GROUPS and among them. E.g., Top layers of the encoder and decoder: ENCODER or DECODER and TOP
+# AND_GROUPS = ['TOP', 'BOTTOM']      # These groups can be AND'ed with the OR_GROUPS and among them

COMET/format_dataset.py

@@ -30,7 +30,7 @@ SEG_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSL
 IMG_NAME_PATTERN = '(.*).nii.gz'
 SEG_NAME_PATTERN = '(.*).nii.gz'
 
-OUT_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128'
+OUT_DIRECTORY = '/mnt/EncryptedData1/Users/javier/ext_datasets/COMET_dataset/OSLO_COMET_CT/Formatted_128x128x128/w_bboxes'
 
 
 if __name__ == '__main__':
@@ -60,6 +60,7 @@ if __name__ == '__main__':
         seg = np.asarray(seg.dataobj)
 
         segs_are_ohe = bool(len(seg.shape) > 3 and seg.shape[3] > 1)
+        bbox = [0]*6
         if args.crop:
             parenchyma = regionprops(seg[..., 0])[0]
             bbox = np.asarray(parenchyma.bbox) + [*[-args.offset]*3, *[args.offset]*3]
@@ -72,10 +73,11 @@ if __name__ == '__main__':
         isot_shape = img.shape
 
         zoom_factors = (np.asarray([128]*3) / np.asarray(img.shape)).tolist()
-
+        img_isotropic = np.copy(img)
         img = zoom(img, zoom_factors, order=3)
         if args.dilate_segmentations:
             seg = binary_dilation(seg, binary_ball, iterations=1)
+        seg_isotropic = np.copy(seg)
         seg = zoom(seg, zoom_factors + [1]*(len(seg.shape) - len(img.shape)), order=0)
         zoom_file = zoom_file.append({'scale_i': zoom_factors[0],
                                       'scale_j': zoom_factors[1],
@@ -92,11 +94,15 @@ if __name__ == '__main__':
         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('image_isotropic', data=img_isotropic[..., np.newaxis], dtype=np.float32)
+        h5_file.create_dataset('segmentation_isotropic', data=seg_isotropic.astype(np.uint8), dtype=np.uint8)
         h5_file.create_dataset('segmentation', data=seg.astype(np.uint8), dtype=np.uint8)
         h5_file.create_dataset('segmentation_expanded', data=seg_expanded.astype(np.uint8), dtype=np.uint8)
         h5_file.create_dataset('segmentation_labels', data=np.unique(seg)[1:])  # Remove the 0 (background label)
         h5_file.create_dataset('isotropic_shape', data=isot_shape)
-
+        if args.crop:
+            h5_file.create_dataset('bounding_box_origin', data=bbox[:3])
+            h5_file.create_dataset('bounding_box_shape', data=bbox[3:] - bbox[:3])
         print('{}: Segmentation labels {}'.format(img_name, np.unique(seg)[1:]))
         h5_file.close()
 

DeepDeformationMapRegistration/losses.py

@@ -626,7 +626,8 @@ class StructuralSimilarityGaussian:
         self.__GF = tf.while_loop(lambda iterator, g_1d: tf.less(iterator, self.dim),
                                   lambda iterator, g_1d: (iterator + 1, tf.expand_dims(g_1d, -1) * tf.transpose(g_1d_expanded)),
                                   [iterator, g_1d],
-                                  [iterator.get_shape(), tf.TensorShape([None]*self.dim)]  # Shape invariants
+                                  [iterator.get_shape(), tf.TensorShape([None]*self.dim)],  # Shape invariants
+                                  back_prop=False,
                                   )[-1]
 
         self.__GF = tf.divide(self.__GF, tf.reduce_sum(self.__GF))  # Normalization
@@ -759,32 +760,41 @@ class GeneralizedDICEScore:
         Learning Los Function for Highly Unbalanced Segmentations" https://arxiv.org/abs/1707.03237
         :param input_shape: Shape of the input image, without the batch dimension, e.g., 2D: [H, W, C], 3D: [H, W, D, C]
         """
+        self.smooth = 1e-10  # If y_pred = y_true = null -> dice should be 1
+        self.num_labels = num_labels
         if input_shape[-1] > 1:
-            self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), input_shape[-1]]
-            self.hot_encode = False
+            try:
+                self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), input_shape[-1]]
+            except TypeError as err:
+                self.flat_shape = [-1, None, input_shape[-1]]
+            self.cardinal_encoded = False
         elif num_labels is not None:
-            self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), num_labels]
-            self.one_hot_enc_shape = [-1, *input_shape[:-1]]
-            self.hot_encode = True
-            warnings.warn('Differentiable one-hot encoding not yet implemented')
+            try:
+                self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), input_shape[-1]]
+            except TypeError as err:
+                self.flat_shape = [-1, None, input_shape[-1]]
+            self.cardinal_enc_shape = [-1, *input_shape[:-1]]
+            self.cardinal_encoded = True
+            warnings.warn('Differentiable cardinal encoding not yet implemented')
         else:
-            raise ValueError('If input_shape is not one hot encoded, then num_labels must be provided')
+            raise ValueError('If input_shape does not correspond to cardinally encoded,'
+                             'then num_labels must be provided')
 
     def one_hot_encoding(self, in_img, name=''):
         # TODO: Test if differentiable!
         labels, indices = tf.unique(tf.reshape(in_img, [-1]), tf.int32, name=name+'_unique')
-        one_hot = tf.one_hot(indices, tf.size(labels), name=name + '_one_hot')
-        one_hot = tf.reshape(one_hot, self.one_hot_enc_shape + [tf.size(labels)], name=name + '_reshape')
-        one_hot = tf.slice(one_hot, [0]*len(self.one_hot_enc_shape) + [1], [-1]*(len(self.one_hot_enc_shape) + 1),
+        one_hot = tf.one_hot(indices, self.num_labels, name=name + '_one_hot')
+        one_hot = tf.reshape(one_hot, self.cardinal_enc_shape + [self.num_labels], name=name + '_reshape')
+        one_hot = tf.slice(one_hot, [0] * len(self.cardinal_enc_shape) + [1], [-1] * (len(self.cardinal_enc_shape) + 1),
                            name=name + '_remove_bg')
         return one_hot
 
     def weigthed_dice(self, y_true, y_pred):
         # y_true = [B, -1, L]
         # y_pred = [B, -1, L]
-        if self.hot_encode:
-            y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
-            y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
+        # if self.cardinal_encoded:
+        #     y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
+        #     y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
         y_true = tf.reshape(y_true, self.flat_shape, name='GDICE_reshape_y_true')    # Flatten along the volume dimensions
         y_pred = tf.reshape(y_pred, self.flat_shape, name='GDICE_reshape_y_pred')    # Flatten along the volume dimensions
 
@@ -793,14 +803,14 @@ class GeneralizedDICEScore:
         w = tf.math.divide_no_nan(1., tf.pow(size_y_true, 2), name='GDICE_weight')
         numerator = w * tf.reduce_sum(y_true * y_pred, axis=1)
         denominator = w * (size_y_true + size_y_pred)
-        return tf.div_no_nan(2 * tf.reduce_sum(numerator, axis=-1), tf.reduce_sum(denominator, axis=-1))
+        return tf.div_no_nan(2 * tf.reduce_sum(numerator, axis=-1) + self.smooth, tf.reduce_sum(denominator, axis=-1) + self.smooth)
 
     def macro_dice(self, y_true, y_pred):
         # y_true = [B, -1, L]
         # y_pred = [B, -1, L]
-        if self.hot_encode:
-            y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
-            y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
+        # if self.cardinal_encoded:
+        #     y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
+        #     y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
         y_true = tf.reshape(y_true, self.flat_shape, name='GDICE_reshape_y_true')    # Flatten along the volume dimensions
         y_pred = tf.reshape(y_pred, self.flat_shape, name='GDICE_reshape_y_pred')    # Flatten along the volume dimensions
 
@@ -808,7 +818,7 @@ class GeneralizedDICEScore:
         size_y_pred = tf.reduce_sum(y_pred, axis=1, name='GDICE_size_y_pred')
         numerator = tf.reduce_sum(y_true * y_pred, axis=1)
         denominator = (size_y_true + size_y_pred)
-        return tf.div_no_nan(2 * numerator, denominator)
+        return tf.div_no_nan(2 * numerator + self.smooth, denominator + self.smooth)
 
     @function_decorator('GeneralizeDICE__loss')
     def loss(self, y_true, y_pred):

SoA_methods/eval_ants.py

@@ -19,6 +19,8 @@ from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
 from DeepDeformationMapRegistration.utils.visualization import save_disp_map_img, plot_predictions
 import DeepDeformationMapRegistration.utils.constants as C
 
+import medpy.metric as medpy_metrics
+
 import voxelmorph as vxm
 
 from argparse import ArgumentParser
@@ -35,22 +37,24 @@ WARPED_FIX = 'warpedfixout'
 FWD_TRFS = 'fwdtransforms'
 INV_TRFS = 'invtransforms'
 
-os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
-os.environ['CUDA_VISIBLE_DEVICES'] = '2'
 
 if __name__ == '__main__':
     parser = ArgumentParser()
     parser.add_argument('--dataset', type=str, help='Directory with the images')
     parser.add_argument('--outdir', type=str, help='Output directory')
+    parser.add_argument('--gpu', type=int, help='GPU')
     args = parser.parse_args()
 
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
     os.makedirs(args.outdir, exist_ok=True)
     os.makedirs(os.path.join(args.outdir, 'SyN'), exist_ok=True)
     os.makedirs(os.path.join(args.outdir, 'SyNCC'), exist_ok=True)
     dataset_files = os.listdir(args.dataset)
     dataset_files.sort()
     dataset_files = [os.path.join(args.dataset, f) for f in dataset_files if re.match(DATASET_NAMES, f)]
-
+    dataset_files.sort()
     dataset_iterator = tqdm(enumerate(dataset_files), desc="Running ANTs")
 
     f = h5py.File(dataset_files[0], 'r')
@@ -67,18 +71,18 @@ if __name__ == '__main__':
                    HausdorffDistanceErosion(3, 10, im_shape=image_shape + [nb_labels]).metric,
                    GeneralizedDICEScore(image_shape + [nb_labels], num_labels=nb_labels).metric_macro]
 
-    fix_img_ph = tf.placeholder(tf.float32, (1, *image_shape, 1), name='fix_img')
-    pred_img_ph = tf.placeholder(tf.float32, (1, *image_shape, 1), name='pred_img')
-    fix_seg_ph = tf.placeholder(tf.float32, (1, *image_shape, nb_labels), name='fix_seg')
-    pred_seg_ph = tf.placeholder(tf.float32, (1, *image_shape, nb_labels), name='pred_seg')
+    fix_img_ph = tf.placeholder(tf.float32, (1, None, None, None, 1), name='fix_img')
+    pred_img_ph = tf.placeholder(tf.float32, (1, None, None, None, 1), name='pred_img')
+    fix_seg_ph = tf.placeholder(tf.float32, (1, None, None, None, nb_labels), name='fix_seg')
+    pred_seg_ph = tf.placeholder(tf.float32, (1, None, None, None, 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)
+    # 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)
 
     config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
     config.gpu_options.allow_growth = True
@@ -90,7 +94,7 @@ if __name__ == '__main__':
     ####
     os.environ["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"] = "{:d}".format(os.cpu_count())  #https://github.com/ANTsX/ANTsPy/issues/261
     print("Running ANTs using {} threads".format(os.environ.get("ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS")))
-    dm_interp = DisplacementMapInterpolator(image_shape, 'griddata')
+    # dm_interp = DisplacementMapInterpolator(image_shape, 'griddata')
     # Header of the metrics csv file
     csv_header = ['File', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'DICE_MACRO', 'HD', 'Time', 'TRE']
 
@@ -147,10 +151,17 @@ if __name__ == '__main__':
 
                 dataset_iterator.set_description('{} ({}): Getting metrics {}'.format(file_num, file_path, reg_method))
                 with sess.as_default():
-                    dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf],
-                                                    {'fix_seg:0': fix_seg[np.newaxis, ...],  # Batch axis
-                                                     'pred_seg:0': pred_seg[np.newaxis, ...]  # Batch axis
-                                                     })
+                    dice = np.mean([medpy_metrics.dc(pred_seg[np.newaxis, ..., l], fix_seg[np.newaxis,..., l]) / np.sum(
+                        fix_seg[np.newaxis,..., l]) for l in range(nb_labels)])
+                    hd = np.mean(
+                        [medpy_metrics.hd(pred_seg[np.newaxis,..., l], fix_seg[np.newaxis,..., l]) for l in range(nb_labels)])
+                    dice_macro = np.mean(
+                        [medpy_metrics.dc(pred_seg[np.newaxis,..., l], fix_seg[np.newaxis,..., l]) for l in range(nb_labels)])
+
+                    # dice, hd, dice_macro = sess.run([dice_tf, hd_tf, dice_macro_tf],
+                    #                                 {'fix_seg:0': fix_seg[np.newaxis, ...],  # Batch axis
+                    #                                  'pred_seg:0': pred_seg[np.newaxis, ...]  # Batch axis
+                    #                                  })
 
                     pred_seg_card = segmentation_ohe_to_cardinal(pred_seg).astype(np.float32)
                     mov_seg_card = segmentation_ohe_to_cardinal(mov_seg).astype(np.float32)
@@ -160,26 +171,28 @@ if __name__ == '__main__':
                                                        {'fix_img:0': fix_img[np.newaxis, ...],  # Batch axis
                                                         'pred_img:0': pred_img[np.newaxis, ...]  # Batch axis
                                                         })
+                    ssim = np.mean(ssim)
                     ms_ssim = ms_ssim[0]
 
                     # TRE
                     disp_map = np.squeeze(np.asarray(nb.load(mov_to_fix_trf_list[0]).dataobj))
+                    dm_interp = DisplacementMapInterpolator(fix_img.shape[:-1], 'griddata', step=2)
                     pred_centroids = dm_interp(disp_map, mov_centroids, backwards=True) + mov_centroids
-                    upsample_scale = 128 / 64
-                    fix_centroids_isotropic = fix_centroids * upsample_scale
-                    pred_centroids_isotropic = pred_centroids * upsample_scale
+                    # upsample_scale = 128 / 64
+                    # fix_centroids_isotropic = fix_centroids * upsample_scale
+                    # pred_centroids_isotropic = pred_centroids * upsample_scale
 
-                    fix_centroids_isotropic = np.divide(fix_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
-                    pred_centroids_isotropic = np.divide(pred_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
-                    tre_array = target_registration_error(fix_centroids_isotropic, pred_centroids_isotropic, False).eval()
+                    # fix_centroids_isotropic = np.divide(fix_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
+                    # pred_centroids_isotropic = np.divide(pred_centroids_isotropic, C.COMET_DATASET_iso_to_cubic_scales)
+                    tre_array = target_registration_error(fix_centroids, pred_centroids, False).eval()
                     tre = np.mean([v for v in tre_array if not np.isnan(v)])
 
-                dataset_iterator.set_description('{} ({}): Saving data {}'.format(file_num, file_path, reg_method))
-                new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd,
-                            t1_syn-t0_syn if reg_method == 'SyN' else t1_syncc-t0_syncc,
-                            tre]
-                with open(metrics_file[reg_method], 'a') as f:
-                    f.write(';'.join(map(str, new_line))+'\n')
+                # dataset_iterator.set_description('{} ({}): Saving data {}'.format(file_num, file_path, reg_method))
+                # new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd,
+                #             t1_syn-t0_syn if reg_method == 'SyN' else t1_syncc-t0_syncc,
+                #             tre]
+                # with open(metrics_file[reg_method], 'a') as f:
+                #     f.write(';'.join(map(str, new_line))+'\n')
 
                 save_nifti(fix_img[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_fix_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
                 save_nifti(mov_img[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_mov_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
@@ -188,12 +201,17 @@ if __name__ == '__main__':
                 save_nifti(mov_seg_card[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
                 save_nifti(pred_seg_card[0, ...], os.path.join(args.outdir, reg_method, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
 
-                plot_predictions(fix_img[np.newaxis, ...], mov_img[np.newaxis, ...], disp_map[np.newaxis, ...], pred_img[np.newaxis, ...], os.path.join(args.outdir, reg_method, '{:03d}_figures_img.png'.format(step)), show=False)
-                plot_predictions(fix_seg[np.newaxis, ...], mov_seg[np.newaxis, ...], disp_map[np.newaxis, ...], pred_seg[np.newaxis, ...], os.path.join(args.outdir, reg_method, '{:03d}_figures_seg.png'.format(step)), show=False)
+                plot_predictions(img_batches=[fix_img[np.newaxis, ...], mov_img[np.newaxis, ...], pred_img[np.newaxis, ...]], disp_map_batch=disp_map[np.newaxis, ...], seg_batches=[fix_seg_card[np.newaxis, ...], mov_seg_card[np.newaxis, ...], pred_seg_card[np.newaxis, ...]], filename=os.path.join(args.outdir, reg_method, '{:03d}_figures_seg.png'.format(step)), show=False)
+                plot_predictions(img_batches=[fix_img[np.newaxis, ...], mov_img[np.newaxis, ...], pred_img[np.newaxis, ...]], disp_map_batch=disp_map[np.newaxis, ...], filename=os.path.join(args.outdir, reg_method, '{:03d}_figures_img.png'.format(step)), show=False)
                 save_disp_map_img(disp_map[np.newaxis, ...], 'Displacement map', os.path.join(args.outdir, reg_method, '{:03d}_disp_map_fig.png'.format(step)), show=False)
 
     for k in metrics_file.keys():
         print('Summary {}\n=======\n'.format(k))
-        print('\nAVG:\n' + str(pd.read_csv(metrics_file[k], sep=';', header=0).mean(axis=0)) + '\nSTD:\n' + str(
-            pd.read_csv(metrics_file[k], sep=';', header=0).std(axis=0)))
+        metrics_df = pd.read_csv(metrics_file[k], sep=';', header=0)
+        print('\nAVG:\n')
+        print(metrics_df.mean(axis=0))
+        print('\nSTD:\n')
+        print(metrics_df.std(axis=0))
+        print('\nHD95perc:\n')
+        print(metrics_df['HD'].describe(percentiles=[.95]))
         print('\n=======\n')