Works when providing the segmentation masks of the images

DeepDeformationMapRegistration/main.py

@@ -1,104 +1,238 @@
-
-# 1. Image files generator
-
-# timer start
-# 2. Preprocess the image
-# 3. Predict the displacement
-# timer stop
-
-# 4. Evaluate the registration: NCC; SSIM; DICE; HD95
-
+import datetime
 import os, sys
-
 import shutil
+import re
+import argparse
+import subprocess
+import logging
 import time
-import tkinter
-
-import h5py
-import matplotlib.pyplot as plt
 
 currentdir = os.path.dirname(os.path.realpath(__file__))
 parentdir = os.path.dirname(currentdir)
 sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
 
 import tensorflow as tf
-# tf.enable_eager_execution(config=config)
 
 import numpy as np
-import pandas as pd
+import nibabel as nib
+from scipy.ndimage import gaussian_filter, zoom
+from skimage.measure import regionprops
+import SimpleITK as sitk
+
 import voxelmorph as vxm
 from voxelmorph.tf.layers import SpatialTransformer
 
 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.losses import StructuralSimilarity_simplified, NCC
 from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
-from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccumulated
-from DeepDeformationMapRegistration.utils.visualization import save_disp_map_img, plot_predictions
-from DeepDeformationMapRegistration.utils.misc import DisplacementMapInterpolator, get_segmentations_centroids, segmentation_ohe_to_cardinal, segmentation_cardinal_to_ohe
-from 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
-import nibabel as nib
-from scipy.ndimage import gaussian_filter, zoom
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.misc import resize_displacement_map
+
+
+MODELS_FILE = {'L': {'BL-N': './models/liver/bl_ncc.h5',
+                     'BL-S': './models/liver/bl_ncc_ssim.h5',
+                     'SG-ND': './models/liver/sg_ncc_dsc.h5',
+                     'SD-NSD': './models/liver/sg_ncc_ssim_dsc.h5',
+                     'UW-NSD': './models/liver/uw_ncc_ssim_dsc.h5',
+                     'UW-NSDH': './models/liver/uw_ncc_ssim_dsc_hd.h5',
+                     },
+               'B': {'BL-N': './models/brain/bl_ncc.h5',
+                     'BL-S': './models/brain/bl_ncc_ssim.h5',
+                     'SG-ND': './models/brain/sg_ncc_dsc.h5',
+                     'SD-NSD': './models/brain/sg_ncc_ssim_dsc.h5',
+                     'UW-NSD': './models/brain/uw_ncc_ssim_dsc.h5',
+                     'UW-NSDH': './models/brain/uw_ncc_ssim_dsc_hd.h5',
+                     }
+               }
 
-import h5py
-import re
-from Brain_study.data_generator import BatchGenerator
+IMAGE_INTPUT_SHAPE = np.asarray([128, 128, 128, 1])
 
-import argparse
 
-from skimage.transform import warp
-import neurite as ne
+def rigidly_align_images(image_1: str, image_2: str) -> nib.Nifti1Image:
+    """
+    Rigidly align the images and resample to the same array size, to the dense displacement map is correct
 
-import tempfile
+    """
+    def resample_to_isotropic(image: sitk.Image) -> sitk.Image:
+        spacing = image.GetSpacing()
+        spacing = min(spacing)
+        resamp_spacing = [spacing] * image.GetDimension()
+        resamp_size = [int(round(or_size*or_space/spacing)) for or_size, or_space in zip(image.GetSize(), image.GetSpacing())]
+        return sitk.Resample(image,
+                             resamp_size, sitk.Transform(), sitk.sitkLinear,image.GetOrigin(),
+                             resamp_spacing, image.GetDirection(), 0, image.GetPixelID())
 
