3D IRCA dataset formatting scripts

Datasets/check_dataset.py

@@ -0,0 +1,62 @@
+import os
+import h5py
+import numpy as np
+from tqdm import tqdm
+import DeepDeformationMapRegistration.utils.constants as C
+
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+LITS_NONE = '/mnt/EncryptedData1/Users/javier/vessel_registration/LiTS/None'
+LITS_TRANS = '/mnt/EncryptedData1/Users/javier/vessel_registration/LiTS/Translation'
+LITS_AFFINE = '/mnt/EncryptedData1/Users/javier/vessel_registration/LiTS/Affine'
+
+
+IMG_SHAPE = (64, 64, 64, 1)
+for dataset in [LITS_NONE, LITS_AFFINE, LITS_TRANS]:
+    dataset_files = [os.path.join(dataset, d) for d in os.listdir(dataset) if os.path.isfile(os.path.join(dataset, d))]
+    f_iter = tqdm(dataset_files)
+    f_iter.set_description('Analyzing ' + dataset)
+    inv_shape_count = 0
+    inv_type_count = 0
+    for i, d in enumerate(f_iter):
+        f = h5py.File(d, 'r')
+        if f[C.H5_FIX_IMG][:].shape != IMG_SHAPE:
+            print(d + ' Invalid FIX IMG. Shape: ' + str(f[C.H5_FIX_IMG][:].shape))
+            inv_shape_count += 1
+        if f[C.H5_MOV_IMG][:].shape != IMG_SHAPE:
+            print(d + ' Invalid MOV IMG. Shape: ' + str(f[C.H5_MOV_IMG][:].shape))
+            inv_shape_count += 1
+        if f[C.H5_FIX_PARENCHYMA_MASK][:].shape != IMG_SHAPE:
+            print(d + ' Invalid FIX PARENCHYMA. Shape: ' + str(f[C.H5_FIX_PARENCHYMA_MASK][:].shape))
+            inv_shape_count += 1
+        if f[C.H5_MOV_PARENCHYMA_MASK][:].shape != IMG_SHAPE:
+            print(d + ' Invalid MOV PARENCHYMA. Shape: ' + str(f[C.H5_MOV_PARENCHYMA_MASK][:].shape))
+            inv_shape_count += 1
+        if f[C.H5_FIX_TUMORS_MASK][:].shape != IMG_SHAPE:
+            print(d + ' Invalid FIX TUMORS. Shape: ' + str(f[C.H5_FIX_TUMORS_MASK][:].shape))
+            inv_shape_count += 1
+        if f[C.H5_MOV_TUMORS_MASK][:].shape != IMG_SHAPE:
+            print(d + ' Invalid MOV TUMORS. Shape: ' + str(f[C.H5_MOV_TUMORS_MASK][:].shape))
+            inv_shape_count += 1
+
+        if f[C.H5_FIX_IMG][:].dtype != np.float32:
+            print(d + ' Invalid FIX IMG. Type: ' + str(f[C.H5_FIX_IMG][:].dtype))
+            inv_type_count += 1
+        if f[C.H5_MOV_IMG][:].dtype != np.float32:
+            print(d + ' Invalid MOV IMG. Type: ' + str(f[C.H5_MOV_IMG][:].dtype))
+            inv_type_count += 1
+        if f[C.H5_FIX_PARENCHYMA_MASK][:].dtype != np.float32:
+            print(d + ' Invalid FIX PARENCHYMA. Type: ' + str(f[C.H5_FIX_PARENCHYMA_MASK][:].dtype))
+            inv_type_count += 1
+        if f[C.H5_MOV_PARENCHYMA_MASK][:].dtype != np.float32:
+            print(d + ' Invalid MOV PARENCHYMA. Type: ' + str(f[C.H5_MOV_PARENCHYMA_MASK][:].dtype))
+            inv_type_count += 1
+        if f[C.H5_FIX_TUMORS_MASK][:].dtype != np.float32:
+            print(d + ' Invalid FIX TUMORS. Type: ' + str(f[C.H5_FIX_TUMORS_MASK][:].dtype))
+            inv_type_count += 1
+        if f[C.H5_MOV_TUMORS_MASK][:].dtype != np.float32:
+            print(d + ' Invalid MOV TUMORS. Type: ' + str(f[C.H5_MOV_TUMORS_MASK][:].dtype))
+            inv_type_count += 1
+
+    print('\n\n>>>>SUMMARY ' + dataset)
+    print('\t\tInvalid shape: ' + str(inv_shape_count) + '\n\t\tInvalid type: ' + str(inv_type_count))

Datasets/irca_pre_processing.py

