Merge pull request #19 from jpdefrutos/clean_repo

DeepDeformationMapRegistration/main.py

@@ -0,0 +1,418 @@
+import datetime
+import os, sys
+import shutil
+import re
+import argparse
+import subprocess
+import logging
+import time
+import warnings
+
+# 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
+
+import numpy as np
+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.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC
+from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.misc import resize_displacement_map
+from DeepDeformationMapRegistration.utils.model_utils import get_models_path, load_model
+from DeepDeformationMapRegistration.utils.logger import LOGGER
+
+from importlib.util import find_spec
+
+
+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
+
+    """
+    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())
+
+    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').astype(np.float32)
+    image_2_padded = np.pad(image_2.dataobj, pad_2, mode='edge').astype(np.float32)
+
+    return image_1_padded, image_2_padded
+
+
+def pad_crop_to_original_shape(crop_image: np.asarray, output_shape: [tuple, np.asarray], top_left_corner: [tuple, np.asarray]):
+    """
+    Pad crop_image so the output image has output_shape with the crop where it originally was found
+    """
+    output_shape = np.asarray(output_shape)
+    top_left_corner = np.asarray(top_left_corner)
+
+    pad = [[c, o - (c + i)] for c, o, i in zip(top_left_corner[:3], output_shape[:3], crop_image.shape[:3])]
+    if len(crop_image.shape) == 4:
+        pad += [[0, 0]]
+    return np.pad(crop_image, pad, mode='constant', constant_values=np.min(crop_image)).astype(crop_image.dtype)
+
+
+def pad_displacement_map(disp_map: np.ndarray, crop_min: np.ndarray, crop_max: np.ndarray, output_shape: (np.ndarray, list)) -> np.ndarray:
+    ret_val = disp_map
+    if np.all([d != i for d, i in zip(disp_map.shape[:3], output_shape)]):
+        padding = [[crop_min[i], max(0, output_shape[i] - crop_max[i])] for i in range(3)] + [[0, 0]]
+        ret_val = np.pad(disp_map, padding, mode='constant')
+    return ret_val
+
+
+def run_livermask(input_image_path, outputdir, filename: str = 'segmentation') -> np.ndarray:
+    assert find_spec('livermask'), 'Livermask is not available'
+    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,
+            compute_segmentation: bool,
+            outputdir: str,
+            filename_filepath: str = 'segmentation',
+            segmentation_file: str = None,
+            debug: bool = False) -> list:
+    segm = None
+    if segmentation_file is None and compute_segmentation:
+        LOGGER.info(f'Computing segmentation using livermask. Only for liver in abdominal CTs')
+        try:
+            segm = run_livermask(image_filepath, outputdir, filename_filepath)
+            LOGGER.info(f'Loaded segmentation using livermask from {os.path.join(outputdir, filename_filepath)}')
+        except (AssertionError, FileNotFoundError) as er:
+            LOGGER.warning(er)
+            LOGGER.warning('No segmentation provided! Using the full volume')
+            pass
+    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
+
+
+def 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 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('-a', '--anatomy', type=str, help='Anatomical structure: liver (L) (Default) or brain (B)',
+                        default='L')
+    parser.add_argument('-s', '--make-segmentation', action='store_true', help='Try to create a segmentation for liver in CT images', default=False)
+    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('-d', '--debug', action='store_true', help='Produce additional debug information', default=False)
+    parser.add_argument('-c', '--clear-outputdir', action='store_true', help='Clear output folder if this has content', default=False)
+    parser.add_argument('--original-resolution', action='store_true',
+                        help='Re-scale the displacement map to the originla resolution and apply it to the original moving image. WARNING: longer processing time.',
+                        default=False)
+    parser.add_argument('--save-displacement-map', action='store_true', help='Save the displacement map. An NPZ file will be created.',
+                        default=False)
+    args = parser.parse_args()
+
+    assert os.path.exists(args.fixed), 'Fixed image not found'
+    assert os.path.exists(args.moving), 'Moving image not found'
+    assert args.model in C.MODEL_TYPES.keys(), 'Invalid model type'
+    assert args.anatomy in C.ANATOMIES.keys(), 'Invalid anatomy option'
+    if os.path.exists(args.outputdir) and len(os.listdir(args.outputdir)):
+        if args.clear_outputdir:
+            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')
+
+    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'
+    LOGGER.setLevel('INFO')
+    if args.debug:
+        LOGGER.setLevel('DEBUG')
+        LOGGER.debug('DEBUG MODE ENABLED')
+
+    if args.original_resolution:
+        LOGGER.info('The results will be rescaled back to the original image resolution. '
+                    'Expect longer post-processing times.')
+    else:
+        LOGGER.info(f'The results will NOT be rescaled. Output shape will be {C.IMG_SHAPE[:3]}.')
+
+    # Load the file and preprocess it
+    LOGGER.info('Loading image files')
+    fixed_image_or = nib.load(args.fixed)
+    moving_image_or = nib.load(args.moving)
+    moving_image_header = moving_image_or.header.copy()
+    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)
+    config.allow_soft_placement = True
+
+    sess = tf.compat.v1.Session(config=config)
+    tf.compat.v1.keras.backend.set_session(sess)
+
+    # Preprocess data
+    # 1. Run Livermask to get the mask around the liver in both the fixed and moving image
+    LOGGER.info('Getting ROI')
+    fixed_segm_bbox = get_roi(args.fixed, args.make_segmentation, args.outputdir,
+                              'fixed_segmentation', args.fixedsegm, args.debug)
+    moving_segm_bbox = get_roi(args.moving, args.make_segmentation, 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 = np.asarray(C.IMG_SHAPE) / np.asarray(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
+    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')
+
+    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'Getting model: {"Brain" if args.anatomy == "B" else "Liver"} -> {args.model}')
+    MODEL_FILE = get_models_path(args.anatomy, args.model, os.getcwd())  # MODELS_FILE[args.anatomy][args.model]
+
+    network, registration_model = load_model(MODEL_FILE, False, True)
+    deb_model = network.apply_transform
+
+    LOGGER.info('Computing registration')
+    with sess.as_default():
+        if args.debug:
+            registration_model.summary(line_length=C.SUMMARY_LINE_LENGTH)
+        LOGGER.info('Computing displacement map...')
+        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()
+        LOGGER.info(f'\t... done ({time_disp_map_end - time_disp_map_start})')
+        disp_map = np.squeeze(disp_map)
+        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, ...]
+
+        LOGGER.info('Applying displacement map...')
+        time_pred_img_start = time.time()
+        pred_image = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]]).eval()
+        time_pred_img_end = time.time()
+        LOGGER.info(f'\t... done ({time_pred_img_end - time_pred_img_start} s)')
+        pred_image = pred_image[0, ...]
+
+        if args.original_resolution:
+            LOGGER.info('Scaling predicted image...')
+            moving_image = moving_image_or
+            fixed_image = fixed_image_or
+            # disp_map = disp_map_or
+            pred_image = zoom(pred_image, 1/zoom_factors)
+            pred_image = pad_crop_to_original_shape(pred_image, fixed_image_or.shape, crop_min)
+            LOGGER.info('Done...')
+
+        LOGGER.info('Computing metrics...')
+        if args.original_resolution:
+            ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
+                                               {'fix_img:0': fixed_image[np.newaxis,
+                                                                         crop_min[0]: crop_max[0],
+                                                                         crop_min[1]: crop_max[1],
+                                                                         crop_min[2]: crop_max[2],
+                                                                         ...],
+                                                'pred_img:0': pred_image[np.newaxis,
+                                                                         crop_min[0]: crop_max[0],
+                                                                         crop_min[1]: crop_max[1],
+                                                                         crop_min[2]: crop_max[2],
+                                                                         ...]})  # to only compare the deformed region!
+        else:
+
+            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': pred_image[np.newaxis, ...]})
+        ssim = np.mean(ssim)
+        ms_ssim = ms_ssim[0]
+
+        if args.original_resolution:
+            save_nifti(pred_image, os.path.join(args.outputdir, 'pred_image.nii.gz'), header=moving_image_header)
+        else:
+            save_nifti(pred_image, os.path.join(args.outputdir, 'pred_image.nii.gz'))
+            save_nifti(fixed_image, os.path.join(args.outputdir, 'fixed_image.nii.gz'))
+            save_nifti(moving_image, os.path.join(args.outputdir, 'moving_image.nii.gz'))
+
+        if args.save_displacement_map or args.debug:
+            if args.original_resolution:
+                # Up sample the displacement map to the full res
+                LOGGER.info('Scaling displacement map...')
+                trf = np.eye(4)
+                np.fill_diagonal(trf, 1 / zoom_factors)
+                disp_map = resize_displacement_map(disp_map, None, trf, moving_image_header.get_zooms())
+                debug_save_image(disp_map, 'disp_map_1_upsampled', args.outputdir, args.debug)
+                disp_map = pad_displacement_map(disp_map, crop_min, crop_max, image_shape_or)
+                debug_save_image(np.squeeze(disp_map), 'disp_map_2_padded', args.outputdir, args.debug)
+                disp_map = gaussian_filter(disp_map, 5)
+                debug_save_image(np.squeeze(disp_map), 'disp_map_3_smoothed', args.outputdir, args.debug)
+                LOGGER.info('\t... done')
+            if args.debug:
+                np.savez_compressed(os.path.join(args.outputdir, 'displacement_map.npz'), disp_map)
+            else:
+                np.savez_compressed(os.path.join(os.path.join(args.outputdir, 'debug'), 'displacement_map.npz'), disp_map)
+        LOGGER.info('Predicted 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\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')
+    exit(0)
+
+
+if __name__ == '__main__':
+    main()

DeepDeformationMapRegistration/networks.py

@@ -1,9 +1,9 @@
 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
-
-PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+# 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
+#
+# PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
 
 import tensorflow as tf
 import voxelmorph as vxm

DeepDeformationMapRegistration/utils/constants.py

@@ -30,7 +30,7 @@ PRED_IMG_GT = 1
 DISP_VECT_GT = 2
 DISP_VECT_LOC_GT = 3
 
-IMG_SIZE = 64  # Assumed a square image
+IMG_SIZE = 128  # Assumed a square image
 IMG_SHAPE = (IMG_SIZE, IMG_SIZE, IMG_SIZE, 1)  # (IMG_SIZE, IMG_SIZE, 1)
 DISP_MAP_SHAPE = (IMG_SIZE, IMG_SIZE, IMG_SIZE, 3)
 BATCH_SHAPE = (None, IMG_SIZE, IMG_SIZE, IMG_SIZE, 2)  # Expected batch shape by the network
@@ -196,8 +196,8 @@ DROPOUT = True
 DROPOUT_RATE = 0.2
 MAX_DATA_SIZE = (1000, 1000, 1)
 PLATEAU_THR = 0.01  # A slope between +-PLATEAU_THR will be considered a plateau for the LR updating function
-ENCODER_FILTERS = [4, 8, 16, 32, 64]
-
+ENCODER_FILTERS = [32, 64, 128, 256, 512, 1024]
+DECODER_FILTERS = ENCODER_FILTERS[::-1] + [16, 16]
 # SSIM
 SSIM_FILTER_SIZE = 11  # Size of Gaussian filter
 SSIM_FILTER_SIGMA = 1.5  # Width of Gaussian filter
@@ -205,7 +205,7 @@ SSIM_K1 = 0.01  # Def. 0.01
 SSIM_K2 = 0.03  # Recommended values 0 < K2 < 0.4
 MAX_VALUE = 1.0  # Maximum intensity values
 
-# Mathematic constants
+# Mathematics constants
 EPS = 1e-8
 EPS_tf = tf.constant(EPS, dtype=tf.float32)
 LOG2 = tf.math.log(tf.constant(2, dtype=tf.float32))
@@ -523,3 +523,11 @@ GAMMA_AUGMENTATION = True
 BRIGHTNESS_AUGMENTATION = False
 NUM_CONTROL_PTS_AUG = 10
 NUM_AUGMENTATIONS = 1
+
+ANATOMIES = {'L': 'liver', 'B': 'brain'}
+MODEL_TYPES = {'BL-N': 'bl_ncc',
+               'BL-NS': 'bl_ncc_ssim',
+               'SG-ND': 'sg_ncc_dsc',
+               'SG-NSD': 'sg_ncc_ssim_dsc',
+               'UW-NSD': 'uw_ncc_ssim_dsc',
+               'UW-NSDH': 'uw_ncc_ssim_dsc_hd'}

DeepDeformationMapRegistration/utils/misc.py

@@ -167,7 +167,7 @@ def segmentation_cardinal_to_ohe(segmentation, labels_list: list = None):
     return cpy
 
 
-def resize_displacement_map(displacement_map: np.ndarray, dest_shape: [list, np.ndarray, tuple], scale_trf: np.ndarray=None):
+def resize_displacement_map(displacement_map: np.ndarray, dest_shape: [list, np.ndarray, tuple], scale_trf: np.ndarray = None, resolution_factors: [tuple, np.ndarray] = np.ones((3,))):
     if scale_trf is None:
         scale_trf = scale_transformation(displacement_map.shape, dest_shape)
     else:
@@ -175,11 +175,12 @@ def resize_displacement_map(displacement_map: np.ndarray, dest_shape: [list, np.
     zoom_factors = scale_trf.diagonal()
     # First scale the values, so we cut down the number of multiplications
     dm_resized = np.copy(displacement_map)
-    dm_resized[..., 0] *= zoom_factors[0]
-    dm_resized[..., 1] *= zoom_factors[1]
-    dm_resized[..., 2] *= zoom_factors[2]
     # Then rescale using zoom
     dm_resized = zoom(dm_resized, zoom_factors)
+    dm_resized *= np.asarray(resolution_factors)
+    # dm_resized[..., 0] *= resolution_factors[0]
+    # dm_resized[..., 1] *= resolution_factors[1]
+    # dm_resized[..., 2] *= resolution_factors[2]
     return dm_resized
 
 

DeepDeformationMapRegistration/utils/model_utils.py

@@ -0,0 +1,65 @@
+import os
+import requests
+from datetime import datetime
+from email.utils import parsedate_to_datetime, formatdate
+from DeepDeformationMapRegistration.utils.constants import ANATOMIES, MODEL_TYPES, ENCODER_FILTERS, DECODER_FILTERS, IMG_SHAPE
+import voxelmorph as vxm
+from DeepDeformationMapRegistration.utils.logger import LOGGER
+
+
+# taken from: https://lenon.dev/blog/downloading-and-caching-large-files-using-python/
+def download(url, destination_file):
+    headers = {}
+
+    if os.path.exists(destination_file):
+        mtime = os.path.getmtime(destination_file)
+        headers["if-modified-since"] = formatdate(mtime, usegmt=True)
+
+    response = requests.get(url, headers=headers, stream=True)
+    response.raise_for_status()
+
+    if response.status_code == requests.codes.not_modified:
+        return
+
+    if response.status_code == requests.codes.ok:
+        with open(destination_file, "wb") as f:
+            for chunk in response.iter_content(chunk_size=1048576):
+                f.write(chunk)
+
+        last_modified = response.headers.get("last-modified")
+        if last_modified:
+            new_mtime = parsedate_to_datetime(last_modified).timestamp()
+            os.utime(destination_file, times=(datetime.now().timestamp(), new_mtime))
+
+
+def get_models_path(anatomy: str, model_type: str, output_root_dir: str):
+    assert anatomy in ANATOMIES.keys(), 'Invalid anatomy'
+    assert model_type in MODEL_TYPES.keys(), 'Invalid model type'
+    anatomy = ANATOMIES[anatomy]
+    model_type = MODEL_TYPES[model_type]
+    url = 'https://github.com/jpdefrutos/DDMR/releases/download/trained_models_v0/' + anatomy + '_' + model_type + '.h5'
+    file_path = os.path.join(output_root_dir, 'models', anatomy, model_type + '.h5')
+    if not os.path.exists(file_path):
+        LOGGER.info(f'Model not found. Downloading from {url}... ')
+        os.makedirs(os.path.split(file_path)[0], exist_ok=True)
+        download(url, file_path)
+        LOGGER.info(f'... downloaded model. Stored in {file_path}')
+    else:
+        LOGGER.info(f'Found model: {file_path}')
+    return file_path
+
+
+def load_model(weights_file_path: str, trainable: bool = False, return_registration_model: bool=True):
+    assert os.path.exists(weights_file_path), f'File {weights_file_path} not found'
+    assert weights_file_path.endswith('h5'), 'Invalid file extension. Expected .h5'
+
+    ret_val = vxm.networks.VxmDense(inshape=IMG_SHAPE[:-1],
+                                    nb_unet_features=[ENCODER_FILTERS, DECODER_FILTERS],
+                                    int_steps=0)
+    ret_val.load_weights(weights_file_path, by_name=True)
+    ret_val.trainable = trainable
+
+    if return_registration_model:
+        ret_val = (ret_val, ret_val.get_registration_model())
+
+    return ret_val

setup.py

@@ -0,0 +1,45 @@
+from setuptools import find_packages, setup
+import os
+
+# with open("README.md", "r", encoding="utf-8") as f:
+#     long_description = f.read()
+
+with open("requirements.txt", "r") as f:
+    reqs = [line.strip('\n') for line in f.readlines()]
+
+setup(
+    name='DeepDeformationMapRegistration',
+    py_modules=['DeepDeformationMapRegistration'],
+    packages=find_packages(include=['DeepDeformationMapRegistration', 'DeepDeformationMapRegistration.*'],
+                           exclude=['test_images', 'test_images.*']),
+    version='1.0',
+    description='Deep-registration training toolkit',
+    # long_description=long_description,
+    author='Javier Pérez de Frutos',
+    classifiers=[
+        'Programming language :: Python :: 3',
+        'License :: OSI Approveed :: MIT License',
+        'Operating System :: OS Independent'
+    ],
+    python_requires='>=3.6',
+    install_requires=[
+        'fastrlock>=0.3',   # required by cupy-cuda110
+        'testresources',    # required by launchpadlib
+        'scipy',
+        'scikit-image',
+        'simpleITK',
+        'voxelmorph==0.1',
+        'pystrum==0.1',
+        'tensorflow-gpu==1.14.0',
+        'tensorflow-addons',
+        'tensorflow-datasets',
+        'tensorflow-metadata',
+        'tensorboard==1.14.0',
+        'nibabel==3.2.1',
+        'numpy==1.18.5',
+        'h5py==2.10'
+    ],
+    entry_points={
+        'console_scripts': ['ddmr=DeepDeformationMapRegistration.main:main']
+    }
+)