-import logging
+    image_1 = sitk.ReadImage(image_1, sitk.sitkFloat32)
+    image_2 = sitk.ReadImage(image_2, sitk.sitkFloat32)
+
+    image_1 = resample_to_isotropic(image_1)
+    image_2 = resample_to_isotropic(image_2)
+
+    rig_reg = sitk.ImageRegistrationMethod()
+    rig_reg.SetMetricAsMeanSquares()
+    rig_reg.SetOptimizerAsRegularStepGradientDescent(4.0, 0.01, 200)
+    rig_reg.SetInitialTransform(sitk.TranslationTransform(image_1.GetDimension()))
+    rig_reg.SetInterpolator(sitk.sitkLinear)
+
+    print('Running rigid registration...')
+    rig_reg_trf = rig_reg.Execute(image_1, image_2)
+    print('Rigid registration completed\n----------------------------')
+    print('Optimizer stop condition: {}'.format(rig_reg.GetOptimizerStopConditionDescription()))
+    print('Iteration: {}'.format(rig_reg.GetOptimizerIteration()))
+    print('Metric value: {}'.format(rig_reg.GetMetricValue()))
+
+    resampler = sitk.ResampleImageFilter()
+    resampler.SetReferenceImage(image_1)
+    resampler.SetInterpolator(sitk.sitkLinear)
+    resampler.SetDefaultPixelValue(100)
+    resampler.SetTransform(rig_reg_trf)
+
+    image_2 = resampler.Execute(image_2)
+
+    # TODO: Build a common image to hold both image_1 and image_2
+
+
+def pad_images(image_1: nib.Nifti1Image, image_2: nib.Nifti1Image):
+    """
+    Align image_1 and image_2 by the top left corner and pad them to the largest dimensions along the three axes
+    """
+    joint_image_shape = np.maximum(image_1.shape, image_2.shape)
+    pad_1 = [[0, p] for p in joint_image_shape - image_1.shape]
+    pad_2 = [[0, p] for p in joint_image_shape - image_2.shape]
+    image_1_padded = np.pad(image_1.dataobj, pad_1, mode='edge')
+    image_2_padded = np.pad(image_2.dataobj, pad_2, mode='edge')
+
+    return image_1_padded, image_2_padded
+
+
+def pad_displacement_map(disp_map: np.ndarray, crop_min: np.ndarray, crop_max: np.ndarray, output_shape: (np.ndarray, list)) -> np.ndarray:
+    padding = [[crop_min[i], image_shape_or[i] - crop_max[i]] for i in range(3)] + [[0, 0]]
+    return np.pad(disp_map, padding, mode='constant')
+
+
+def run_livermask(input_image_path, outputdir, filename: str = 'segmentation') -> np.ndarray:
+    logger.info('Getting parenchyma segmentations...')
+    shutil.copy2(input_image_path, os.path.join(outputdir, f'{filename}.nii.gz'))
+    livermask_cmd = "{} -m livermask.livermask --input {} --output {}".format(sys.executable,
+                                                                              input_image_path,
+                                                                              os.path.join(outputdir,
+                                                                                           f'{filename}.nii.gz'))
+    subprocess.run(livermask_cmd)
+    logger.info('done!')
+    segmentation_path = os.path.join(outputdir, f'{filename}.nii.gz')
+    return np.asarray(nib.load(segmentation_path).dataobj, dtype=int)
+
+
+def debug_save_image(image: (np.ndarray, nib.Nifti1Image), filename: str, outputdir: str, debug: bool = True):
+    def disp_map_modulus(disp_map, scale: float = None):
+        disp_map_mod = np.sqrt(np.sum(np.power(disp_map, 2), -1))
+        if scale:
+            min_disp = np.min(disp_map_mod)
+            max_disp = np.max(disp_map_mod)
+            disp_map_mod = disp_map_mod - min_disp / (max_disp - min_disp)
+            disp_map_mod *= scale
+            logger.debug('Scaled displacement map to [0., 1.] range')
+        return disp_map_mod
+
+    if debug:
+        os.makedirs(os.path.join(outputdir, 'debug'), exist_ok=True)
+        if image.shape[-1] > 1:
+            image = disp_map_modulus(image, 1.)
+        save_nifti(image, os.path.join(outputdir, 'debug', filename+'.nii.gz'), verbose=False)
+        logger.debug(f'Saved {filename} at {os.path.join(outputdir, filename+".nii.gz")}')
+
+
+def get_roi(image_filepath: str,
+            anatomy: str,
+            outputdir: str,
+            filename_filepath: str = 'segmentation',
+            segmentation_file: str = None,
+            debug: bool = False) -> list:
+    segm = None
+    if segmentation_file is None and anatomy == 'L':
+        segm = run_livermask(image_filepath, outputdir, filename_filepath)
+        logger.info(f'Loaded segmentation using livermask from {os.path.join(outputdir, filename_filepath)}')
+    elif segmentation_file is not None:
+        segm = np.asarray(nib.load(segmentation_file).dataobj, dtype=int)
+        logger.info(f'Loaded fixed segmentation from {segmentation_file}')
+    else:
+        logger.warning('No segmentation provided! Using the full volume')
+    if segm is not None:
+        segm[segm > 0] = 1
+        ret_val = regionprops(segm)[0].bbox
+        debug_save_image(segm, f'img_1_{filename_filepath}', outputdir, debug)
+    else:
+        ret_val = [0, 0, 0] + list(nib.load(image_filepath).shape[:3])
+    logger.debug(f'ROI found at coordinates {ret_val}')
+    return ret_val
 
