diff --git a/pyproject.toml b/pyproject.toml index b8c95e5b..4a096105 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,6 @@ plotting = ["nilearn"] resmon = ["psutil >=5.4"] -popylar = ["popylar >= 0.2"] - test = [ "coverage", "pytest >= 4.4", @@ -76,10 +74,15 @@ test = [ "pytest-xdist >= 1.28" ] +antsopt = [ + "ConfigSpace", + "smac", +] + # Aliases docs = ["eddymotion[doc]"] tests = ["eddymotion[test]"] -all = ["eddymotion[doc,test,dev,plotting,resmon,popylar]"] +all = ["eddymotion[doc,test,dev,plotting,resmon,antsopt]"] [project.scripts] eddymotion = "eddymotion.cli.run:main" @@ -164,7 +167,7 @@ known-first-party=["eddymotion"] pythonpath = "src/ test/" norecursedirs = [".*", "_*"] addopts = "-v --doctest-modules" -doctest_optionflags = "ALLOW_UNICODE NORMALIZE_WHITESPACE" +doctest_optionflags = "ALLOW_UNICODE NORMALIZE_WHITESPACE ELLIPSIS" env = "PYTHONHASHSEED=0" filterwarnings = ["ignore::DeprecationWarning"] diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py new file mode 100644 index 00000000..25b6e2a5 --- /dev/null +++ b/scripts/optimize_registration.py @@ -0,0 +1,222 @@ +# 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/ +# +"""Optimize ANTs' configurations.""" + +import asyncio +import logging +from os import getenv +from pathlib import Path +from shutil import rmtree + +import nibabel as nb +import nitransforms as nt +import numpy as np +from ConfigSpace import Configuration, ConfigurationSpace +from smac import HyperparameterOptimizationFacade, Scenario +from smac.utils.configspace import get_config_hash + +from eddymotion.registration import ants as erants +from eddymotion.registration import utils + +logger = logging.getLogger("ants-optimization") + +# When inside ipython / jupyter +# import nest_asyncio +# nest_asyncio.apply() + +TIME_PENALTY_WEIGHT = 0.1 +SEED = 2390232 + +## Generate config dictionary +configdict = { + # "convergence_threshold": (1e-5, 1e-6), + # "winsorize_lower_quantile": (0.001, 0.2), + # "winsorize_upper_quantile": (0.9, 0.999), + "transform_parameters": (0.01, 2.0), + "smoothing_sigmas": (0.0, 1.0), + "shrink_factors": (1, 2), + "radius_or_number_of_bins": (3, 5), + "sampling_percentage": (0.1, 0.4), + # "metric": ["GC"], + "sampling_strategy": ["Random", "Regular"], +} +paramspace = ConfigurationSpace(configdict) + + +async def ants(cmd, cwd, stdout, stderr, semaphore): + async with semaphore: + proc = await asyncio.create_subprocess_shell( + cmd, + cwd=cwd, + stdout=stdout, + stderr=stderr, + ) + returncode = await proc.wait() + return returncode + + +DATASET_PATH = Path(getenv("TEST_DATA_HOME", f"{getenv('HOME')}/.cache/eddymotion-tests")) + +WORKDIR = Path.home() / "tmp" / "eddymotiondev" +WORKDIR.mkdir(parents=True, exist_ok=True) + +EXPERIMENTDIR = WORKDIR / "smac" +if EXPERIMENTDIR.exists(): + rmtree(EXPERIMENTDIR, ignore_errors=True) + +EXPERIMENTDIR.mkdir(parents=True, exist_ok=True) + +rng = np.random.default_rng(SEED) +MOTION_PARAMETERS = np.hstack( + (rng.uniform(-0.4, 0.4, size=(60, 3)), rng.uniform(-2.0, 2.0, size=(60, 3))) +) +CONVERSIONS = [ + nb.affines.from_matvec(nb.eulerangles.euler2mat(*parameters[:3]), parameters[3:]) + for parameters in MOTION_PARAMETERS +] + +REFERENCES = ( + DATASET_PATH / "dwi-b0_desc-avg.nii.gz", + DATASET_PATH / "hcph-b0_desc-avg.nii.gz", +) + + +async def train_coro( + config: Configuration, + seed: int = 0, + verbose: bool = False, +) -> float: + tmp_folder = EXPERIMENTDIR / get_config_hash(config) + tmp_folder.mkdir(parents=True, exist_ok=True) + align_kwargs = {k: config[k] for k in configdict.keys()} + + ref_xfms = [] + tasks = [] + semaphore = asyncio.Semaphore(18) + nconv = len(CONVERSIONS) + for i, T in enumerate(CONVERSIONS): + for j in (0, 1): + fixed_path = REFERENCES[j] + brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") + refnii = nb.load(fixed_path) + xfm = nt.linear.Affine(T, reference=refnii) + ref_xfms.append(xfm) + + index = i * len(REFERENCES) + j + moving_path = tmp_folder / f"test-{index:04d}.nii.gz" + (~xfm).apply(refnii, reference=refnii).to_filename(moving_path) + + cmdline = erants.generate_command( + fixed_path, + moving_path, + fixedmask_path=brainmask_path, + output_transform_prefix=f"conversion-{index:04d}", + **align_kwargs, + ) + + tasks.append( + ants( + cmdline, + cwd=str(tmp_folder), + stdout=Path(tmp_folder / f"ants-{index:04d}.out").open("w+"), + stderr=Path(tmp_folder / f"ants-{index:04d}.err").open("w+"), + semaphore=semaphore, + ) + ) + + results = await asyncio.gather(*tasks, return_exceptions=True) + + diff = [] + times = [] + start = [] + for i, r in enumerate(results): + if r: + return 1e6 + + j = i % 2 + fixed_path = REFERENCES[j] + brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") + + fixednii = nb.load(fixed_path) + movingnii = nb.load(tmp_folder / f"test-{i:04d}.nii.gz") + xform = nt.linear.Affine( + nt.io.itk.ITKLinearTransform.from_filename( + tmp_folder / f"conversion-{i:04d}0GenericAffine.mat" + ).to_ras( + reference=fixednii, + moving=movingnii, + ), + ) + + masknii = nb.load(brainmask_path) + initial_error = utils.displacements_within_mask( + masknii, + ref_xfms[i], + ) + + disps = utils.displacements_within_mask( + masknii, + xform, + ref_xfms[i], + ) + diff.append(np.percentile(disps, 95)) + start.append(np.percentile(initial_error, 95)) + + # Parse log -- Total elapsed time: 1.0047e+00 + for line in reversed(Path(tmp_folder / f"ants-{i:04d}.out").read_text().splitlines()): + if line.strip().startswith("Total elapsed time:"): + times.append(float(line.strip().split(" ")[-1])) + + meandiff = np.mean(diff) + meantime = np.mean(times) + error = ((1.0 - TIME_PENALTY_WEIGHT) * meandiff + TIME_PENALTY_WEIGHT * meantime) / np.mean( + start + ) + + logger.info( + f"Normalized objective ({nconv} it.): {error:0.3f} " + f"({meandiff:0.2f} mm | {meantime:0.2f} s). " + f"Avg. p95 initial error: {np.mean(start):0.2f} mm." + ) + if verbose: + logger.info(f"\n\nParameters:\n{align_kwargs}" f"\n\nConversions folder: {tmp_folder}.") + + return error + + +def train(config: Configuration, seed: int = 0) -> float: + loop = asyncio.get_event_loop() + return loop.run_until_complete(train_coro(config, seed)) + + +# Scenario object specifying the optimization environment +scenario = Scenario(paramspace, n_trials=200) + +# Use SMAC to find the best configuration/hyperparameters +smac = HyperparameterOptimizationFacade(scenario, train) +incumbent = smac.optimize() + +print(incumbent) + +loop = asyncio.get_event_loop() +loop.run_until_complete(train_coro(incumbent, verbose=True)) diff --git a/src/eddymotion/conftest.py b/src/eddymotion/conftest.py new file mode 100644 index 00000000..5d47bfc3 --- /dev/null +++ b/src/eddymotion/conftest.py @@ -0,0 +1,54 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2021 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/ +# +"""py.test configuration.""" + +import os +from pathlib import Path + +import nibabel +import numpy +import pytest + +test_data_env = os.getenv("TEST_DATA_HOME", str(Path.home() / "eddymotion-tests")) +test_output_dir = os.getenv("TEST_OUTPUT_DIR") +test_workdir = os.getenv("TEST_WORK_DIR") + +_datadir = (Path(__file__).parent.parent / "eddymotion" / "data").absolute() + + +def pytest_report_header(config): + return f"""\ +TEST_DATA_HOME={test_data_env}. +TEST_OUTPUT_DIR={test_output_dir or ' (output files will be discarded)'}. +TEST_WORK_DIR={test_workdir or ' (intermediate files will be discarded)'}. +""" + + +@pytest.fixture(autouse=True) +def doctest_imports(doctest_namespace): + """Populates doctests with some conveniency imports.""" + doctest_namespace["np"] = numpy + doctest_namespace["nb"] = nibabel + doctest_namespace["os"] = os + doctest_namespace["Path"] = Path + doctest_namespace["repodata"] = _datadir diff --git a/src/eddymotion/data/fileA.nii.gz b/src/eddymotion/data/fileA.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/data/fileB.nii.gz b/src/eddymotion/data/fileB.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/data/maskA.nii.gz b/src/eddymotion/data/maskA.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index f238174d..925ed1a3 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -22,7 +22,12 @@ # """Using ANTs for image registration.""" +from __future__ import annotations + from collections import namedtuple +from json import loads +from pathlib import Path +from warnings import warn import nibabel as nb import nitransforms as nt @@ -31,8 +36,46 @@ from nitransforms.linear import Affine from pkg_resources import resource_filename as pkg_fn +PARAMETERS_SINGLE_VALUE = { + "collapse_output_transforms", + "dimension", + "initial_moving_transform", + "initialize_transforms_per_stage", + "interpolation", + "output_transform_prefix", + "verbose", + "winsorize_lower_quantile", + "winsorize_upper_quantile", + "write_composite_transform", +} + +PARAMETERS_SINGLE_LIST = { + "radius_or_number_of_bins", + "sampling_percentage", + "metric", + "sampling_strategy", +} +PARAMETERS_DOUBLE_LIST = {"shrink_factors", "smoothing_sigmas", "transform_parameters"} + + +def _to_nifti( + data: np.ndarray, affine: np.ndarray, filename: str | Path, clip: bool = True +) -> None: + """ + Save data as a NIfTI file, optionally applying clipping. -def _to_nifti(data, affine, filename, clip=True): + Parameters + ---------- + data : :obj:`~numpy.ndarray` + The image data to be saved. + affine : :obj:`~numpy.ndarray` + The affine transformation matrix. + filename : :obj:`os.pathlike` + The file path where the NIfTI file will be saved. + clip : :obj:`bool`, optional + Whether to apply clipping to the data before saving, by default True. + + """ data = np.squeeze(data) if clip: from eddymotion.data.filtering import advanced_clip @@ -48,34 +91,42 @@ def _to_nifti(data, affine, filename, clip=True): 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. +def _prepare_registration_data( + dwframe: np.ndarray, + predicted: np.ndarray, + affine: np.ndarray, + vol_idx: int, + dirname: Path | str, + reg_target_type: str, +) -> tuple[Path, Path]: + """ + Prepare the registration data: save the fixed and moving images to disk. Parameters ---------- - dwframe : :obj:`numpy.ndarray` + dwframe : :obj:`~numpy.ndarray` DWI data object. - predicted : :obj:`numpy.ndarray` + predicted : :obj:`~numpy.ndarray` Predicted data. affine : :obj:`numpy.ndarray` Affine transformation matrix. - vol_idx : :obj:`int + vol_idx : :obj:`int` DWI volume index. - dirname : :obj:`Path` + dirname : :obj:`os.pathlike` Directory name where the data is saved. reg_target_type : :obj:`str` Target registration type. Returns ------- - fixed : :obj:`Path` + fixed : :obj:`~pathlib.Path` Fixed image filename. - moving : :obj:`Path` + moving : :obj:`~pathlib.Path` Moving image filename. """ - moving = dirname / f"moving{vol_idx:05d}.nii.gz" - fixed = dirname / f"fixed{vol_idx:05d}.nii.gz" + moving = Path(dirname) / f"moving{vol_idx:05d}.nii.gz" + fixed = Path(dirname) / f"fixed{vol_idx:05d}.nii.gz" _to_nifti(dwframe, affine, moving) _to_nifti( predicted, @@ -86,22 +137,287 @@ def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg return fixed, moving +def _get_ants_settings(settings: str = "b0-to-b0_level0") -> Path: + """ + Retrieve the path to ANTs settings configuration file. + + Parameters + ---------- + settings : :obj:`str`, optional + Name of the settings configuration, by default ``"b0-to-b0_level0"``. + + Returns + ------- + :obj:`~pathlib.Path` + The path to the configuration file. + + Examples + -------- + >>> _get_ants_settings() + PosixPath('.../config/b0-to-b0_level0.json') + + >>> _get_ants_settings("b0-to-b0_level1") + PosixPath('.../config/b0-to-b0_level1.json') + + """ + return Path( + pkg_fn( + "eddymotion.registration", + f"config/{settings}.json", + ) + ) + + +def _massage_mask_path(mask_path: str | Path | list[str], nlevels: int) -> list[str]: + """ + Generate nipype-compatible masks paths. + + Parameters + ---------- + mask_path : :obj:`os.pathlike` or :obj:`list` + Path(s) to the mask file(s). + nlevels : :obj:`int` + Number of registration levels. + + Returns + ------- + :obj:`list` + A list of mask paths formatted for *Nipype*. + + Examples + -------- + >>> _massage_mask_path("/some/path", 2) + ['/some/path', '/some/path'] + + >>> _massage_mask_path(["/some/path"] * 2, 2) + ['/some/path', '/some/path'] + + >>> _massage_mask_path(["/some/path"] * 2, 4) + ['NULL', 'NULL', '/some/path', '/some/path'] + + """ + if isinstance(mask_path, (str, Path)): + return [str(mask_path)] * nlevels + if len(mask_path) < nlevels: + return ["NULL"] * (nlevels - len(mask_path)) + mask_path + if len(mask_path) > nlevels: + warn("More mask paths than levels", stacklevel=1) + return mask_path[:nlevels] + return mask_path + + +def generate_command( + fixed_path: str | Path, + moving_path: str | Path, + fixedmask_path: str | Path | list[str] | None = None, + movingmask_path: str | Path | list[str] | None = None, + init_affine: str | Path | None = None, + default: str = "b0-to-b0_level0", + **kwargs: dict, +) -> str: + """ + Generate an ANTs' command line. + + Parameters + ---------- + fixed_path : :obj:`os.pathlike` + Path to the fixed image. + moving_path : :obj:`os.pathlike` + Path to the moving image. + fixedmask_path : :obj:`os.pathlike` or :obj:`list`, optional + Path to the fixed image mask, by default None. + movingmask_path : :obj:`os.pathlike` or :obj:`list`, optional + Path to the moving image mask, by default None. + init_affine : :obj:`os.pathlike`, optional + Initial affine transformation, by default None. + default : :obj:`str`, optional + Default settings configuration, by default "b0-to-b0_level0". + **kwargs : :obj:`dict` + Additional parameters for ANTs registration. + + Returns + ------- + :obj:`str` + The ANTs registration command line string. + + Examples + -------- + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 12.0 ] \ + --metric GC[ .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 3, Random, 0.4 ] \ + --convergence [ 20, 1e-06, 4 ] --smoothing-sigmas 2.71vox --shrink-factors 3 \ + --use-histogram-matching 1 \ + --transform Rigid[ 1.96 ] \ + --metric GC[ .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 4, Random, 0.18 ] \ + --convergence [ 10, 1e-07, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 2 \ + --use-histogram-matching 1 \ + -v --winsorize-image-intensities [ 0.063, 0.991 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Regular, 0.2 \ + ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ + --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ + --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Random, 0.1 \ + ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... fixedmask_path=repodata / 'maskA.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Regular, 0.2 ] \ + --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ + --use-histogram-matching 1 --masks [ \ + .../maskA.nii.gz, NULL ] \ + --transform Rigid[ 0.001 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Random, 0.1 ] \ + --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 --masks [ \ + .../maskA.nii.gz, NULL ] \ + -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Regular, 0.2 \ + ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ + --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ + --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Random, 0.1 \ + ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... fixedmask_path=[repodata / 'maskA.nii.gz'], + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Regular, 0.2 ] \ + --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ + --use-histogram-matching 1 --masks [ NULL, NULL ] \ + --transform Rigid[ 0.001 ] --metric Mattes[ \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ + 1, 32, Random, 0.1 ] \ + --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 --masks [ \ + .../maskA.nii.gz, NULL ] \ + -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' + + """ + + # Bootstrap settings from defaults file and override with single-valued parameters in args + settings = loads(_get_ants_settings(default).read_text()) | { + k: kwargs.pop(k) for k in PARAMETERS_SINGLE_VALUE if k in kwargs + } + + # Determine number of levels and assert consistency of levels + levels = {len(settings[p]) for p in PARAMETERS_SINGLE_LIST if p in settings} + nlevels = levels.pop() + if levels: + raise RuntimeError(f"Malformed settings file (levels: {levels})") + + # Override list (and nested-list) parameters + for key, value in kwargs.items(): + if key in PARAMETERS_DOUBLE_LIST: + value = [value] + elif key not in PARAMETERS_SINGLE_LIST: + continue + + if levels == 1: + settings[key] = [value] + else: + settings[key][-1] = value + + # Set fixed masks if provided + if fixedmask_path is not None: + settings["fixed_image_masks"] = [ + str(p) for p in _massage_mask_path(fixedmask_path, nlevels) + ] + + # Set moving masks if provided + if movingmask_path is not None: + settings["moving_image_masks"] = [ + str(p) for p in _massage_mask_path(movingmask_path, nlevels) + ] + + # Set initalizing affine if provided + if init_affine is not None: + settings["initial_moving_transform"] = str(init_affine) + + # Generate command line with nipype and return + return Registration( + fixed_image=str(Path(fixed_path).absolute()), + moving_image=str(Path(moving_path).absolute()), + **settings, + ).cmdline + + 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. + fixed: Path, + moving: Path, + bmask_img: nb.spatialimages.SpatialImage, + em_affines: np.ndarray, + affine: np.ndarray, + shape: tuple[int, int, int], + bval: int, + fieldmap: nb.spatialimages.SpatialImage, + i_iter: int, + vol_idx: int, + dirname: Path, + reg_target_type: str, + align_kwargs: dict, +) -> nt.base.BaseTransform: + """ + Register the moving image to the fixed image. Parameters ---------- @@ -109,7 +425,7 @@ def _run_registration( Fixed image filename. moving : :obj:`Path` Moving image filename. - bmask_img : :class:`~nibabel.nifti1.Nifti1Image` + bmask_img : :class:`~nibabel.spatialimages.SpatialImage` Brainmask image. em_affines : :obj:`numpy.ndarray` Estimated eddy motion affine transformation matrices. @@ -119,7 +435,7 @@ def _run_registration( Shape of the DWI frame. bval : :obj:`int` b-value of the corresponding DWI volume. - fieldmap : :class:`~nibabel.nifti1.Nifti1Image` + fieldmap : :class:`~nibabel.spatialimages.SpatialImage` Fieldmap. i_iter : :obj:`int` Iteration number. @@ -134,8 +450,9 @@ def _run_registration( Returns ------- - xform : :class:`~nitransforms.linear.Affine` + xform : :obj:`~nitransforms.base.BaseTransform` Registration transformation. + """ if isinstance(reg_target_type, str): diff --git a/src/eddymotion/registration/config/b0-to-b0_level0.json b/src/eddymotion/registration/config/b0-to-b0_level0.json new file mode 100644 index 00000000..5035c4c2 --- /dev/null +++ b/src/eddymotion/registration/config/b0-to-b0_level0.json @@ -0,0 +1,37 @@ +{ + "collapse_output_transforms": true, + "convergence_threshold": [ 1E-6, 1E-7 ], + "convergence_window_size": [ 4, 2 ], + "dimension": 3, + "initialize_transforms_per_stage": false, + "interpolation": "Linear", + "metric": [ "GC", "GC" ], + "metric_weight": [ 1.0, 1.0 ], + "number_of_iterations": [ + [ 20 ], + [ 10 ] + ], + "radius_or_number_of_bins": [ 3, 4 ], + "sampling_percentage": [ 0.4, 0.18 ], + "sampling_strategy": [ "Random", "Random" ], + "shrink_factors": [ + [ 3 ], + [ 2 ] + ], + "sigma_units": [ "vox", "vox" ], + "smoothing_sigmas": [ + [ 2.71 ], + [ 0.0 ] + ], + "transform_parameters": [ + [ 12.0 ], + [ 1.96 ] + ], + "transforms": [ "Rigid", "Rigid" ], + "use_histogram_matching": [ true, true ], + "verbose": true, + "winsorize_lower_quantile": 0.063, + "winsorize_upper_quantile": 0.991, + "write_composite_transform": false + } + \ No newline at end of file diff --git a/src/eddymotion/registration/utils.py b/src/eddymotion/registration/utils.py new file mode 100644 index 00000000..0ac54240 --- /dev/null +++ b/src/eddymotion/registration/utils.py @@ -0,0 +1,114 @@ +# 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/ +# +""" +Utilities to aid in performing and evaluating image registration. + +This module provides functions to compute displacements of image coordinates +under a transformation, useful for assessing the accuracy of image registration +processes. + +""" + +from __future__ import annotations + +from itertools import product + +import nibabel as nb +import nitransforms as nt +import numpy as np + + +def displacements_within_mask( + mask_img: nb.spatialimages.SpatialImage, + test_xfm: nt.base.BaseTransform, + reference_xfm: nt.base.BaseTransform | None = None, +) -> np.ndarray: + """ + Compute the distance between voxel coordinates mapped through two transforms. + + Parameters + ---------- + mask_img : :obj:`~nibabel.spatialimages.SpatialImage` + A mask image that defines the region of interest. Voxel coordinates + within the mask are transformed. + test_xfm : :obj:`~nitransforms.base.BaseTransform` + The transformation to test. This transformation is applied to the + voxel coordinates. + reference_xfm : :obj:`~nitransforms.base.BaseTransform`, optional + A reference transformation to compare with. If ``None``, the identity + transformation is assumed (no transformation). + + Returns + ------- + :obj:`~numpy.ndarray` + An array of displacements (in mm) for each voxel within the mask. + + """ + # Mask data as boolean (True for voxels inside the mask) + maskdata = np.asanyarray(mask_img.dataobj) > 0 + # Convert voxel coordinates to world coordinates using affine transform + xyz = nb.affines.apply_affine( + mask_img.affine, + np.argwhere(maskdata), + ) + # Apply the test transformation + targets = test_xfm.map(xyz) + + # Compute the difference (displacement) between the test and reference transformations + diffs = targets - xyz if reference_xfm is None else targets - reference_xfm.map(xyz) + return np.linalg.norm(diffs, axis=-1) + + +def displacement_framewise( + img: nb.spatialimages.SpatialImage, + test_xfm: nt.base.BaseTransform, + radius: float = 50.0, +): + """ + Compute the framewise displacement (FD) for a given transformation. + + Parameters + ---------- + img : :obj:`~nibabel.spatialimages.SpatialImage` + The reference image. Used to extract the center coordinates. + test_xfm : :obj:`~nitransforms.base.BaseTransform` + The transformation to test. Applied to coordinates around the image center. + radius : :obj:`float`, optional + The radius (in mm) of the spherical neighborhood around the center of the image. + Default is 50.0 mm. + + Returns + ------- + :obj:`float` + The average framewise displacement (FD) for the test transformation. + + """ + affine = img.affine + # Compute the center of the image in voxel space + center_ijk = 0.5 * (np.array(img.shape[:3]) - 1) + # Convert to world coordinates + center_xyz = nb.affines.apply_affine(affine, center_ijk) + # Generate coordinates of points at radius distance from center + fd_coords = np.array(list(product(*((radius, -radius),) * 3))) + center_xyz + # Compute the average displacement from the test transformation + return np.mean(np.linalg.norm(test_xfm.map(fd_coords) - fd_coords, axis=-1)) diff --git a/test/test_integration.py b/test/test_integration.py index 485a99cf..1e5f1a4f 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -74,7 +74,7 @@ def test_proximity_estimator_trivial_model(datadir): models=("b0",), seed=None, align_kwargs={ - "fixed_modality": "dwi", + "fixed_modality": "b0", "moving_modality": "b0", "num_threads": min(cpu_count(), 8), }, diff --git a/test/test_registration.py b/test/test_registration.py index a7415859..4f193911 100644 --- a/test/test_registration.py +++ b/test/test_registration.py @@ -33,7 +33,8 @@ from nipype.interfaces.ants.registration import Registration from pkg_resources import resource_filename as pkg_fn -from eddymotion.data.dmri import DWI +from eddymotion.registration.ants import _massage_mask_path +from eddymotion.registration.utils import displacements_within_mask @pytest.mark.parametrize("r_x", [0.0, 0.1, 0.3]) @@ -42,19 +43,17 @@ @pytest.mark.parametrize("t_x", [0.0, 1.0]) @pytest.mark.parametrize("t_y", [0.0, 1.0]) @pytest.mark.parametrize("t_z", [0.0, 1.0]) -def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): +@pytest.mark.parametrize("dataset", ["hcph", "dwi"]) +# @pytest.mark.parametrize("dataset", ["dwi"]) +def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z): """Check that the registration parameters for b=0 gives a good estimate of known affine""" - fixed = tmp_path / "b0.nii.gz" + fixed = datadir / f"{dataset}-b0_desc-avg.nii.gz" + fixed_mask = datadir / f"{dataset}-b0_desc-brain.nii.gz" moving = tmp_path / "moving.nii.gz" - dwdata = DWI.from_filename(datadir / "dwi.h5") - b0nii = nb.Nifti1Image(dwdata.bzero, dwdata.affine, None) - b0nii.header.set_qform(dwdata.affine, code=1) - b0nii.header.set_sform(dwdata.affine, code=1) - b0nii.to_filename(fixed) - + b0nii = nb.load(fixed) T = from_matvec(euler2mat(x=r_x, y=r_y, z=r_z), (t_x, t_y, t_z)) xfm = nt.linear.Affine(T, reference=b0nii) @@ -64,10 +63,11 @@ def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): terminal_output="file", from_file=pkg_fn( "eddymotion.registration", - "config/dwi-to-b0_level0.json", + "config/b0-to-b0_level0.json", ), fixed_image=str(fixed.absolute()), moving_image=str(moving.absolute()), + fixed_image_masks=[str(fixed_mask)], random_seed=1234, num_threads=cpu_count(), ) @@ -78,6 +78,15 @@ def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): reference=b0nii, ) - coords = xfm.reference.ndcoords.T - rms = np.sqrt(((xfm.map(coords) - xform.map(coords)) ** 2).sum(1)).mean() - assert rms < 0.8 + masknii = nb.load(fixed_mask) + assert displacements_within_mask(masknii, xform, xfm).mean() < ( + 0.6 * np.mean(b0nii.header.get_zooms()[:3]) + ) + + +def test_massage_mask_path(): + """Test the case where a warning must be issued.""" + with pytest.warns(UserWarning, match="More mask paths than levels"): + maskpath = _massage_mask_path(["/some/path"] * 2, 1) + + assert maskpath == ["/some/path"]