ENH: insert the gradunwarp workflow in anat template workflow

pyproject.toml

@@ -25,7 +25,8 @@ dependencies = [
     "matplotlib >= 2.2.0",
     "nibabel >= 4.0.1",
     "nipype >= 1.7.0",
-    "niworkflows >= 1.10.1",
+#    "niworkflows >= 1.10.1",
+    "niworkflows@git+https://github.com/bpinsard/niworkflows.git@enh/gradunwarp",
     "numpy",
     "packaging",
     "pybids >= 0.11.1",

smriprep/cli/run.py

@@ -173,6 +173,12 @@ def get_parser():
         action='store_true',
         help='treat dataset as longitudinal - may increase runtime',
     )
+    g_conf.add_argument(
+        '--gradunwarp-file',
+        metavar='PATH',
+        type=Path,
+        help='Path to vendor file for gradunwarp gradient distortion ' 'correction.',
+    )
 
     #  ANTs options
     g_ants = parser.add_argument_group('Specific options for ANTs registrations')

smriprep/workflows/anatomical.py

@@ -24,6 +24,7 @@
 
 import typing as ty
 
+import bids
 from nipype import logging
 from nipype.interfaces import (
     freesurfer as fs,
@@ -47,12 +48,14 @@
 from niworkflows.interfaces.freesurfer import (
     StructuralReference,
 )
+from niworkflows.interfaces.gradunwarp import GradUnwarp
 from niworkflows.interfaces.header import ValidateImage
 from niworkflows.interfaces.images import Conform, TemplateDimensions
 from niworkflows.interfaces.nibabel import ApplyMask, Binarize
 from niworkflows.interfaces.nitransforms import ConcatenateXFMs
 from niworkflows.utils.misc import add_suffix
 from niworkflows.utils.spaces import Reference, SpatialReferences
+from niworkflows.workflows.gradunwarp import init_gradunwarp_wf
 
 from ..data import load_resource
 from ..interfaces import DerivativesDataSink
@@ -94,6 +97,7 @@
 def init_anat_preproc_wf(
     *,
     bids_root: str,
+    layout: bids.BIDSLayout,
     output_dir: str,
     freesurfer: bool,
     hires: bool,
@@ -113,6 +117,7 @@ def init_anat_preproc_wf(
     name: str = 'anat_preproc_wf',
     skull_strip_fixed_seed: bool = False,
     fs_no_resume: bool = False,
+    gradunwarp_file: str | None = None,
 ):
     """
     Stage the anatomical preprocessing steps of *sMRIPrep*.
@@ -150,6 +155,8 @@ def init_anat_preproc_wf(
     ----------
     bids_root : :obj:`str`
         Path of the input BIDS dataset root
+    layout : BIDSLayout object
+        BIDS dataset layout
     output_dir : :obj:`str`
         Directory in which to save derivatives
     freesurfer : :obj:`bool`
@@ -189,6 +196,8 @@ def init_anat_preproc_wf(
         EXPERT: Import pre-computed FreeSurfer reconstruction without resuming.
         The user is responsible for ensuring that all necessary files are present.
         (default: ``False``).
+    gradunwarp_file : :obj:`str`, optional
+        Gradient unwarping filename (default: None)
 
     Inputs
     ------
@@ -265,6 +274,7 @@ def init_anat_preproc_wf(
 
     anat_fit_wf = init_anat_fit_wf(
         bids_root=bids_root,
+        layout=layout,
         output_dir=output_dir,
         freesurfer=freesurfer,
         hires=hires,
@@ -282,6 +292,7 @@ def init_anat_preproc_wf(
         omp_nthreads=omp_nthreads,
         skull_strip_fixed_seed=skull_strip_fixed_seed,
         fs_no_resume=fs_no_resume,
+        gradunwarp_file=gradunwarp_file,
     )
     template_iterator_wf = init_template_iterator_wf(spaces=spaces, sloppy=sloppy)
     ds_std_volumes_wf = init_ds_anat_volumes_wf(
@@ -448,6 +459,7 @@ def init_anat_preproc_wf(
 def init_anat_fit_wf(
     *,
     bids_root: str,
+    layout: bids.BIDSLayout,
     output_dir: str,
     freesurfer: bool,
     hires: bool,
@@ -466,6 +478,7 @@ def init_anat_fit_wf(
     name='anat_fit_wf',
     skull_strip_fixed_seed: bool = False,
     fs_no_resume: bool = False,
+    gradunwarp_file: str | None = None,
 ):
     """
     Stage the anatomical preprocessing steps of *sMRIPrep*.
@@ -511,6 +524,8 @@ def init_anat_fit_wf(
     ----------
     bids_root : :obj:`str`
         Path of the input BIDS dataset root
+    layout : BIDSLayout object
+        BIDS dataset layout
     output_dir : :obj:`str`
         Directory in which to save derivatives
     freesurfer : :obj:`bool`
@@ -546,6 +561,12 @@ def init_anat_fit_wf(
         Do not use a random seed for skull-stripping - will ensure
         run-to-run replicability when used with --omp-nthreads 1
         (default: ``False``).
+    fs_no_resume : bool
+        EXPERT: Import pre-computed FreeSurfer reconstruction without resuming.
+        The user is responsible for ensuring that all necessary files are present.
+        (default: ``False``).
+    gradunwarp_file : :obj:`str`, optional
+        Gradient unwarping filename (default: None)
 
     Inputs
     ------
@@ -760,12 +781,14 @@ def init_anat_fit_wf(
 non-uniformity (INU) with `N4BiasFieldCorrection` [@n4], distributed with ANTs {ants_ver}
 [@ants, RRID:SCR_004757]"""
         desc += '.\n' if num_t1w > 1 else ', and used as T1w-reference throughout the workflow.\n'
-
+        t1w_metas = [layout.get_file(t).get_metadata() for t in t1w]
         anat_template_wf = init_anat_template_wf(
             longitudinal=longitudinal,
             omp_nthreads=omp_nthreads,
             num_files=num_t1w,
             contrast='T1w',
+            gradunwarp_file=gradunwarp_file,
+            metadata=t1w_metas,
             name='anat_template_wf',
         )
         ds_template_wf = init_ds_template_wf(output_dir=output_dir, num_t1w=num_t1w)
@@ -1131,11 +1154,14 @@ def init_anat_fit_wf(
 
     if t2w and not have_t2w:
         LOGGER.info('ANAT Stage 7: Creating T2w template')
+        t2w_metas = [layout.get_file(t).get_metadata() for t in t1w]
         t2w_template_wf = init_anat_template_wf(
             longitudinal=longitudinal,
             omp_nthreads=omp_nthreads,
             num_files=len(t2w),
             contrast='T2w',
+            metadata=t2w_metas,
+            gradunwarp_file=gradunwarp_file,
             name='t2w_template_wf',
         )
         bbreg = pe.Node(
@@ -1376,6 +1402,8 @@ def init_anat_template_wf(
     omp_nthreads: int,
     num_files: int,
     contrast: str,
+    metadata: dict,
+    gradunwarp_file: str | None = None,
     name: str = 'anat_template_wf',
 ):
     """
@@ -1388,7 +1416,8 @@ def init_anat_template_wf(
 
             from smriprep.workflows.anatomical import init_anat_template_wf
             wf = init_anat_template_wf(
-                longitudinal=False, omp_nthreads=1, num_files=1, contrast="T1w"
+                longitudinal=False, omp_nthreads=1, num_files=1, contrast="T1w",
+                gradunwarp_file=None,
             )
 
     Parameters
@@ -1402,6 +1431,8 @@ def init_anat_template_wf(
         Number of images
     contrast : :obj:`str`
         Name of contrast, for reporting purposes, e.g., T1w, T2w, PDw
+    gradunwarp_file : :obj:`str`, optional
+        Gradient unwarping filename (default: None)
     name : :obj:`str`, optional
         Workflow name (default: anat_template_wf)
 
@@ -1457,14 +1488,50 @@ def init_anat_template_wf(
             ('target_zooms', 'target_zooms'),
             ('target_shape', 'target_shape'),
         ]),
-        (denoise, anat_conform, [('output_image', 'in_file')]),
         (anat_ref_dimensions, outputnode, [
             ('out_report', 'out_report'),
             ('t1w_valid_list', 'anat_valid_list'),
         ]),
     ])
     # fmt:on
 
+    # 0.5 Gradient unwarping (optional)
+    if gradunwarp_file:
+        nds = [
+            (
+                not meta.get('NonlinearGradientCorrection', None)
+                or 'ND' in meta.get('ImageType', [])
+                or False
+            )
+            for meta in metadata
+        ]
+        if any(nds) and not all(nds):
+            raise RuntimeError(
+                f'Inconsistent distortion correction metadata across {contrast} images.'
+            )
+        if not any(nds):
+            # gradient unwarping not needed for that contrast
+            gradunwarp_file = None
+
+    if gradunwarp_file:
+        gradunwarp_ver = GradUnwarp.version()
+        workflow.__desc__ = (
+            (workflow.__desc__ or '')
+            + f"""\
+ {"Each" if num_files > 1 else "The"} {contrast} image was corrected for gradient
+ non-linearity with `gradunwarp` [@gradunwarp] {gradunwarp_ver} [@gradunwarp]\n"""
+        )
+        gradunwarp_wf = init_gradunwarp_wf('gradunward_T1w')
+        gradunwarp_wf.inputs.inputnode.grad_file = gradunwarp_file
+        # fmt:off
+        workflow.connect([
+            (denoise, gradunwarp_wf, [('output_image', 'inputnode.input_file')]),
+            (gradunwarp_wf, anat_conform, [('outputnode.corrected_file', 'in_file')]),
+        ])
+    else:
+        workflow.connect([(denoise, anat_conform, [('output_image', 'in_file')])])
+        # fmt:on
+
     if num_files == 1:
         get1st = pe.Node(niu.Select(index=[0]), name='get1st')
         outputnode.inputs.anat_realign_xfm = [str(load_resource('itkIdentityTransform.txt'))]

smriprep/workflows/base.py

@@ -422,6 +422,7 @@ def init_single_subject_wf(
     # Preprocessing of T1w (includes registration to MNI)
     anat_preproc_wf = init_anat_preproc_wf(
         bids_root=layout.root,
+        layout=layout,
         sloppy=sloppy,
         debug=debug,
         precomputed=deriv_cache,

smriprep/workflows/tests/test_anatomical.py

@@ -1,18 +1,26 @@
 from pathlib import Path
 
+import bids
 import nibabel as nb
 import numpy as np
 import pytest
 from nipype.pipeline.engine.utils import generate_expanded_graph
+from niworkflows.tests.data import load_test_data
 from niworkflows.utils.spaces import Reference, SpatialReferences
 from niworkflows.utils.testing import generate_bids_skeleton
 
 from ..anatomical import init_anat_fit_wf, init_anat_preproc_wf
 
+gradunwarp_file_params = [None, load_test_data('gradunwarp_coeffs.grad')]
+
 BASE_LAYOUT = {
     '01': {
         'anat': [
-            {'run': 1, 'suffix': 'T1w'},
+            {
+                'run': 1,
+                'suffix': 'T1w',
+                'metadata': {'ImageType': ['ND']},
+            },
             {'run': 2, 'suffix': 'T1w'},
             {'suffix': 'T2w'},
         ],
@@ -73,14 +81,17 @@ def test_init_anat_preproc_wf(
     output_dir = tmp_path / 'output'
     output_dir.mkdir()
 
+    bids_layout = bids.BIDSLayout(bids_root)
+
     init_anat_preproc_wf(
         bids_root=str(bids_root),
+        layout=bids_layout,
         output_dir=str(output_dir),
         freesurfer=freesurfer,
         hires=False,
         longitudinal=False,
         msm_sulc=False,
-        t1w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_T1w.nii.gz')],
+        t1w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_run-1_T1w.nii.gz')],
         t2w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_T2w.nii.gz')],
         skull_strip_mode='force',
         skull_strip_template=Reference('OASIS30ANTs'),
@@ -96,23 +107,28 @@ def test_init_anat_preproc_wf(
 
 @pytest.mark.parametrize('msm_sulc', [True, False])
 @pytest.mark.parametrize('skull_strip_mode', ['skip', 'force'])
+@pytest.mark.parametrize('gradunwarp_file', gradunwarp_file_params)
 def test_anat_fit_wf(
     bids_root: Path,
     tmp_path: Path,
     msm_sulc: bool,
     skull_strip_mode: str,
+    gradunwarp_file: str,
 ):
     output_dir = tmp_path / 'output'
     output_dir.mkdir()
 
+    bids_layout = bids.BIDSLayout(bids_root)
+
     init_anat_fit_wf(
         bids_root=str(bids_root),
+        layout=bids_layout,
         output_dir=str(output_dir),
         freesurfer=True,
         hires=False,
         longitudinal=False,
         msm_sulc=msm_sulc,
-        t1w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_T1w.nii.gz')],
+        t1w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_run-1_T1w.nii.gz')],
         t2w=[str(bids_root / 'sub-01' / 'anat' / 'sub-01_T2w.nii.gz')],
         skull_strip_mode=skull_strip_mode,
         skull_strip_template=Reference('OASIS30ANTs'),
@@ -122,6 +138,7 @@ def test_anat_fit_wf(
         ),
         precomputed={},
         omp_nthreads=1,
+        gradunwarp_file=gradunwarp_file,
     )
 
 
@@ -200,9 +217,12 @@ def test_anat_fit_precomputes(
             for path in xfm.values():
                 Path(path).touch()
 
+    bids_layout = bids.BIDSLayout(bids_root)
+
     # Create workflow
     wf = init_anat_fit_wf(
         bids_root=str(bids_root),
+        layout=bids_layout,
         output_dir=str(output_dir),
         freesurfer=True,
         hires=False,