-MODELS_FILE = {'liver': {'BL-N': './models/liver/bl_ncc.h5',
-                         'BL-S': './models/liver/bl_ssim.h5',
-                         'SG-ND': './models/liver/sg_ncc_dsc.h5',
-                         'SD-NSD': './models/liver/sg_ncc_ssim_dsc.h5',
-                         'UW-NSD': './models/liver/uw_ncc_ssim_dsc.h5',
-                         'UW-NSDH': './models/liver/uw_ncc_ssim_dsc_hd.h5',
-                         },
-               'brain': {'BL-N': './models/brain/bl_ncc.h5',
-                         'BL-S': './models/brain/bl_ssim.h5',
-                         'SG-ND': './models/brain/sg_ncc_dsc.h5',
-                         'SD-NSD': './models/brain/sg_ncc_ssim_dsc.h5',
-                         'UW-NSD': './models/brain/uw_ncc_ssim_dsc.h5',
-                         'UW-NSDH': './models/brain/uw_ncc_ssim_dsc_hd.h5',
-                         }
-               }
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('-f', '--fixed', type=str, help='Path to fixed image file (NIfTI)')
-    parser.add_argument('-m', '--moving', type=str, help='Path to oving image file (NIfTI)')
+    parser.add_argument('-m', '--moving', type=str, help='Path to moving segmentation image file (NIfTI)', default=None)
+    parser.add_argument('-F', '--fixedsegm', type=str, help='Path to fixed image segmentation file(NIfTI)',
+                        default=None)
+    parser.add_argument('-M', '--movingsegm', type=str, help='Path to moving image file (NIfTI)')
     parser.add_argument('-o', '--outputdir', type=str, help='Output directory', default='./Registration_output')
-    parser.add_argument('--gpu', type=int, help='In case of multi-GPU systems, limits the execution to the defined GPU number', default=None)
-    parser.add_argument('--model', type=str, help='Which model to use: BL-N, BL-S, SG-ND, SG-NSD, UW-NSD, UW-NSDH', default='UW-NSD')
+    parser.add_argument('-a', '--anatomy', type=str, help='Anatomical structure: liver (L) (Default) or brain (B)',
+                        default='L')
+    parser.add_argument('--gpu', type=int,
+                        help='In case of multi-GPU systems, limits the execution to the defined GPU number',
+                        default=None)
+    parser.add_argument('--model', type=str, help='Which model to use: BL-N, BL-S, SG-ND, SG-NSD, UW-NSD, UW-NSDH',
+                        default='UW-NSD')
+    parser.add_argument('--debug', '-d', action='store_true', help='Produce additional debug information', default=False)
+    parser.add_argument('-y', action='store_true', help='Erase output folder if this has content', default=False)
     # parser.add_argument('--brain', type=bool, action='store_true', help='Perform brain MRi registration', default=False)
     args = parser.parse_args()
-    logger = logging.getLogger(__name__)
 
     assert os.path.exists(args.fixed), 'Fixed image not found'
     assert os.path.exists(args.moving), 'Moving image not found'
     assert args.model in ['BL-N', 'BL-S', 'SG-ND', 'SG-NSD', 'UW-NSD', 'UW-NSDH'], 'Invalid model type'
+    assert args.anatomy in ['L', 'B'], 'Invalid anatomy option'
+
+    if os.path.exists(args.outputdir) and len(os.listdir(args.outputdir)):
+        if args.y:
+            erase = 'y'
+        else:
+            erase = input('Output directory is not empty, erase content? (y/n)')
+        if erase.lower() in ['y', 'yes']:
+            shutil.rmtree(args.outputdir, ignore_errors=True)
+            print('Erased directory: ' + args.outputdir)
+        elif erase.lower() in ['n', 'no']:
+            args.outputdir = os.path.join(args.outputdir, datetime.datetime.now().strftime('%H%M%S_%Y%m%d'))
+            print('New output directory: ' + args.outputdir)
+    os.makedirs(args.outputdir, exist_ok=True)
 
