Merge pull request #425 from mgxd/fix/composite-xfms

FIX: Generate composite transforms with `--multi-step-reg`

Author: mgxd

Dockerfile

@@ -72,12 +72,12 @@ RUN mkdir -p /opt/afni-latest \
         -name "3dAutomask" -or \
         -name "3dvolreg" \) -delete
 
-# ANTs 2.4.4
+# ANTs 2.5.4
 FROM downloader as ants
 RUN mkdir -p /opt && \
-    curl -sSLO "https://github.com/ANTsX/ANTs/releases/download/v2.4.4/ants-2.4.4-ubuntu-22.04-X64-gcc.zip" && \
-    unzip ants-2.4.4-ubuntu-22.04-X64-gcc.zip -d /opt && \
-    rm ants-2.4.4-ubuntu-22.04-X64-gcc.zip
+    curl -sSLO "https://github.com/ANTsX/ANTs/releases/download/v2.5.4/ants-2.5.4-ubuntu-22.04-X64-gcc.zip" && \
+    unzip ants-2.5.4-ubuntu-22.04-X64-gcc.zip -d /opt && \
+    rm ants-2.5.4-ubuntu-22.04-X64-gcc.zip
 
 # Connectome Workbench 1.5.0
 FROM downloader as workbench
@@ -180,7 +180,7 @@ RUN apt-get update -qq \
     && ldconfig
 
 COPY --from=afni /opt/afni-latest /opt/afni-latest
-COPY --from=ants /opt/ants-2.4.4 /opt/ants
+COPY --from=ants /opt/ants-2.5.4 /opt/ants
 COPY --from=workbench /opt/workbench /opt/workbench
 
 # AFNI config

nibabies/interfaces/patches.py

@@ -4,6 +4,15 @@
     freesurfer as fs,
 )
 from nipype.interfaces.ants.base import ANTSCommand, ANTSCommandInputSpec
+from nipype.interfaces.ants.registration import (
+    CompositeTransformUtil as _CompositeTransformUtil,
+)
+from nipype.interfaces.ants.registration import (
+    CompositeTransformUtilInputSpec as _CompositeTransformUtilInputSpec,
+)
+from nipype.interfaces.ants.registration import (
+    CompositeTransformUtilOutputSpec as _CompositeTransformUtilOutputSpec,
+)
 from nipype.interfaces.base import File, InputMultiObject, TraitedSpec, traits
 
 
@@ -105,3 +114,72 @@ def _list_outputs(self):
         outputs = self._outputs().get()
         outputs['out_xfm'] = Path(self.inputs.out_xfm).absolute()
         return outputs
+
+
+class CompositeTransformUtilInputSpec(_CompositeTransformUtilInputSpec):
+    order_transforms = traits.Bool(
+        True,
+        usedefault=True,
+        desc='Order disassembled transforms into [Affine, Displacement] pairs.',
+    )
+
+
+class CompositeTransformUtilOutputSpec(_CompositeTransformUtilOutputSpec):
+    out_transforms = traits.List(desc='list of transform components')
+
+
+class CompositeTransformUtil(_CompositeTransformUtil):
+    """Outputs have changed in newer versions of ANTs."""
+
+    input_spec = CompositeTransformUtilInputSpec
+    output_spec = CompositeTransformUtilOutputSpec
+
+    def _list_outputs(self):
+        outputs = self.output_spec().get()
+
+        # Ordering may change depending on forward/inverse transform
+        # Forward: <prefix>_00_AffineTransform.mat, <prefix>_01_DisplacementFieldTransform.nii.gz
+        # Inverse: <prefix>_01_AffineTransform.mat, <prefix>_00_DisplacementFieldTransform.nii.gz
+        if self.inputs.process == 'disassemble':
+            transforms = [
+                str(Path(x).absolute())
+                for x in sorted(Path().glob(f'{self.inputs.output_prefix}_*'))
+            ]
+
+            if self.inputs.order_transforms:
+                transforms = _order_xfms(transforms)
+            outputs['out_transforms'] = transforms
+
+            # Potentially could be more than one affine / displacement per composite transform...
+            outputs['affine_transform'] = [
+                x for x in transforms if 'AffineTransform' in Path(x).name
+            ][0]
+            outputs['displacement_field'] = [
+                x for x in transforms if 'DisplacementFieldTransform' in Path(x).name
+            ][0]
+        elif self.inputs.process == 'assemble':
+            outputs['out_file'] = Path(self.inputs.out_file).absolute()
+        return outputs
+
+
+def _order_xfms(vals):
+    """
+    Assumes [affine, displacement] or [displacement, affine] transform pairs.
+
+    >>> _order_xfms(['DisplacementFieldTransform.nii.gz', 'AffineTransform.mat'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz']
+
+    >>> _order_xfms(['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz']
+
+    >>> _order_xfms(['DisplacementFieldTransform.nii.gz', 'AffineTransform.mat', \
+        'AffineTransform.mat'])
+    ['AffineTransform.mat', 'DisplacementFieldTransform.nii.gz', 'AffineTransform.mat']
+    """
+    for i in range(0, len(vals) - 1, 2):
+        if (
+            'DisplacementFieldTransform' in Path(vals[i]).name
+            and 'AffineTransform' in Path(vals[i + 1]).name
+        ):
+            vals[i], vals[i + 1] = vals[i + 1], vals[i]
+    return vals

