Repository: uncbiag/ICON (language: python)

Issue:
Add multichanel support
Support multichannel image as input to the network and similarity measure.

It is useful when there are multiple modalities such as MRI-T1, MRI-T2 and MRI-TI enhanced images. 

This pull request contains the three following major changes:
1. Expose an argument called input_channels for tallUNet2 function. This argument indicates the number of channels of one input image tensor.
2. Refactor all the similarity measures to inherit from SimilarityBase class. And SimilarityBase class contains a variable member indicating whether this Similarity Measure requires inbounds_tag for evaluation.
3. Add itk interface for multi-modality registration.


Discussion / hints:
I really like the getattr trick to maintain backwards compatibility. 
It looks good to me!


Fix patch:
diff --git a/training_scripts/cvpr_clean/HCP_ants_eval.py b/training_scripts/cvpr_clean/HCP_ants_eval.py
index 55f61a5b..68d74cfb 100644
--- a/training_scripts/cvpr_clean/HCP_ants_eval.py
+++ b/training_scripts/cvpr_clean/HCP_ants_eval.py
@@ -14,23 +14,19 @@ def preprocess(image):
     return image
 
 dices = []
-
-import glob
-paths = glob.glob("/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/*WarpedLabels*")
-atlas_registered = [p.split("/malf3")[-1].split("_")[0] for p in paths]
-
-def get_sub_seg(n):
-    path = f"/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/{n}_label.nii.gz"
-    return itk.imread(path)
+flips = []
+pair_list = []
 
 def mean_dice_f(sA, sB):
     sA, sB = [itk.image_from_array(s.numpy()) for s in (sA, sB)]
     return utils.itk_mean_dice(sA, sB)
 
+from synthmorph.HCP_segs import (atlas_registered, get_sub_seg, get_brain_image)
+
 random.seed(1)
-for _ in range(20):
+for _ in range(100):
     n_A, n_B = (random.choice(atlas_registered) for _ in range(2))