+    log_format = '%(asctime)s [%(levelname)s]:\t%(message)s'
+    logging.basicConfig(filename=os.path.join(args.outputdir, 'log.log'), filemode='w',
+                        format=log_format, datefmt='%Y-%m-%d %H:%M:%S')
+    logger = logging.getLogger(__name__)
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    stdout_handler.setFormatter(logging.Formatter(log_format, datefmt='%Y-%m-%d %H:%M:%S'))
+    logger.addHandler(stdout_handler)
     if isinstance(args.gpu, int):
         os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
         os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)  # Check availability before running using 'nvidia-smi'
+    if args.debug:
+        logger.setLevel('DEBUG')
+        logger.debug('DEBUG MODE ENABLED')
 
     # Load the file and preprocess it
-    fixed_image = nib.load(args.fixed)
-    moving_image = nib.load(args.moving)
+    logger.info('Loading image files')
+    fixed_image_or = nib.load(args.fixed)
+    moving_image_or = nib.load(args.moving)
+    image_shape_or = np.asarray(fixed_image_or.shape)
+    fixed_image_or, moving_image_or = pad_images(fixed_image_or, moving_image_or)
+    fixed_image_or = fixed_image_or[..., np.newaxis]  # add channel dim
+    moving_image_or = moving_image_or[..., np.newaxis]  # add channel dim
+    debug_save_image(fixed_image_or, 'img_0_loaded_fix_image', args.outputdir, args.debug)
+    debug_save_image(moving_image_or, 'img_0_loaded_moving_image', args.outputdir, args.debug)
 
     # TF stuff
+    logger.info('Setting up configuration')
     config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
     config.gpu_options.allow_growth = True
     config.log_device_placement = False  ## to log device placement (on which device the operation ran)
@@ -108,303 +242,140 @@ if __name__ == '__main__':
     tf.compat.v1.keras.backend.set_session(sess)
 
     # Preprocess data
-    if args.erase:
-        shutil.rmtree(args.outputdir, ignore_errors=True)
-    os.makedirs(args.outputdir, exist_ok=True)
-    lm_output_dir = os.path.join(args.outputdir, 'livermask')
-    os.makedirs(lm_output_dir, exist_ok=True)
-
     # 1. Run Livermask to get the mask around the liver in both the fixed and moving image
-    logger.info('Getting parenchyma segmentations...')
-    livermask_cmd = "python -m livermaks.livermask --input {} --output {}".format(args.fixed, os.path.join(lm_output_dir, 'fixed.nii.gz'))
-    os.system(livermask_cmd)
-    logger.info('... fixed image done...')
-    livermask_cmd = "python -m livermaks.livermask --input {} --output {}".format(args.moving, os.path.join(lm_output_dir, 'moving.nii.gz'))
-    os.system(livermask_cmd)
-    logger.info('... moving image done.')
-
-    # 2. Crop around the liver
-    # 2.1 Load the segmentations
-    # 2.2 Find the outermost box containing both boxes
-    # 2.3 Crop the fixed and moving images using such boxes
-    # 2.4 Resize the images to the expected input size
+    logger.info('Getting ROI')
+    fixed_segm_bbox = get_roi(args.fixed, args.anatomy, args.outputdir,
+                              'fixed_segmentation', args.fixedsegm, args.debug)
+    moving_segm_bbox = get_roi(args.moving, args.anatomy, args.outputdir,
+                               'moving_segmentation', args.movingsegm, args.debug)
+
+    crop_min = np.amin(np.vstack([fixed_segm_bbox[:3], moving_segm_bbox[:3]]), axis=0)
+    crop_max = np.amax(np.vstack([fixed_segm_bbox[3:], moving_segm_bbox[3:]]), axis=0)
+
+    # 2.2 Crop the fixed and moving images using such boxes
+    fixed_image = fixed_image_or[crop_min[0]: crop_max[0],
+                                 crop_min[1]: crop_max[1],
+                                 crop_min[2]: crop_max[2], ...]
+    debug_save_image(fixed_image, 'img_2_cropped_fixed_image', args.outputdir, args.debug)
+
+    moving_image = moving_image_or[crop_min[0]: crop_max[0],
+                                   crop_min[1]: crop_max[1],
+                                   crop_min[2]: crop_max[2], ...]
+    debug_save_image(moving_image, 'img_2_cropped_moving_image', args.outputdir, args.debug)
+
+    image_shape_crop = fixed_image.shape
+    # 2.3 Resize the images to the expected input size
+    zoom_factors = IMAGE_INTPUT_SHAPE / image_shape_crop
+    fixed_image = zoom(fixed_image, zoom_factors)
+    moving_image = zoom(moving_image, zoom_factors)
+    fixed_image = min_max_norm(fixed_image)
+    moving_image = min_max_norm(moving_image)
+    debug_save_image(fixed_image, 'img_3_preproc_fixed_image', args.outputdir, args.debug)
+    debug_save_image(moving_image, 'img_3_preproc_moving_image', args.outputdir, args.debug)
 
     # 3. Build the whole graph