nibabies/workflows/anatomical/fit.py

@@ -991,7 +991,6 @@ def init_infant_anat_fit_wf(
                 ('anat2std_xfm', 'inputnode.anat2std_xfm'),
                 ('std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
-            (anat_buffer, concat_reg_wf, [('anat_preproc', 'inputnode.anat_preproc')]),
             (sourcefile_buffer, ds_concat_reg_wf, [
                 ('anat_source_files', 'inputnode.source_files')
             ]),
@@ -1909,7 +1908,6 @@ def init_infant_single_anat_fit_wf(
                 ('anat2std_xfm', 'inputnode.anat2std_xfm'),
                 ('std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
-            (anat_buffer, concat_reg_wf, [('anat_preproc', 'inputnode.anat_preproc')]),
             (sourcefile_buffer, ds_concat_reg_wf, [
                 ('anat_source_files', 'inputnode.source_files')
             ]),

nibabies/workflows/anatomical/registration.py

@@ -15,15 +15,11 @@
 from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
 from smriprep.workflows.fit.registration import (
     TemplateDesc,
-    TemplateFlowSelect,
     _fmt_cohort,
     get_metadata,
     tf_ver,
 )
 
-from nibabies.config import DEFAULT_MEMORY_MIN_GB
-from nibabies.interfaces.patches import ConcatXFM
-
 
 def init_coregistration_wf(
     *,
@@ -312,7 +308,7 @@ def init_concat_registrations_wf(
     name='concat_registrations_wf',
 ):
     """
-    Concatenate two transforms to produce a single transform, from native to ``template``.
+    Concatenate two transforms to produce a single composite transform from native to template.
 
     Parameters
     ----------
@@ -347,6 +343,8 @@ def init_concat_registrations_wf(
         further use in downstream nodes.
 
     """
+    from nibabies.interfaces.patches import CompositeTransformUtil
+
     ntpls = len(templates)
     workflow = Workflow(name=name)
 
@@ -384,9 +382,7 @@ def init_concat_registrations_wf(
             workflow.__desc__ += '.\n' if template == templates[-1] else ', '
 
     inputnode = pe.Node(
-        niu.IdentityInterface(
-            fields=['template', 'anat_preproc', 'anat2std_xfm', 'intermediate', 'std2anat_xfm']
-        ),
+        niu.IdentityInterface(fields=['template', 'intermediate', 'anat2std_xfm', 'std2anat_xfm']),
         name='inputnode',
     )
     inputnode.inputs.template = templates
@@ -413,29 +409,37 @@ def init_concat_registrations_wf(
         TemplateDesc(), run_without_submitting=True, iterfield='template', name='split_desc'
     )
 
-    tf_select = pe.MapNode(
-        TemplateFlowSelect(resolution=1),
-        name='tf_select',
-        run_without_submitting=True,
-        iterfield=['template', 'template_spec'],
+    merge_anat2std = pe.Node(niu.Merge(2), name='merge_anat2std', run_without_submitting=True)
+    merge_std2anat = merge_anat2std.clone('merge_std2anat')
+
+    disassemble_anat2std = pe.MapNode(
+        CompositeTransformUtil(process='disassemble', output_prefix='anat2std'),
+        iterfield=['in_file'],
+        name='disassemble_anat2std',
     )
 
-    merge_anat2std = pe.MapNode(
-        niu.Merge(2), name='merge_anat2std', iterfield=['in1', 'in2'], run_without_submitting=True
+    disassemble_std2anat = pe.MapNode(
+        CompositeTransformUtil(process='disassemble', output_prefix='std2anat'),
+        iterfield=['in_file'],
+        name='disassemble_std2anat',
     )
-    merge_std2anat = merge_anat2std.clone('merge_std2anat')
 
-    concat_anat2std = pe.MapNode(
-        ConcatXFM(),
-        name='concat_anat2std',
-        mem_gb=DEFAULT_MEMORY_MIN_GB,
-        iterfield=['transforms', 'reference_image'],
+    merge_anat2std_composites = pe.Node(
+        niu.Merge(1, ravel_inputs=True),
+        name='merge_anat2std_composites',
     )
-    concat_std2anat = pe.MapNode(
-        ConcatXFM(),
-        name='concat_std2anat',
-        mem_gb=DEFAULT_MEMORY_MIN_GB,
-        iterfield=['transforms', 'reference_image'],
+    merge_std2anat_composites = pe.Node(
+        niu.Merge(1, ravel_inputs=True),
+        name='merge_std2anat_composites',
+    )
+
+    assemble_anat2std = pe.Node(
+        CompositeTransformUtil(process='assemble', out_file='anat2std.h5'),
+        name='assemble_anat2std',
+    )
+    assemble_std2anat = pe.Node(
+        CompositeTransformUtil(process='assemble', out_file='std2anat.h5'),
+        name='assemble_std2anat',
     )
 
     fmt_cohort = pe.MapNode(
@@ -446,24 +450,24 @@ def init_concat_registrations_wf(
     )
 
     workflow.connect([
+        # Template concatenation
         (inputnode, merge_anat2std, [('anat2std_xfm', 'in2')]),
         (inputnode, merge_std2anat, [('std2anat_xfm', 'in2')]),
-        (inputnode, concat_std2anat, [('anat_preproc', 'reference_image')]),
         (inputnode, intermed_xfms, [('intermediate', 'intermediate')]),
         (inputnode, intermed_xfms, [('template', 'std')]),
-
         (intermed_xfms, merge_anat2std, [('int2std_xfm', 'in1')]),
         (intermed_xfms, merge_std2anat, [('std2int_xfm', 'in1')]),
-
-        (merge_anat2std, concat_anat2std, [('out', 'transforms')]),
-        (merge_std2anat, concat_std2anat, [('out', 'transforms')]),
-
+        (merge_anat2std, disassemble_anat2std, [('out', 'in_file')]),
+        (merge_std2anat, disassemble_std2anat, [('out', 'in_file')]),
+        (disassemble_anat2std, merge_anat2std_composites, [('out_transforms', 'in1')]),
+        (disassemble_std2anat, merge_std2anat_composites, [('out_transforms', 'in1')]),
+        (merge_anat2std_composites, assemble_anat2std, [('out', 'in_file')]),
+        (merge_std2anat_composites, assemble_std2anat, [('out', 'in_file')]),
+        (assemble_anat2std, outputnode, [('out_file', 'anat2std_xfm')]),
+        (assemble_std2anat, outputnode, [('out_file', 'std2anat_xfm')]),
+
+        # Template name wrangling
         (inputnode, split_desc, [('template', 'template')]),
-        (split_desc, tf_select, [
-            ('name', 'template'),
-            ('spec', 'template_spec'),
-        ]),
-        (tf_select, concat_anat2std, [('t1w_file', 'reference_image')]),
         (split_desc, fmt_cohort, [
             ('name', 'template'),
             ('spec', 'spec'),
@@ -472,8 +476,6 @@ def init_concat_registrations_wf(
             ('template', 'template'),
             ('spec', 'template_spec'),
         ]),
-        (concat_anat2std, outputnode, [('out_xfm', 'anat2std_xfm')]),
-        (concat_std2anat, outputnode, [('out_xfm', 'std2anat_xfm')]),
     ])  # fmt:skip
 
     return workflow

pyproject.toml

@@ -24,8 +24,7 @@ dependencies = [
     "nipype >= 1.8.5",
     "nireports >= 23.2.0",
     "nitime",
-#    "nitransforms >= 24.1.1",
-    "nitransforms @ git+https://github.com/nipy/nitransforms.git@enh/itk-displacement-field",
+    "nitransforms >= 24.1.1",
     "niworkflows >= 1.12.1",
     "numpy >= 1.21.0",
     "packaging",