@@ -0,0 +1,145 @@
+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 h5py
+import numpy as np
+import zipfile
+import re
+import dicom2nifti as d2n
+import nibabel as nib
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from DeepDeformationMapRegistration.utils.misc import try_mkdir
+from tqdm import tqdm
+import shutil
+
+
+IRCA_PATH = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/'
+
+SEGMENTATIONS = ('venoussystem', 'venacava', 'portalvein', 'liver', 'livertumor')
+
+SEGS_VESSELS = ('venoussystem', 'venacava', 'portalvein')
+SEGS_PARENCH = ('liver',)
+SEGS_TUMOR = ('livertumor', 'tumor', 'liverkyst', 'livercyst')
+
+SEGMENTATIONS = (SEGS_PARENCH + SEGS_VESSELS + SEGS_TUMOR)
+
+DEST_FOLDER = os.path.join(IRCA_PATH, 'nifti3')
+
+ZIP_EXT = '.zip'
+NIFTI_EXT = '.nii.gz'
+H5_EXT = '.hd5'
+
+PATIENT_DICOM = 'PATIENT_DICOM'
+TEMP = 'temp'
+SEGS_DICOM = 'MASKS_DICOM'
+
+CONVERTED_FILE = 'none' + NIFTI_EXT
+VOL_FILE = 'volume-{:04d}' + NIFTI_EXT
+SEG_FILE = 'segmentation-{:04d}' + NIFTI_EXT
+
+SEG_TO_CT_ORIENTATION_MAT = np.eye(4)
+SEG_TO_CT_ORIENTATION_MAT[0] = -1
+
+
+def merge_segmentations(file_list):
+    nib_file = nib.concat_images(file_list)
+    np_file = np.asarray(nib_file.dataobj)
+    np_file = np.sign(np.sum(np_file, -1)) * np.max(np_file)
+    return nib.Nifti1Image(np_file, nib_file.affine)
+
+
+if __name__ == '__main__':
+    # 1. List of folders
+    folder_list = [os.path.join(IRCA_PATH, d) for d in os.listdir(IRCA_PATH) if d.lower().startswith('3dircadb1.')]
+    folder_list.sort()
+
+    try_mkdir(DEST_FOLDER)
+    folder_iter = tqdm(folder_list)
+    for pat_dir in folder_iter:
+        pat_dir = folder_list[13]
+        i = int(pat_dir.split('.')[-1])
+        # 2. Unzip PATIENT_DICOM.zip
+        temp_folder = os.path.join(pat_dir, TEMP)
+        folder_iter.set_description('Volume DICOM: Unzipping PATIENT_DICOM.zip')
+        zipfile.ZipFile(os.path.join(pat_dir, PATIENT_DICOM + ZIP_EXT)).extractall(temp_folder)
+
+        folder_iter.set_description('Volume DICOM: Converting DICOM to Nifti')
+        d2n.convert_directory(os.path.join(temp_folder, PATIENT_DICOM), os.path.join(temp_folder, PATIENT_DICOM))
+        os.rename(os.path.join(temp_folder, PATIENT_DICOM, CONVERTED_FILE), os.path.join(DEST_FOLDER, VOL_FILE.format(i)))
+
+        folder_iter.set_description('Volume DICOM: CT stored in: ' + os.path.join(DEST_FOLDER, VOL_FILE.format(i)))
+        # os.rename also moves the file to the destination path. So the original one ceases to exist
+
+        # 3. Unzip MASKS_DICOM.zip
+        folder_iter.set_description('Segmentations DICOM: Unzipping MASKS_DICOM.zip')
+        zipfile.ZipFile(os.path.join(pat_dir, SEGS_DICOM + ZIP_EXT)).extractall(temp_folder)
+        seg_nib = list()
+        seg_ves = list()
+        seg_par = list()
+        seg_tumor = list()
+        seg_dirs = list()
+        for root, dir_list, file_list in os.walk(os.path.join(temp_folder, SEGS_DICOM)):
+            for fold in dir_list:
+                if fold.startswith(SEGMENTATIONS):
+                    # if 'liverkyst' in fold:
+                    #     continue
+                    # else:
+                    seg_dirs.append(fold)
+
+        seg_dirs.sort()
+        for fold in seg_dirs:
+            folder_iter.set_description('Segmentations DICOM: Converting ' + fold)
+            d2n.convert_directory(os.path.join(temp_folder, SEGS_DICOM, fold),
+                                  os.path.join(temp_folder, SEGS_DICOM))
+            os.rename(os.path.join(temp_folder, SEGS_DICOM, CONVERTED_FILE),
+                      os.path.join(temp_folder, SEGS_DICOM, fold + '_nifti_' + NIFTI_EXT))
+            if fold.startswith(SEGS_VESSELS):
+                seg_ves.append(os.path.join(temp_folder, SEGS_DICOM, fold + '_nifti_' + NIFTI_EXT))
+            elif fold.startswith(SEGS_TUMOR):
+                seg_tumor.append(os.path.join(temp_folder, SEGS_DICOM, fold + '_nifti_' + NIFTI_EXT))
+            elif fold.startswith(SEGS_PARENCH):
+                seg_par.append(os.path.join(temp_folder, SEGS_DICOM, fold + '_nifti_' + NIFTI_EXT))
+            else:
+                continue
+
+        folder_iter.set_description('Segmentations DICOM: Concatenating segmentations')
+        # Merge the vessel segmentations
+
+        segs_to_merge = tuple()
+        if len(seg_par) > 1:
+            segs_to_merge += (merge_segmentations(seg_par),)
+        else:
+            segs_to_merge += tuple(seg_par)  # seg_par is a list
+
+        if len(seg_ves) > 1:
+            segs_to_merge += (merge_segmentations(seg_ves),)
+        else:
+            segs_to_merge += tuple(seg_ves)  # seg_ves is a list
+
+        if len(seg_tumor) > 1:
+            segs_to_merge += (merge_segmentations(seg_tumor),)
+        else:
+            segs_to_merge += tuple(seg_tumor)  # seg_tumor is a list
+
+        # # Merge the tumors segmentations
+        # if len(seg_tumor):
+        #     segs_to_merge.append(merge_segmentations(seg_tumor))
+        # else:
+        #     print('No tumors found in ' + pat_dir)
+
+
+        # Merge with the parenchyma and save
+        folder_iter.set_description('Segmentations DICOM: Saving segmentations')
+        if len(segs_to_merge) > 1:
+            nib.save(nib.concat_images(segs_to_merge, check_affines=True), os.path.join(DEST_FOLDER, SEG_FILE.format(i)))
+        else:
+            nib.save(segs_to_merge[0], os.path.join(DEST_FOLDER, SEG_FILE.format(i)))
+
+        folder_iter.set_description('Segmentations DICOM: Segmentation stored in ' + os.path.join(DEST_FOLDER, SEG_FILE.format(i)))
+
+        shutil.rmtree(temp_folder)
+        folder_iter.set_description('Temporal file deleted')
+        # 4. Load DICOM and transform to nifti
+        # 5. Store as nifty in hd5

