CPD scripts

Centerline/centerline.py

@@ -0,0 +1,241 @@
+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 tensorflow as tf
+# tf.enable_eager_execution()
+# import neurite.py.utils as neurite_utils
+
+from skimage.morphology import skeletonize_3d, ball
+from skimage.morphology import binary_closing, binary_opening
+from skimage.filters import median
+from skimage.measure import regionprops, label
+from skimage.transform import warp
+
+from scipy.ndimage import zoom
+from scipy.interpolate import LinearNDInterpolator, Rbf
+
+import h5py
+import numpy as np
+from tqdm import tqdm
+import re
+import nibabel as nib
+from nilearn.image import resample_img
+
+from Centerline.graph_utils import graph_to_ndarray, deform_graph, get_bifurcation_nodes, subsample_graph, \
+    apply_displacement
+from Centerline.skeleton_to_graph import get_graph_from_skeleton
+from Centerline.visualization_utils import plot_skeleton, compare_graphs
+
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils import constants as C
+
+import cupy
+from cupyx.scipy.ndimage import zoom as zoom_gpu
+from cupyx.scipy.ndimage import map_coordinates
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['Affine', 'None', 'Translation']
+DATASET_FILENAME = 'volume'
+IMGS_FOLDER = '/home/jpdefrutos/workspace/DeepDeformationMapRegistration/Centerline/centerlines'
+
+DATASTE_RAW_FILES = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/nifti3'
+LITS_SEGMENTATION_FILE = 'segmentation'
+LITS_CT_FILE = 'volume'
+
+
+def warp_volume(volume, disp_map, indexing='ij'):
+    assert indexing is 'ij' or 'xy', 'Invalid indexing option. Only "ij" or "xy"'
+    grid_i = np.linspace(0, disp_map.shape[0], disp_map.shape[0], endpoint=False)
+    grid_j = np.linspace(0, disp_map.shape[1], disp_map.shape[1], endpoint=False)
+    grid_k = np.linspace(0, disp_map.shape[2], disp_map.shape[2], endpoint=False)
+    grid_i, grid_j, grid_k = np.meshgrid(grid_i, grid_j, grid_k, indexing=indexing)
+    grid_i = (grid_i.flatten() + disp_map[..., 0].flatten())[..., np.newaxis]
+    grid_j = (grid_j.flatten() + disp_map[..., 1].flatten())[..., np.newaxis]
+    grid_k = (grid_k.flatten() + disp_map[..., 2].flatten())[..., np.newaxis]
+    coords = np.hstack([grid_i, grid_j, grid_k]).reshape([*disp_map.shape[:-1], -1])
+    coords = coords.transpose((-1, 0, 1, 2))
+    # The returned volume has indexing xy
+    return warp(volume, coords)
+
+
+def keep_largest_segmentation(img):
+    label_img = label(img)
+    rp = regionprops(label_img)  # Regions labeled with 0 (bg) are ignored
+    biggest_area = (0, 0)
+    for l in range(0, label_img.max()):
+        if rp[l].area > biggest_area[1]:
+            biggest_area = (l + 1, rp[l].area)
+    img[label_img != biggest_area[0]] = 0.
+    return img
+
+
+def preprocess_image(img, keep_largest=False):
+    ret = binary_closing(img, ball(1))
+    ret = binary_opening(ret, ball(1))
+    #ret = median(ret, ball(1), mode='constant')
+    if keep_largest:
+        ret = keep_largest_segmentation(ret)
+    return ret.astype(np.float)
+
+
+def build_displacement_map_interpolator(disp_map, backwards=False, indexing='ij'):
+    grid_i = np.linspace(0, disp_map.shape[0], disp_map.shape[0], endpoint=False)
+    grid_j = np.linspace(0, disp_map.shape[1], disp_map.shape[1], endpoint=False)
+    grid_k = np.linspace(0, disp_map.shape[2], disp_map.shape[2], endpoint=False)
+    grid_i, grid_j, grid_k = np.meshgrid(grid_i, grid_j, grid_k, indexing=indexing)
+    grid_i = grid_i.flatten()
+    grid_j = grid_j.flatten()
+    grid_k = grid_k.flatten()
+    # To generate the moving image, we used backwards mapping were the input was the fix image
+    # Now we are doing direct mapping from the fix graph coordinates to the moving coordinates
+    # The application points of the displacement map are thus the transformed "moving image"-grid
+    #   and the displacement vectors are reversed
+    if backwards:
+        coords = np.hstack([grid_i[..., np.newaxis], grid_j[..., np.newaxis], grid_k[..., np.newaxis]])
+        return LinearNDInterpolator(coords, np.reshape(disp_map, [-1, 3]))
+    else:
+        grid_i = (grid_i + disp_map[..., 0].flatten())
+        grid_j = (grid_j + disp_map[..., 1].flatten())
+        grid_k = (grid_k + disp_map[..., 2].flatten())
+
+        coords = np.hstack([grid_i[..., np.newaxis], grid_j[..., np.newaxis], grid_k[..., np.newaxis]])
+        return LinearNDInterpolator(coords, -np.reshape(disp_map, [-1, 3]))
+
+
+def resample_segmentation(img, output_shape, preserve_range, threshold=None, gpu=True):
+    # Preserve range can be a bool (keep or not the original dyn. range) or a list with a new dyn. range
+    zoom_f = np.divide(np.asarray(output_shape), np.asarray(img.shape))
+
+    if gpu:
+        out_img = zoom_gpu(cupy.asarray(img), zoom_f, order=1)  # order = 0 or 1
+    else:
+        out_img = zoom(img, zoom_f)
+    if isinstance(preserve_range, bool):
+        if preserve_range:
+            range_min, range_max = np.amin(img), np.amax(img)
+            out_img = min_max_norm(out_img)
+            out_img = out_img * (range_max - range_min) + range_min
+    elif isinstance(preserve_range, list):
+        range_min, range_max = preserve_range
+        out_img = min_max_norm(out_img)
+        out_img = out_img * (range_max - range_min) + range_min
+
+    if threshold is not None and out_img.min() < threshold < out_img.max():
+        range_min, range_max = np.amin(out_img), np.amax(out_img)
+        out_img[out_img > threshold] = range_max
+        out_img[out_img < range_max] = range_min
+    return cupy.asnumpy(out_img) if gpu else out_img
+
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        for file_path in iterator:
+            file_num = int(re.findall('(\d+)', os.path.split(file_path)[-1])[0])
+
+            iterator.set_description('{} ({}): laoding data'.format(file_num, dataset_name))
+            vol_file = h5py.File(file_path, 'r')
+            # fix_vessels = vol_file[C.H5_FIX_VESSELS_MASK][..., 0]
+            disp_map = vol_file[C.H5_GT_DISP][:]
+            bbox = vol_file['parameters/bbox'][:]
+            bbox_min = bbox[:3]
+            bbox_max = bbox[3:] + bbox_min
+
+            # Load vessel segmentation mask and resize to 64^3
+            fix_labels = nib.load(os.path.join(DATASTE_RAW_FILES, 'segmentation-{:04d}.nii.gz'.format(file_num)))
+            fix_vessels = fix_labels.slicer[..., 1]
+            fix_vessels = resample_img(fix_vessels, np.eye(3))
+            fix_vessels = np.asarray(fix_vessels.dataobj)
+            fix_vessels = preprocess_image(fix_vessels)
+            fix_vessels = resample_segmentation(fix_vessels, vol_file['parameters/first_reshape'][:], [0, 1], 0.3,
+                                                gpu=True)
+            fix_vessels = fix_vessels[bbox_min[0]:bbox_max[0], bbox_min[1]:bbox_max[1], bbox_min[2]:bbox_max[2]]
+            fix_vessels = resample_segmentation(fix_vessels, [64] * 3, [0, 1], 0.3, gpu=True)
+            fix_vessels = preprocess_image(fix_vessels)
+
+            mov_vessels = preprocess_image(warp_volume(fix_vessels, disp_map))
+            mov_skel = skeletonize_3d(mov_vessels)
+            ### Fix the incorrect scaling ###
+            disp_map *= 2
+            bbox_size = np.asarray(bbox[3:])  # Only load the bbox size
+            rescale_factors = 64 / bbox_size
+
+            disp_map[..., 0] = np.multiply(disp_map[..., 0], rescale_factors[0])
+            disp_map[..., 1] = np.multiply(disp_map[..., 1], rescale_factors[1])
+            disp_map[..., 2] = np.multiply(disp_map[..., 2], rescale_factors[2])
+            #################################
+
+            iterator.set_description('{} ({}): getting graphs'.format(file_num, dataset_name))
+            # Prepare displacement map
+            disp_map_interpolator = build_displacement_map_interpolator(disp_map, backwards=False)
+
+            # Get skeleton and graph
+            fix_skel = skeletonize_3d(fix_vessels)
+            fix_graph = get_graph_from_skeleton(fix_skel, subsample=True)
+            mov_graph = get_graph_from_skeleton(mov_skel, subsample=True) # deform_graph(fix_graph, disp_map_interpolator)
+
+            ##### TODO: ERASE Check the mov graph ######
+            # check_mov_vessels = vol_file[C.H5_MOV_VESSELS_MASK][..., 0]
+            # check_mov_vessels = preprocess_image(check_mov_vessels)
+            # check_mov_skel = skeletonize_3d(check_mov_vessels)
+            # check_mov_graph = get_graph_from_skeleton(check_mov_skel, subsample=True)
+            ###########
+            fix_pts, fix_nodes, fix_edges = graph_to_ndarray(fix_graph)
+            mov_pts, mov_nodes, mov_edges = graph_to_ndarray(mov_graph)
+
+            fix_bifur_loc, fix_bifur_id = get_bifurcation_nodes(fix_graph)
+            mov_bifur_loc, mov_bifur_id = get_bifurcation_nodes(mov_graph)
+
+            iterator.set_description('{} ({}): saving data'.format(file_num, dataset_name))
+            pts_file_path, pts_file_name = os.path.split(file_path)
+            pts_file_name = pts_file_name.replace(DATASET_FILENAME, 'points')
+            pts_file_path = os.path.join(pts_file_path, pts_file_name)
+            pts_file = h5py.File(pts_file_path, 'w')
+
+            pts_file.create_dataset('fix/points', data=fix_pts)
+            pts_file.create_dataset('fix/nodes', data=fix_nodes)
+            pts_file.create_dataset('fix/edges', data=fix_edges)
+            pts_file.create_dataset('fix/bifurcations', data=fix_bifur_loc)
+            pts_file.create_dataset('fix/graph', data=fix_graph)
+            pts_file.create_dataset('fix/img', data=fix_vessels)
+            pts_file.create_dataset('fix/skeleton', data=fix_skel)
+            pts_file.create_dataset('fix/centroid', data=vol_file[C.H5_FIX_CENTROID][:])
+
+            pts_file.create_dataset('mov/points', data=mov_pts)
+            pts_file.create_dataset('mov/nodes', data=mov_nodes)
+            pts_file.create_dataset('mov/edges', data=mov_edges)
+            pts_file.create_dataset('mov/bifurcations', data=mov_bifur_loc)
+            pts_file.create_dataset('mov/graph', data=mov_graph)
+            pts_file.create_dataset('mov/img', data=mov_vessels)
+            pts_file.create_dataset('mov/skeleton', data=mov_skel)
+            pts_file.create_dataset('mov/centroid', data=vol_file[C.H5_MOV_CENTROID][:])
+
+            pts_file.create_dataset('parameters/voxel_size', data=vol_file['parameters/voxel_size'][:])
+            pts_file.create_dataset('parameters/original_affine', data=vol_file['parameters/original_affine'][:])
+            pts_file.create_dataset('parameters/isotropic_affine', data=vol_file['parameters/isotropic_affine'][:])
+            pts_file.create_dataset('parameters/original_shape', data=vol_file['parameters/original_shape'][:])
+            pts_file.create_dataset('parameters/isotropic_shape', data=vol_file['parameters/isotropic_shape'][:])
+            pts_file.create_dataset('parameters/first_reshape', data=vol_file['parameters/first_reshape'][:])
+            pts_file.create_dataset('parameters/bbox', data=vol_file['parameters/bbox'][:])
+            pts_file.create_dataset('parameters/last_reshape', data=vol_file['parameters/last_reshape'][:])
+
+            pts_file.create_dataset('displacement_map', data=disp_map)
+
+            vol_file.close()
+            pts_file.close()
+
+            iterator.set_description('{} ({}): drawing plots'.format(file_num, dataset_name))
+            num = pts_file_name.split('-')[-1].split('.hd5')[0]
+            imgs_folder = os.path.join(IMGS_FOLDER, dataset_name, num)
+            os.makedirs(imgs_folder, exist_ok=True)
+            plot_skeleton(fix_vessels, fix_skel, fix_graph, os.path.join(imgs_folder, 'fix'), ['.pdf', '.png'])
+            plot_skeleton(mov_vessels, mov_skel, mov_graph, os.path.join(imgs_folder, 'mov'), ['.pdf', '.png'])
+            iterator.set_description('{} ({})'.format(file_num, dataset_name))

Centerline/cpd_utils.py

