Add test for affine_registration and add interfaces to use for emc

Author: dPys

dmriprep/interfaces/registration.py

@@ -0,0 +1,147 @@
+"""Register tools interfaces."""
+import numpy as np
+import nibabel as nb
+import dmriprep
+from nipype import logging
+from pathlib import Path
+from nipype.utils.filemanip import fname_presuffix
+from nipype.interfaces.base import (
+    traits,
+    TraitedSpec,
+    BaseInterfaceInputSpec,
+    InputMultiObject,
+    SimpleInterface,
+    File,
+)
+
+
+LOGGER = logging.getLogger("nipype.interface")
+
+
+class _ApplyAffineInputSpec(BaseInterfaceInputSpec):
+    moving_image = File(
+        exists=True, mandatory=True, desc="image to apply transformation from"
+    )
+    fixed_image = File(
+        exists=True, mandatory=True, desc="image to apply transformation to"
+    )
+    transform_affine = InputMultiObject(
+        File(exists=True), mandatory=True, desc="transformation affine"
+    )
+    invert_transform = traits.Bool(False, usedefault=True)
+
+
+class _ApplyAffineOutputSpec(TraitedSpec):
+    warped_image = File(exists=True, desc="Outputs warped image")
+
+
+class ApplyAffine(SimpleInterface):
+    """
+    Interface to apply an affine transformation to an image.
+    """
+
+    input_spec = _ApplyAffineInputSpec
+    output_spec = _ApplyAffineOutputSpec
+
+    def _run_interface(self, runtime):
+        from dmriprep.utils.registration import apply_affine
+
+        warped_image_nifti = apply_affine(
+            nb.load(self.inputs.moving_image),
+            nb.load(self.inputs.fixed_image),
+            np.load(self.inputs.transform_affine[0]),
+            self.inputs.invert_transform,
+        )
+        cwd = Path(runtime.cwd).absolute()
+        warped_file = fname_presuffix(
+            self.inputs.moving_image,
+            use_ext=False,
+            suffix="_warped.nii.gz",
+            newpath=str(cwd),
+        )
+
+        warped_image_nifti.to_filename(warped_file)
+
+        self._results["warped_image"] = warped_file
+        return runtime
+
+
+class _RegisterInputSpec(BaseInterfaceInputSpec):
+    moving_image = File(
+        exists=True, mandatory=True, desc="image to apply transformation from"
+    )
+    fixed_image = File(
+        exists=True, mandatory=True, desc="image to apply transformation to"
+    )
+    nbins = traits.Int(default_value=32, usedefault=True)
+    sampling_prop = traits.Float(defualt_value=1, usedefault=True)
+    metric = traits.Str(default_value="MI", usedefault=True)
+    level_iters = traits.List(
+        trait=traits.Any(), value=[10000, 1000, 100], usedefault=True
+    )
+    sigmas = traits.List(trait=traits.Any(), value=[5.0, 2.5, 0.0], usedefault=True)
+    factors = traits.List(trait=traits.Any(), value=[4, 2, 1], usedefault=True)
+    params0 = traits.ArrayOrNone(value=None, usedefault=True)
+    pipeline = traits.List(
+        trait=traits.Any(),
+        value=["c_of_mass", "translation", "rigid", "affine"],
+        usedefault=True,
+    )
+
+
+class _RegisterOutputSpec(TraitedSpec):
+    forward_transforms = traits.List(
+        File(exists=True), desc="List of output transforms for forward registration"
+    )
+    warped_image = File(exists=True, desc="Outputs warped image")
+
+
+class Register(SimpleInterface):
+    """
+    Interface to perform affine registration.
+    """
+
+    input_spec = _RegisterInputSpec
+    output_spec = _RegisterOutputSpec
+
+    def _run_interface(self, runtime):
+        from dmriprep.utils.registration import affine_registration
+
+        reg_types = ["c_of_mass", "translation", "rigid", "affine"]
+        pipeline = [
+            getattr(dmriprep.utils.register, i)
+            for i in self.inputs.pipeline
+            if i in reg_types
+        ]
+
+        warped_image_nifti, forward_transform_mat = affine_registration(
+            nb.load(self.inputs.moving_image),
+            nb.load(self.inputs.fixed_image),
+            self.inputs.nbins,
+            self.inputs.sampling_prop,
+            self.inputs.metric,
+            pipeline,
+            self.inputs.level_iters,
+            self.inputs.sigmas,
+            self.inputs.factors,
+            self.inputs.params0,
+        )
+        cwd = Path(runtime.cwd).absolute()
+        warped_file = fname_presuffix(
+            self.inputs.moving_image,
+            use_ext=False,
+            suffix="_warped.nii.gz",
+            newpath=str(cwd),
+        )
+        forward_transform_file = fname_presuffix(
+            self.inputs.moving_image,
+            use_ext=False,
+            suffix="_forward_transform.npy",
+            newpath=str(cwd),
+        )
+        warped_image_nifti.to_filename(warped_file)
+
+        np.save(forward_transform_file, forward_transform_mat)
+        self._results["warped_image"] = warped_file
+        self._results["forward_transforms"] = [forward_transform_file]
+        return runtime

