From 8b2075814cf94a63a8faeb2f38de4ecf4e87053c Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 29 Aug 2024 16:25:50 +0200 Subject: [PATCH] rfc: move image registration into its own submodule --- src/eddymotion/estimator.py | 166 +--------------- src/eddymotion/registration/__init__.py | 0 src/eddymotion/registration/ants.py | 184 ++++++++++++++++++ .../config/dwi-to-b0_level0.json | 0 .../config/dwi-to-dwi_level0.json | 0 .../config/dwi-to-dwi_level1.json | 0 test/test_integration.py | 8 +- ...test_estimator.py => test_registration.py} | 5 +- 8 files changed, 202 insertions(+), 161 deletions(-) create mode 100644 src/eddymotion/registration/__init__.py create mode 100644 src/eddymotion/registration/ants.py rename src/eddymotion/{ => registration}/config/dwi-to-b0_level0.json (100%) rename src/eddymotion/{ => registration}/config/dwi-to-dwi_level0.json (100%) rename src/eddymotion/{ => registration}/config/dwi-to-dwi_level1.json (100%) rename test/{test_estimator.py => test_registration.py} (96%) diff --git a/src/eddymotion/estimator.py b/src/eddymotion/estimator.py index d274cfd8..0322fc05 100644 --- a/src/eddymotion/estimator.py +++ b/src/eddymotion/estimator.py @@ -22,21 +22,16 @@ # """A model-based algorithm for the realignment of dMRI data.""" -from collections import namedtuple from pathlib import Path from tempfile import TemporaryDirectory, mkstemp import nibabel as nb -import nitransforms as nt -import numpy as np -from nipype.interfaces.ants.registration import Registration -from nitransforms.linear import Affine -from pkg_resources import resource_filename as pkg_fn from tqdm import tqdm from eddymotion import utils as eutils from eddymotion.data.splitting import lovo_split from eddymotion.model import ModelFactory +from eddymotion.registration.ants import _prepare_registration_data, _run_registration class EddyMotionEstimator: @@ -104,11 +99,13 @@ def estimate( align_kwargs["num_threads"] = omp_nthreads n_iter = len(models) - for i_iter, model in enumerate(models): - reg_target_type = ( - "dwi" if model.lower() not in ("b0", "s0", "avg", "average", "mean") else "b0" - ) + reg_target_type = ( + align_kwargs.pop("fixed_modality", None), + align_kwargs.pop("moving_modality", None), + ) + + for i_iter, model in enumerate(models): # When downsampling these need to be set per-level bmask_img = _prepare_brainmask_data(dwdata.brainmask, dwdata.affine) @@ -199,22 +196,6 @@ def estimate( return dwdata.em_affines -def _to_nifti(data, affine, filename, clip=True): - data = np.squeeze(data) - if clip: - from eddymotion.data.filtering import advanced_clip - - data = advanced_clip(data) - nii = nb.Nifti1Image( - data, - affine, - None, - ) - nii.header.set_sform(affine, code=1) - nii.header.set_qform(affine, code=1) - nii.to_filename(filename) - - def _prepare_brainmask_data(brainmask, affine): """Prepare the brainmask data: save the data to disk. @@ -268,136 +249,3 @@ def _prepare_kwargs(dwdata, kwargs): if hasattr(dwdata, "total_duration"): kwargs["xlim"] = dwdata.total_duration - - -def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg_target_type): - """Prepare the registration data: save the fixed and moving images to disk. - - Parameters - ---------- - dwframe : :obj:`numpy.ndarray` - DWI data object. - predicted : :obj:`numpy.ndarray` - Predicted data. - affine : :obj:`numpy.ndarray` - Affine transformation matrix. - vol_idx : :obj:`int - DWI volume index. - dirname : :obj:`Path` - Directory name where the data is saved. - reg_target_type : :obj:`str` - Target registration type. - - Returns - ------- - fixed : :obj:`Path` - Fixed image filename. - moving : :obj:`Path` - Moving image filename. - """ - - moving = dirname / f"moving{vol_idx:05d}.nii.gz" - fixed = dirname / f"fixed{vol_idx:05d}.nii.gz" - _to_nifti(dwframe, affine, moving) - _to_nifti( - predicted, - affine, - fixed, - clip=reg_target_type == "dwi", - ) - return fixed, moving - - -def _run_registration( - fixed, - moving, - bmask_img, - em_affines, - affine, - shape, - bval, - fieldmap, - i_iter, - vol_idx, - dirname, - reg_target_type, - align_kwargs, -): - """Register the moving image to the fixed image. - - Parameters - ---------- - fixed : :obj:`Path` - Fixed image filename. - moving : :obj:`Path` - Moving image filename. - bmask_img : :class:`~nibabel.nifti1.Nifti1Image` - Brainmask image. - em_affines : :obj:`numpy.ndarray` - Estimated eddy motion affine transformation matrices. - affine : :obj:`numpy.ndarray` - Affine transformation matrix. - shape : :obj:`tuple` - Shape of the DWI frame. - bval : :obj:`int` - b-value of the corresponding DWI volume. - fieldmap : :class:`~nibabel.nifti1.Nifti1Image` - Fieldmap. - i_iter : :obj:`int` - Iteration number. - vol_idx : :obj:`int` - DWI frame index. - dirname : :obj:`Path` - Directory name where the transformation is saved. - reg_target_type : :obj:`str` - Target registration type. - align_kwargs : :obj:`dict` - Parameters to configure the image registration process. - - Returns - ------- - xform : :class:`~nitransforms.linear.Affine` - Registration transformation. - """ - - registration = Registration( - terminal_output="file", - from_file=pkg_fn( - "eddymotion", - f"config/dwi-to-{reg_target_type}_level{i_iter}.json", - ), - fixed_image=str(fixed.absolute()), - moving_image=str(moving.absolute()), - **align_kwargs, - ) - if bmask_img: - registration.inputs.fixed_image_masks = ["NULL", bmask_img] - - if em_affines is not None and np.any(em_affines[vol_idx, ...]): - reference = namedtuple("ImageGrid", ("shape", "affine"))(shape=shape, affine=affine) - - # create a nitransforms object - if fieldmap: - # compose fieldmap into transform - raise NotImplementedError - else: - initial_xform = Affine(matrix=em_affines[vol_idx], reference=reference) - mat_file = dirname / f"init_{i_iter}_{vol_idx:05d}.mat" - initial_xform.to_filename(mat_file, fmt="itk") - registration.inputs.initial_moving_transform = str(mat_file) - - # execute ants command line - result = registration.run(cwd=str(dirname)).outputs - - # read output transform - xform = nt.linear.Affine( - nt.io.itk.ITKLinearTransform.from_filename(result.forward_transforms[0]).to_ras( - reference=fixed, moving=moving - ), - ) - # debugging: generate aligned file for testing - xform.apply(moving, reference=fixed).to_filename( - dirname / f"aligned{vol_idx:05d}_{int(bval):04d}.nii.gz" - ) - - return xform diff --git a/src/eddymotion/registration/__init__.py b/src/eddymotion/registration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py new file mode 100644 index 00000000..f238174d --- /dev/null +++ b/src/eddymotion/registration/ants.py @@ -0,0 +1,184 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2024 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +"""Using ANTs for image registration.""" + +from collections import namedtuple + +import nibabel as nb +import nitransforms as nt +import numpy as np +from nipype.interfaces.ants.registration import Registration +from nitransforms.linear import Affine +from pkg_resources import resource_filename as pkg_fn + + +def _to_nifti(data, affine, filename, clip=True): + data = np.squeeze(data) + if clip: + from eddymotion.data.filtering import advanced_clip + + data = advanced_clip(data) + nii = nb.Nifti1Image( + data, + affine, + None, + ) + nii.header.set_sform(affine, code=1) + nii.header.set_qform(affine, code=1) + nii.to_filename(filename) + + +def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg_target_type): + """Prepare the registration data: save the fixed and moving images to disk. + + Parameters + ---------- + dwframe : :obj:`numpy.ndarray` + DWI data object. + predicted : :obj:`numpy.ndarray` + Predicted data. + affine : :obj:`numpy.ndarray` + Affine transformation matrix. + vol_idx : :obj:`int + DWI volume index. + dirname : :obj:`Path` + Directory name where the data is saved. + reg_target_type : :obj:`str` + Target registration type. + + Returns + ------- + fixed : :obj:`Path` + Fixed image filename. + moving : :obj:`Path` + Moving image filename. + """ + + moving = dirname / f"moving{vol_idx:05d}.nii.gz" + fixed = dirname / f"fixed{vol_idx:05d}.nii.gz" + _to_nifti(dwframe, affine, moving) + _to_nifti( + predicted, + affine, + fixed, + clip=reg_target_type == "dwi", + ) + return fixed, moving + + +def _run_registration( + fixed, + moving, + bmask_img, + em_affines, + affine, + shape, + bval, + fieldmap, + i_iter, + vol_idx, + dirname, + reg_target_type, + align_kwargs, +): + """Register the moving image to the fixed image. + + Parameters + ---------- + fixed : :obj:`Path` + Fixed image filename. + moving : :obj:`Path` + Moving image filename. + bmask_img : :class:`~nibabel.nifti1.Nifti1Image` + Brainmask image. + em_affines : :obj:`numpy.ndarray` + Estimated eddy motion affine transformation matrices. + affine : :obj:`numpy.ndarray` + Affine transformation matrix. + shape : :obj:`tuple` + Shape of the DWI frame. + bval : :obj:`int` + b-value of the corresponding DWI volume. + fieldmap : :class:`~nibabel.nifti1.Nifti1Image` + Fieldmap. + i_iter : :obj:`int` + Iteration number. + vol_idx : :obj:`int` + DWI frame index. + dirname : :obj:`Path` + Directory name where the transformation is saved. + reg_target_type : :obj:`str` + Target registration type. + align_kwargs : :obj:`dict` + Parameters to configure the image registration process. + + Returns + ------- + xform : :class:`~nitransforms.linear.Affine` + Registration transformation. + """ + + if isinstance(reg_target_type, str): + reg_target_type = (reg_target_type, reg_target_type) + + registration = Registration( + terminal_output="file", + from_file=pkg_fn( + "eddymotion.registration", + f"config/{reg_target_type[0]}-to-{reg_target_type[1]}_level{i_iter}.json", + ), + fixed_image=str(fixed.absolute()), + moving_image=str(moving.absolute()), + **align_kwargs, + ) + if bmask_img: + registration.inputs.fixed_image_masks = ["NULL", bmask_img] + + if em_affines is not None and np.any(em_affines[vol_idx, ...]): + reference = namedtuple("ImageGrid", ("shape", "affine"))(shape=shape, affine=affine) + + # create a nitransforms object + if fieldmap: + # compose fieldmap into transform + raise NotImplementedError + else: + initial_xform = Affine(matrix=em_affines[vol_idx], reference=reference) + mat_file = dirname / f"init_{i_iter}_{vol_idx:05d}.mat" + initial_xform.to_filename(mat_file, fmt="itk") + registration.inputs.initial_moving_transform = str(mat_file) + + # execute ants command line + result = registration.run(cwd=str(dirname)).outputs + + # read output transform + xform = nt.linear.Affine( + nt.io.itk.ITKLinearTransform.from_filename(result.forward_transforms[0]).to_ras( + reference=fixed, moving=moving + ), + ) + # debugging: generate aligned file for testing + xform.apply(moving, reference=fixed).to_filename( + dirname / f"aligned{vol_idx:05d}_{int(bval):04d}.nii.gz" + ) + + return xform diff --git a/src/eddymotion/config/dwi-to-b0_level0.json b/src/eddymotion/registration/config/dwi-to-b0_level0.json similarity index 100% rename from src/eddymotion/config/dwi-to-b0_level0.json rename to src/eddymotion/registration/config/dwi-to-b0_level0.json diff --git a/src/eddymotion/config/dwi-to-dwi_level0.json b/src/eddymotion/registration/config/dwi-to-dwi_level0.json similarity index 100% rename from src/eddymotion/config/dwi-to-dwi_level0.json rename to src/eddymotion/registration/config/dwi-to-dwi_level0.json diff --git a/src/eddymotion/config/dwi-to-dwi_level1.json b/src/eddymotion/registration/config/dwi-to-dwi_level1.json similarity index 100% rename from src/eddymotion/config/dwi-to-dwi_level1.json rename to src/eddymotion/registration/config/dwi-to-dwi_level1.json diff --git a/test/test_integration.py b/test/test_integration.py index 081fda32..a3e436ec 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -68,7 +68,13 @@ def test_proximity_estimator_trivial_model(datadir): estimator = EddyMotionEstimator() em_affines = estimator.estimate( - dwdata=dwi_motion, models=("b0",), align_kwargs=None, seed=None + dwdata=dwi_motion, + models=("b0",), + seed=None, + align_kwargs={ + "fixed_modality": "dwi", + "moving_modality": "b0", + }, ) # Uncomment to see the realigned dataset diff --git a/test/test_estimator.py b/test/test_registration.py similarity index 96% rename from test/test_estimator.py rename to test/test_registration.py index c9024cfa..a7415859 100644 --- a/test/test_estimator.py +++ b/test/test_registration.py @@ -22,6 +22,8 @@ # """Unit tests exercising the estimator.""" +from os import cpu_count + import nibabel as nb import nitransforms as nt import numpy as np @@ -61,12 +63,13 @@ def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): registration = Registration( terminal_output="file", from_file=pkg_fn( - "eddymotion", + "eddymotion.registration", "config/dwi-to-b0_level0.json", ), fixed_image=str(fixed.absolute()), moving_image=str(moving.absolute()), random_seed=1234, + num_threads=cpu_count(), ) result = registration.run(cwd=str(tmp_path)).outputs