diff --git a/dmriprep/utils/register.py b/dmriprep/utils/register.py new file mode 100644 index 00000000..7e6ecd3b --- /dev/null +++ b/dmriprep/utils/register.py @@ -0,0 +1,189 @@ +""" +Linear affine registration tools for motion correction. +""" +import attr + +import numpy as np +import nibabel as nb +from dipy.align.metrics import CCMetric, EMMetric, SSDMetric + +from dipy.align.imaffine import ( + transform_centers_of_mass, + AffineMap, + MutualInformationMetric, + AffineRegistration, +) + +from dipy.align.transforms import ( + TranslationTransform3D, + RigidTransform3D, + AffineTransform3D, +) +from nipype.utils.filemanip import fname_presuffix + +syn_metric_dict = {"CC": CCMetric, "EM": EMMetric, "SSD": SSDMetric} + + +__all__ = [ + "c_of_mass", + "translation", + "rigid", + "affine", + "AffineRegistration", +] + + +def apply_affine(moving, static, transform_affine, invert=False): + """Apply an affine to transform an image from one space to another. + + Parameters + ---------- + moving : array + The image to be resampled + + static : array + + Returns + ------- + warped_img : the moving array warped into the static array's space. + + """ + affine_map = AffineMap( + transform_affine, static.shape, static.affine, moving.shape, moving.affine + ) + if invert is True: + warped_arr = affine_map.transform_inverse(np.asarray(moving.dataobj)) + else: + warped_arr = affine_map.transform(np.asarray(moving.dataobj)) + + return nb.Nifti1Image(warped_arr, static.affine) + + +def average_affines(transforms): + affine_list = [np.load(aff) for aff in transforms] + average_affine_file = fname_presuffix( + transforms[0], use_ext=False, suffix="_average.npy" + ) + np.save(average_affine_file, np.mean(affine_list, axis=0)) + return average_affine_file + + +# Affine registration pipeline: +affine_metric_dict = {"MI": MutualInformationMetric, "CC": CCMetric} + + +def c_of_mass( + moving, static, static_affine, moving_affine, reg, starting_affine, params0=None +): + transform = transform_centers_of_mass(static, static_affine, moving, moving_affine) + transformed = transform.transform(moving) + return transform + + +def translation( + moving, static, static_affine, moving_affine, reg, starting_affine, params0=None +): + transform = TranslationTransform3D() + translation = reg.optimize( + static, + moving, + transform, + params0, + static_affine, + moving_affine, + starting_affine=starting_affine, + ) + + return translation + + +def rigid( + moving, static, static_affine, moving_affine, reg, starting_affine, params0=None +): + transform = RigidTransform3D() + rigid = reg.optimize( + static, + moving, + transform, + params0, + static_affine, + moving_affine, + starting_affine=starting_affine, + ) + return rigid + + +def affine(moving, static, static_affine, moving_affine, reg, starting_affine, + params0=None): + """ + """ + transform = AffineTransform3D() + affine = reg.optimize( + static, + moving, + transform, + params0, + static_affine, + moving_affine, + starting_affine=starting_affine, + ) + + return affine + + +@attr.s(slots=True, frozen=True) +class AffineRegistration(): + def __init__(self): + nbins = attr.ib(default=32) + sampling_prop = attr.ib(default=1.0) + metric = attr.ib(default="MI") + level_iters = attr.ib(default=[10000, 1000, 100]) + sigmas = attr.ib(defaults=[3, 1, 0.0]) + factors = attr.ib(defaults=[4, 2, 1]) + pipeline = attr.ib(defaults=[c_of_mass, translation, rigid, affine]) + + def fit(self, static, moving, params0=None): + """ + static, moving : nib.Nifti1Image class images + """ + if params0 is None: + starting_affine = np.eye(4) + else: + starting_affine = params0 + + use_metric = affine_metric_dict[self.metric](self.nbins, + self.sampling_prop) + affreg = AffineRegistration( + metric=use_metric, + level_iters=self.level_iters, + sigmas=self.sigmas, + factors=self.factors) + + # Go through the selected transformation: + for func in self.pipeline: + transform = func( + np.asarray(moving.dataobj), + np.asarray(static.dataobj), + static.affine, + moving.affine, + affreg, + starting_affine, + params0, + ) + starting_affine = transform.affine + + self.static_affine_ = static.affine + self.moving_affine_ = moving.affine + self.affine_ = starting_affine + self.reg_ = AffineMap(starting_affine, + static.shape, static.affine, + moving.shape, moving.affine) + + def apply(self, moving): + """ + + """ + data = moving.get_fdata() + assert np.all(moving.affine, self.moving_affine_) + return nb.Nifti1Image(np.array(self.reg_.transform(data)), + self.static_affine_) diff --git a/setup.cfg b/setup.cfg index 61b04423..a9be9d6f 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,6 +26,7 @@ install_requires = nibabel ~= 3.0 nipype ~= 1.4 niworkflows @ git+https://github.com/nipreps/niworkflows.git@master + nitransforms numpy pybids >=0.10.2 pyyaml