-
-
-    # Loss and metric functions. Common to all models
-    # loss_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).loss,
-    #              NCC(image_input_shape).loss,
-    #              vxm.losses.MSE().loss,
-    #              MultiScaleStructuralSimilarity(max_val=1., filter_size=3).loss]
-    #
-    # metric_fncs = [StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric,
-    #                NCC(image_input_shape).metric,
-    #                vxm.losses.MSE().loss,
-    #                MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric,
-    #                GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric,
-    #                HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).metric,
-    #                GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).metric_macro]
-
+    logger.info('Building TF graph')
     ### METRICS GRAPH ###
-    # 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)
-    # 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', step=4)
-
-    for MODEL_FILE, DATA_ROOT_DIR in zip(MODEL_FILE_LIST, DATA_ROOT_DIR_LIST):
-        print('MODEL LOCATION: ', MODEL_FILE)
-
-        # data_folder = '/mnt/EncryptedData1/Users/javier/train_output/DDMR/THESIS/BASELINE_Affine_ncc___mse_ncc_160606-25022021'
-        output_folder = os.path.join(DATA_ROOT_DIR, args.outdirname)  # '/mnt/EncryptedData1/Users/javier/train_output/DDMR/THESIS/eval/BASELINE_TRAIN_Affine_ncc_EVAL_Affine'
-        # os.makedirs(os.path.join(output_folder, 'images'), exist_ok=True)
-
-        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,
-                                                              'AdamAccumulated': AdamAccumulated,
-                                                              'loss': loss_fncs,
-                                                              'metric': metric_fncs},
-                                                 compile=False)
-        except ValueError as e:
-            enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
-            dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
-            nb_features = [enc_features, dec_features]
-            if re.search('^UW|SEGGUIDED_', MODEL_FILE):
-                network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
-                                                                 nb_labels=nb_labels,
-                                                                 nb_unet_features=nb_features,
-                                                                 int_steps=0,
-                                                                 int_downsize=1,
-                                                                 seg_downsize=1)
-            else:
-                network = vxm.networks.VxmDense(inshape=image_output_shape,
-                                                nb_unet_features=nb_features,
-                                                int_steps=0)
-            network.load_weights(MODEL_FILE, by_name=True)
-        # Record metrics
-        metrics_file = os.path.join(output_folder, 'metrics.csv')
-        with open(metrics_file, 'w') as f:
-            f.write(';'.join(csv_header)+'\n')
-
-        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:
-                    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'][:]
-                    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()
-                    pred_img, disp_map, pred_seg = network.predict([mov_img, fix_img, mov_seg, fix_seg])    # predict([source, target])
-                    t1 = time.time()
-                else:
-                    t0 = time.time()
-                    pred_img, disp_map = network.predict([mov_img, fix_img])
-                    pred_seg = warp_segmentation.predict([mov_seg, disp_map])
-                    t1 = time.time()
-
-                pred_img = min_max_norm(pred_img)
-                mov_centroids, missing_lbls = get_segmentations_centroids(mov_seg[0, ...], ohe=True, expected_lbls=range(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 = 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
-                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)])
-                # ['File', 'SSIM', 'MS-SSIM', 'NCC', 'MSE', 'DICE', 'HD', 'Time', 'TRE', 'No_missing_lbls', 'Missing_lbls']
-
-                new_line = [step, ssim, ms_ssim, ncc, mse, dice, dice_macro, hd, t1-t0, tre, len(missing_lbls), missing_lbls]
-                with open(metrics_file, 'a') as f:
-                    f.write(';'.join(map(str, new_line))+'\n')
-
-                save_nifti(fix_img[0, ...], os.path.join(output_folder, '{:03d}_fix_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(mov_img[0, ...], os.path.join(output_folder, '{:03d}_mov_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(pred_img[0, ...], os.path.join(output_folder, '{:03d}_pred_img_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(fix_seg_card[0, ...], os.path.join(output_folder, '{:03d}_fix_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(mov_seg_card[0, ...], os.path.join(output_folder, '{:03d}_mov_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-                save_nifti(pred_seg_card[0, ...], os.path.join(output_folder, '{:03d}_pred_seg_ssim_{:.03f}_dice_{:.03f}.nii.gz'.format(step, ssim, dice)), verbose=False)
-
-                # with h5py.File(os.path.join(output_folder, '{:03d}_centroids.h5'.format(step)), 'w') as f:
-                    # f.create_dataset('fix_centroids', dtype=np.float32, data=fix_centroids)
-                    # f.create_dataset('mov_centroids', dtype=np.float32, data=mov_centroids)
-                    # f.create_dataset('pred_centroids', dtype=np.float32, data=pred_centroids)
-                    # f.create_dataset('fix_centroids_isotropic', dtype=np.float32, data=fix_centroids_isotropic)
-                    # f.create_dataset('mov_centroids_isotropic', dtype=np.float32, data=mov_centroids_isotropic)
-
-                # magnitude = np.sqrt(np.sum(disp_map[0, ...] ** 2, axis=-1))
-                # _ = plt.hist(magnitude.flatten())
-                # plt.title('Histogram of disp. magnitudes')
-                # plt.show(block=False)
-                # plt.savefig(os.path.join(output_folder, '{:03d}_hist_mag_ssim_{:.03f}_dice_{:.03f}.png'.format(step, ssim, dice)))
-                # plt.close()
-
-                plot_predictions(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('[FR] SSIM {:.04f}\tM_DICE: {:.04f}'.format(ssim, dice_macro))
-
-        print('Summary\n=======\n')
-        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
-    print('Done')
+    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')
+
+    ssim_tf = StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric(fix_img_ph, pred_img_ph)
+    ncc_tf = NCC(image_shape_or).metric(fix_img_ph, pred_img_ph)
+    mse_tf = vxm.losses.MSE().loss(fix_img_ph, pred_img_ph)
+    ms_ssim_tf = MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric(fix_img_ph, pred_img_ph)
+
+    logger.info(f'Using model: {"Brain" if args.anatomy == "B" else "Liver"} -> {args.model}')
+    MODEL_FILE = MODELS_FILE[args.anatomy][args.model]
+
+    # try:
+    #     network = tf.keras.models.load_model(MODEL_FILE,
+    #                                          {'VxmDense': vxm.networks.VxmDense,
+    #                                           # 'VxmDenseSemiSupervisedSeg': vxm.networks.VxmDenseSemiSupervisedSeg,
+    #                                           'AdamAccumulated': AdamAccumulated
+    #                                           },
+    #                                          compile=False)
+    # except ValueError as e:
+    #     enc_features = [32, 64, 128, 256, 512, 1024]  # const.ENCODER_FILTERS
+    #     dec_features = enc_features[::-1] + [16, 16]  # const.ENCODER_FILTERS[::-1]
+    #     nb_features = [enc_features, dec_features]
+    #     if re.search('^UW|SEGGUIDED_', MODEL_FILE):
+    #         network = vxm.networks.VxmDense(inshape=IMAGE_INTPUT_SHAPE[:-1],
+    #                                                    nb_unet_features=nb_features,
+    #                                                    int_steps=0,
+    #                                                    int_downsize=1,
+    #                                                    seg_downsize=1)
+    #     else:
+    #         network = vxm.networks.VxmDense(inshape=IMAGE_INTPUT_SHAPE[:-1],
+    #                                                    nb_unet_features=nb_features,
+    #                                                    int_steps=0)
+    #     network.load_weights(MODEL_FILE, by_name=True)
+
+    enc_features = [32, 64, 128, 256, 512, 1024]  # const.ENCODER_FILTERS
+    dec_features = enc_features[::-1] + [16, 16]  # const.ENCODER_FILTERS[::-1]
+    nb_features = [enc_features, dec_features]
+    network = vxm.networks.VxmDense(inshape=IMAGE_INTPUT_SHAPE[:-1],
+                                    nb_unet_features=nb_features,
+                                    int_steps=0)
+    network.load_weights(MODEL_FILE, by_name=True)
+    network.trainable = False
+
+    registration_model = network.get_registration_model()
+    deb_model = network.apply_transform
+
+    logger.info('Performing registration')
+    with sess.as_default():
+        if args.debug:
+            registration_model.summary(line_length=C.SUMMARY_LINE_LENGTH)
+        time_disp_map_start = time.time()
+        # disp_map = registration_model.predict([moving_image[np.newaxis, ...], fixed_image[np.newaxis, ...]])
+        p, disp_map = network.predict([moving_image[np.newaxis, ...], fixed_image[np.newaxis, ...]])
+        time_disp_map_end = time.time()
+        debug_save_image(np.squeeze(disp_map), 'disp_map_0_raw', args.outputdir, args.debug)
+        debug_save_image(p[0, ...], 'img_4_net_pred_image', args.outputdir, args.debug)
+        # pred_image = min_max_norm(pred_image)
+        # pred_image_isot = zoom(pred_image[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+        # fixed_image_isot = zoom(fixed_image[0, ...], zoom_factors, order=3)[np.newaxis, ...]
+
+        # Up sample the displacement map to the full res
+        trf = np.eye(4)
+        np.fill_diagonal(trf, 1/zoom_factors)
+        disp_map = resize_displacement_map(np.squeeze(disp_map), None, trf)
+        debug_save_image(np.squeeze(disp_map), 'disp_map_1_upsampled', args.outputdir, args.debug)
+        disp_map_or = pad_displacement_map(disp_map, crop_min, crop_max, image_shape_or)
+        debug_save_image(np.squeeze(disp_map_or), 'disp_map_2_padded', args.outputdir, args.debug)
+        disp_map_or = gaussian_filter(disp_map_or, 5)
+        debug_save_image(np.squeeze(disp_map_or), 'disp_map_3_smoothed', args.outputdir, args.debug)
+
+        time_pred_img_start = time.time()
+        pred_image = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([moving_image_or[np.newaxis, ...], disp_map_or[np.newaxis, ...]]).eval()
+        time_pred_img_end = time.time()
+        ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
+                                           {'fix_img:0': fixed_image_or[np.newaxis, ...], 'pred_img:0': pred_image})
+        ssim = np.mean(ssim)
+        ms_ssim = ms_ssim[0]
+        pred_image = pred_image[0, ...]
+
+        save_nifti(pred_image, os.path.join(args.outputdir, 'pred_image.nii.gz'))
+        np.savez_compressed(os.path.join(args.outputdir, 'displacement_map.npz'), disp_map_or)
+        logger.info('Predicted image (full image) and displacement map saved in: '.format(args.outputdir))
+        logger.info(f'Displacement map prediction time: {time_disp_map_end - time_disp_map_start} s')
+        logger.info(f'Predicted image time: {time_pred_img_end - time_pred_img_start} s')
+
+        logger.info('Similarity metrics (Full image)\n------------------')
+        logger.info('SSIM: {:.03f}'.format(ssim))
+        logger.info('NCC: {:.03f}'.format(ncc))
+        logger.info('MSE: {:.03f}'.format(mse))
+        logger.info('MS SSIM: {:.03f}'.format(ms_ssim))
+
+        ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
+                                           {'fix_img:0': fixed_image[np.newaxis, ...], 'pred_img:0': p})
+        ssim = np.mean(ssim)
+        ms_ssim = ms_ssim[0]
+        logger.info('\nSimilarity metrics (ROI)\n------------------')
+        logger.info('SSIM: {:.03f}'.format(ssim))
+        logger.info('NCC: {:.03f}'.format(ncc))
+        logger.info('MSE: {:.03f}'.format(mse))
+        logger.info('MS SSIM: {:.03f}'.format(ms_ssim))
+
+    del registration_model
+    logger.info('Done')