Datasets/ircad_dataset.py

@@ -0,0 +1,426 @@
+import os, sys
+
+os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
+os.environ['CUDA_VISIBLE_DEVICES'] = "1"  # Check availability before running using 'nvidia-smi'
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)
+import multiprocessing as mp
+
+mp.set_start_method('spawn')
+
+import tensorflow as tf
+
+# tf.enable_eager_execution()
+
+import numpy as np
+import nibabel as nib
+from skimage.transform import resize
+from skimage.filters import median
+from scipy.ndimage import binary_dilation, generate_binary_structure
+from nilearn.image import math_img
+import h5py
+from tqdm import tqdm
+import re
+
+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 DeepDeformationMapRegistration.utils.cmd_args_parser import parse_arguments
+import DeepDeformationMapRegistration.utils.constants as const
+from tools.thinPlateSplines_tf import ThinPlateSplines
+from keras_model.ext.neuron.layers import SpatialTransformer
+from tools.voxelMorph import interpn
+from generate_dataset.utils import plot_central_slices, plot_def_map, single_img_gif, two_img_gif, plot_slices, \
+    crop_images, plot_displacement_map, bbox_3D
+from generate_dataset import utils
+from tools.misc import try_mkdir
+from generate_dataset.utils import unzip_file, delete_temp
+
+DATASTE_RAW_FILES = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/nifti'
+LITS_SEGMENTATION_FILE = 'segmentations'
+LITS_CT_FILE = 'volume'
+
+IMG_SIZE_LARGE = const.IMG_SHAPE[:-1]
+IMG_SIZE_LARGE_x2 = [2 * x for x in const.IMG_SHAPE[:-1]]
+FINE_GRID_SHAPE = tuple(x // 1 for x in IMG_SIZE_LARGE_x2)  # tuple(np.asarray(IMG_SIZE_LARGE) // 10)
+CTRL_GRID = const.CoordinatesGrid()
+CTRL_GRID.set_coords_grid(IMG_SIZE_LARGE_x2, [const.TPS_NUM_CTRL_PTS_PER_AXIS, const.TPS_NUM_CTRL_PTS_PER_AXIS,
+                                              const.TPS_NUM_CTRL_PTS_PER_AXIS], batches=False, norm=False,
+                          img_type=tf.float32)
+
+FULL_FINE_GRID = const.CoordinatesGrid()
+FULL_FINE_GRID.set_coords_grid(IMG_SIZE_LARGE_x2, FINE_GRID_SHAPE, batches=False, norm=False)
+
+OFFSET_NAME_NUM = 0
+
+TH_BIN = 0.50
+
+DILATION_STRUCT = generate_binary_structure(3, 1)
+
+LARGE_PT_DIM = CTRL_GRID.shape_grid_flat + np.asarray([9, 0])
+SINGLE_PT_DIM = CTRL_GRID.shape_grid_flat + np.asarray([1, 0])
+USE_LARGE_PT = False
+ADD_AFFINE_TRF = False
+
+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)
+
+
+
+def tf_graph_translation():
+    # Place holders
+    fix_img = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_img')
+    fix_tumors = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_tumors')
+    fix_parenchyma = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_parenchyma')
+
+    # Apply Affine translation
+    w = tf.constant(np.random.uniform(-1, 1, 3) * const.MAX_DISP_DM_PERC * IMG_SIZE_LARGE_x2[0], dtype=tf.float32)
+    pad = tf.cast(tf.abs(w) + 1., tf.int32)
+    padding = tf.stack([pad, pad], 1)
+## PURE TRANSLATION
+    # Shift the target grid 'w' units
+    #control_grid = tf.identity(CTRL_GRID.grid_flat())
+    #trg_grid = tf.add(control_grid, w)
+
+    #tps = ThinPlateSplines(control_grid, trg_grid)
+    #def_grid = tps.interpolate(FULL_FINE_GRID.grid_flat())
+## PURE TRANSLATION
+
+    def_grid = tf.add(FULL_FINE_GRID.grid_flat(), w)
+    disp_map = def_grid - FULL_FINE_GRID.grid_flat()
+    disp_map = tf.reshape(disp_map, (*FINE_GRID_SHAPE, -1))
+    # disp_map = interpn(disp_map, FULL_FINE_GRID.grid)
+
+    # add the batch and channel dimensions
+    fix_img = tf.pad(fix_img, padding, "CONSTANT", constant_values=0.)
+    fix_tumors = tf.pad(fix_tumors, padding, "CONSTANT", constant_values=0.)
+    fix_parenchyma = tf.pad(fix_parenchyma, padding, "CONSTANT", constant_values=0.)
+
+    sampl_grid = tf.add(def_grid, tf.cast(pad, def_grid.dtype)) # Because of the padding, the sampling points are now translated 'pad' units
+    fix_img = tf.expand_dims(fix_img, -1)
+    fix_tumors = tf.expand_dims(fix_tumors, -1)
+    fix_parenchyma = tf.expand_dims(fix_parenchyma, -1)
+
+    mov_img = interpn(fix_img, sampl_grid, interp_method='linear')
+    mov_img = tf.squeeze(tf.reshape(mov_img, IMG_SIZE_LARGE_x2))
+
+    mov_tumors = interpn(fix_tumors, sampl_grid, interp_method='linear')
+    mov_tumors = tf.squeeze(tf.reshape(mov_tumors, IMG_SIZE_LARGE_x2))
+
+    mov_parenchyma = interpn(fix_parenchyma, sampl_grid, interp_method='linear')
+    mov_parenchyma = tf.squeeze(tf.reshape(mov_parenchyma, IMG_SIZE_LARGE_x2))
+
+    disp_map = tf.cast(disp_map, tf.float32)
+    return mov_img, mov_parenchyma, mov_tumors, disp_map, w  # , w, trg_grid, def_grid
+
+
+def build_affine_trf(img_size, alpha, beta, gamma, ti, tj, tk):
+    img_centre = tf.expand_dims(tf.divide(img_size, 2.), -1)
+
+    # Rotation matrix around the image centre
+    # R* = T(p) R(ang) T(-p)
+    # tf.cos and tf.sin expect radians
+    zero = tf.zeros((1,))
+    one = tf.ones((1,))
+    R = tf.convert_to_tensor([[tf.math.cos(gamma) * tf.math.cos(beta),
+                               tf.math.cos(gamma) * tf.math.sin(beta) * tf.math.sin(alpha) - tf.math.sin(gamma) * tf.math.cos(alpha),
+                               tf.math.cos(gamma) * tf.math.sin(beta) * tf.math.cos(alpha) + tf.math.sin(gamma) * tf.math.sin(alpha),
+                               zero],
+                              [tf.math.sin(gamma) * tf.math.cos(beta),
+                               tf.math.sin(gamma) * tf.math.sin(beta) * tf.math.sin(gamma) + tf.math.cos(gamma) * tf.math.cos(alpha),
+                               tf.math.sin(gamma) * tf.math.sin(beta) * tf.math.cos(gamma) - tf.math.cos(gamma) * tf.math.sin(gamma),
+                               zero],
+                              [-tf.math.sin(beta),
+                               tf.math.cos(beta) * tf.math.sin(alpha),
+                               tf.math.cos(beta) * tf.math.cos(alpha),
+                               zero],
+                              [zero, zero, zero, one]], tf.float32)
+    R = tf.squeeze(R)
+
+    Tc = tf.convert_to_tensor([[one, zero, zero, img_centre[0]],
+                               [zero, one, zero, img_centre[1]],
+                               [zero, zero, one, img_centre[2]],
+                               [zero, zero, zero, one]], tf.float32)
+    Tc = tf.squeeze(Tc)
+    Tc_ = tf.convert_to_tensor([[one, zero, zero, -img_centre[0]],
+                                [zero, one, zero, -img_centre[1]],
+                                [zero, zero, one, -img_centre[2]],
+                                [zero, zero, zero, one]], tf.float32)
+    Tc_ = tf.squeeze(Tc_)
+
+    T = tf.convert_to_tensor([[one, zero, zero, ti],
+                              [zero, one, zero, tj],
+                              [zero, zero, one, tk],
+                              [zero, zero, zero, one]], tf.float32)
+    T = tf.squeeze(T)
+
+    return tf.matmul(T, tf.matmul(Tc, tf.matmul(R, Tc_)))
+
+
+def transform_points(points: tf.Tensor):
+    alpha = tf.random.uniform((1,), -const.MAX_ANGLE_RAD, const.MAX_ANGLE_RAD)
+    beta = tf.random.uniform((1,), -const.MAX_ANGLE_RAD, const.MAX_ANGLE_RAD)
+    gamma = tf.random.uniform((1,), -const.MAX_ANGLE_RAD, const.MAX_ANGLE_RAD)
+
+    ti = tf.constant(np.random.uniform(-1, 1, 1) * const.MAX_DISP_DM / 2, dtype=tf.float32)
+    tj = tf.constant(np.random.uniform(-1, 1, 1) * const.MAX_DISP_DM / 2, dtype=tf.float32)
+    tk = tf.constant(np.random.uniform(-1, 1, 1) * const.MAX_DISP_DM / 2, dtype=tf.float32)
+
+    M = build_affine_trf(tf.convert_to_tensor(IMG_SIZE_LARGE_x2, tf.float32), alpha, beta, gamma, ti, tj, tk)
+    if points.shape.as_list()[-1] == 3:
+        points = tf.transpose(points)
+    new_pts = tf.matmul(M[:3, :3], points)
+    new_pts = tf.expand_dims(M[:3, -1], -1) + new_pts
+    return tf.transpose(new_pts), M  # Remove the last row of ones
+
+
+def tf_graph_deform():
+    # Place holders
+    fix_img = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_img')
+    fix_tumors = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_tumors')
+    fix_vessels = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_vessels')
+    fix_parenchyma = tf.placeholder(tf.float32, IMG_SIZE_LARGE_x2, 'fix_parenchyma')
+    large_point = tf.placeholder_with_default(input=tf.constant(False, tf.bool), shape=(), name='large_point')
+    add_affine = tf.placeholder_with_default(input=tf.constant(False, tf.bool), shape=(), name='add_affine')
+
+    search_voxels = tf.cond(tf.equal(tf.reduce_sum(fix_tumors), 0.0),
+                            lambda: fix_parenchyma,
+                            lambda: fix_tumors)
+
+    # Apply TPS deformation
+    # 1. get a point in the label img and add it to the control grid and target grid
+    idx_points_in_label = tf.where(tf.greater(search_voxels, 0.0))  # Indices of the points in the label image with intensity greater than 0
+
+    random_idx = tf.random.uniform([], minval=0, maxval=tf.shape(idx_points_in_label)[0],
+                                  dtype=tf.int32)  # Randomly select one of the points
+    disp_location = tf.gather_nd(idx_points_in_label, tf.expand_dims(random_idx, 0))  # And get the coordinates
+    disp_location = tf.cast(disp_location, tf.float32)
+    # Get the coordinates of the control point displaces
+    rand_disp = tf.constant(np.random.uniform(-1, 1, 3) * const.MAX_DISP_DM, dtype=tf.float32)
+    warped_location = disp_location + rand_disp
+
+    def get_box_neighbours(location, radius=3):
+        n1 = tf.add(rand_disp, tf.constant(np.asarray([radius, radius, radius]), location.dtype))
+        n2 = tf.add(rand_disp, tf.constant(np.asarray([-radius, radius, radius]), location.dtype))
+        n3 = tf.add(rand_disp, tf.constant(np.asarray([radius, -radius, radius]), location.dtype))
+        n4 = tf.add(rand_disp, tf.constant(np.asarray([-radius, -radius, radius]), location.dtype))
+        n5 = tf.add(rand_disp, tf.constant(np.asarray([radius, radius, -radius]), location.dtype))
+        n6 = tf.add(rand_disp, tf.constant(np.asarray([-radius, radius, -radius]), location.dtype))
+        n7 = tf.add(rand_disp, tf.constant(np.asarray([radius, -radius, -radius]), location.dtype))
+        n8 = tf.add(rand_disp, tf.constant(np.asarray([-radius, -radius, -radius]), location.dtype))
+
+        return tf.stack([location, n1, n2, n3, n4, n5, n6, n7, n8], 0)
+
+    disp_location, warped_location = tf.cond(large_point,
+                               lambda: (get_box_neighbours(disp_location, 3), get_box_neighbours(warped_location, 3)),
+                               lambda: (tf.expand_dims(rand_disp, 0), tf.expand_dims(warped_location, 0)))
+
+    # 2. Add the new point to the control grid and the target grid
+    control_grid = tf.concat([CTRL_GRID.grid_flat(), disp_location], axis=0)
+    trg_grid = tf.concat([CTRL_GRID.grid_flat(), warped_location], axis=0)
+
+    trg_grid, aff = tf.cond(add_affine,
+                            lambda: transform_points(trg_grid),
+                            lambda: (trg_grid, tf.eye(4, 4)))
+
+    # I need to know the shape before running TPS
+    control_grid.set_shape([73, 3] if USE_LARGE_PT else [65, 3])
+    trg_grid.set_shape([73, 3] if USE_LARGE_PT else [65, 3])
+
+    tps = ThinPlateSplines(control_grid, trg_grid)
+    def_grid = tps.interpolate(FULL_FINE_GRID.grid_flat())
+
+    disp_map = def_grid - FULL_FINE_GRID.grid_flat()
+    disp_map = tf.reshape(disp_map, (*FINE_GRID_SHAPE, -1))
+    # disp_map = interpn(disp_map, FULL_FINE_GRID.grid)
+
+    # add the batch and channel dimensions
+    fix_img = tf.expand_dims(tf.expand_dims(fix_img, -1), 0)
+    fix_tumors = tf.expand_dims(tf.expand_dims(fix_tumors, -1), 0)
+    fix_vessels = tf.expand_dims(tf.expand_dims(fix_vessels, -1), 0)
+    fix_parenchyma = tf.expand_dims(tf.expand_dims(fix_parenchyma, -1), 0)
+    disp_map = tf.cast(tf.expand_dims(disp_map, 0), tf.float32)
+
+    mov_tumors = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_tumors, disp_map])
+    mov_vessels = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_vessels, disp_map])
+    mov_parenchyma = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_parenchyma, disp_map])
+    mov_img = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_img, disp_map])
+
+    return tf.squeeze(mov_img),\
+           tf.squeeze(mov_parenchyma),\
+           tf.squeeze(mov_tumors),\
+           tf.squeeze(mov_vessels),\
+           tf.squeeze(disp_map),\
+           disp_location,\
+           rand_disp,\
+           aff #, w, trg_grid, def_grid
+
+
+if __name__ == '__main__':
+    parse_arguments(sys.argv[1:])
+    volume_list = [os.path.join(DATASTE_RAW_FILES, f) for f in os.listdir(DATASTE_RAW_FILES) if f.startswith(LITS_CT_FILE)]
+    volume_list.sort()
+    segmentation_list = [os.path.join(DATASTE_RAW_FILES, f) for f in os.listdir(DATASTE_RAW_FILES) if
+                         f.startswith(LITS_SEGMENTATION_FILE)]
+    segmentation_list.sort()
+
+    file_path_pairs = [[v, s] for v, s in zip(volume_list, segmentation_list)]
+
+    print('Generating HD5 files at {} ...', format(const.DESTINATION_FOLDER))
+    # with Pool(10) as p, tf.Session(config=config) as sess:
+    #     tqdm(p.map(generate_training_sample, file_path_pairs))
+    intensity_window_w = 350
+    intensity_window_l = 40
+    intensity_clipping_range = intensity_window_l + np.asarray([-intensity_window_w // 2, intensity_window_w // 2],
+                                                               np.int)  # Slicer range for abdominal CT
+
+    try_mkdir(const.DESTINATION_FOLDER)
+
+    print('PART 1: Deformation')
+    # Then do the fancy stuff
+    init = tf.initialize_all_variables()
+    get_mov_img = tf_graph_deform()
+    sess = tf.Session(config=config)
+    with sess.as_default():
+        sess.run(init)
+        sess.graph.finalize()
+        for img_path, labels_path in tqdm(file_path_pairs):
+            if img_path is not None and labels_path is not None:
+                #img_path = unzip_file(img_path)
+                #labels_path = unzip_file(labels_path)
+
+                fix_img = nib.load(img_path)  # By convention, nibabel world axes are always in RAS+ orientation
+                img_header = fix_img.header
+                fix_labels = nib.load(labels_path)
+                fix_img = np.asarray(fix_img.dataobj)
+                fix_labels = np.asarray(fix_labels.dataobj)
+                if fix_labels.shape[-1] < 4:
+                    print('[INF] ' + img_path + ' has no tumor segmentations')
+                    continue
+                # fix_artery = fix_labels[..., 0]
+                fix_vessels = fix_labels[..., 1]
+                fix_parenchyma = fix_labels[..., 2]
+                fix_tumors = fix_labels[..., 3]
+
+                # Clip intensity values
+                fix_img = utils.intesity_clipping(fix_img, intensity_clipping_range, augment=True)
+
+                # Reshape
+                fix_img = resize(fix_img, IMG_SIZE_LARGE_x2)
+                fix_parenchyma = resize(fix_parenchyma, IMG_SIZE_LARGE_x2)
+                fix_tumors = resize(fix_tumors, IMG_SIZE_LARGE_x2)
+                fix_vessels = resize(fix_vessels, IMG_SIZE_LARGE_x2)
+
+                fix_parenchyma = median(fix_parenchyma, np.ones((5, 5, 5)))
+
+                # Compute deformation
+                mov_img, mov_parenchyma, mov_tumors, mov_vessels, disp_map, disp_loc, disp_vec, aff = sess.run(get_mov_img,
+                                                                                   feed_dict={
+                                                                                       'fix_img:0': fix_img,
+                                                                                       'fix_tumors:0': fix_tumors,
+                                                                                       'fix_vessels:0': fix_vessels,
+                                                                                       'fix_parenchyma:0': fix_parenchyma,
+                                                                                       'large_point:0': USE_LARGE_PT,
+                                                                                       'add_affine:0': ADD_AFFINE_TRF})
+                # Cleaning
+                mov_img = utils.intesity_clipping(mov_img, intensity_clipping_range)
+
+                if USE_LARGE_PT:
+                    disp_loc = disp_loc[0, ...]
+
+                # Define the bbox around the union of the parenchyma of both volumes, so none falls outside
+                bbox_mask = np.sign(mov_parenchyma + fix_parenchyma)
+                bbox_mask = binary_dilation(bbox_mask, DILATION_STRUCT)
+                bbox_mask = binary_dilation(bbox_mask, DILATION_STRUCT).astype(np.float32)
+
+                # The point of application is referred to the whole image coordinate, not to the local BB
+                min_i, _, min_j, _, min_k, _ = bbox_3D(bbox_mask)
+                disp_loc = (disp_loc - np.asarray([min_i, min_j, min_k])) / 2
+                # Crop the image to only contain the liver
+                # The origin moved according to the mask information. And the images will be resized in a factor of 2!!
+
+                fix_img, _ = crop_images(fix_img, bbox_mask, IMG_SIZE_LARGE)
+                fix_tumors, _ = crop_images(fix_tumors, bbox_mask, IMG_SIZE_LARGE)
+                fix_vessels, _ = crop_images(fix_vessels, bbox_mask, IMG_SIZE_LARGE)
+                disp_map, _ = crop_images(disp_map, bbox_mask, IMG_SIZE_LARGE)
+                fix_parenchyma, _ = crop_images(fix_parenchyma, bbox_mask, IMG_SIZE_LARGE)
+
+                # We will later crop even further, so we don't want to downsample too much
+                # Crop the image to only contain the liver
+                mov_img, _ = crop_images(mov_img, bbox_mask, IMG_SIZE_LARGE)
+                mov_tumors, _ = crop_images(mov_tumors, bbox_mask, IMG_SIZE_LARGE)
+                mov_vessels, _ = crop_images(mov_vessels, bbox_mask, IMG_SIZE_LARGE)
+                mov_parenchyma, _ = crop_images(mov_parenchyma, bbox_mask, IMG_SIZE_LARGE)
+
+                # Just to be sure we have binary masks
+                fix_tumors[fix_tumors > TH_BIN] = 1.0
+                fix_tumors[fix_tumors < 1.0] = 0.0
+
+                fix_vessels[fix_vessels > TH_BIN] = 1.0
+                fix_vessels[fix_vessels < 1.0] = 0.0
+
+                fix_parenchyma[fix_parenchyma > TH_BIN] = 1.0
+                fix_parenchyma[fix_parenchyma < 1.0] = 0.0
+
+                mov_tumors[mov_tumors > TH_BIN] = 1.0
+                mov_tumors[mov_tumors < 1.0] = 0.0
+
+                mov_vessels[mov_vessels > TH_BIN] = 1.0
+                mov_vessels[mov_vessels < 1.0] = 0.0
+
+                mov_parenchyma[mov_parenchyma > TH_BIN] = 1.0
+                mov_parenchyma[mov_parenchyma < 1.0] = 0.0
+
+                # Save everything
+                fix_img = np.expand_dims(fix_img, -1)
+                fix_tumors = np.expand_dims(fix_tumors, -1)
+                fix_vessels = np.expand_dims(fix_vessels, -1)
+                fix_parenchyma = np.expand_dims(fix_parenchyma, -1)
+                fix_segmentations = np.stack([fix_parenchyma, fix_vessels, fix_tumors], -1)
+
+                mov_img = np.expand_dims(mov_img, -1)
+                mov_tumors = np.expand_dims(mov_tumors, -1)
+                mov_vessels = np.expand_dims(mov_vessels, -1)
+                mov_parenchyma = np.expand_dims(mov_parenchyma, -1)
+
+                # Save everything
+                file_name = os.path.split(img_path)[-1].split('.')[0]
+                vol_num = int(re.split('-|_', file_name)[-1])
+                hd5_filename = 'volume-{:04d}'.format(vol_num + OFFSET_NAME_NUM)
+                hd5_filename = os.path.join(const.DESTINATION_FOLDER, hd5_filename + '.hd5')
+                hd5_file = h5py.File(hd5_filename, 'w')
+
+                hd5_file.create_dataset(const.H5_FIX_IMG, data=fix_img, dtype='float32')
+                hd5_file.create_dataset(const.H5_FIX_PARENCHYMA_MASK, data=fix_parenchyma, dtype='float32')
+                hd5_file.create_dataset(const.H5_FIX_VESSELS_MASK, data=fix_vessels, dtype='float32')
+                hd5_file.create_dataset(const.H5_FIX_TUMORS_MASK, data=fix_tumors, dtype='float32')
+                hd5_file.create_dataset(const.H5_FIX_SEGMENTATIONS, data=fix_segmentations, dtype='float32')
+
+                hd5_file.create_dataset(const.H5_PARAMS_INTENSITY_RANGE, (2,), data=intensity_clipping_range,
+                                        dtype='float32')
+
+                hd5_file.create_dataset(const.H5_MOV_IMG, const.IMG_SHAPE, data=mov_img, dtype='float32')
+                hd5_file.create_dataset(const.H5_MOV_PARENCHYMA_MASK, const.IMG_SHAPE, data=mov_parenchyma,
+                                        dtype='float32')
+                hd5_file.create_dataset(const.H5_MOV_VESSELS_MASK, const.IMG_SHAPE, data=mov_vessels, dtype='float32')
+                hd5_file.create_dataset(const.H5_MOV_TUMORS_MASK, const.IMG_SHAPE, data=mov_tumors, dtype='float32')
+                hd5_file.create_dataset(const.H5_MOV_SEGMENTATIONS, data=fix_segmentations, dtype='float32')
+
+                hd5_file.create_dataset(const.H5_GT_DISP, const.DISP_MAP_SHAPE, data=disp_map, dtype='float32')
+                hd5_file.create_dataset(const.H5_GT_DISP_VECT_LOC, data=disp_loc, dtype='float32')
+                hd5_file.create_dataset(const.H5_GT_DISP_VECT, data=disp_vec, dtype='float32')
+                hd5_file.create_dataset(const.H5_GT_AFFINE_M, data=aff, dtype='float32')
+
+                hd5_file.create_dataset('params/voxel_size', data=img_header.get_zooms()[:3])
+                hd5_file.create_dataset('params/original_shape', data=img_header.get_data_shape())
+                hd5_file.create_dataset('params/bbox_origin', data=[min_i, min_j, min_k])
+                hd5_file.create_dataset('params/first_reshape', data=IMG_SIZE_LARGE_x2)
+
+                # delete_temp(img_path)
+                # delete_temp(labels_path)
+
+                hd5_file.close()
+        sess.close()
+    print('...Done generating HD5 files')