dmriprep/utils/register.py renamed to dmriprep/utils/registration.py

@@ -0,0 +1,56 @@
+import pytest
+import numpy as np
+import numpy.testing as npt
+import nibabel as nb
+import dipy.data as dpd
+from dmriprep.utils.registration import affine_registration
+
+
+def setup_module():
+    global subset_b0, subset_dwi_data, subset_t2, subset_b0_img, \
+           subset_t2_img, gtab, hardi_affine, MNI_T2_affine
+    MNI_T2 = dpd.read_mni_template()
+    hardi_img, gtab = dpd.read_stanford_hardi()
+    MNI_T2_data = MNI_T2.get_fdata()
+    MNI_T2_affine = MNI_T2.affine
+    hardi_data = hardi_img.get_fdata()
+    hardi_affine = hardi_img.affine
+    b0 = hardi_data[..., gtab.b0s_mask]
+    mean_b0 = np.mean(b0, -1)
+
+    # Select some arbitrary chunks of data so this goes quicker
+    subset_b0 = mean_b0[40:50, 40:50, 40:50]
+    subset_dwi_data = nb.Nifti1Image(hardi_data[40:50, 40:50, 40:50],
+                                     hardi_affine)
+    subset_t2 = MNI_T2_data[40:60, 40:60, 40:60]
+    subset_b0_img = nb.Nifti1Image(subset_b0, hardi_affine)
+    subset_t2_img = nb.Nifti1Image(subset_t2, MNI_T2_affine)
+
+
+def test_affine_registration():
+    moving = subset_b0
+    static = subset_b0
+    moving_affine = static_affine = np.eye(4)
+    xformed, affine = affine_registration(moving, static)
+    # We don't ask for much:
+    npt.assert_almost_equal(affine[:3, :3], np.eye(3), decimal=1)
+
+    with pytest.raises(ValueError):
+        # For array input, must provide affines:
+        xformed, affine = affine_registration(moving, static)
+
+    # If providing nifti image objects, don't need to provide affines:
+    moving_img = nb.Nifti1Image(moving, moving_affine)
+    static_img = nb.Nifti1Image(static, static_affine)
+    xformed, affine = affine_registration(moving_img, static_img)
+    npt.assert_almost_equal(affine[:3, :3], np.eye(3), decimal=1)
+
+    # Using strings with full paths as inputs also works:
+    t1_name, b0_name = dpd.get_fnames('syn_data')
+    moving = b0_name
+    static = t1_name
+    xformed, affine = affine_registration(moving, static,
+                                          level_iters=[5, 5],
+                                          sigmas=[3, 1],
+                                          factors=[4, 2])
+    npt.assert_almost_equal(affine[:3, :3], np.eye(3), decimal=1)