@@ -0,0 +1,48 @@
+from pycpd import DeformableRegistration, RigidRegistration
+import numpy as np
+import time
+from scipy.interpolate import Rbf
+import warnings
+
+def cpd_non_rigid_transform_pt(pt, Y, G, W):
+    from scipy.interpolate import LinearNDInterpolator
+    interp = LinearNDInterpolator(points=Y, values=np.dot(G, W), fill_value=0.)
+    return interp(pt)
+
+
+def radial_basis_function(pts, vals, function='thin-plate'):
+    # The Rbf function does not handle n-D hyper-surfaces, so we need an interpolator per displacements. Actually it does mode='N-D'
+    pts_unique, idxs = np.unique(pts, return_index=True, axis=0)  # Prevent singular matrices
+    ill_conditioned = False
+    with warnings.catch_warnings(record=True) as caught_warns:
+        warnings.simplefilter('always')
+        dx = Rbf(pts_unique[:, 0], pts_unique[:, 1], pts_unique[:, 2], vals[idxs][:, 0], function=function)
+        dy = Rbf(pts_unique[:, 0], pts_unique[:, 1], pts_unique[:, 2], vals[idxs][:, 1], function=function)
+        dz = Rbf(pts_unique[:, 0], pts_unique[:, 1], pts_unique[:, 2], vals[idxs][:, 2], function=function)
+        for w in caught_warns:
+            print(w)
+            ill_conditioned = ill_conditioned or 'ill-conditioned matrix' in str(w).lower()
+    return lambda int_pt: np.asarray([dx(*int_pt), dy(*int_pt), dz(*int_pt)]), ill_conditioned
+
+
+def deform_registration(fix_pts, mov_pts, callback_fnc=None, time_it=False, max_iterations=100, tolerance=1e-8, alpha=None, beta=None):
+    deform_reg = DeformableRegistration(**{'Y': mov_pts, 'X': fix_pts},
+                                        alpha=alpha, beta=beta, tolerance=tolerance, max_iterations=max_iterations)
+    start_t = time.time()
+    trf_mov_pts, deform_p = deform_reg.register(callback_fnc)
+    end_t = time.time()
+    if time_it:
+        return end_t - start_t, deform_reg
+    else:
+        return trf_mov_pts, deform_p, deform_reg
+
+
+def rigid_registration(fix_pts, mov_pts, callback_fnc=None, time_it=False):
+    rigid_reg = RigidRegistration(**{'Y': mov_pts, 'X': fix_pts})
+    start_t = time.time()
+    trf_mov_pts, trf_p = rigid_reg.register(callback_fnc)
+    end_t = time.time()
+    if time_it:
+        return end_t - start_t, rigid_reg
+    else:
+        return trf_mov_pts, trf_p, rigid_reg

Centerline/evaluate_BayesianCPD_skeleton.py