-    image_A, image_B = (preprocess(itk.imread(f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz")) for n in (n_A, n_B))
+    image_A, image_B = (preprocess(get_brain_image(n)) for n in (n_A, n_B))
 
     segmentation_A, segmentation_B = (get_sub_seg(n) for n in (n_A, n_B))
     segmentation_A, segmentation_B = [
@@ -51,11 +47,19 @@ def mean_dice_f(sA, sB):
     warped_segmentation_B = ants.apply_transforms(fixed=segmentation_A, moving=segmentation_B, transformlist=reg_res['fwdtransforms'], interpolator="nearestNeighbor")
 
 
-    mean_dice = mean_dice_f(segmentation_A,warped_segmentation_B)
-
-    utils.log(mean_dice)
+    mean_dice = mean_dice_f(segmentation_A, warped_segmentation_B)
 
+    jac_np = np.array(itk.displacement_field_jacobian_determinant_filter(itk.imread(reg_res['fwdtransforms'][0])))
+    flip = np.mean(jac_np<0) * 100.0
+    
+    flips.append(flip)
     dices.append(mean_dice)
+    utils.log(f"{_}/100 mean DICE {n_A} to {n_B}: {mean_dice} | running DICE: {np.mean(dices)} | Percentage: {flip} | running Per: {np.mean(flips)} ")
+    pair_list.append([n_A, n_B])
 
 utils.log("Mean DICE")
-utils.log(np.mean(dices))
+utils.log(f"Final DICE: {np.mean(dices)} | final percentage of negative jacobian: {np.mean(flips)}")
+
+with open(f'{footsteps.output_dir}/pair_list.txt', 'w') as f:
+    for p in pair_list:
+        f.write(",".join(p)+'\n')
diff --git a/training_scripts/cvpr_clean/HCP_eval.py b/training_scripts/cvpr_clean/HCP_eval.py
index 9c93e5a9..c97a82e9 100644
--- a/training_scripts/cvpr_clean/HCP_eval.py
+++ b/training_scripts/cvpr_clean/HCP_eval.py
@@ -21,34 +21,52 @@ def preprocess(image):
     #image = itk.CastImageFilter[itk.Image[itk.SS, 3], itk.Image[itk.F, 3]].New()(image)
     max_ = np.max(np.array(image))
     image = itk.shift_scale_image_filter(image, shift=0., scale = .9 / max_)
+    
     #image = itk.clamp_image_filter(image, bounds=(0, 1))
     return image
 
 
 input_shape = [1, 1, 130, 155, 130]
-net = cvpr_network.make_network(input_shape, include_last_step=True)
+net = cvpr_network.make_network(input_shape, include_last_step=True)#, framework="svf")
 
 
 utils.log(net.regis_net.load_state_dict(torch.load(weights_path), strict=False))
 net.eval()
 
 dices = []
+flips = []
+shape_prod = np.prod(input_shape)
 
-import glob
-paths = glob.glob("/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/*WarpedLabels*")
-atlas_registered = [p.split("/malf3")[-1].split("_")[0] for p in paths]
-
-def get_sub_seg(n):
-    path = f"/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/{n}_label.nii.gz"
-    return itk.imread(path)
+from synthmorph.HCP_segs import (atlas_registered, get_sub_seg, get_brain_image)
 
+pair_list = []
 random.seed(1)
-for _ in range(20):
+for _ in range(100):
     n_A, n_B = (random.choice(atlas_registered) for _ in range(2))
-    image_A, image_B = (preprocess(itk.imread(f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz")) for n in (n_A, n_B))
+    image_A, image_B = (preprocess(get_brain_image(n)) for n in (n_A, n_B))
 
     #import pdb; pdb.set_trace()
-    phi_AB, phi_BA = itk_wrapper.register_pair(net, image_A, image_B, finetune_steps=None)
+    phi_AB, phi_BA, loss= itk_wrapper.register_pair(net, image_A, image_B, finetune_steps=None,
+        return_artifacts=True,
+    )
+
+    # print(net.warped_image_A.shape)
+    
+    # import shutil
+    # shutil.copyfile(f"/playpen-raid2/Data/HCP/HCP_1200/{n_A}/T1w/T1w_acpc_dc_restore_brain.nii.gz", f"{footsteps.output_dir}/{n_A}_T1w_acpc_dc_restore_brain.nii.gz")
+    # shutil.copyfile(f"/playpen-raid2/Data/HCP/HCP_1200/{n_B}/T1w/T1w_acpc_dc_restore_brain.nii.gz", f"{footsteps.output_dir}/{n_B}_T1w_acpc_dc_restore_brain.nii.gz")
+
+    # interpolator = itk.LinearInterpolateImageFunction.New(image_A)
+    # warped_image_A = itk.resample_image_filter(
+    #         image_A,
+    #         transform=phi_AB,
+    #         interpolator=interpolator,
+    #         size=itk.size(image_B),
+    #         output_spacing=itk.spacing(image_B),
+    #         output_direction=image_B.GetDirection(),
+    #         output_origin=image_B.GetOrigin(),
+    #     )
+    # np.save(f"{footsteps.output_dir}/{n_A}_and_{n_B}_warped.np", np.array(warped_image_A))
 
     segmentation_A, segmentation_B = (get_sub_seg(n) for n in (n_A, n_B))
 
@@ -62,10 +80,16 @@ def get_sub_seg(n):
             reference_image=segmentation_B
             )
     mean_dice = utils.itk_mean_dice(segmentation_B, warped_segmentation_A)
-
-    utils.log(mean_dice)
-
+    
+    flip = loss.flips / shape_prod * 100.
+    flips.append(flip)
     dices.append(mean_dice)
+    utils.log(f"{_}/100 mean DICE {n_A} to {n_B}: {mean_dice} | running DICE: {np.mean(dices)} | Percentage: {flip} | running Per: {np.mean(flips)} ")
+    pair_list.append([n_A, n_B])
 
 utils.log("Mean DICE")
-utils.log(np.mean(dices))
+utils.log(f"Final DICE: {np.mean(dices)} | final percentage of negative jacobian: {np.mean(flips)}")
+
+with open(f'{footsteps.output_dir}/pair_list.txt', 'w') as f:
+    for p in pair_list:
+        f.write(",".join(p)+'\n')
diff --git a/training_scripts/cvpr_clean/HCP_synthmorph_eval.py b/training_scripts/cvpr_clean/HCP_synthmorph_eval.py
deleted file mode 100644
index 17c1f98d..00000000
--- a/training_scripts/cvpr_clean/HCP_synthmorph_eval.py
+++ /dev/null
@@ -1,177 +0,0 @@
-import footsteps
-import subprocess
-import glob
-import sys
-import random
-import itk
-
-# footsteps.initialize(output_root="evaluation_results/", run_name="asdf")
-import numpy as np
-import utils
-
-
-def itk_rotate_scale_image(img, label=True):
-    if not label:
-        a = np.array(img)
-
-        b = a[a != 0]
-
-        print(np.min(a), np.percentile(b, 5) , np.max(a))
-
-        max_ = np.max(np.array(img))
-        img = itk.shift_scale_image_filter(img, shift=0.0, scale=1.0 / max_)
-    scale = [0.618, 0.618, 0.618]
-    input_size = itk.size(img)
-    input_spacing = itk.spacing(img)
-    input_origin = itk.origin(img)
-    dimension = img.GetImageDimension()
-
-    output_size = [int(input_size[d] * scale[d]) for d in range(dimension)]
-    output_spacing = [input_spacing[d] / scale[d] for d in range(dimension)]
-    output_origin = [
-        input_origin[d] + 0.5 * (output_spacing[d] - input_spacing[d])
-        for d in range(dimension)
-    ]
-
-    if label:
-        interpolator = itk.NearestNeighborInterpolateImageFunction.New(img)
-    else:
-        interpolator = itk.LinearInterpolateImageFunction.New(img)
-
-    transform = itk.CenteredEuler3DTransform[itk.D].New()
-
-    params = transform.GetParameters()
-
-    params[0] = 0.5
-
-    transform.SetParameters(params)
-
-    transform.SetCenter([0.0, 65.0, -10.0])
-
-    resampled = itk.resample_image_filter(
-        img,
-        transform=transform,
-        interpolator=interpolator,
-        size=output_size,
-        output_spacing=output_spacing,
-        output_origin=output_origin,
-        output_direction=img.GetDirection(),
-    )
-    #if not label:
-
-    #    max_ = np.max(np.array(resampled))
-    #    resampled = itk.shift_scale_image_filter(resampled, shift=0.0, scale=1.0 / max_)
-
-    print(resampled)
-    resamplednt = np.array(resampled)
-
-
-    iref = itk.imread("ref.nii.gz")
-    iref = itk.CastImageFilter[itk.Image[itk.UC, 3], itk.Image[itk.D, 3]].New()(iref)
-    iref = itk.shift_scale_image_filter(iref, shift=0.0, scale = 1.0 / 255)
-
-    resamplednt = np.transpose(resamplednt, (1, 0, 2))
-    resamplednt = np.flip(resamplednt, axis=1)
-
-
-    npresampleditk = itk.image_from_array(resampled)#.astype(np.float64)
-
-    transform2 = itk.MatrixTransform[itk.D, 3]
-
-    
-
-    sys.exit()
-
-    #resampled.SetSpacing(iref.GetSpacing())
-    #resampled.SetDirection(iref.GetDirection())
-    #resampled.SetOrigin(iref.GetOrigin())
-    #resampled = itk.checker_board_image_filter(
-    #    resampled, iref
-
-    #)
-    # print(img)
-    # print(resampled)
-    # exit()
-
-    return resampled
-
-
-def preprocess(image):
-    # image = itk.CastImageFilter[itk.Image[itk.SS, 3], itk.Image[itk.F, 3]].New()(image)
-    # image = itk.clamp_image_filter(image, bounds=(0, 1))
-    return image
-
-
-dices = []
-
-
-paths = glob.glob("/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/*WarpedLabels*")
-atlas_registered = [p.split("/malf3")[-1].split("_")[0] for p in paths]
-
-
-def get_sub_seg(n):
-    path = f"/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/{n}_label.nii.gz"
-    return itk.imread(path)
-
-
-def mean_dice_f(sA, sB):
-    sA, sB = [itk.image_from_array(s.numpy()) for s in (sA, sB)]
-    return utils.itk_mean_dice(sA, sB)
-
-
-random.seed(1)
-for _ in range(20):
-    n_A, n_B = (random.choice(atlas_registered) for _ in range(2))
-    image_A, image_B = (
-        preprocess(
-            itk.imread(
-                f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz"
-            )
-        )
-        for n in (n_A, n_B)
-    )
-
-    segmentation_A, segmentation_B = (get_sub_seg(n) for n in (n_A, n_B))
-
-    segmentation_A, segmentation_B = [
-        itk.CastImageFilter[itk.Image[itk.SS, 3], itk.Image[itk.UC, 3]].New()(image)
-        for image in (segmentation_A, segmentation_B)
-    ]
-
-    image_A = itk_rotate_scale_image(image_A, label=False)
-    image_B = itk_rotate_scale_image(image_B, label=False)
-    segmentation_A = itk_rotate_scale_image(segmentation_A, label=True)
-    segmentation_B = itk_rotate_scale_image(segmentation_B, label=True)
-
-    itk.imwrite(image_A, "A.nii.gz")
-    itk.imwrite(image_B, "B.nii.gz")
-    itk.imwrite(segmentation_A, "segA.nii.gz")
-    itk.imwrite(segmentation_B, "segB.nii.gz")
-
-    #subprocess.run("vshow B.nii.gz -z", shell=True)
-    #sys.exit()
-
-    cmd = """python /playpen-raid1/tgreer/voxelmorph/voxelmorph/scripts/tf/register.py --fixed A.nii.gz --moving B.nii.gz --moved out.nii.gz --model /playpen-raid1/tgreer/voxelmorph/brains-dice-vel-0.5-res-16-256f.h5 --warp warp.nii.gz"""
-    # cmd = """python /playpen-raid1/tgreer/voxelmorph/voxelmorph/scripts/tf/register.py --fixed A.nii.gz --moving B.nii.gz --moved out.nii.gz --model shapes-dice-vel-3-res-8-16-32-256f.h5 --warp warp.nii.gz"""
-    subprocess.run(cmd, shell=True)
-
-    import voxelmorph
-    import voxelmorph as vxm
-
-    vsegfix = voxelmorph.py.utils.load_labels("segA.nii.gz")[1][0][None, :, :, :, None]
-    vsegmov = voxelmorph.py.utils.load_labels("segB.nii.gz")[1][0][None, :, :, :, None]
-    warp = voxelmorph.py.utils.load_volfile("warp.nii.gz", add_batch_axis=True)
-    warped_seg = vxm.networks.Transform(
-        vsegfix.shape[1:-1], interp_method="nearest"
-    ).predict([vsegmov, warp])
-    overlap = vxm.py.utils.dice(vsegfix, warped_seg, labels=list(range(1, 29)))
-
-    mean_dice = np.mean(overlap)
-    dices.append(np.mean(overlap))
-
-    utils.log(mean_dice)
-
-    dices.append(mean_dice)
-
-utils.log("Mean DICE")
-utils.log(np.mean(dices))
diff --git a/training_scripts/cvpr_clean/cvpr_network.py b/training_scripts/cvpr_clean/cvpr_network.py
index 76f607b0..85a5ff2c 100644
--- a/training_scripts/cvpr_clean/cvpr_network.py
+++ b/training_scripts/cvpr_clean/cvpr_network.py
@@ -39,22 +39,28 @@ def forward(self, image_A, image_B):
         # can use information about whether a sample is interpolated
         # or extrapolated
 
-        inbounds_tag = torch.zeros(tuple(image_A.shape), device=image_A.device)
-        if len(self.input_shape) - 2 == 3:
-            inbounds_tag[:, :, 1:-1, 1:-1, 1:-1] = 1.0
-        elif len(self.input_shape) - 2 == 2:
-            inbounds_tag[:, :, 1:-1, 1:-1] = 1.0
+        if getattr(self.similarity, "isInterpolated", False):
+            # tag images during warping so that the similarity measure
+            # can use information about whether a sample is interpolated
+            # or extrapolated
+            inbounds_tag = torch.zeros([image_A.shape[0]] + [1] + list(image_A.shape[2:]), device=image_A.device)
+            if len(self.input_shape) - 2 == 3:
+                inbounds_tag[:, :, 1:-1, 1:-1, 1:-1] = 1.0
+            elif len(self.input_shape) - 2 == 2:
+                inbounds_tag[:, :, 1:-1, 1:-1] = 1.0
+            else:
+                inbounds_tag[:, :, 1:-1] = 1.0
         else:
-            inbounds_tag[:, :, 1:-1] = 1.0
+            inbounds_tag = None
 
         self.warped_image_A = compute_warped_image_multiNC(
-            torch.cat([image_A, inbounds_tag], axis=1),
+            torch.cat([image_A, inbounds_tag], axis=1) if inbounds_tag is not None else image_A,
             self.phi_AB_vectorfield,
             self.spacing,
             1,
         )
         self.warped_image_B = compute_warped_image_multiNC(
-            torch.cat([image_B, inbounds_tag], axis=1),
+            torch.cat([image_B, inbounds_tag], axis=1) if inbounds_tag is not None else image_B,
             self.phi_BA_vectorfield,
             self.spacing,
             1,
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_freesurfer_affine_eval.py b/training_scripts/cvpr_clean/synthmorph/HCP_freesurfer_affine_eval.py
new file mode 100644
index 00000000..45767f0c
--- /dev/null
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_freesurfer_affine_eval.py
@@ -0,0 +1,141 @@
+import os
+import random
+
+import footsteps
+footsteps.initialize(output_root="evaluation_results/")
+import nibabel as nib
+import numpy as np
+import tensorflow as tf
+import voxelmorph as vxm
+
+import utils
+from HCP_synthmorph_data_prepare import process_image
+from HCP_segs import atlas_registered, get_brain_image_path, get_sub_seg_path
+from HCP_synthmorph_helper import net_to_vox, transform, read_affine
+
+# This script borrows some functions from https://github.com/freesurfer/freesurfer/blob/a810044ae08a24402436c1d43472b3b3df06592a/mri_synthmorph/mri_synthmorph
+
+if __name__ == "__main__":
+    
+    prealign_folder = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/evaluation_results/synthmorph_preprocessed"
+    ref_path = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/ref.nii.gz"
+    ref = vxm.py.utils.load_volfile(ref_path, add_batch_axis = True, add_feat_axis = True)
+
+    in_shape = ref.shape[1:-1]
+    nb_feats = ref.shape[-1]
+
+    dices = []
+    flips = []
+    pair_list = []
+
+    random.seed(1)
+    for _ in range(100):
+        n_A, n_B = (random.choice(atlas_registered) for _ in range(2))
+        if prealign_folder == "" or not os.path.exists(prealign_folder):
+            prealign_folder = f"{footsteps.output_dir}/synthmorph_preprocessed"
+            if not os.path.exists(prealign_folder):
+                os.mkdir(prealign_folder)
+            
+            process_image(n_A, prealign_folder)
+            process_image(n_B, prealign_folder)
+        
+        # Input data.
+        mov = nib.load(f"{prealign_folder}/{n_A}_affine.nii.gz")
+        fix = nib.load(f"{prealign_folder}/{n_B}_affine.nii.gz")
+
+        # ################################# Prepare transformation from image to network ###########################################
+        # Coordinate transforms. We will need these to take the images from their
+        # native voxel spaces to network space. Voxel and network spaces are different
+        # for each image. Network space is an isotropic 1-mm space centered on the
+        # original image. Its axes are aligned with the original voxel data but flipped
+        # and swapped to gross LIA orientation, which the network will expect.
+        net_to_mov = net_to_vox(mov, out_shape=in_shape)
+        net_to_fix = net_to_vox(fix, out_shape=in_shape)
+        mov_to_net = np.linalg.inv(net_to_mov)
+        fix_to_net = np.linalg.inv(net_to_fix)
+
+        # Transforms from and to world space (RAS). There is only one world.
+        mov_to_ras = mov.affine
+        fix_to_ras = fix.affine
+        ras_to_mov = np.linalg.inv(mov_to_ras)
+        ras_to_fix = np.linalg.inv(fix_to_ras)
+
+        # Transforms between zero-centered and zero-based voxel spaces.
+        ind_to_cen = np.eye(4)
+        ind_to_cen[:-1, -1] = -0.5 * (np.asarray(in_shape) - 1)
+        cen_to_ind = np.eye(4)
+        cen_to_ind[:-1, -1] = +0.5 * (np.asarray(in_shape) - 1)
+
+        # # Incorporate an initial linear transform operating in RAS. It goes from fixed
+        # # to moving coordinates, so we start with fixed network space on the right.
+        # if arg.init:
+        #     aff = np.loadtxt(arg.init)
+        #     net_to_mov = ras_to_mov @ aff @ fix_to_ras @ net_to_fix
+
+        ################################# Run networks on the transformed image ###########################################
+        fix_affine_ras = read_affine(f"{prealign_folder}/{n_B}_Affine.lta")
+        mov_affine_ras = read_affine(f"{prealign_folder}/{n_A}_Affine.lta")
+        pre_fix = nib.load(get_brain_image_path(n_B))
+        pre_mov = nib.load(get_brain_image_path(n_A))
+
+        ################################# Compose all the transformations  ###########################################
+        # Construct grid of zero-based index coordinates and shape (3, N) in native
+        # fixed voxel space, where N is the number of voxels.
+        x_fix = (tf.range(x, dtype=tf.float32) for x in pre_fix.shape)
+        x_fix = tf.meshgrid(*x_fix, indexing='ij')
+        x_fix = tf.stack(x_fix)
+        x_fix = tf.reshape(x_fix, shape=(3, -1))
+
+        # Transform x_fix from previous fix to after affine fix
+        pre_fix_to_fix = np.linalg.inv(fix.affine)@fix_affine_ras@pre_fix.affine
+        x_out = pre_fix_to_fix[:-1, -1:] + (pre_fix_to_fix[:-1, :-1] @ x_fix)
+
+        # Transform fixed voxel coordinates to the fixed network space.
+        x_out = fix_to_net[:-1, -1:] + (fix_to_net[:-1, :-1] @ x_out)
+        # x_out = tf.transpose(x_out)
+
+        # # Add predicted warp to coordinates to go to the moving network space.
+        # trans = tf.squeeze(trans)
+        # x_out += ne.utils.interpn(trans, x_out, fill_value=0)
+        # x_out = tf.transpose(x_out)
+
+        # Transform coordinates to the native moving voxel space. Subtract fixed
+        # coordinates to obtain displacement from fixed to moving voxel space.
+        x_out = net_to_mov[:-1, -1:] + (net_to_mov[:-1, :-1] @ x_out)
+
+        # Transform to before affine move
+        mov_to_pre_mov = np.linalg.inv(pre_mov.affine)@np.linalg.inv(mov_affine_ras)@mov.affine
+        x_out = mov_to_pre_mov[:-1, -1:] + (mov_to_pre_mov[:-1, :-1] @ x_out)
+
+        trans_vox = tf.transpose(x_out - x_fix)
+        trans_vox = tf.reshape(trans_vox, shape=(*pre_fix.shape, -1))
+
+        # Displacement from fixed to moving RAS coordinates.
+        x_ras = fix_to_ras[:-1, -1:] + (fix_to_ras[:-1, :-1] @ x_fix)
+        x_out = mov_to_ras[:-1, -1:] + (mov_to_ras[:-1, :-1] @ x_out)
+        trans_ras = tf.transpose(x_out - x_ras)
+        trans_ras = tf.reshape(trans_ras, shape=(*pre_fix.shape, -1))
+
+        ################################# Evaluate the segmentations  ###########################################
+        mov = nib.load(get_sub_seg_path(n_A))
+        fixed = nib.load(get_sub_seg_path(n_B))
+        fixed_np = fixed.get_fdata(dtype=np.float32)
+
+        warped_seg = transform(mov, trans=trans_vox, shape=fixed_np.shape, interp_method='nearest').numpy()[0,:,:,:,0]
+        
+        dice = utils.mean_dice(fixed_np, warped_seg)
+
+        jd = vxm.py.utils.jacobian_determinant(trans_vox)
+        flip = np.mean(jd<0) * 100.0
+
+        flips.append(flip)
+        dices.append(dice)
+        utils.log(f"{_}/100 mean DICE {n_A} to {n_B}: {dice} | running DICE: {np.mean(dices)} | Percentage: {flip} | running Per: {np.mean(flips)} ")
+        pair_list.append([n_A, n_B])
+
+    utils.log("Mean DICE")
+    utils.log(f"Final DICE: {np.mean(dices)} | final percentage of negative jacobian: {np.mean(flips)}")
+
+    with open(f'{footsteps.output_dir}/pair_list.txt', 'w') as f:
+        for p in pair_list:
+            f.write(",".join(p)+'\n')
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_segs.py b/training_scripts/cvpr_clean/synthmorph/HCP_segs.py
index 154fdb51..1d7812d7 100644
--- a/training_scripts/cvpr_clean/synthmorph/HCP_segs.py
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_segs.py
@@ -2,14 +2,18 @@
 import itk
 
 import glob
-paths = glob.glob("/playpen-raid2/Data/HCP/manual_subcortical_segmentations_BWH/*/*_*_novdc-label.nii.gz")
-atlas_registered = [p.split("/")[-1].split("_")[0] for p in paths]
+# paths = glob.glob("/playpen-raid2/Data/HCP/manual_subcortical_segmentations_BWH/*/*_*_novdc-label.nii.gz")
+# atlas_registered = [p.split("/")[-1].split("_")[0] for p in paths]
+
+paths = glob.glob("/playpen-raid2/Data/HCP/manual_subcortical_segmentations_BWH/*/*_*_*_restore.*")
+atlas_registered = [p.split("/")[-2] for p in paths]
+
+
 
 #paths = glob.glob("/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/*WarpedLabels*")
 #atlas_registered = [p.split("/malf3")[-1].split("_")[0] for p in paths]
 
-
-def get_sub_seg(n):
+def get_sub_seg_path(n):
     path = f"/playpen-raid2/Data/HCP/manual_subcortical_segmentations_BWH/{n}/{n}_*_novdc-label.nii.gz"
     path = glob.glob(path)
     print(path)
@@ -17,8 +21,15 @@ def get_sub_seg(n):
         path = path[0]
     else:
         path = f"/playpen-raid1/tgreer/Subcortical_Atlas_Fusion2/{n}_label.nii.gz"
+    return path
+
+def get_brain_image_path(n):
+    return f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz"
+
+def get_sub_seg(n):
+    path = get_sub_seg_path(n)
     return itk.imread(path)
 
 def get_brain_image(n):
-    return itk.imread(f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz")
+    return itk.imread(get_brain_image_path(n))
 
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_data_prepare.py b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_data_prepare.py
new file mode 100644
index 00000000..ab16b417
--- /dev/null
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_data_prepare.py
@@ -0,0 +1,32 @@
+import subprocess
+from HCP_segs import atlas_registered, get_brain_image_path, get_sub_seg_path
+import os
+
+'''
+To run mri_robust_register, you need first source /playpen-raid1/tgreer/freesurfer/activate.sh
+'''
+
+def process_image(n, output_folder):
+    ref_path = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/ref.nii.gz"
+
+    image_path = get_brain_image_path(n)
+    label_path = get_sub_seg_path(n)
+
+    subprocess.run(f"mri_robust_register --mov {image_path} --dst {ref_path} -lta {output_folder}/{n}_Affine.lta --satit --iscale --verbose 0", shell=True)
+    subprocess.run(f"mri_robust_register --mov {image_path} --dst {ref_path} -lta {output_folder}/{n}_Affine.lta --satit --iscale --ixform {output_folder}/{n}_Affine.lta --affine --verbose 0", shell=True)
+    subprocess.run(f"mri_vol2vol --mov {image_path} --o {output_folder}/{n}_affine.nii.gz --lta {output_folder}/{n}_Affine.lta --targ {ref_path}", shell=True)
+    subprocess.run(f"mri_vol2vol --mov {label_path} --o {output_folder}/{n}_label_affine.nii.gz --lta {output_folder}/{n}_Affine.lta  --targ {ref_path} --nearest --keep-precision", shell=True)
+
+if __name__ == "__main__":
+    prealign_folder = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/evaluation_results/synthmorph_preprocessed"
+    override = False
+    
+    if not os.path.exists(prealign_folder):
+        os.makedirs(prealign_folder)
+    
+    for n in atlas_registered:
+        if not override and os.path.exists(f"{prealign_folder}/{n}_affine.nii.gz"):
+            continue
+        process_image(n, prealign_folder)
+            
+
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval.py b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval.py
index c73a059e..2dcb2339 100644
--- a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval.py
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval.py
@@ -1,190 +1,27 @@
-import argparse
-import glob
 import os
 import random
-import subprocess
-import sys
-import time
-
-import footsteps
-import itk
 import neurite as ne
 import nibabel as nib
-# footsteps.initialize(output_root="evaluation_results/", run_name="asdf")
 import numpy as np
 import tensorflow as tf
 import voxelmorph as vxm
 
+import footsteps
+footsteps.initialize(output_root="evaluation_results/")
+
 import utils
-import voxelmorph as vxm
+from HCP_synthmorph_data_prepare import process_image
+from HCP_synthmorph_helper import net_to_vox, transform, read_affine
+from HCP_segs import atlas_registered, get_brain_image_path, get_sub_seg_path
 
 # This script borrows some functions from https://github.com/freesurfer/freesurfer/blob/a810044ae08a24402436c1d43472b3b3df06592a/mri_synthmorph/mri_synthmorph
 
 
-def save(path, dat, affine, dtype=None):
-    # Usi NiBabel's caching functionality to avoid re-reading from disk.
-    if isinstance(dat, nib.filebasedimages.FileBasedImage):
-        if dtype is None:
-            dtype = dat.dataobj.dtype
-        dat = dat.get_fdata(dtype=np.float32)
-
-    dat = np.squeeze(dat)
-    dat = np.asarray(dat, dtype)
-
-    # Avoid warning about missing units when reading with FS.
-    out = nib.Nifti1Image(dat, affine)
-    out.header.set_xyzt_units(xyz='mm', t='sec')
-    nib.save(out, filename=path)
-
-
-def ori_to_ori(old, new='LIA', old_shape=None, zero_center=False):
-    '''Construct matrix transforming coordinates from a voxel space with a new
-    predominant anatomical axis orientation to an old orientation, by swapping
-    and flipping axes. Operates in zero-based index space unless the space is
-    to be zero-centered. The old shape must be specified if the old image is
-    not a NiBabel object.'''
-    def extract_ori(x):
-        if isinstance(x, nib.filebasedimages.FileBasedImage):
-            x = x.affine
-        if isinstance(x, np.ndarray):
-            return nib.orientations.io_orientation(x)
-        if isinstance(x, str):
-            return nib.orientations.axcodes2ornt(x)
-
-    # Old shape.
-    if zero_center:
-        old_shape = (1, 1, 1)
-    if old_shape is None:
-        old_shape = old.shape
-
-    # Transform from new to old index coordinates.
-    old = extract_ori(old)
-    new = extract_ori(new)
-    new_to_old = nib.orientations.ornt_transform(old, new)
-    return nib.orientations.inv_ornt_aff(new_to_old, old_shape)
-
-
-def net_to_vox(im, out_shape):
-    '''Construct coordinate transform from isotropic 1-mm voxel space with
-    gross LIA orentiation centered on the FOV - to the original image index
-    space. The target space is a scaled and shifted voxel space, not world
-    space.'''
-    if isinstance(im, str):
-        im = nib.load(im)
-
-    # Gross LIA to predominant anatomical orientation of input image.
-    assert isinstance(im, nib.filebasedimages.FileBasedImage) 
-    lia_to_ori = ori_to_ori(im, new='LIA', old_shape=out_shape)
-
-    # Scaling from millimeter to input voxels.
-    vox_size = np.sqrt(np.sum(im.affine[:-1, :-1] ** 2, axis=0))
-    scale = np.diag((*1 / vox_size, 1))
-
-    # Shift from cen
-    shift = np.eye(4)
-    shift[:-1, -1] = 0.5 * (im.shape - out_shape / vox_size)
-
-    # Total transform.
-    return shift @ scale @ lia_to_ori
-
-
-def transform(im, trans, shape, normalize=False, interp_method='linear'):
-    '''Apply transformation matrix or field operating in zero-based index space
-    to an image.'''
-    if isinstance(im, nib.filebasedimages.FileBasedImage):
-        im = im.get_fdata(dtype=np.float32)
-
-    # Add singleton feature dimension if needed.
-    if tf.rank(im) == 3:
-        im = im[..., tf.newaxis]
-
-    # Remove last row of matrix transforms.
-    if tf.rank(trans) == 2 and trans.shape[0] == trans.shape[1]:
-        trans = trans[:-1, :]
-
-    out = vxm.utils.transform(
-        im, trans, interp_method=interp_method, fill_value=0, shift_center=False, shape=shape,
-    )
-
-    if normalize:
-        out -= tf.reduce_min(out)
-        out /= tf.reduce_max(out)
-    return out[tf.newaxis, ...]
-
-
-
-def vm_dense(
-    in_shape=None,
-    input_model=None,
-    enc_nf=[256] * 4,
-    dec_nf=[256] * 4,
-    add_nf=[256] * 4,
-    int_steps=5,
-    upsample=True,
-    half_res=True,
-):
-    '''Deformable registration network.'''
-    if input_model is None:
-        source = tf.keras.Input(shape=(*in_shape, 1))
-        target = tf.keras.Input(shape=(*in_shape, 1))
-        input_model = tf.keras.Model(*[(source, target)] * 2)
-    source, target = input_model.outputs[:2]
-
-    in_shape = np.asarray(source.shape[1:-1])
-    num_dim = len(in_shape)
-    assert num_dim in (2, 3), 'only 2D and 3D supported'
-
-    down = getattr(tf.keras.layers, f'MaxPool{num_dim}D')()
-    up = getattr(tf.keras.layers, f'UpSampling{num_dim}D')()
-    act = tf.keras.layers.LeakyReLU(0.2)
-    conv = getattr(tf.keras.layers, f'Conv{num_dim}D')
-    prop = dict(kernel_size=3, padding='same')
-
-    # Encoder.
-    x = tf.keras.layers.concatenate((source, target))
-    if half_res:
-        x = down(x)
-    enc = [x]
-    for n in enc_nf:
-        x = conv(n, **prop)(x)
-        x = act(x)
-        enc.append(x)
-        x = down(x)
-
-    # Decoder.
-    for n in dec_nf:
-        x = conv(n, **prop)(x)
-        x = act(x)
-        x = tf.keras.layers.concatenate([up(x), enc.pop()])
-
-    # Additional convolutions.
-    for n in add_nf:
-        x = conv(n, **prop)(x)
-        x = act(x)
-
-    # Transform.
-    x = conv(num_dim, **prop)(x)
-    if int_steps > 0:
-        x = vxm.layers.VecInt(method='ss', int_steps=int_steps)(x)
-
-    # Rescaling.
-    zoom = source.shape[1] // x.shape[1]
-    if upsample and zoom > 1:
-        x = vxm.layers.RescaleTransform(zoom)(x)
-
-    return tf.keras.Model(input_model.inputs, outputs=x)
-
-def read_affine(lta_path=""):
-    with open(lta_path, 'rb') as f:
-        lines = f.readlines()
-        affine = lines[8:12]
-        affine = [str(i).split("'")[1].split(' ')[:-1] for i in affine]
-        
-    return np.array(affine, np.float)
-
 if __name__ == "__main__":
-    prealign_folder = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/results/debug/eval_HCP-43"
-    ref = vxm.py.utils.load_volfile("ref.nii.gz", add_batch_axis = True, add_feat_axis = True)
+     
+    prealign_folder = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/evaluation_results/synthmorph_preprocessed"
+    ref_path = "/playpen-raid2/lin.tian/projects/icon_lung/ICON/training_scripts/cvpr_clean/ref.nii.gz"
+    ref = vxm.py.utils.load_volfile(ref_path, add_batch_axis = True, add_feat_axis = True)
 
     in_shape = ref.shape[1:-1]
     nb_feats = ref.shape[-1]
@@ -199,61 +36,22 @@ def read_affine(lta_path=""):
 
         dices = []
         flips = []
-
-
-        from HCP_segs import atlas_registered, get_brain_image, get_sub_seg
-
-        def mean_dice_f(sA, sB):
-            sA, sB = [itk.image_from_array(s.numpy()) for s in (sA, sB)]
-            return utils.itk_mean_dice(sA, sB)
-
-        
+        pair_list = []
 
         random.seed(1)
         for _ in range(100):
             n_A, n_B = (random.choice(atlas_registered) for _ in range(2))
-            if prealign_folder != "" and os.path.exists(prealign_folder):
-                pair_dir = f"{prealign_folder}/{n_A}_{n_B}"
-            else:
-                image_A, image_B = (
-                        itk.imread(
-                            f"/playpen-raid2/Data/HCP/HCP_1200/{n}/T1w/T1w_acpc_dc_restore_brain.nii.gz"
-                        )
-                    for n in (n_A, n_B)
-                )
-
-                segmentation_A, segmentation_B = (get_sub_seg(n) for n in (n_A, n_B))
-
-                pair_dir = f"{footsteps.output_dir}/{n_A}_{n_B}"
-                if not os.path.exists(pair_dir):
-                    os.mkdir(pair_dir)
-                itk.imwrite(segmentation_A, f"{pair_dir}/segA_orig.nii.gz")
-                itk.imwrite(image_A, f"{pair_dir}/imageA_orig.nii.gz")
-                itk.imwrite(segmentation_B, f"{pair_dir}/segB_orig.nii.gz")
-                itk.imwrite(image_B, f"{pair_dir}/imageB_orig.nii.gz")
-
-                # Affine pre-aling to reference image
-                subprocess.run(f"mri_robust_register --mov {pair_dir}/imageA_orig.nii.gz --dst ref.nii.gz -lta {pair_dir}/A_Affine.lta --satit --iscale --verbose 0", shell=True)
-                subprocess.run(f"mri_robust_register --mov {pair_dir}/imageA_orig.nii.gz --dst ref.nii.gz -lta {pair_dir}/A_Affine.lta --satit --iscale --ixform {pair_dir}/A_Affine.lta --affine --verbose 0", shell=True)
-                subprocess.run(f"mri_vol2vol --mov {pair_dir}/imageA_orig.nii.gz --o {pair_dir}/A_affine.nii.gz --lta {pair_dir}/A_Affine.lta --targ ref.nii.gz", shell=True)
-                subprocess.run(f"mri_vol2vol --mov {pair_dir}/segA_orig.nii.gz --o {pair_dir}/Aseg_affine.nii.gz --lta {pair_dir}/A_Affine.lta --targ ref.nii.gz --nearest --keep-precision", shell=True)
-
-                subprocess.run(f"mri_robust_register --mov {pair_dir}/imageB_orig.nii.gz --dst ref.nii.gz -lta {pair_dir}/B_Affine.lta --satit --iscale --verbose 0", shell=True)
-                subprocess.run(f"mri_robust_register --mov {pair_dir}/imageB_orig.nii.gz --dst ref.nii.gz -lta {pair_dir}/B_Affine.lta --satit --iscale --ixform {pair_dir}/B_Affine.lta --affine --verbose 0", shell=True)
-                subprocess.run(f"mri_vol2vol --mov {pair_dir}/imageB_orig.nii.gz --o {pair_dir}/B_affine.nii.gz --lta {pair_dir}/B_Affine.lta --targ ref.nii.gz", shell=True)
-                subprocess.run(f"mri_vol2vol --mov {pair_dir}/segB_orig.nii.gz --o {pair_dir}/Bseg_affine.nii.gz --lta {pair_dir}/B_Affine.lta --targ ref.nii.gz --nearest --keep-precision", shell=True)
-
-            #cmd = """python /playpen-raid1/tgreer/voxelmorph/voxelmorph/scripts/tf/register.py --fixed A.nii.gz --moving B.nii.gz --moved out.nii.gz --model /playpen-raid1/tgreer/voxelmorph/brains-dice-vel-0.5-res-16-256f.h5 --warp warp.nii.gz"""
-            #cmd = """python /playpen-raid1/tgreer/voxelmorph/voxelmorph/scripts/tf/register.py --fixed A_affine.nii.gz --moving B_affine.nii.gz --moved out.nii.gz --model shapes-dice-vel-3-res-8-16-32-256f.h5 --warp warp.nii.gz"""
-            #subprocess.run(cmd, shell=True)
-            
-            output_dir = f"{footsteps.output_dir}/{n_A}_{n_B}"
-            if not os.path.exists(output_dir):
-                os.mkdir(output_dir)
+            if prealign_folder == "" or not os.path.exists(prealign_folder):
+                prealign_folder = f"{footsteps.output_dir}/synthmorph_preprocessed"
+                if not os.path.exists(prealign_folder):
+                    os.mkdir(prealign_folder)
+                
+                process_image(n_A, prealign_folder)
+                process_image(n_B, prealign_folder)
             
             # Input data.
-            mov = nib.load(f"{pair_dir}/A_affine.nii.gz")
-            fix = nib.load(f"{pair_dir}/B_affine.nii.gz")
+            mov = nib.load(f"{prealign_folder}/{n_A}_affine.nii.gz")
+            fix = nib.load(f"{prealign_folder}/{n_B}_affine.nii.gz")
 
             # ################################# Prepare transformation from image to network ###########################################
             # Coordinate transforms. We will need these to take the images from their
@@ -292,10 +90,10 @@ def mean_dice_f(sA, sB):
 
             trans = regis_net.register(*inputs) 
 
-            fix_affine_ras = read_affine(f"{pair_dir}/B_Affine.lta")
-            mov_affine_ras = read_affine(f"{pair_dir}/A_Affine.lta")
-            pre_fix = nib.load(f"{pair_dir}/imageB_orig.nii.gz")
-            pre_mov = nib.load(f"{pair_dir}/imageA_orig.nii.gz")
+            fix_affine_ras = read_affine(f"{prealign_folder}/{n_B}_Affine.lta")
+            mov_affine_ras = read_affine(f"{prealign_folder}/{n_A}_Affine.lta")
+            pre_fix = nib.load(get_brain_image_path(n_B))
+            pre_mov = nib.load(get_brain_image_path(n_A))
 
             ################################# Compose all the transformations  ###########################################
             # Construct grid of zero-based index coordinates and shape (3, N) in native
@@ -336,22 +134,25 @@ def mean_dice_f(sA, sB):
             trans_ras = tf.reshape(trans_ras, shape=(*pre_fix.shape, -1))
 
             ################################# Evaluate the segmentations  ###########################################
-            mov = nib.load(f"{pair_dir}/segA_orig.nii.gz")
-            fixed = nib.load(f"{pair_dir}/segB_orig.nii.gz")
-            fixed_np = fixed.get_data().squeeze()
+            mov = nib.load(get_sub_seg_path(n_A))
+            fixed = nib.load(get_sub_seg_path(n_B))
+            fixed_np = fixed.get_fdata(dtype=np.float32)
 
             warped_seg = transform(mov, trans=trans_vox, shape=fixed_np.shape, interp_method='nearest').numpy()[0,:,:,:,0] 
-            save(f"{output_dir}/warped_segA_orig.nii.gz", warped_seg, fixed.affine, int)
-            
-            overlap = vxm.py.utils.dice(fixed_np, warped_seg, labels=list(range(1, 29)))
-            mean_dice = np.mean(overlap)
+
+            dice = utils.mean_dice(fixed_np, warped_seg)
 
             jd = vxm.py.utils.jacobian_determinant(trans_vox)
             flip = np.mean(jd<0) * 100.0
 
             flips.append(flip)
-            dices.append(mean_dice)
-            utils.log(f"{_}/100 mean DICE {n_A} to {n_B}: {mean_dice} | running DICE: {np.mean(dices)} | Percentage: {flip} | running Per: {np.mean(flips)} ")
+            dices.append(dice)
+            utils.log(f"{_}/100 mean DICE {n_A} to {n_B}: {dice} | running DICE: {np.mean(dices)} | Percentage: {flip} | running Per: {np.mean(flips)} ")
+            pair_list.append([n_A, n_B])
 
         utils.log("Mean DICE")
         utils.log(f"Final DICE: {np.mean(dices)} | final percentage of negative jacobian: {np.mean(flips)}")
+
+        with open(f'{footsteps.output_dir}/pair_list.txt', 'w') as f:
+            for p in pair_list:
+                f.write(",".join(p)+'\n')
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_helper.py b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_helper.py
new file mode 100644
index 00000000..a9f34b23
--- /dev/null
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_helper.py
@@ -0,0 +1,167 @@
+import nibabel as nib
+import numpy as np
+import tensorflow as tf
+import voxelmorph as vxm
+
+# This script borrows some functions from https://github.com/freesurfer/freesurfer/blob/a810044ae08a24402436c1d43472b3b3df06592a/mri_synthmorph/mri_synthmorph
+
+def save(path, dat, affine, dtype=None):
+    # Usi NiBabel's caching functionality to avoid re-reading from disk.
+    if isinstance(dat, nib.filebasedimages.FileBasedImage):
+        if dtype is None:
+            dtype = dat.dataobj.dtype
+        dat = dat.get_fdata(dtype=np.float32)
+
+    dat = np.squeeze(dat)
+    dat = np.asarray(dat, dtype)
+
+    # Avoid warning about missing units when reading with FS.
+    out = nib.Nifti1Image(dat, affine)
+    out.header.set_xyzt_units(xyz='mm', t='sec')
+    nib.save(out, filename=path)
+
+
+def ori_to_ori(old, new='LIA', old_shape=None, zero_center=False):
+    '''Construct matrix transforming coordinates from a voxel space with a new
+    predominant anatomical axis orientation to an old orientation, by swapping
+    and flipping axes. Operates in zero-based index space unless the space is
+    to be zero-centered. The old shape must be specified if the old image is
+    not a NiBabel object.'''
+    def extract_ori(x):
+        if isinstance(x, nib.filebasedimages.FileBasedImage):
+            x = x.affine
+        if isinstance(x, np.ndarray):
+            return nib.orientations.io_orientation(x)
+        if isinstance(x, str):
+            return nib.orientations.axcodes2ornt(x)
+
+    # Old shape.
+    if zero_center:
+        old_shape = (1, 1, 1)
+    if old_shape is None:
+        old_shape = old.shape
+
+    # Transform from new to old index coordinates.
+    old = extract_ori(old)
+    new = extract_ori(new)
+    new_to_old = nib.orientations.ornt_transform(old, new)
+    return nib.orientations.inv_ornt_aff(new_to_old, old_shape)
+
+
+def net_to_vox(im, out_shape):
+    '''Construct coordinate transform from isotropic 1-mm voxel space with
+    gross LIA orentiation centered on the FOV - to the original image index
+    space. The target space is a scaled and shifted voxel space, not world
+    space.'''
+    if isinstance(im, str):
+        im = nib.load(im)
+
+    # Gross LIA to predominant anatomical orientation of input image.
+    assert isinstance(im, nib.filebasedimages.FileBasedImage) 
+    lia_to_ori = ori_to_ori(im, new='LIA', old_shape=out_shape)
+
+    # Scaling from millimeter to input voxels.
+    vox_size = np.sqrt(np.sum(im.affine[:-1, :-1] ** 2, axis=0))
+    scale = np.diag((*1 / vox_size, 1))
+
+    # Shift from cen
+    shift = np.eye(4)
+    shift[:-1, -1] = 0.5 * (im.shape - out_shape / vox_size)
+
+    # Total transform.
+    return shift @ scale @ lia_to_ori
+
+
+def transform(im, trans, shape, normalize=False, interp_method='linear'):
+    '''Apply transformation matrix or field operating in zero-based index space
+    to an image.'''
+    if isinstance(im, nib.filebasedimages.FileBasedImage):
+        im = im.get_fdata(dtype=np.float32)
+
+    # Add singleton feature dimension if needed.
+    if tf.rank(im) == 3:
+        im = im[..., tf.newaxis]
+
+    # Remove last row of matrix transforms.
+    if tf.rank(trans) == 2 and trans.shape[0] == trans.shape[1]:
+        trans = trans[:-1, :]
+
+    out = vxm.utils.transform(
+        im, trans, interp_method=interp_method, fill_value=0, shift_center=False, shape=shape,
+    )
+
+    if normalize:
+        out -= tf.reduce_min(out)
+        out /= tf.reduce_max(out)
+    return out[tf.newaxis, ...]
+
+
+
+def vm_dense(
+    in_shape=None,
+    input_model=None,
+    enc_nf=[256] * 4,
+    dec_nf=[256] * 4,
+    add_nf=[256] * 4,
+    int_steps=5,
+    upsample=True,
+    half_res=True,
+):
+    '''Deformable registration network.'''
+    if input_model is None:
+        source = tf.keras.Input(shape=(*in_shape, 1))
+        target = tf.keras.Input(shape=(*in_shape, 1))
+        input_model = tf.keras.Model(*[(source, target)] * 2)
+    source, target = input_model.outputs[:2]
+
+    in_shape = np.asarray(source.shape[1:-1])
+    num_dim = len(in_shape)
+    assert num_dim in (2, 3), 'only 2D and 3D supported'
+
+    down = getattr(tf.keras.layers, f'MaxPool{num_dim}D')()
+    up = getattr(tf.keras.layers, f'UpSampling{num_dim}D')()
+    act = tf.keras.layers.LeakyReLU(0.2)
+    conv = getattr(tf.keras.layers, f'Conv{num_dim}D')
+    prop = dict(kernel_size=3, padding='same')
+
+    # Encoder.
+    x = tf.keras.layers.concatenate((source, target))
+    if half_res:
+        x = down(x)
+    enc = [x]
+    for n in enc_nf:
+        x = conv(n, **prop)(x)
+        x = act(x)
+        enc.append(x)
+        x = down(x)
+
+    # Decoder.
+    for n in dec_nf:
+        x = conv(n, **prop)(x)
+        x = act(x)
+        x = tf.keras.layers.concatenate([up(x), enc.pop()])
+
+    # Additional convolutions.
+    for n in add_nf:
+        x = conv(n, **prop)(x)
+        x = act(x)
+
+    # Transform.
+    x = conv(num_dim, **prop)(x)
+    if int_steps > 0:
+        x = vxm.layers.VecInt(method='ss', int_steps=int_steps)(x)
+
+    # Rescaling.
+    zoom = source.shape[1] // x.shape[1]
+    if upsample and zoom > 1:
+        x = vxm.layers.RescaleTransform(zoom)(x)
+
+    return tf.keras.Model(input_model.inputs, outputs=x)
+
+def read_affine(lta_path=""):
+    with open(lta_path, 'rb') as f:
+        lines = f.readlines()
+        affine = lines[8:12]
+        affine = [str(i).split("'")[1].split(' ')[:-1] for i in affine]
+        
+    return np.array(affine, np.float)
\ No newline at end of file
diff --git a/training_scripts/cvpr_clean/synthmorph/utils.py b/training_scripts/cvpr_clean/synthmorph/utils.py
index cb91f76c..1c3ed387 100644
--- a/training_scripts/cvpr_clean/synthmorph/utils.py
+++ b/training_scripts/cvpr_clean/synthmorph/utils.py
@@ -7,11 +7,11 @@ def log(*val):
     print(*val)
     print(*val, file=logfile)
 
+# Keep this interface the same for all the experiments.
 def itk_mean_dice(im1, im2):
     array1 = itk.array_from_image(im1)
     array2 = itk.array_from_image(im2)
     dices = []
-    print(np.max(array1), np.max(array2))
     for index in range(1, max(np.max(array1), np.max(array2)) + 1):
         m1 = array1 == index
         m2 = array2 == index
@@ -21,3 +21,15 @@ def itk_mean_dice(im1, im2):
         d = 2 * np.sum(intersection) / (np.sum(m1) + np.sum(m2))
         dices.append(d)
     return np.mean(dices)
+
+def mean_dice(array1, array2):
+    dices = []
+    for index in range(1, max(int(np.max(array1)), int(np.max(array2))) + 1):
+        m1 = array1 == index
+        m2 = array2 == index
+        
+        intersection = np.logical_and(m1, m2)
+        
+        d = 2 * np.sum(intersection) / (np.sum(m1) + np.sum(m2))
+        dices.append(d)
+    return np.mean(dices)


Test patch:
diff --git a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval_test.py b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval_test.py
index 5d9b8fef..24d6f7aa 100644
--- a/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval_test.py
+++ b/training_scripts/cvpr_clean/synthmorph/HCP_synthmorph_eval_test.py
@@ -6,7 +6,7 @@
 import tensorflow as tf
 import voxelmorph as vxm
 
-from HCP_synthmorph_eval import save, transform, read_affine
+from HCP_synthmorph_helper import save, transform, read_affine
 
 # This script borrows some functions from https://github.com/freesurfer/freesurfer/blob/a810044ae08a24402436c1d43472b3b3df06592a/mri_synthmorph/mri_synthmorph