@@ -0,0 +1,160 @@
+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
+from tqdm import tqdm
+from functools import partial
+import numpy as np
+from scipy.spatial.distance import euclidean
+import pandas as pd
+from EvaluationScripts.Evaluate_class import resize_img_to_original_space, resize_pts_to_original_space
+from Centerline.visualization_utils import plot_cpd_registration_step, plot_cpd
+from Centerline.cpd_utils import cpd_non_rigid_transform_pt, radial_basis_function, deform_registration, rigid_registration
+from scipy.spatial.distance import cdist
+from skimage.morphology import skeletonize_3d
+import re
+from probreg import bcpd
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['Affine', 'None', 'Translation']
+DATASET_FILENAME = 'points'
+
+OUT_IMG_FOLDER = '/mnt/EncryptedData1/Users/javier/vessel_registration/Centerline/cpd/skeleton'
+
+SCALE = 1e-2  # mm to cm
+# CPD PARAMS (deform)
+MAX_ITER = 200
+ALPHA = 0.1
+BETA = 1.0  # None = Use default
+TOLERANCE = 1e-8
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        df = pd.DataFrame(columns=['DATASET',
+                                   'ITERATIONS_DEF', 'ITERATIONS_R_DEF__R', 'ITERATIONS_R_DEF__DEF',
+                                   'TIME_DEF', 'TIME_R_DEF',
+                                   'Q_DEF', 'Q_R_DEF__R', 'Q_R_DEF__DEF',
+                                   'TRE_DEF', 'TRE_R_DEF',
+                                   'DS_DISP',
+                                   'DATA_PATH',
+                                   'DIST_CENTR', 'DIST_CENTR_DEF_95', 'SAMPLE_NUM'])
+        for i, file_path in enumerate(iterator):
+            fn = os.path.split(file_path)[-1].split('.hd5')[0]
+            fnum = int(re.findall('(\d+)', fn)[0])
+            iterator.set_description('{}: start'.format(fn))
+            pts_file = h5py.File(file_path, 'r')
+            # fix_pts = pts_file['fix/points'][:]
+            # fix_nodes = pts_file['fix/nodes'][:]
+            fix_skel = pts_file['fix/skeleton'][:]
+            fix_centroid = pts_file['fix/centroid'][:]
+
+            # mov_pts = pts_file['mov/points'][:]
+            # mov_nodes = pts_file['mov/nodes'][:]
+            mov_skel = pts_file['mov/skeleton'][:]
+            mov_centroid = pts_file['mov/centroid'][:]
+
+            bbox = pts_file['parameters/bbox'][:]
+            first_reshape = pts_file['parameters/first_reshape'][:]
+            isotropic_shape = pts_file['parameters/isotropic_shape'][:]
+            iterator.set_description('{}: Loaded data'.format(fn))
+            # TODO: bring back to original shape!
+            # Reshape to original_shape
+            # fix_nodes = resize_pts_to_original_space(fix_nodes, bbox, [64]*3, first_reshape, original_shape)
+            # fix_pts = resize_pts_to_original_space(fix_pts, bbox, [64] * 3, first_reshape, original_shape)
+            fix_centroid = resize_pts_to_original_space(fix_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            fix_skel = resize_img_to_original_space(fix_skel, bbox, first_reshape, isotropic_shape)
+            fix_skel = skeletonize_3d(fix_skel)
+            fix_skel_pts = np.argwhere(fix_skel)
+            # mov_nodes = resize_pts_to_original_space(mov_nodes, bbox, [64] * 3, first_reshape, original_shape)
+            # mov_pts = resize_pts_to_original_space(mov_pts, bbox, [64] * 3, first_reshape, original_shape)
+            mov_centroid = resize_pts_to_original_space(mov_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            mov_skel = resize_img_to_original_space(mov_skel, bbox, first_reshape, isotropic_shape)
+            mov_skel = skeletonize_3d(mov_skel)
+            mov_skel_pts = np.argwhere(mov_skel)
+            iterator.set_description('{}: reshaped data'.format(fn))
+
+            ill_cond_def = False
+            ill_cond_r_def = False
+            # Deformable only
+            iterator.set_description('{}: Computing only deformable reg.'.format(fn))
+
+            tf_param = bcpd.registration_bcpd(mov_skel_pts*SCALE, fix_skel_pts*SCALE)
+
+            if np.isnan(deform_reg_def.diff):
+                tre_def = np.nan
+                pred_mov_centroid = mov_centroid
+            else:
+                tps, ill_cond_def = radial_basis_function(mov_skel_pts, np.dot(*deform_reg_def.get_registration_parameters()) / SCALE)
+                displacement_mov_centroid = tps(mov_centroid)
+                pred_mov_centroid = mov_centroid + displacement_mov_centroid
+
+                tre_def = euclidean(pred_mov_centroid, fix_centroid)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_skel_pts, mov_skel_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_skel_pts, deform_reg_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            # Rigid followed by deformable
+            iterator.set_description('{}: Computing rigid and deform. reg.'.format(fn))
+
+            rigid_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/rigid'.format(dataset_name, fnum)))
+            deform_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/deform'.format(dataset_name, fnum)))
+
+            # rigid_yt, rigid_p, rigid_reg_r_def = rigid_registration(fix_pts, mov_pts, rigid_cb)
+            # deform_yt, deform_p, deform_reg_r_def = deform_registration(fix_pts, rigid_yt, deform_cb)
+
+            time_r_def__r, rigid_reg_r_def = rigid_registration(fix_skel_pts*SCALE, mov_skel_pts*SCALE, time_it=True)
+            rigid_yt = rigid_reg_r_def.TY
+            time_r_def__def, deform_reg_r_def = deform_registration(fix_skel_pts*SCALE, rigid_yt, time_it=True,
+                                                                    tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                                    alpha=ALPHA, beta=BETA)
+
+            if np.isnan(deform_reg_r_def.diff):
+                pred_mov_centroid = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+            else:
+                mov_centroid_t = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+                tps, ill_cond_r_def = radial_basis_function(rigid_yt / SCALE,
+                                                            np.dot(*deform_reg_r_def.get_registration_parameters()) / SCALE)
+                displacement_mov_centroid_t = tps(mov_centroid_t)
+                pred_mov_centroid = mov_centroid_t + displacement_mov_centroid_t
+
+            tre_r_def = euclidean(pred_mov_centroid, fix_centroid)
+            dist_centroid_to_pts = cdist(mov_centroid[np.newaxis, ...], mov_skel_pts)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_skel_pts, mov_skel_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_skel_pts, deform_reg_r_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+
+            iterator.set_description('{}: Saving data'.format(fn))
+            df = df.append({'DATASET': dataset_name,
+                       'ITERATIONS_DEF': deform_reg_def.iteration,
+                       'ITERATIONS_R_DEF__R': rigid_reg_r_def.iteration,
+                       'ITERATIONS_R_DEF__DEF': deform_reg_r_def.iteration,
+                       'TIME_DEF': time_def,
+                       'TIME_R_DEF': time_r_def__r + time_r_def__def,
+                       'Q_DEF': deform_reg_def.diff,
+                       'Q_R_DEF__R': rigid_reg_r_def.q,
+                       'Q_R_DEF__DEF': deform_reg_r_def.diff,
+                       'ILL_COND_DEF': ill_cond_def,
+                       'ILL_COND_R_DEF': ill_cond_r_def,
+                       'TRE_DEF': tre_def, 'TRE_R_DEF': tre_r_def,
+                       'DS_DISP':euclidean(mov_centroid, fix_centroid),
+                       'DATA_PATH': file_path,
+                       'DIST_CENTR': np.min(dist_centroid_to_pts),
+                       'DIST_CENTR_DEF_95': np.percentile(dist_centroid_to_pts, 95),
+                       'SAMPLE_NUM':fnum}, ignore_index=True)
+            pts_file.close()
+
+        df.to_csv(os.path.join(OUT_IMG_FOLDER, 'cpd_{}.csv'.format(dataset_name)))

Centerline/evaluate_CPD_dense.py

@@ -0,0 +1,156 @@
+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
+from tqdm import tqdm
+from functools import partial
+import numpy as np
+from scipy.spatial.distance import euclidean
+import pandas as pd
+from EvaluationScripts.Evaluate_class import resize_img_to_original_space, resize_pts_to_original_space
+from Centerline.visualization_utils import plot_cpd_registration_step, plot_cpd
+from Centerline.cpd_utils import cpd_non_rigid_transform_pt, radial_basis_function, deform_registration, rigid_registration
+from scipy.spatial.distance import cdist
+import re
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['None', 'Affine', 'None', 'Translation']
+DATASET_FILENAME = 'points'
+
+OUT_IMG_FOLDER = '/mnt/EncryptedData1/Users/javier/vessel_registration/Centerline/cpd/dense_final'
+
+SCALE = 1e-2  # mm to cm
+
+# CPD PARAMS (deform)
+MAX_ITER = 200
+ALPHA = 2.
+BETA = 2.  # None = Use default
+TOLERANCE = 1e-8
+RBF_FUNCTION='thin-plate'
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        df = pd.DataFrame(columns=['DATASET',
+                                   'ITERATIONS_DEF', 'ITERATIONS_R_DEF__R', 'ITERATIONS_R_DEF__DEF',
+                                   'TIME_DEF', 'TIME_R_DEF',
+                                   'Q_DEF', 'Q_R_DEF__R', 'Q_R_DEF__DEF',
+                                   'TRE_DEF', 'TRE_R_DEF',
+                                   'DS_DISP',
+                                   'DATA_PATH',
+                                   'DIST_CENTR', 'DIST_CENTR_DEF_95', 'SAMPLE_NUM'])
+        for i, file_path in enumerate(iterator):
+            fn = os.path.split(file_path)[-1].split('.hd5')[0]
+            fnum = int(re.findall('(\d+)', fn)[0])
+            iterator.set_description('{}: start'.format(fn))
+            pts_file = h5py.File(file_path, 'r')
+            fix_pts = pts_file['fix/points'][:]
+            # fix_nodes = pts_file['fix/nodes'][:]
+            fix_centroid = pts_file['fix/centroid'][:]
+
+            mov_pts = pts_file['mov/points'][:]
+            # mov_nodes = pts_file['mov/nodes'][:]
+            mov_centroid = pts_file['mov/centroid'][:]
+
+            bbox = pts_file['parameters/bbox'][:]
+            first_reshape = pts_file['parameters/first_reshape'][:]
+            original_shape = pts_file['parameters/isotropic_shape'][:]
+            iterator.set_description('{}: Loaded data'.format(fn))
+            # TODO: bring back to original shape!
+            # Reshape to original_shape
+            # fix_nodes = resize_pts_to_original_space(fix_nodes, bbox, [64]*3, first_reshape, original_shape)
+            fix_pts = resize_pts_to_original_space(fix_pts, bbox, [64] * 3, first_reshape, original_shape)
+            fix_centroid = resize_pts_to_original_space(fix_centroid, bbox, [64] * 3, first_reshape, original_shape)
+            # mov_nodes = resize_pts_to_original_space(mov_nodes, bbox, [64] * 3, first_reshape, original_shape)
+            mov_pts = resize_pts_to_original_space(mov_pts, bbox, [64] * 3, first_reshape, original_shape)
+            mov_centroid = resize_pts_to_original_space(mov_centroid, bbox, [64] * 3, first_reshape, original_shape)
+            iterator.set_description('{}: reshaped data'.format(fn))
+
+            ill_cond_def = False
+            ill_cond_r_def = False
+            # Deformable only
+            iterator.set_description('{}: Computing only deformable reg.'.format(fn))
+
+            # deform_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum)))
+
+            # _, _, deform_reg_def = deform_registration(fix_pts, mov_pts, deform_cb)
+            time_def, deform_reg_def = deform_registration(fix_pts*SCALE, mov_pts*SCALE, time_it=True,
+                                                           tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                           alpha=ALPHA, beta=BETA)
+            if np.isnan(deform_reg_def.diff):
+                tre_def = np.nan
+                pred_mov_centroid = np.zeros((3,))
+            else:
+                tps, ill_cond_def = radial_basis_function(mov_pts, np.dot(*deform_reg_def.get_registration_parameters()) / SCALE, RBF_FUNCTION)
+                displacement_mov_centroid = tps(mov_centroid)
+                pred_mov_centroid = mov_centroid + displacement_mov_centroid
+
+                tre_def = euclidean(pred_mov_centroid, fix_centroid)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_pts, mov_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_pts, deform_reg_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            # Rigid followed by deformable
+            iterator.set_description('{}: Computing rigid and deform. reg.'.format(fn))
+
+            # rigid_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER,
+            #                                                                        '{}/{:04d}/RIGID_DEF/rigid'.format(
+            #                                                                            dataset_name, fnum)))
+            # deform_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER,
+            #                                                                         '{}/{:04d}/RIGID_DEF/deform'.format(
+            #                                                                             dataset_name, fnum)))
+            # rigid_yt, rigid_p, rigid_reg_r_def = rigid_registration(fix_pts, mov_pts, rigid_cb)
+            # deform_yt, deform_p, deform_reg_r_def = deform_registration(fix_pts, rigid_yt, deform_cb)
+
+            time_r_def__r, rigid_reg_r_def = rigid_registration(fix_pts*SCALE, mov_pts*SCALE, time_it=True)
+            rigid_yt = rigid_reg_r_def.TY
+            time_r_def__def, deform_reg_r_def = deform_registration(fix_pts*SCALE, rigid_yt, time_it=True,
+                                                                    tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                                    alpha=ALPHA, beta=BETA)
+
+            if np.isnan(deform_reg_r_def.diff):
+                pred_mov_centroid = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+            else:
+                mov_centroid_t = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+                tps, ill_cond_r_def = radial_basis_function(rigid_yt / SCALE, np.dot(*deform_reg_r_def.get_registration_parameters()) / SCALE, RBF_FUNCTION)
+                displacement_mov_centroid_t = tps(mov_centroid_t)
+                pred_mov_centroid = mov_centroid_t + displacement_mov_centroid_t
+
+            tre_r_def = euclidean(pred_mov_centroid, fix_centroid)
+            dist_centroid_to_pts = cdist(mov_centroid[np.newaxis, ...], mov_pts)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_pts, mov_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_pts, deform_reg_r_def.TY, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            iterator.set_description('{}: Saving data'.format(fn))
+            df = df.append({'DATASET': dataset_name,
+                       'ITERATIONS_DEF': deform_reg_def.iteration,
+                       'ITERATIONS_R_DEF__R': rigid_reg_r_def.iteration,
+                       'ITERATIONS_R_DEF__DEF': deform_reg_r_def.iteration,
+                       'TIME_DEF': time_def,
+                       'TIME_R_DEF': time_r_def__r + time_r_def__def,
+                       'Q_DEF': deform_reg_def.diff,
+                       'Q_R_DEF__R': rigid_reg_r_def.q,
+                       'Q_R_DEF__DEF': deform_reg_r_def.diff,
+                       'ILL_COND_DEF': ill_cond_def,
+                       'ILL_COND_R_DEF': ill_cond_r_def,
+                       'TRE_DEF': tre_def, 'TRE_R_DEF': tre_r_def,
+                       'DS_DISP':euclidean(mov_centroid, fix_centroid),
+                       'DATA_PATH': file_path,
+                       'DIST_CENTR': np.min(dist_centroid_to_pts),
+                       'DIST_CENTR_DEF_95': np.percentile(dist_centroid_to_pts, 95),
+                       'SAMPLE_NUM': fnum}, ignore_index=True)
+            pts_file.close()
+
+        df.to_csv(os.path.join(OUT_IMG_FOLDER, 'cpd_{}.csv'.format(dataset_name)))

Centerline/evaluate_CPD_nodes.py

@@ -0,0 +1,158 @@
+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
+from tqdm import tqdm
+from functools import partial
+import numpy as np
+from scipy.spatial.distance import euclidean
+import pandas as pd
+from EvaluationScripts.Evaluate_class import resize_img_to_original_space, resize_pts_to_original_space
+from Centerline.visualization_utils import plot_cpd_registration_step, plot_cpd
+from Centerline.cpd_utils import cpd_non_rigid_transform_pt, radial_basis_function, deform_registration, rigid_registration
+from scipy.spatial.distance import cdist
+import re
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['Affine', 'None', 'Translation']
+DATASET_FILENAME = 'points'
+
+OUT_IMG_FOLDER = '/mnt/EncryptedData1/Users/javier/vessel_registration/Centerline/cpd/nodes_final'
+
+SCALE = 1e-2  # mm to cm
+
+# CPD PARAMS (deform)
+MAX_ITER = 200
+ALPHA = 2.
+BETA = 2.  # None = Use default
+TOLERANCE = 1e-8
+RBF_FUNCTION='thin-plate'
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        df = pd.DataFrame(columns=['DATASET',
+                                   'ITERATIONS_DEF', 'ITERATIONS_R_DEF__R', 'ITERATIONS_R_DEF__DEF',
+                                   'TIME_DEF', 'TIME_R_DEF',
+                                   'Q_DEF', 'Q_R_DEF__R', 'Q_R_DEF__DEF',
+                                   'TRE_DEF', 'TRE_R_DEF',
+                                   'DS_DISP',
+                                   'DATA_PATH',
+                                   'DIST_CENTR', 'DIST_CENTR_DEF_95', 'SAMPLE_NUM'])
+        for i, file_path in enumerate(iterator):
+            fn = os.path.split(file_path)[-1].split('.hd5')[0]
+            fnum = int(re.findall('(\d+)', fn)[0])
+            iterator.set_description('{}: start'.format(fn))
+            pts_file = h5py.File(file_path, 'r')
+            fix_pts = pts_file['fix/points'][:]
+            fix_nodes = pts_file['fix/nodes'][:]
+            fix_centroid = pts_file['fix/centroid'][:]
+
+            mov_pts = pts_file['mov/points'][:]
+            mov_nodes = pts_file['mov/nodes'][:]
+            mov_centroid = pts_file['mov/centroid'][:]
+
+            bbox = pts_file['parameters/bbox'][:]
+            first_reshape = pts_file['parameters/first_reshape'][:]
+            isotropic_shape = pts_file['parameters/isotropic_shape'][:]
+            iterator.set_description('{}: Loaded data'.format(fn))
+            # TODO: bring back to original shape!
+            # Reshape to original_shape
+            fix_nodes = resize_pts_to_original_space(fix_nodes, bbox, [64] * 3, first_reshape, isotropic_shape)
+            fix_pts = resize_pts_to_original_space(fix_pts, bbox, [64] * 3, first_reshape, isotropic_shape)
+            fix_centroid = resize_pts_to_original_space(fix_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            mov_nodes = resize_pts_to_original_space(mov_nodes, bbox, [64] * 3, first_reshape, isotropic_shape)
+            mov_pts = resize_pts_to_original_space(mov_pts, bbox, [64] * 3, first_reshape, isotropic_shape)
+            mov_centroid = resize_pts_to_original_space(mov_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            iterator.set_description('{}: reshaped data'.format(fn))
+
+            if mov_nodes.shape[0] == 1:
+                # Otherwise we only have a point, and CPD can't handle that... absurd!
+                fix_nodes = fix_pts
+                mov_nodes = mov_pts
+
+            ill_cond_def = False
+            ill_cond_r_def = False
+            # Deformable only
+            iterator.set_description('{}: Computing only deformable reg.'.format(fn))
+
+            # deform_cb = partial(plot_cpd_registration_step,
+            # out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum)))
+
+            # _, _, deform_reg_def = deform_registration(fix_nodes, mov_nodes, deform_cb)
+            time_def, deform_reg_def = deform_registration(fix_nodes*SCALE, mov_nodes*SCALE, time_it=True,
+                                                           tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                           alpha=ALPHA, beta=BETA)
+            if np.isnan(deform_reg_def.diff):
+                tre_def = np.nan
+                pred_mov_centroid = np.zeros((3,))
+            else:
+                tps, ill_cond_def = radial_basis_function(mov_nodes, np.dot(*deform_reg_def.get_registration_parameters()) / SCALE, RBF_FUNCTION)
+                displacement_mov_centroid = tps(mov_centroid)
+                pred_mov_centroid = mov_centroid + displacement_mov_centroid
+
+                tre_def = euclidean(pred_mov_centroid, fix_centroid)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_nodes, mov_nodes, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_nodes, deform_reg_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            # Rigid followed by deformable
+            iterator.set_description('{}: Computing rigid and deform. reg.'.format(fn))
+
+            # rigid_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/rigid'.format(dataset_name, fnum)))
+            # deform_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/deform'.format(dataset_name, fnum)))
+            # rigid_yt, rigid_p, rigid_reg_r_def = rigid_registration(fix_nodes, mov_nodes, rigid_cb)
+            # deform_yt, deform_p, deform_reg_r_def = deform_registration(fix_nodes, rigid_yt, deform_cb)
+
+            time_r_def__r, rigid_reg_r_def = rigid_registration(fix_nodes*SCALE, mov_nodes*SCALE, time_it=True)
+            rigid_yt = rigid_reg_r_def.TY
+            time_r_def__def, deform_reg_r_def = deform_registration(fix_nodes*SCALE, rigid_yt, time_it=True,
+                                                                    tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                                    alpha=ALPHA, beta=BETA)
+
+            if np.isnan(deform_reg_r_def.diff):
+                pred_mov_centroid = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+            else:
+                mov_centroid_t = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+                tps, ill_cond_r_def = radial_basis_function(rigid_yt / SCALE, np.dot(*deform_reg_r_def.get_registration_parameters()) / SCALE, RBF_FUNCTION)
+                displacement_mov_centroid_t = tps(mov_centroid_t)
+                pred_mov_centroid = mov_centroid_t + displacement_mov_centroid_t
+
+            tre_r_def = euclidean(pred_mov_centroid, fix_centroid)
+            dist_centroid_to_pts = cdist(mov_centroid[np.newaxis, ...], mov_nodes)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_nodes, mov_nodes, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_nodes, deform_reg_r_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            iterator.set_description('{}: Saving data'.format(fn))
+            df = df.append({'DATASET':dataset_name,
+                       'ITERATIONS_DEF': deform_reg_def.iteration,
+                       'ITERATIONS_R_DEF__R': rigid_reg_r_def.iteration,
+                       'ITERATIONS_R_DEF__DEF': deform_reg_r_def.iteration,
+                       'TIME_DEF': time_def,
+                       'TIME_R_DEF': time_r_def__r + time_r_def__def,
+                       'Q_DEF': deform_reg_def.diff,
+                       'Q_R_DEF__R': rigid_reg_r_def.q,
+                       'Q_R_DEF__DEF': deform_reg_r_def.diff,
+                       'ILL_COND_DEF': ill_cond_def,
+                       'ILL_COND_R_DEF': ill_cond_r_def,
+                       'TRE_DEF':tre_def, 'TRE_R_DEF':tre_r_def,
+                       'DS_DISP':euclidean(mov_centroid, fix_centroid),
+                       'DATA_PATH':file_path,
+                       'DIST_CENTR':np.min(dist_centroid_to_pts),
+                       'DIST_CENTR_DEF_95':np.percentile(dist_centroid_to_pts, 95),
+                       'SAMPLE_NUM':fnum}, ignore_index=True)
+            pts_file.close()
+
+        df.to_csv(os.path.join(OUT_IMG_FOLDER, 'cpd_{}.csv'.format(dataset_name)))

Centerline/evaluate_CPD_skeleton.py

@@ -0,0 +1,164 @@
+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
+from tqdm import tqdm
+from functools import partial
+import numpy as np
+from scipy.spatial.distance import euclidean
+import pandas as pd
+from EvaluationScripts.Evaluate_class import resize_img_to_original_space, resize_pts_to_original_space
+from Centerline.visualization_utils import plot_cpd_registration_step, plot_cpd
+from Centerline.cpd_utils import cpd_non_rigid_transform_pt, radial_basis_function, deform_registration, rigid_registration
+from scipy.spatial.distance import cdist
+from skimage.morphology import skeletonize_3d
+import re
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['Affine', 'None', 'Translation']
+DATASET_FILENAME = 'points'
+
+OUT_IMG_FOLDER = '/mnt/EncryptedData1/Users/javier/vessel_registration/Centerline/cpd/skeleton'
+
+SCALE = 1e-2  # mm to cm
+# CPD PARAMS (deform)
+MAX_ITER = 200
+ALPHA = 0.1
+BETA = 1.0  # None = Use default
+TOLERANCE = 1e-8
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        df = pd.DataFrame(columns=['DATASET',
+                                   'ITERATIONS_DEF', 'ITERATIONS_R_DEF__R', 'ITERATIONS_R_DEF__DEF',
+                                   'TIME_DEF', 'TIME_R_DEF',
+                                   'Q_DEF', 'Q_R_DEF__R', 'Q_R_DEF__DEF',
+                                   'TRE_DEF', 'TRE_R_DEF',
+                                   'DS_DISP',
+                                   'DATA_PATH',
+                                   'DIST_CENTR', 'DIST_CENTR_DEF_95', 'SAMPLE_NUM'])
+        for i, file_path in enumerate(iterator):
+            fn = os.path.split(file_path)[-1].split('.hd5')[0]
+            fnum = int(re.findall('(\d+)', fn)[0])
+            iterator.set_description('{}: start'.format(fn))
+            pts_file = h5py.File(file_path, 'r')
+            # fix_pts = pts_file['fix/points'][:]
+            # fix_nodes = pts_file['fix/nodes'][:]
+            fix_skel = pts_file['fix/skeleton'][:]
+            fix_centroid = pts_file['fix/centroid'][:]
+
+            # mov_pts = pts_file['mov/points'][:]
+            # mov_nodes = pts_file['mov/nodes'][:]
+            mov_skel = pts_file['mov/skeleton'][:]
+            mov_centroid = pts_file['mov/centroid'][:]
+
+            bbox = pts_file['parameters/bbox'][:]
+            first_reshape = pts_file['parameters/first_reshape'][:]
+            isotropic_shape = pts_file['parameters/isotropic_shape'][:]
+            iterator.set_description('{}: Loaded data'.format(fn))
+            # TODO: bring back to original shape!
+            # Reshape to original_shape
+            # fix_nodes = resize_pts_to_original_space(fix_nodes, bbox, [64]*3, first_reshape, original_shape)
+            # fix_pts = resize_pts_to_original_space(fix_pts, bbox, [64] * 3, first_reshape, original_shape)
+            fix_centroid = resize_pts_to_original_space(fix_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            fix_skel = resize_img_to_original_space(fix_skel, bbox, first_reshape, isotropic_shape)
+            fix_skel = skeletonize_3d(fix_skel)
+            fix_skel_pts = np.argwhere(fix_skel)
+            # mov_nodes = resize_pts_to_original_space(mov_nodes, bbox, [64] * 3, first_reshape, original_shape)
+            # mov_pts = resize_pts_to_original_space(mov_pts, bbox, [64] * 3, first_reshape, original_shape)
+            mov_centroid = resize_pts_to_original_space(mov_centroid, bbox, [64] * 3, first_reshape, isotropic_shape)
+            mov_skel = resize_img_to_original_space(mov_skel, bbox, first_reshape, isotropic_shape)
+            mov_skel = skeletonize_3d(mov_skel)
+            mov_skel_pts = np.argwhere(mov_skel)
+            iterator.set_description('{}: reshaped data'.format(fn))
+
+            ill_cond_def = False
+            ill_cond_r_def = False
+            # Deformable only
+            iterator.set_description('{}: Computing only deformable reg.'.format(fn))
+
+            deform_cb = partial(plot_cpd_registration_step,
+                                out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum)))
+
+            # _, _, deform_reg_def = deform_registration(fix_pts, mov_pts, deform_cb)
+            time_def, deform_reg_def = deform_registration(fix_skel_pts*SCALE, mov_skel_pts*SCALE, time_it=True,
+                                                           tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                           alpha=ALPHA, beta=BETA)
+            if np.isnan(deform_reg_def.diff):
+                tre_def = np.nan
+                pred_mov_centroid = mov_centroid
+            else:
+                tps, ill_cond_def = radial_basis_function(mov_skel_pts, np.dot(*deform_reg_def.get_registration_parameters()) / SCALE)
+                displacement_mov_centroid = tps(mov_centroid)
+                pred_mov_centroid = mov_centroid + displacement_mov_centroid
+
+                tre_def = euclidean(pred_mov_centroid, fix_centroid)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_skel_pts, mov_skel_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_skel_pts, deform_reg_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+            # Rigid followed by deformable
+            iterator.set_description('{}: Computing rigid and deform. reg.'.format(fn))
+
+            rigid_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/rigid'.format(dataset_name, fnum)))
+            deform_cb = partial(plot_cpd_registration_step, out_folder=os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF/deform'.format(dataset_name, fnum)))
+
+            # rigid_yt, rigid_p, rigid_reg_r_def = rigid_registration(fix_pts, mov_pts, rigid_cb)
+            # deform_yt, deform_p, deform_reg_r_def = deform_registration(fix_pts, rigid_yt, deform_cb)
+
+            time_r_def__r, rigid_reg_r_def = rigid_registration(fix_skel_pts*SCALE, mov_skel_pts*SCALE, time_it=True)
+            rigid_yt = rigid_reg_r_def.TY
+            time_r_def__def, deform_reg_r_def = deform_registration(fix_skel_pts*SCALE, rigid_yt, time_it=True,
+                                                                    tolerance=TOLERANCE, max_iterations=MAX_ITER,
+                                                                    alpha=ALPHA, beta=BETA)
+
+            if np.isnan(deform_reg_r_def.diff):
+                pred_mov_centroid = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+            else:
+                mov_centroid_t = rigid_reg_r_def.transform_point_cloud(mov_centroid*SCALE)/SCALE
+                tps, ill_cond_r_def = radial_basis_function(rigid_yt / SCALE,
+                                                            np.dot(*deform_reg_r_def.get_registration_parameters()) / SCALE)
+                displacement_mov_centroid_t = tps(mov_centroid_t)
+                pred_mov_centroid = mov_centroid_t + displacement_mov_centroid_t
+
+            tre_r_def = euclidean(pred_mov_centroid, fix_centroid)
+            dist_centroid_to_pts = cdist(mov_centroid[np.newaxis, ...], mov_skel_pts)
+
+            plot_file = os.path.join(OUT_IMG_FOLDER, '{}/{:04d}/RIGID_DEF'.format(dataset_name, fnum))
+            os.makedirs(plot_file, exist_ok=True)
+            plot_cpd(fix_skel_pts, mov_skel_pts, fix_centroid, mov_centroid, plot_file + '/before_registration')
+            plot_cpd(fix_skel_pts, deform_reg_r_def.TY/SCALE, fix_centroid, pred_mov_centroid, plot_file + '/after_registration')
+
+
+            iterator.set_description('{}: Saving data'.format(fn))
+            df = df.append({'DATASET': dataset_name,
+                       'ITERATIONS_DEF': deform_reg_def.iteration,
+                       'ITERATIONS_R_DEF__R': rigid_reg_r_def.iteration,
+                       'ITERATIONS_R_DEF__DEF': deform_reg_r_def.iteration,
+                       'TIME_DEF': time_def,
+                       'TIME_R_DEF': time_r_def__r + time_r_def__def,
+                       'Q_DEF': deform_reg_def.diff,
+                       'Q_R_DEF__R': rigid_reg_r_def.q,
+                       'Q_R_DEF__DEF': deform_reg_r_def.diff,
+                       'ILL_COND_DEF': ill_cond_def,
+                       'ILL_COND_R_DEF': ill_cond_r_def,
+                       'TRE_DEF': tre_def, 'TRE_R_DEF': tre_r_def,
+                       'DS_DISP':euclidean(mov_centroid, fix_centroid),
+                       'DATA_PATH': file_path,
+                       'DIST_CENTR': np.min(dist_centroid_to_pts),
+                       'DIST_CENTR_DEF_95': np.percentile(dist_centroid_to_pts, 95),
+                       'SAMPLE_NUM':fnum}, ignore_index=True)
+            pts_file.close()
+
+        df.to_csv(os.path.join(OUT_IMG_FOLDER, 'cpd_{}.csv'.format(dataset_name)))

Centerline/get_vessels.py

@@ -0,0 +1,30 @@
+from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
+from tqdm import tqdm
+import os
+import h5py
+import DeepDeformationMapRegistration.utils.constants as C
+
+DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL'
+DATASET_NAMES = ['Affine', 'None', 'Translation']
+DATASET_FILENAME = 'volume'
+
+
+if __name__ == '__main__':
+    for dataset_name in DATASET_NAMES:
+        dataset_loc = os.path.join(DATASET_LOCATION, dataset_name)
+        dataset_files = os.listdir(dataset_loc)
+        dataset_files.sort()
+        dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f]
+
+        iterator = tqdm(dataset_files)
+        for fn in iterator:
+            f = os.path.split(fn)[-1].split('.hd5')[0]
+            vol_file = h5py.File(fn, 'r')
+            fix_vessels = vol_file[C.H5_FIX_VESSELS_MASK][..., 0]
+            mov_vessels = vol_file[C.H5_MOV_VESSELS_MASK][..., 0]
+
+            dst_folder = os.path.join(os.getcwd(), 'VESSELS', dataset_name)
+            os.makedirs(dst_folder, exist_ok=True)
+            save_nifti(fix_vessels, os.path.join(dst_folder, f+'_fix.nii.gz'))
+            save_nifti(mov_vessels, os.path.join(dst_folder, f+'_mov.nii.gz'))
+            vol_file.close()

Centerline/graph_utils.py

@@ -0,0 +1,85 @@
+import networkx as nx
+import numpy as np
+from scipy.interpolate import RegularGridInterpolator, LinearNDInterpolator
+
+
+def graph_to_ndarray(graph):
+    out_nodes = np.empty((1, 3))
+    out_edges = np.empty((1, 3))
+    visited_nodes = list()
+    visited_node_pairs = list()
+    for (start_node, end_node) in graph.edges():
+        if (not (start_node, end_node) in visited_node_pairs) and (not (end_node, start_node) in visited_node_pairs):
+            edge = graph[start_node][end_node]['pts']
+            out_edges = np.vstack([out_edges, edge])
+
+            # Avoid duplicates
+            if not (start_node in visited_nodes):
+                out_nodes = np.vstack([out_nodes, graph.nodes[start_node]['o']])
+                visited_nodes.append(start_node)
+            if not (end_node in visited_nodes):
+                out_nodes = np.vstack([out_nodes, graph.nodes[end_node]['o']])
+                visited_nodes.append(end_node)
+
+            visited_node_pairs.append((start_node, end_node))
+
+    return np.vstack([out_edges, out_nodes]), out_nodes, out_edges
+
+
+def get_bifurcation_nodes(graph: nx.Graph):
+    # Vertex degree relates to the number of branches connected to a given node
+    out_nodes = np.empty((1, 3))
+    bif_nodes_id = list()
+    for node_num, deg in graph.degree:
+        if deg > 1:
+            bif_nodes_id.append(node_num)
+            out_nodes = np.vstack([out_nodes, graph.nodes[node_num]['o']])
+
+    return out_nodes, bif_nodes_id
+
+
+def apply_displacement(pts_list: np.ndarray, interpolator: [RegularGridInterpolator, LinearNDInterpolator]):
+    pts_list = pts_list.astype(np.float)
+    ret_val = pts_list + interpolator(pts_list).squeeze()
+    return ret_val
+
+
+def deform_graph(graph, dm_interpolator: [RegularGridInterpolator, LinearNDInterpolator]):
+    def_graph = nx.Graph()
+    for (start_node, end_node) in graph.edges():
+        edge = graph[start_node][end_node]['pts']
+        def_edge = apply_displacement(edge, dm_interpolator)
+
+        def_start_node_pts = apply_displacement(graph.nodes[start_node]['pts'], dm_interpolator)
+        def_end_node_pts = apply_displacement(graph.nodes[end_node]['pts'], dm_interpolator)
+
+        def_start_node_o = apply_displacement(graph.nodes[start_node]['o'], dm_interpolator)
+        def_end_node_o = apply_displacement(graph.nodes[end_node]['o'], dm_interpolator)
+
+        def_graph.add_node(start_node, pts=def_start_node_pts, o=def_start_node_o)
+        def_graph.add_node(end_node, pts=def_end_node_pts, o=def_end_node_o)
+        def_graph.add_edge(start_node, end_node, pts=def_edge, weight=len(def_edge))
+    return def_graph
+
+
+def subsample_graph(graph: nx.Graph, num_samples=3):
+    sub_graph = nx.Graph()
+    for (start_node, end_node) in graph.edges():
+        edge = graph[start_node][end_node]['pts']
+        edge_len = edge.shape[0]
+        sub_edge_len = (edge_len - 2) // num_samples # Do not count the pts corresponding to the nodes (-2)
+
+        sub_edge = [edge[0]]
+        include_last = bool((edge_len - 2) % num_samples)  # Skip the last point, as this is too close to the node
+        if sub_edge_len:
+            idxs = np.arange(0, edge_len, num_samples)[1:] if include_last else np.arange(0, edge_len, num_samples)[1:-1]
+            for i in idxs:
+                sub_edge.append(edge[i])
+
+        sub_edge.append(edge[-1])
+        sub_edge = np.asarray(sub_edge)
+        sub_graph.add_node(start_node, pts=graph.nodes[start_node]['pts'], o=graph.nodes[start_node]['o'])
+        sub_graph.add_node(end_node, pts=graph.nodes[end_node]['pts'], o=graph.nodes[end_node]['o'])
+        sub_graph.add_edge(start_node, end_node, pts=sub_edge, weight=len(sub_edge))
+    return sub_graph
+

Centerline/skeleton_to_graph.py

@@ -0,0 +1,167 @@
+# SRC: https://github.com/Image-Py/sknw/blob/master/sknw/sknw.py
+import numpy as np
+import networkx as nx
+from Centerline.graph_utils import subsample_graph
+
+
+def neighbors(shape):
+    dim = len(shape)
+    block = np.ones([3] * dim)
+    block[tuple([1] * dim)] = 0
+    idx = np.where(block > 0)
+    idx = np.array(idx, dtype=np.uint8).T
+    idx = np.array(idx - [1] * dim)
+    acc = np.cumprod((1,) + shape[::-1][:-1])
+    return np.dot(idx, acc[::-1])
+
+
+ # my mark
+def mark(img, nbs):  # mark the array use (0, 1, 2)
+    img = img.ravel()
+    for p in range(len(img)):
+        if img[p] == 0: continue
+        s = 0
+        for dp in nbs:
+            if img[p + dp] != 0: s += 1
+        if s == 2:
+            img[p] = 1
+        else:
+            img[p] = 2
+
+
+ # trans index to r, c...
+def idx2rc(idx, acc):
+    rst = np.zeros((len(idx), len(acc)), dtype=np.int16)
+    for i in range(len(idx)):
+        for j in range(len(acc)):
+            rst[i, j] = idx[i] // acc[j]
+            idx[i] -= rst[i, j] * acc[j]
+    rst -= 1
+    return rst
+
+
+ # fill a node (may be two or more points)
+def fill(img, p, num, nbs, acc, buf):
+    back = img[p]
+    img[p] = num
+    buf[0] = p
+    cur = 0;
+    s = 1;
+
+    while True:
+        p = buf[cur]
+        for dp in nbs:
+            cp = p + dp
+            if img[cp] == back:
+                img[cp] = num
+                buf[s] = cp
+                s += 1
+        cur += 1
+        if cur == s: break
+    return idx2rc(buf[:s], acc)
+
+
+ # trace the edge and use a buffer, then buf.copy, if use [] numba not works
+def trace(img, p, nbs, acc, buf):
+    c1 = 0;
+    c2 = 0;
+    newp = 0
+    cur = 0
+
+    while True:
+        buf[cur] = p
+        img[p] = 0
+        cur += 1
+        for dp in nbs:
+            cp = p + dp
+            if img[cp] >= 10:
+                if c1 == 0:
+                    c1 = img[cp]
+                else:
+                    c2 = img[cp]
+            if img[cp] == 1:
+                newp = cp
+        p = newp
+        if c2 != 0: break
+    return (c1 - 10, c2 - 10, idx2rc(buf[:cur], acc))
+
+
+ # parse the image then get the nodes and edges
+def parse_struc(img, pts, nbs, acc):
+    img = img.ravel()
+    buf = np.zeros(131072, dtype=np.int64)
+    num = 10
+    nodes = []
+    for p in pts:
+        if img[p] == 2:
+            nds = fill(img, p, num, nbs, acc, buf)
+            num += 1
+            nodes.append(nds)
+
+    edges = []
+    for p in pts:
+        for dp in nbs:
+            if img[p + dp] == 1:
+                edge = trace(img, p + dp, nbs, acc, buf)
+                edges.append(edge)
+    return nodes, edges
+
+
+# use nodes and edges build a networkx graph
+def build_graph(nodes, edges, multi=False):
+    graph = nx.MultiGraph() if multi else nx.Graph()
+    for i in range(len(nodes)):
+        graph.add_node(i, pts=nodes[i], o=nodes[i].mean(axis=0))
+    for s, e, pts in edges:
+        l = np.linalg.norm(pts[1:] - pts[:-1], axis=1).sum()
+        graph.add_edge(s, e, pts=pts, weight=l)
+    return graph
+
+
+def buffer(ske):
+    buf = np.zeros(tuple(np.array(ske.shape) + 2), dtype=np.uint16)
+    buf[tuple([slice(1, -1)] * buf.ndim)] = ske
+    return buf
+
+
+def build_sknw(ske, multi=False):
+    buf = buffer(ske)
+    nbs = neighbors(buf.shape)
+    acc = np.cumprod((1,) + buf.shape[::-1][:-1])[::-1]
+    mark(buf, nbs)
+    pts = np.array(np.where(buf.ravel() == 2))[0]
+    nodes, edges = parse_struc(buf, pts, nbs, acc)
+    return build_graph(nodes, edges, multi)
+
+
+# draw the graph
+def draw_graph(img, graph, cn=255, ce=128):
+    acc = np.cumprod((1,) + img.shape[::-1][:-1])[::-1]
+    img = img.ravel()
+    for idx in graph.nodes():
+        pts = graph.nodes[idx]['pts']
+        img[np.dot(pts, acc)] = cn
+    for (s, e) in graph.edges():
+        eds = graph[s][e]
+        for i in eds:
+            pts = eds[i]['pts']
+            img[np.dot(pts, acc)] = ce
+
+
+def get_graph_from_skeleton(mask, subsample=False):
+    graph = build_sknw(mask, False)
+    if len(graph.nodes) > 1 and len(graph.edges) and subsample:
+        graph = subsample_graph(graph, 3)
+    return graph
+
+
+if __name__ == '__main__':
+    g = nx.MultiGraph()
+    g.add_nodes_from([1, 2, 3, 4, 5])
+    g.add_edges_from([(1, 2), (1, 3), (2, 3), (4, 5), (5, 4)])
+    print(g.nodes())
+    print(g.edges())
+    a = g.subgraph(1)
+    print('d')
+    print(a)
+    print('d')

Centerline/skeletonization.py

@@ -0,0 +1,817 @@
+import sys, os
+import numpy as np
+import nibabel as nib
+from scipy import ndimage as ndi
+from scipy.signal import convolve
+from numpy.linalg import norm
+import networkx as nx
+import logging
+import traceback
+import timeit
+import time
+import math
+from ast import literal_eval as make_tuple
+from skimage.measure import label
+import subprocess
+import platform
+import glob
+
+
+def loadVolume(volumeFolderPath, volumeName):
+    """
+    Load nifti files (*.nii or *.nii.gz).
+    Parameters
+    ----------
+    volumeFolderPath : str
+        Folder of the volume file.
+    volumeName : str
+        Name of the volume file.
+
+    Returns
+    -------
+    volume : ndarray
+        Volume data in the form of numpy ndarray.
+    affine : ndarray
+        Associated affine transformation matrix in the form of numpy ndarray.
+    """
+    volumeFilePath = os.path.join(volumeFolderPath, volumeName)
+    volumeImg = nib.load(volumeFilePath)
+    volume = volumeImg.get_data()
+    shape = volume.shape
+    affine = volumeImg.affine
+    print('Volume loaded from {} with shape = {}.'.format(volumeFilePath, shape))
+
+    return volume, affine
+
+
+def saveVolume(volume, affine, path, astype=None):
+    """
+    Save the given volume to the specified location in specified data type.
+    Parameters
+    ----------
+    volume : ndarray
+        Volume data to be saved.
+    affine : ndarray
+        The affine transformation matrix associated with the volume.
+    path : str
+        The absolute path where the volume is going to be saved.
+    astype : numpy dtype, optional
+        The desired data type of the volume data.
+    """
+    if astype is None:
+        astype = np.uint8
+
+    nib.save(nib.Nifti1Image(volume.astype(astype), affine), path)
+    print('Volume saved to {} as type {}.'.format(path, astype))
+
+
+def labelVolume(volume, minSize=1, maxHop=3):
+    """
+    Partition the volume into several connected components and attach labels.
+    Parameters
+    ----------
+    volume : ndarray
+        Volume to be partitioned.
+    minSize : int, optional
+        The connected component that is less than this size will be disgarded.
+    maxHop : int, optional
+        Controls how neighboring voxels are defined. See `label` doc for details.
+
+    Returns
+    -------
+    labeled : ndarray
+        The partitioned and labeled volume. Each connected component has a label (a positive integer) and the background
+        is labeled as 0.
+    labelResult : list
+        In the form of [[label1, size1], [label2, size2], ...]
+    """
+    labeled, maxNum = label(volume, return_num=True, connectivity=maxHop)
+    counts = np.bincount(labeled.ravel())
+    countLoc = np.nonzero(counts)[0]
+    sizeList = counts[countLoc]
+    labelResult = list(zip(countLoc[sizeList >= minSize], sizeList[sizeList >= minSize]))
+    # print(labelResult)
+    # print('Total segments: {}'.format(np.count_nonzero(sizeList >= minSize)))
+    return labeled, labelResult
+
+
+def analyze(vesselVolumeMask, baseFolder):
+    """
+    Main function to provoke the skeletonization process. Note that here I am using the docker version of the code. If
+    you have already downloaded the original C++ code and successfully compiled it, then you can run that compiled code
+    instead of this one.
+    """
+    vesselVolumeMask = vesselVolumeMask.astype(np.uint8)
+    vesselVolumeMask[vesselVolumeMask != 0] = 1
+    vesselVolumeMask = np.swapaxes(vesselVolumeMask, 0, 2)
+    shape = vesselVolumeMask.shape
+
+    vesselVolumeMaskLabeled, vesselVolumeMaskLabelResult = labelVolume(vesselVolumeMask, minSize=1)
+    directory = os.path.join(baseFolder, 'skeletonizationResult')
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+        print('Directory {} created.'.format(directory))
+
+    vesselVolumeMaskLabelInfoFilename = 'vesselVolumeMaskLabelInfo.npz'
+    vesselVolumeMaskLabelInfoFilePath = os.path.join(directory, vesselVolumeMaskLabelInfoFilename)
+    np.savez_compressed(vesselVolumeMaskLabelInfoFilePath, vesselVolumeMaskLabeled=vesselVolumeMaskLabeled,
+                        vesselVolumeMaskLabelResult=vesselVolumeMaskLabelResult)
+    print('{} saved to {}.'.format(vesselVolumeMaskLabelInfoFilename, vesselVolumeMaskLabelInfoFilePath))
+
+    # directory2 = directory + 'labelNum=' + str(labelNum) + '/'
+    # if not os.path.exists(directory2):
+    #     os.makedirs(directory2)
+    # with open(directory2 + 'BB.txt', 'w') as f1:
+    #     f1.write('1\n')
+    #     f1.write('{} {} {}\n'.format(0, 0, 0))
+    #     f1.write('{} {} {}'.format(*shape))
+    # '''
+    BBFilePath = os.path.join(directory, 'BB.txt')
+    f1 = open(BBFilePath, 'w')
+    f1.write('1\n')
+    f1.write('{} {} {}\n'.format(0, 0, 0))
+    f1.write('{} {} {}'.format(*shape))
+    f1.close()
+
+    vesselCoords = np.array(np.where(vesselVolumeMask)).T
+    xyzFilePath = os.path.join(directory, 'xyz.txt')
+    np.savetxt(xyzFilePath, vesselCoords, fmt='%1u')
+    f2 = open(xyzFilePath, "r")
+    contents = f2.readlines()
+    f2.close()
+
+    contents.insert(0, '{}\n'.format(len(vesselCoords)))
+
+    f2 = open(xyzFilePath, "w")
+    contents = "".join(contents)
+    f2.write(contents)
+    f2.close()
+    # '''
+
+    # '''
+    currentPlatform = platform.system()
+    print('Current platform is {}.'.format(currentPlatform))
+    if currentPlatform == 'Windows':
+        cmd = '"C:/Program Files/Docker/Docker/Resources/bin/docker.exe" run -v ' + '"' + directory + '"' + ':/write_directory -e THRESH=1e-12 -e CC_FLAG=1 -e CONVERSION_TYPE=1 amytabb/curveskel-tabb-medeiros-2018-docker'
+    elif currentPlatform == 'Darwin':
+        cmd = 'docker run -v ' + '"' + directory + '"' + ':/write_directory -e THRESH=1e-12 -e CC_FLAG=1 -e CONVERSION_TYPE=1 amytabb/curveskel-tabb-nih-aug2018-docker2'
+    elif currentPlatform == 'Linux':
+        cmd = '/usr/local/bin/docker run -v ' + '"' + directory + '"' + ':/write_directory -e THRESH=1e-12 -e CC_FLAG=1 -e CONVERSION_TYPE=1 amytabb/curveskel-tabb-medeiros-2018-docker'
+        cmd = 'sudo docker run -v ' + '"' + directory + '"' + ':/write_directory -e THRESH=1e-12 -e CC_FLAG=1 -e CONVERSION_TYPE=1 amytabb/curveskel-tabb-medeiros-2018-docker'
+        cmd = 'sudo docker run -v ' + '"' + directory + '"' + ':/write_directory -e THRESH=1e-12 -e CC_FLAG=1 -e CONVERSION_TYPE=1 amytabb/curveskel-tabb-nih-aug2018-docker2'
+
+    print('cmd={}'.format(cmd))
+    subprocess.call(cmd, shell=True)
+    # '''
+
+
+def combineSkeletonSegments(skeletonSegmentFolderPath):
+    """
+    Collect and combine the results from the skeletonization.
+    Parameters
+    ----------
+    skeletonSegmentFolderPath : str
+        The folder that contains the segments information (result_segments_xyz*.txt).
+
+    Returns
+    -------
+    segmentList : list
+        A list containing the segment information. Each sublist represents a segment and each element in the sublist
+        represents a centerpoint coordinates.
+    """
+    segmentList = []
+    files = glob.glob(os.path.join(skeletonSegmentFolderPath, 'result_segments_xyz*.txt'))
+    for segmentFile in files:
+        result = readSegmentFile(segmentFile)
+        segmentList += result
+
+    return segmentList
+
+
+def readSegmentFile(segmentFile):
+    """
+    Parse the segment files (result_segments_xyz*.txt) and return segments information in a list.
+    Parameters
+    ----------
+    segmentFile : str
+        Path to the segment file.
+
+    Returns
+    -------
+    segmentList : list
+        A list containing the segment information. Each sublist represents a segment and each element in the sublist
+        represents a centerpoint coordinates.
+    """
+    isFirstLine = True
+    isSegmentLength = True
+    segmentList = []
+    with open(segmentFile) as f:
+        for line in f:
+            if isFirstLine:
+                numOfSegments = int(line)
+                isFirstLine = False
+            else:
+                if isSegmentLength:
+                    segmentLength = int(line)
+                    isSegmentLength = False
+                    segmentCounter = 1
+                    segment = []
+                else:
+                    if segmentCounter <= segmentLength:
+                        voxel = tuple([int(x) for x in line.split(' ')])
+                        segment.append(voxel[::-1])
+                        segmentCounter += 1
+                    else:
+                        segmentCounter += 1
+                        isSegmentLength = True
+                        segmentList.append(segment)
+                        assert (len(segment) == segmentLength)
+
+    return segmentList
+
+
+# def drawSegments(segmentList):
+#     pass
+
+def processSegments(segmentList, shape):
+    """
+    Re-partition the segments so that each segment is a simple branch, i.e., it does not contain bifurcation point
+    unless at the two ends.
+    Note that this function might be replaced by another more concise function `getSegmentList`.
+    Parameters
+    ----------
+    segmentList : list
+        A list containing the segment information. Each sublist represents a segment and each element in the sublist
+        represents a centerpoint coordinates.
+    shape : tuple
+        Shape of the vessel volume (used for ploting).
+
+    Returns
+    -------
+    G : NetworkX graph
+        A graph in which each node represents a centerpoint and each edge represents a portion of a vessel branch.
+    segmentList : list
+        A list containing the segment information. Each sublist represents a segment and each element in the sublist
+        represents a centerpoint coordinates.
+    errorSegments : list
+        A list that contains segments that cannot be fixed.
+    """
+    ## Import pyqtgraph ##
+    from pyqtgraph.Qt import QtCore, QtGui
+    import pyqtgraph as pg
+    import pyqtgraph.opengl as gl
+
+    ## Init ##
+    app = pg.QtGui.QApplication([])
+    w = gl.GLViewWidget()
+    w.opts['distance'] = 800
+    w.setGeometry(0, 110, 1600, 900)
+    offset = np.array(shape) / (-2.0)
+
+    G = nx.Graph()
+    colorList = [pg.glColor('r'), pg.glColor('g'), pg.glColor('b'), pg.glColor('c'), pg.glColor('m'), pg.glColor('y')]
+    colorPointer = 0
+    skeleton = np.full(shape, 0)
+    for segment in segmentList:
+        # G.add_path(list(map(tuple, segment)))
+        G.add_path(segment)
+        segmentCoords = np.array(segment)
+        skeleton[tuple(segmentCoords.T)] = 1
+        # segmentCoordsView = segmentCoords + offset
+        # aa = gl.GLLinePlotItem(pos=segmentCoordsView, color=colorList[colorPointer], width=3)
+        # w.addItem(aa)
+        # colorPointer = colorPointer + 1 if colorPointer < len(colorList) - 1 else 0
+
+    # skeletonCoords = np.array(np.where(skeleton)).T
+    # skeletonCoordsView = (skeletonCoords + offset) * affineTransform
+    # aa = gl.GLScatterPlotItem(pos=skeletonCoordsView, size=5)
+    # w.addItem(aa)
+
+    # w.show()
+
+    voxelDegrees = np.array([v for _, v in G.degree(G.nodes())])
+    maxVoxelDegree = np.amax(voxelDegrees)
+    voxelDegreesZippedResult = list(zip(np.arange(maxVoxelDegree + 1), np.bincount(voxelDegrees)))
+    print('Voxel degree distribution is \n{}'.format(voxelDegreesZippedResult))
+    print('Number of cycles is {}'.format(len(nx.cycle_basis(G))))
+
+    # Remove duplicate segments
+    keepList = np.full((len(segmentList),), True)
+    duplicateCounter = 0
+    for idx, seg in enumerate(segmentList):
+        for idx2, seg2 in enumerate(segmentList[idx + 1:]):
+            if seg == seg2 or seg == seg2[::-1]:
+                keepList[idx + idx2] = False
+                duplicateCounter += 1
+
+    segmentList = [seg for idx, seg in enumerate(segmentList) if keepList[idx]]
+    print('{} duplicate segments removed!'.format(duplicateCounter))
+
+    # Cut segments into sub-pieces if there are bifurcation points in the middle
+    extraSegments = []
+    keepList = np.full((len(segmentList),), True)
+    for idx, segment in enumerate(segmentList):
+        voxelDegrees = np.array([v for _, v in G.degree(segment)])
+        if len(voxelDegrees) >= 3:
+            if voxelDegrees[0] == 2 or voxelDegrees[-1] == 2 or (not np.all(voxelDegrees[1:-1] == 2)):
+                keepList[idx] = False
+                locs = np.nonzero(voxelDegrees != 2)[0]
+                if voxelDegrees[0] == 2:
+                    locs = np.hstack((0, locs))
+
+                if voxelDegrees[-1] == 2:
+                    locs = np.hstack((locs, len(voxelDegrees)))
+
+                newSegments = []
+                for ii in range(len(locs) - 1):
+                    newSegments.append(segment[locs[ii]:(locs[ii + 1] + 1)])
+
+                extraSegments += newSegments
+
+    segmentList = [seg for idx, seg in enumerate(segmentList) if keepList[idx]]
+    segmentList += extraSegments
+
+    # Remove duplicate segments again
+    keepList = np.full((len(segmentList),), True)
+    duplicateCounter = 0
+    for idx, seg in enumerate(segmentList):
+        for idx2, seg2 in enumerate(segmentList[idx + 1:]):
+            if seg == seg2 or seg == seg2[::-1]:
+                keepList[idx + idx2] = False
+                duplicateCounter += 1
+
+    segmentList = [seg for idx, seg in enumerate(segmentList) if keepList[idx]]
+    print('{} duplicate segments removed in the second stage!'.format(duplicateCounter))
+
+    # Remove segment if it is completely contained in another segment
+    # keepList = np.full((len(segmentList),), True)
+    # sublistCounter = 0
+    # for idx, seg in enumerate(segmentList):
+    #     for idx2, seg2 in enumerate(segmentList[idx + 1:]):
+    #         if contains(seg, seg2):
+    #             keepList[idx] = False
+    #             sublistCounter += 1
+    #         elif contains(seg2, seg):
+    #             keepList[idx + idx2] = False
+    #             sublistCounter += 1
+
+    # segmentList = [seg for idx, seg in enumerate(segmentList) if keepList[idx]]
+    # print('{} sublist segments removed!'.format(sublistCounter))
+
+    # Treat the segment if either end is not correct
+    hasInvalidSegments = False
+    for idx, segment in enumerate(segmentList):
+        voxelDegrees = np.array([v for _, v in G.degree(segment)])
+        if len(voxelDegrees) == 2:
+            if voxelDegrees[0] == 2 or voxelDegrees[-1] == 2:
+                # print('Degrees on either end is 2: {}'.format(voxelDegrees))
+                hasInvalidSegments = True
+        elif len(voxelDegrees) > 2:
+            if voxelDegrees[0] == 2 or voxelDegrees[-1] == 2 or np.any(voxelDegrees[1:-1] != 2):
+                # print('Degrees not correct: {}'.format(voxelDegrees))
+                hasInvalidSegments = True
+
+    if not hasInvalidSegments:
+        drawSegments(segmentList, shape)
+        print('No errors!')
+        errorSegments = []
+        return G, segmentList, errorSegments
+
+    iterCounter = 1
+    while hasInvalidSegments:
+        print('\n\nIter={}'.format(iterCounter))
+        keepList = np.full((len(segmentList),), True)
+        extraSegments = []
+        for idx, segment in enumerate(segmentList):
+            if keepList[idx]:
+                voxelDegrees = np.array([v for _, v in G.degree(segment)])
+                if voxelDegrees[0] == 2 and voxelDegrees[-1] == 2:
+                    print('Both end have 2 neighbours')
+                elif voxelDegrees[0] == 2 or voxelDegrees[-1] == 2:
+                    # print('Degrees on either end is 2: {}'.format(voxelDegrees))
+                    # pass
+                    # segmentCoords = np.array(segment)
+                    if voxelDegrees[0] == 2:
+                        otherSegmentInfo = [(idx2, seg) for idx2, seg in enumerate(segmentList) if
+                                            (seg[0] == segment[0] or seg[-1] == segment[0]) and keepList[
+                                                idx2] and idx != idx2]
+                        if len(otherSegmentInfo) != 0:
+                            if len(otherSegmentInfo) > 1:
+                                # print(contains(segment, otherSegmentInfo[0][1]), contains(otherSegmentInfo[1][1], segment))
+                                otherSegmentInfoTemp = []
+                                for idx2, seg in otherSegmentInfo:
+                                    if contains(segment, seg) or contains(segment[::-1], seg):
+                                        keepList[idx] = False
+                                        continue
+                                    elif contains(seg, segment) or contains(seg[::-1], segment):
+                                        keepList[idx2] = False
+                                        otherSegmentInfoTemp.append((idx2, seg))
+
+                                otherSegmentInfo = otherSegmentInfoTemp
+                                # otherSegmentInfo = [segInfo for segInfo in otherSegmentInfo if not (contains(segment, segInfo[1]) or contains(segInfo[1], segment))]
+                                if len(otherSegmentInfo) > 1:
+                                    print('More than one other segments found!')
+                                    print('Current segment ({}) is {} ({})'.format(idx, segment, voxelDegrees))
+                                    for otherSegmentIdx, otherSegment in otherSegmentInfo:
+                                        otherSegmentVoxelDegrees = np.array([v for _, v in G.degree(otherSegment)])
+                                        print('Idx = {}: {} ({})'.format(otherSegmentIdx, otherSegment,
+                                                                         otherSegmentVoxelDegrees))
+                                elif len(otherSegmentInfo) == 1:
+                                    otherSegmentIdx, otherSegment = otherSegmentInfo[0]
+                                else:
+                                    print('No valid other segments found!')
+                                    continue
+                            else:
+                                otherSegmentIdx, otherSegment = otherSegmentInfo[0]
+                                if contains(segment, otherSegment) or contains(segment[::-1], otherSegment):
+                                    keepList[idx] = False
+                                    continue
+                                elif contains(otherSegment, segment) or contains(otherSegment[::-1], segment):
+                                    keepList[otherSegmentIdx] = False
+                                    continue
+
+                            newSegment = otherSegment + segment[1:] if otherSegment[-1] == segment[0] else otherSegment[
+                                                                                                           ::-1] + segment[
+                                                                                                                   1:]
+                            if not validateSegment(G, newSegment):
+                                newSegmentVoxelDegrees = np.array([v for _, v in G.degree(newSegment)])
+                                print('Old degree is {} () and new degree is {} ()'.format(voxelDegrees,
+                                                                                           newSegmentVoxelDegrees))
+                            else:
+                                print('Two segments ({} and {}) merged together!'.format(idx, otherSegmentIdx))
+
+                            extraSegments.append(newSegment)
+                            keepList[idx] = False
+                            keepList[otherSegmentIdx] = False
+                        else:
+                            print(
+                                'Could not find other segments for segment({}) {} with degrees {}'.format(idx, segment,
+                                                                                                          voxelDegrees))
+                            possibleSegmentsInfo = [(idx2, seg) for idx2, seg in enumerate(segmentList) if
+                                                    (seg[0] == segment[0] or seg[-1] == segment[0]) and idx != idx2]
+                            print('Possible segments: {}'.format(len(possibleSegmentsInfo)))
+
+                    elif voxelDegrees[-1] == 2:
+                        otherSegmentInfo = [(idx2, seg) for idx2, seg in enumerate(segmentList) if
+                                            (seg[0] == segment[-1] or seg[-1] == segment[-1]) and keepList[
+                                                idx2] and idx != idx2]
+                        if len(otherSegmentInfo) != 0:
+                            if len(otherSegmentInfo) > 1:
+                                # print(contains(segment, otherSegmentInfo[0][1]), contains(otherSegmentInfo[1][1], segment))
+                                otherSegmentInfoTemp = []
+                                for idx2, seg in otherSegmentInfo:
+                                    if contains(segment, seg) or contains(segment[::-1], seg):
+                                        keepList[idx] = False
+                                        continue
+                                    elif contains(seg, segment) or contains(seg[::-1], segment):
+                                        keepList[idx2] = False
+                                        otherSegmentInfoTemp.append((idx2, seg))
+
+                                otherSegmentInfo = otherSegmentInfoTemp
+                                # otherSegmentInfo = [segInfo for segInfo in otherSegmentInfo if not (contains(segment, segInfo[1]) or contains(segInfo[1], segment))]
+                                if len(otherSegmentInfo) > 1:
+                                    print('More than one other segments found!')
+                                    print('Current segment ({}) is {} ({})'.format(idx, segment, voxelDegrees))
+                                    for otherSegmentIdx, otherSegment in otherSegmentInfo:
+                                        otherSegmentVoxelDegrees = np.array([v for _, v in G.degree(otherSegment)])
+                                        print('Idx = {}: {} ({})'.format(otherSegmentIdx, otherSegment,
+                                                                         otherSegmentVoxelDegrees))
+                                elif len(otherSegmentInfo) == 1:
+                                    otherSegmentIdx, otherSegment = otherSegmentInfo[0]
+                                else:
+                                    print('No valid other segments found!')
+                                    continue
+                            else:
+                                otherSegmentIdx, otherSegment = otherSegmentInfo[0]
+                                if contains(segment, otherSegment) or contains(segment[::-1], otherSegment):
+                                    keepList[idx] = False
+                                    continue
+                                elif contains(otherSegment, segment) or contains(otherSegment[::-1], segment):
+                                    keepList[otherSegmentIdx] = False
+                                    continue
+
+                            newSegment = segment[:-1] + otherSegment if otherSegment[0] == segment[-1] else segment[
+                                                                                                            :-1] + otherSegment[
+                                                                                                                   ::-1]
+                            if not validateSegment(G, newSegment):
+                                newSegmentVoxelDegrees = np.array([v for _, v in G.degree(newSegment)])
+                                print('Old degree is {} () and new degree is {} ()'.format(voxelDegrees,
+                                                                                           newSegmentVoxelDegrees))
+                            else:
+                                print('Two segments ({} and {}) merged together!'.format(idx, otherSegmentIdx))
+
+                            extraSegments.append(newSegment)
+                            keepList[idx] = False
+                            keepList[otherSegmentIdx] = False
+                        else:
+                            print(
+                                'Could not find other segments for segment({}) {} with degrees {}'.format(idx, segment,
+                                                                                                          voxelDegrees))
+                            possibleSegmentsInfo = [(idx2, seg) for idx2, seg in enumerate(segmentList) if
+                                                    (seg[0] == segment[-1] or seg[-1] == segment[-1]) and idx != idx2]
+                            print('Possible segments: {}'.format(len(possibleSegmentsInfo)))
+
+        segmentList = [segment for idx, segment in enumerate(segmentList) if keepList[idx]]
+        segmentList += extraSegments
+        hasInvalidSegments = False
+        errorSegments = []
+        for idx, segment in enumerate(segmentList):
+            voxelDegrees = np.array([v for _, v in G.degree(segment)])
+            if len(voxelDegrees) == 2:
+                if voxelDegrees[0] == 2 or voxelDegrees[-1] == 2:
+                    print('Degrees on either end is 2: {}'.format(voxelDegrees))
+                    hasInvalidSegments = True
+                    errorSegments.append(segment)
+            elif len(voxelDegrees) > 2:
+                if voxelDegrees[0] == 2 or voxelDegrees[-1] == 2 or np.any(voxelDegrees[1:-1] != 2):
+                    print('Degrees not correct: {}'.format(voxelDegrees))
+                    hasInvalidSegments = True
+                    errorSegments.append(segment)
+
+        print('hasInvalidSegments = {}'.format(hasInvalidSegments))
+        iterCounter += 1
+        if len(extraSegments) == 0:
+            hasInvalidSegments = False
+            print('While loop aborted because there is no change in segments!')
+
+    for errorSegment in errorSegments:
+        segmentList.remove(errorSegment)
+
+    # np.savez_compressed(directory + 'segmentList.npz', segmentList=segmentList)
+    # if partIdx != 10:
+    #     nib.save(nib.Nifti1Image(skeleton.astype(np.int16), vesselImg.affine), directory + skeletonNamePartial + str(partIdx) + '.nii.gz')
+    # else:
+    #     nib.save(nib.Nifti1Image(skeleton.astype(np.int16), vesselImg.affine), directory + skeletonNameTotal + '.nii.gz')
+
+    # nx.write_graphml(G, directory + 'graphRepresentation.graphml')
+
+    # drawAbstractGraph(offset, segmentList)
+    # drawAbstractGraph(offset, errorSegments)
+
+    print(errorSegments)
+
+    return G, segmentList, errorSegments
+
+
+def getSegmentList(G, nodeInfoDict):
+    """
+    Generate segmentList from graph and nodeInfoDict.
+    Parameters
+    ----------
+    G : NetworkX graph
+        The graph representation of the network.
+    nodeInfoDict : dict
+        A dictionary containing the information about nodes.
+
+    Returns
+    -------
+    segmentList : list
+        A list of segments in which each segment is a simple branch.
+    """
+    startNodeIDList = [nodeID for nodeID in nodeInfoDict.keys() if nodeInfoDict[nodeID]['parentNodeID'] == -1]
+    print('startNodeIDList = {}'.format(startNodeIDList))
+    segmentList = []
+    for startNodeID in startNodeIDList:
+        segmentList = getSegmentListDetail(G, nodeInfoDict, segmentList, startNodeID)
+
+    print('There are {} segments in segmentList'.format(len(segmentList)))
+    print(segmentList)
+    return segmentList
+
+
+def getSegmentListDetail(G, nodeInfoDict, segmentList, startNodeID):
+    """
+    Implementation of `getSegmentList`. Use DFS to traverse all the segments.
+    Parameters
+    ----------
+    G : NetworkX graph
+        The graph representation of the network.
+    nodeInfoDict : dict
+        A dictionary containing the information about nodes.
+    segmentList : list
+        A list of segments in which each segment is a simple branch.
+    startNodeID : int
+        The index of the start point of a segment.
+
+    Returns
+    -------
+    segmentList : list
+        A list of segments in which each segment is a simple branch.
+    """
+    neighborNodeIDList = [nodeID for nodeID in list(G[startNodeID].keys()) if
+                          'visited' not in G[startNodeID][nodeID]]  # use adjacency dict to find neighbors
+    newSegmentList = []
+    for neighborNodeID in neighborNodeIDList:
+        newSegment = [startNodeID, neighborNodeID]
+        G[startNodeID][neighborNodeID]['visited'] = True
+        currentNodeID = neighborNodeID
+        while G.degree(currentNodeID) == 2:
+            newNodeID = [nodeID for nodeID in G[currentNodeID].keys() if 'visited' not in G[currentNodeID][nodeID]][0]
+            G[currentNodeID][newNodeID]['visited'] = True
+            newSegment.append(newNodeID)
+            currentNodeID = newNodeID
+
+        newSegmentList.append(newSegment)
+        segmentList.append(newSegment)
+        segmentList = getSegmentListDetail(G, nodeInfoDict, segmentList, currentNodeID)
+
+    return segmentList
+
+
+def sublist(ls1, ls2):
+    '''
+    >>> sublist([], [1,2,3])
+    True
+    >>> sublist([1,2,3,4], [2,5,3])
+    True
+    >>> sublist([1,2,3,4], [0,3,2])
+    False
+    >>> sublist([1,2,3,4], [1,2,5,6,7,8,5,76,4,3])
+    False
+    '''
+
+    def get_all_in(one, another):
+        for element in one:
+            if element in another:
+                yield element
+
+    for x1, x2 in zip(get_all_in(ls1, ls2), get_all_in(ls2, ls1)):
+        if x1 != x2:
+            return False
+
+    return True
+
+
+def contains(lst1, lst2):
+    lst1, lst2 = (lst2, lst1) if len(lst1) > len(lst2) else (lst1, lst2)
+    if lst1[0] in lst2:
+        startLoc = lst2.index(lst1[0])
+    else:
+        return False
+
+    if lst1[-1] in lst2:
+        endLoc = lst2.index(lst1[-1])
+    else:
+        return False
+
+    if startLoc < endLoc:
+        if lst1 == lst2[startLoc:(endLoc + 1)]:
+            return True
+        else:
+            return False
+    else:
+        if lst1 == lst2[endLoc:(startLoc + 1)][::-1]:
+            return True
+        else:
+            return False
+
+
+def validateSegment(G, segment):
+    """
+    Check whether a segment is a simple branch.
+    Parameters
+    ----------
+    G : NetworkX graph
+        A graph in which each node represents a centerpoint and each edge represents a portion of a vessel branch.
+    segment : list
+        A list containing the coordinates of the centerpoints of a segment.
+
+    Returns
+    -------
+    result : bool
+        If True, the segment is a simple branch.
+    """
+    voxelDegrees = np.array([v for _, v in G.degree(segment)])
+    if voxelDegrees[0] != 2 and voxelDegrees[-1] != 2:
+        if len(voxelDegrees) == 2:
+            result = True
+        elif len(voxelDegrees) > 2:
+            if np.all(voxelDegrees[1:-1] == 2):
+                result = True
+            else:
+                result = False
+        else:
+            print('Error! Segment with length 1 found!')
+            result = False
+    else:
+        result = False
+
+    return result
+
+
+def drawSegments(segmentList, shape):
+    """
+    Plot all the segments in `segmentList`. Try to assign different colors to the segments connected to the same node.
+    Parameters
+    ----------
+    segmentList : list
+        A list containing the segment information. Each sublist represents a segment and each element in the sublist
+        represents a centerpoint coordinates.
+    shape : tuple
+        Shape of the vessel volume (used for ploting).
+    """
+    ## Import pyqtgraph ##
+    from pyqtgraph.Qt import QtCore, QtGui
+    import pyqtgraph as pg
+    import pyqtgraph.opengl as gl
+
+    ## Init ##
+    app = pg.QtGui.QApplication([])
+    w = gl.GLViewWidget()
+    w.opts['distance'] = 800
+    w.setGeometry(0, 110, 1600, 900)
+    offset = np.array(shape) / (-2.0)
+
+    colorList = [pg.glColor('r'), pg.glColor('g'), pg.glColor('b'), pg.glColor('c'), pg.glColor('m'), pg.glColor('y')]
+    colorNames = ['Red', 'Green', 'Blue', 'Cyan', 'Magneta', 'Yellow']
+    numOfColors = len(colorList)
+    nodeColorDict = {}
+    for segment in segmentList:
+        startVoxel = segment[0]
+        endVoxel = segment[-1]
+        if startVoxel in nodeColorDict and endVoxel in nodeColorDict:  # and endVoxel in [voxel for voxel, _ in nodeColorDict[startVoxel]]:
+            nodeColorDict[startVoxel].append([endVoxel, -1])
+            nodeColorDict[endVoxel].append([startVoxel, -1])
+        else:
+            if startVoxel not in nodeColorDict:
+                nodeColorDict[startVoxel] = [[endVoxel, -1]]
+            else:
+                nodeColorDict[startVoxel].append([endVoxel, -1])
+
+            if endVoxel not in nodeColorDict:
+                nodeColorDict[endVoxel] = [[startVoxel, -1]]
+            else:
+                nodeColorDict[endVoxel].append([startVoxel, -1])
+
+        existingColorsInStart = [colorCode for _, colorCode in nodeColorDict[startVoxel]]
+        existingColorsInEnd = [colorCode for _, colorCode in nodeColorDict[endVoxel]]
+        availableColors = [colorCode for colorCode in range(numOfColors) if
+                           colorCode not in existingColorsInStart and colorCode not in existingColorsInEnd]
+        # print('color in start: {} and color in end: {}'.format(existingColorsInStart, existingColorsInEnd))
+        chosenColor = availableColors[0] if len(availableColors) != 0 else 0
+        nodeColorDict[startVoxel][-1][1] = chosenColor
+        nodeColorDict[endVoxel][-1][1] = chosenColor
+
+        segmentCoords = np.array(segment)
+        aa = gl.GLLinePlotItem(pos=segmentCoords, color=colorList[chosenColor], width=3)
+        aa.translate(*offset)
+        w.addItem(aa)
+
+    w.show()
+    pg.QtGui.QApplication.exec_()
+    # sys.exit(app.exec_())
+
+
+def main():
+    start_time = timeit.default_timer()
+    baseFolder = os.path.abspath(os.path.dirname(__file__))
+
+    ## Load existing volume ##
+    vesselVolumeMaskFolderPath = baseFolder
+    vesselVolumeMaskFileName = 'vesselVolumeMask.nii.gz'
+    vesselVolumeMask, vesselVolumeMaskAffine = loadVolume(vesselVolumeMaskFolderPath, vesselVolumeMaskFileName)
+
+    ## Skeletonization ##
+    # analyze(vesselVolumeMask, baseFolder)
+
+    skeletonSegmentFolderPath = os.path.join(baseFolder, 'skeletonizationResult/segments_by_cc')
+    segmentListRough = combineSkeletonSegments(skeletonSegmentFolderPath)
+
+    shape = vesselVolumeMask.shape
+    # drawSegments(segmentListRough, shape)
+
+    G, segmentList, errorSegments = processSegments(segmentListRough, shape=shape)
+    # drawSegments(segmentList, shape)
+    G = nx.Graph()
+    segmentIndex = 0
+    for segment in segmentList:
+        G.add_path(segment, segmentIndex=segmentIndex)
+        segmentIndex += 1
+
+    ## Save graph representation ##
+    graphFileName = 'graphRepresentation.graphml'
+    graphFilePath = os.path.join(baseFolder, graphFileName)
+    nx.write_graphml(G, graphFilePath)
+    print('{} saved to {}.'.format(graphFileName, graphFilePath))
+
+    ## Save segmentList ##
+    segmentListFileName = 'segmentList.npz'
+    segmentListFilePath = os.path.join(baseFolder, segmentListFileName)
+    np.savez_compressed(segmentListFilePath, segmentList=segmentList)
+    print('{} saved to {}.'.format(segmentListFileName, segmentListFilePath))
+
+    ## Save skeleton.nii.gz ##
+    skeleton = np.zeros_like(vesselVolumeMask)
+    for segment in segmentList:
+        skeleton[tuple(np.array(segment).T)] = 1
+
+    skeletonFileName = 'skeleton.nii.gz'
+    skeletonFilePath = os.path.join(baseFolder, skeletonFileName)
+    saveVolume(skeleton, vesselVolumeMaskAffine, skeletonFilePath, astype=np.uint8)
+
+    elapsed = timeit.default_timer() - start_time
+    print('Elapsed: {} sec'.format(elapsed))
+
+
+if __name__ == "__main__":
+    main() 

Centerline/thinPlateSplines.py

@@ -0,0 +1,82 @@
+import numpy as np
+from scipy.spatial.distance import pdist, cdist, squareform
+from sklearn.metrics import pairwise_distances
+
+class ThinPlateSplines:
+    def __init__(self, ctrl_pts: np.ndarray, target_pts: np.ndarray, reg=0.0):
+        """
+
+        :param ctrl_pts: [N, d] tensor of control d-dimensional points
+        :param target_pts: [N, d] tensor of target d-dimensional points
+        :param reg: regularization coefficient
+        """
+        self.__ctrl_pts = ctrl_pts
+        self.__target_pts = target_pts
+        self.__reg = reg
+        self.__num_ctrl_pts = ctrl_pts.shape[0]
+        self.__dim = ctrl_pts.shape[1]
+
+        self.__K = None
+        self.__compute_coeffs()
+        self.__aff_params = self.__coeffs[self.__num_ctrl_pts:, ...]   # Affine  parameters of the TPS
+        self.__non_aff_paramms = self.__coeffs[:self.__num_ctrl_pts, ...]      # Non-affine parameters of  he TPS
+
+    def __compute_coeffs(self):
+        target_pts_aug = np.vstack([self.__target_pts,
+                                    np.zeros([self.__dim + 1, self.__dim])]).astype(self.__target_pts.dtype)
+
+        T_i = np.linalg.inv(self.__make_T()).astype(self.__target_pts.dtype)
+        self.__coeffs = np.matmul(T_i, target_pts_aug).astype(self.__target_pts.dtype)
+
+    def __make_T(self):
+        # cp: [K x 2] control points
+        # T: [(K+3) x (K+3)]
+        P = np.hstack([np.ones([self.__num_ctrl_pts, 1], dtype=np.float), self.__ctrl_pts])
+        zeros = np.zeros([self.__dim + 1, self.__dim + 1], dtype=np.float)
+        self.__K = self.__U_dist(self.__ctrl_pts)
+        alfa = np.mean(self.__K)
+
+        self.__K = self.__K + np.ones_like(self.__K) * np.power(alfa, 2) * self.__reg
+
+        top = np.hstack([P, self.__K])
+        bottom = np.hstack([P.transpose(), zeros])
+
+        return np.vstack([top, bottom])
+
+    def __U_dist(self, ctrl_pts, int_pts=None):
+        dist = pairwise_distances(ctrl_pts, int_pts)
+
+        if ctrl_pts.shape[-1] == 2:
+            u_dist = np.square(dist) * np.log(dist + 1e-6)
+        else:
+            u_dist = np.sqrt(dist)
+
+        return u_dist
+
+    def __lift_pts(self, int_pts: np.ndarray, num_pts):
+        # int_pts: [N x 2], input points
+        # cp: [K x 2], control points
+        # pLift: [N x (3+K)], lifted input points
+
+        int_pts_lift = np.hstack([self.__U_dist(self.__ctrl_pts, int_pts),
+                                  np.ones([num_pts, 1], dtype=np.float),
+                                  int_pts])
+        return int_pts_lift
+
+    def _get_coefficients(self):
+        return self.__coeffs
+
+    def interpolate(self, int_points):
+        """
+
+        :param int_points: [K, d] flattened d-points of a mesh
+        :return:
+        """
+        num_pts = int_points.shape[0]
+        int_points_lift = self.__lift_pts(int_points, num_pts)
+        return np.dot(int_points_lift, self.__coeffs)
+
+    @property
+    def bending_energy(self):
+        aux = tf.matmul(self.__non_aff_paramms, self.__K, transpose_a=True)
+        return tf.matmul(aux, self.__non_aff_paramms)

Centerline/visualization_utils.py

@@ -0,0 +1,161 @@
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.lines import Line2D
+import numpy as np
+from DeepDeformationMapRegistration.utils.visualization import add_axes_arrows_3d, remove_tick_labels, set_axes_size
+import os
+
+
+def _plot_graph(graph, ax, nodes_colour='C3', edges_colour='C1', plot_nodes=True, plot_edges=True, add_axes=True):
+    if plot_edges:
+        for (start_node, end_node) in graph.edges():
+            edge_pts = graph[start_node][end_node]['pts']
+            edge_pts = np.vstack([graph.nodes[start_node]['o'], edge_pts])
+            edge_pts = np.vstack([edge_pts, graph.nodes[end_node]['o']])
+            ax.plot(edge_pts[:, 0], edge_pts[:, 1], edge_pts[:, 2], edges_colour)
+
+    if plot_nodes:
+        nodes = graph.nodes()
+        ps = np.array([nodes[i]['o'] for i in nodes])
+        if len(ps.shape) > 1:
+            ax.scatter(ps[:, 0], ps[:, 1], ps[:, 2], nodes_colour)
+        else:
+            ax.scatter(ps[0], ps[1], ps[2], nodes_colour)
+    ax.set_xlim(0, 63)
+    ax.set_ylim(0, 63)
+    ax.set_zlim(0, 63)
+    remove_tick_labels(ax, True)
+    if add_axes:
+        add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b')
+    ax.view_init(None, 45)
+
+    return ax
+
+
+def plot_skeleton(img, skeleton, graph, filename='skeleton', extension=['.png']):
+    if not isinstance(extension, list):
+        extension = [extension]
+    # Skeleton
+    f = plt.figure(figsize=(5, 5))
+    ax = f.add_subplot(111, projection='3d')
+
+    coords = np.argwhere(skeleton)
+    i = coords[:, 0]
+    j = coords[:, 1]
+    k = coords[:, 2]
+
+    seg = ax.voxels(img, facecolors=(0., 0., 1., 0.3), label='image')
+    ske = ax.scatter(i, j, k, color='C1', label='skeleton', s=1)
+    ax.set_xlim(0, 63)
+    ax.set_ylim(0, 63)
+    ax.set_zlim(0, 63)
+    remove_tick_labels(ax, True)
+    add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b')
+    ax.view_init(None, 45)
+    for ex in extension:
+        f.savefig(filename + '_segmentation_skeleton' + ex)
+
+    # Combined
+    ax = _plot_graph(graph, ax, 'r', 'r')
+
+    for ex in extension:
+        f.savefig(filename + '_combined' + ex)
+    plt.close()
+
+    # Graph
+    f = plt.figure(figsize=(5, 5))
+    ax = f.add_subplot(111, projection='3d')
+
+    ax = _plot_graph(graph, ax)
+
+    for ex in extension:
+        f.savefig(filename + '_graph' + ex)
+    plt.close()
+
+
+
+
+def compare_graphs(graph_0, graph_1, graph_names=None, filename='compare_graphs'):
+    f = plt.figure(figsize=(12, 5))
+    if graph_names is None:
+        graph_names =['graph_0', 'graph_1']
+    else:
+        assert len(graph_names) == 2, 'A different name is expected for each graph'
+    ax = f.add_subplot(131, projection='3d')
+    ax = _plot_graph(graph_0, ax)
+    ax.set_title(graph_names[0], y=-0.01)
+
+    ax = f.add_subplot(132, projection='3d')
+    ax = _plot_graph(graph_1, ax)
+    ax.set_title(graph_names[1])
+
+    ax = f.add_subplot(133, projection='3d')
+    ax = _plot_graph(graph_0, ax, 'C2', 'C2', plot_nodes=False)
+    ax = _plot_graph(graph_1, ax, 'C4', 'C4', plot_nodes=False)
+    legend_elements = [Line2D([0], [0], color='C2', lw=2, label=graph_names[0]),
+                       Line2D([0], [0], color='C4', lw=2, label=graph_names[1])]
+    ax.legend(handles=legend_elements)
+
+    f.savefig(filename + '_compare_graphs.png')
+    plt.close()
+
+
+def plot_cpd_registration_step(iteration, error, X, Y, out_folder, add_axes=True, pdf=True):
+    fig = plt.figure(figsize=(8, 8))
+    ax = fig.add_axes([0, 0, .9, .9], projection='3d')
+    ax.scatter(X[:, 0],  X[:, 1], X[:, 2], color='C1', label='Fixed')
+    ax.scatter(Y[:, 0],  Y[:, 1], Y[:, 2], color='C2', label='Moving')
+
+    ax.text2D(0.95, 0.98, 'Iteration: {:d}'.format(
+        iteration), horizontalalignment='right', verticalalignment='center', transform=ax.transAxes, fontsize='x-large')
+    #ax.text2D(0.95, 0.90, 'Error: {:10.4f}'.format(
+    #    error), horizontalalignment='right', verticalalignment='center', transform=ax.transAxes, fontsize='x-large')
+    ax.legend(loc='upper left', fontsize='x-large')
+
+    if add_axes:
+        x_range = [np.min(np.hstack([X[:, 0], Y[:, 0]])), np.max(np.hstack([X[:, 0], Y[:, 0]]))]
+        y_range = [np.min(np.hstack([X[:, 1], Y[:, 1]])), np.max(np.hstack([X[:, 1], Y[:, 1]]))]
+        z_range = [np.min(np.hstack([X[:, 2], Y[:, 2]])), np.max(np.hstack([X[:, 2], Y[:, 2]]))]
+        ax.set_xlim(x_range[0], x_range[1])
+        ax.set_ylim(y_range[0], y_range[1])
+        ax.set_zlim(z_range[0], z_range[1])
+
+        remove_tick_labels(ax, True)
+        add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b', arrow_length=25, dist_arrow_text=3)
+    ax.view_init(None, 45)
+
+    os.makedirs(out_folder, exist_ok=True)
+    fig.savefig(os.path.join(out_folder, '{:04d}.png'.format(iteration)))
+    if pdf:
+        fig.savefig(os.path.join(out_folder, '{:04d}.pdf'.format(iteration)))
+    plt.close()
+
+
+def plot_cpd(fix_pts, mov_pts, fix_centroid, mov_centroid, file_name):
+    fig = plt.figure(figsize=(8, 8))
+    ax = fig.add_axes([0, 0, .9, .9], projection='3d')
+    ax.scatter(fix_pts[:, 0],  fix_pts[:, 1], fix_pts[:, 2], color='C1', label='Fixed')
+    ax.scatter(mov_pts[:, 0],  mov_pts[:, 1], mov_pts[:, 2], color='C2', label='Moving')
+    ax.scatter(fix_centroid[0],  fix_centroid[1], fix_centroid[2], color='none', s=100, edgecolor='b', label='Centroid')
+    ax.scatter(mov_centroid[0],  mov_centroid[1], mov_centroid[2], color='none', s=100, edgecolor='b')
+    ax.scatter(fix_centroid[0],  fix_centroid[1], fix_centroid[2], color='C1')
+    ax.scatter(mov_centroid[0],  mov_centroid[1], mov_centroid[2], color='C2')
+
+    x_range = [np.min(np.hstack([fix_pts[:, 0], mov_pts[:, 0], fix_centroid[0], mov_centroid[0]])),
+               np.max(np.hstack([fix_pts[:, 0], mov_pts[:, 0], fix_centroid[0], mov_centroid[0]]))]
+    y_range = [np.min(np.hstack([fix_pts[:, 1], mov_pts[:, 1], fix_centroid[1], mov_centroid[1]])),
+               np.max(np.hstack([fix_pts[:, 1], mov_pts[:, 1], fix_centroid[1], mov_centroid[1]]))]
+    z_range = [np.min(np.hstack([fix_pts[:, 2], mov_pts[:, 2], fix_centroid[2], mov_centroid[2]])),
+               np.max(np.hstack([fix_pts[:, 2], mov_pts[:, 2], fix_centroid[2], mov_centroid[2]]))]
+    ax.set_xlim(x_range[0], x_range[1])
+    ax.set_ylim(y_range[0], y_range[1])
+    ax.set_zlim(z_range[0], z_range[1])
+
+    remove_tick_labels(ax, True)
+    add_axes_arrows_3d(ax, x_color='r', y_color='g', z_color='b', arrow_length=25, dist_arrow_text=3)
+    ax.view_init(None, 45)
+    ax.legend(fontsize='xx-large')
+    fig.savefig(file_name + '.png')
+    fig.savefig(file_name + '.pdf')
+    plt.close()
+