ENH: Use 2 step registration for adult templates

nibabies/interfaces/patches.py

@@ -1,5 +1,10 @@
-import nipype.interfaces.freesurfer as fs
-from nipype.interfaces.base import File, traits
+from pathlib import Path
+
+from nipype.interfaces import (
+    freesurfer as fs,
+)
+from nipype.interfaces.ants.base import ANTSCommand, ANTSCommandInputSpec
+from nipype.interfaces.base import File, InputMultiObject, TraitedSpec, traits
 
 
 class _MRICoregInputSpec(fs.registration.MRICoregInputSpec):
@@ -23,3 +28,80 @@ class MRICoreg(fs.MRICoreg):
     """
 
     input_spec = _MRICoregInputSpec
+
+
+class ConcatXFMInputSpec(ANTSCommandInputSpec):
+    transforms = InputMultiObject(
+        traits.Either(File(exists=True), 'identity'),
+        argstr='%s',
+        mandatory=True,
+        desc='transform files: will be applied in reverse order. For '
+        'example, the last specified transform will be applied first.',
+    )
+    out_xfm = traits.File(
+        'concat_xfm.h5',
+        usedefault=True,
+        argstr='--output [ %s, 1 ]',
+        desc='output file name',
+    )
+    reference_image = File(
+        argstr='--reference-image %s',
+        mandatory=True,
+        desc='reference image space that you wish to warp INTO',
+        exists=True,
+    )
+    invert_transform_flags = InputMultiObject(traits.Bool())
+
+
+class ConcatXFMOutputSpec(TraitedSpec):
+    out_xfm = File(desc='Combined transform')
+
+
+class ConcatXFM(ANTSCommand):
+    """
+    Streamed use of antsApplyTransforms to combine nonlinear xfms into a single file
+
+    Examples
+    --------
+
+    >>> from nibabies.interfaces.patches import ConcatXFM
+    >>> cxfm = ConcatXFM()
+    >>> cxfm.inputs.transforms = ['xfm1.h5', 'xfm0.h5']
+    >>> cxfm.inputs.reference_image = 'sub-01_T1w.nii.gz'
+    >>> cxfm.cmdline
+    'antsApplyTransforms --output [ concat_xfm.h5, 1 ] --transform .../xfm1.h5 \
+--transform .../xfm0.h5 --reference_image .../sub-01_T1w.nii.gz'
+
+    """
+
+    _cmd = 'antsApplyTransforms'
+    input_spec = ConcatXFMInputSpec
+    output_spec = ConcatXFMOutputSpec
+
+    def _get_transform_filenames(self):
+        retval = []
+        invert_flags = self.inputs.invert_transform_flags
+        if not invert_flags:
+            invert_flags = [False] * len(self.inputs.transforms)
+        elif len(self.inputs.transforms) != len(invert_flags):
+            raise ValueError(
+                'ERROR: The invert_transform_flags list must have the same number '
+                'of entries as the transforms list.'
+            )
+
+        for transform, invert in zip(self.inputs.transforms, invert_flags, strict=False):
+            if invert:
+                retval.append(f'--transform [ {transform}, 1 ]')
+            else:
+                retval.append(f'--transform {transform}')
+        return ' '.join(retval)
+
+    def _format_arg(self, opt, spec, val):
+        if opt == 'transforms':
+            return self._get_transform_filenames()
+        return super()._format_arg(opt, spec, val)
+
+    def _list_outputs(self):
+        outputs = self._outputs().get()
+        outputs['out_xfm'] = Path(self.inputs.out_xfm).absolute()
+        return outputs

nibabies/workflows/anatomical/fit.py

@@ -9,6 +9,7 @@
 from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
 from niworkflows.interfaces.header import ValidateImage
 from niworkflows.interfaces.nibabel import ApplyMask, Binarize
+from niworkflows.interfaces.utility import KeySelect
 from niworkflows.utils.connections import pop_file
 from smriprep.workflows.anatomical import (
     _is_skull_stripped,
@@ -37,7 +38,10 @@
 from nibabies.workflows.anatomical.brain_extraction import init_infant_brain_extraction_wf
 from nibabies.workflows.anatomical.outputs import init_anat_reports_wf
 from nibabies.workflows.anatomical.preproc import init_anat_preproc_wf
-from nibabies.workflows.anatomical.registration import init_coregistration_wf
+from nibabies.workflows.anatomical.registration import (
+    init_concat_registrations_wf,
+    init_coregistration_wf,
+)
 from nibabies.workflows.anatomical.segmentation import init_segmentation_wf
 from nibabies.workflows.anatomical.surfaces import init_mcribs_dhcp_wf
 
@@ -220,9 +224,9 @@ def init_infant_anat_fit_wf(
         name='seg_buffer',
     )
     # Stage 5 - collated template names, forward and reverse transforms
-    template_buffer = pe.Node(niu.Merge(2), name='template_buffer')
-    anat2std_buffer = pe.Node(niu.Merge(2), name='anat2std_buffer')
-    std2anat_buffer = pe.Node(niu.Merge(2), name='std2anat_buffer')
+    template_buffer = pe.Node(niu.Merge(3), name='template_buffer')
+    anat2std_buffer = pe.Node(niu.Merge(3), name='anat2std_buffer')
+    std2anat_buffer = pe.Node(niu.Merge(3), name='std2anat_buffer')
 
     # Stage 6 results: Refined stage 2 results; may be direct copy if no refinement
     refined_buffer = pe.Node(
@@ -885,21 +889,48 @@ def init_infant_anat_fit_wf(
         seg_buffer.inputs.anat_tpms = anat_tpms
 
     # Stage 5: Normalization
+
+    # If selected adult templates are requested (MNI152 6th Gen or 2009)
+    # opt to concatenate transforms first from native -> infant template (MNIInfant),
+    # and then use a previously computed MNIInfant<cohort> -> MNI transform
+    # this minimizes the chance of a bad registration.
+
     templates = []
+    concat_xfms = []
     found_xfms = {}
+    intermediate = None  # The intermediate space when concatenating xfms - includes cohort
+    intermediate_targets = {
+        'MNI152NLin6Asym',
+    }  # TODO: 'MNI152NLin2009cAsym'
+
     for template in spaces.get_spaces(nonstandard=False, dim=(3,)):
+        # resolution / spec will not differentiate here
+        if template.startswith('MNIInfant'):
+            intermediate = template
         xfms = precomputed.get('transforms', {}).get(template, {})
         if set(xfms) != {'forward', 'reverse'}:
-            templates.append(template)
+            if template in intermediate_targets:
+                concat_xfms.append(template)
+            else:
+                templates.append(template)
         else:
             found_xfms[template] = xfms
 
+    # Create another set of buffers to handle the case where we aggregate found and generated
+    # xfms to be concatenated
+    concat_template_buffer = pe.Node(niu.Merge(2), name='concat_template_buffer')
+    concat_template_buffer.inputs.in1 = list(found_xfms)
+    concat_anat2std_buffer = pe.Node(niu.Merge(2), name='concat_anat2std_buffer')
+    concat_anat2std_buffer.inputs.in1 = [xfm['forward'] for xfm in found_xfms.values()]
+    concat_std2anat_buffer = pe.Node(niu.Merge(2), name='concat_std2anat_buffer')
+    concat_std2anat_buffer.inputs.in1 = [xfm['reverse'] for xfm in found_xfms.values()]
+
     template_buffer.inputs.in1 = list(found_xfms)
     anat2std_buffer.inputs.in1 = [xfm['forward'] for xfm in found_xfms.values()]
     std2anat_buffer.inputs.in1 = [xfm['reverse'] for xfm in found_xfms.values()]
 
     if templates:
-        LOGGER.info(f'ANAT Stage 5: Preparing normalization workflow for {templates}')
+        LOGGER.info(f'ANAT Stage 5a: Preparing normalization workflow for {templates}')
         register_template_wf = init_register_template_wf(
             sloppy=sloppy,
             omp_nthreads=omp_nthreads,
@@ -923,11 +954,53 @@ def init_infant_anat_fit_wf(
                 ('outputnode.std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
             (register_template_wf, template_buffer, [('outputnode.template', 'in2')]),
+            (register_template_wf, concat_template_buffer, [('outputnode.template', 'in2')]),
             (register_template_wf, std2anat_buffer, [('outputnode.std2anat_xfm', 'in2')]),
+            (register_template_wf, concat_std2anat_buffer, [('outputnode.std2anat_xfm', 'in2')]),
             (register_template_wf, anat2std_buffer, [('outputnode.anat2std_xfm', 'in2')]),
+            (register_template_wf, concat_anat2std_buffer, [('outputnode.anat2std_xfm', 'in2')]),
         ])  # fmt:skip
+
+    if concat_xfms:
+        LOGGER.info(f'ANAT Stage 5b: Concatenating normalization for {concat_xfms}')
+        # 1. Select intermediate's transforms
+        select_infant_mni = pe.Node(
+            KeySelect(fields=['template', 'anat2std_xfm', 'std2anat_xfm'], key=intermediate),
+            name='select_infant_mni',
+            run_without_submitting=True,
+        )
+        concat_reg_wf = init_concat_registrations_wf(templates=concat_xfms)
+        ds_concat_reg_wf = init_ds_template_registration_wf(
+            output_dir=str(output_dir),
+            image_type=reference_anat,
+            name='ds_concat_registration_wf',
+        )
+
+        workflow.connect([
+            (concat_template_buffer, select_infant_mni, [('out', 'template')]),
+            (concat_anat2std_buffer, select_infant_mni, [('out', 'anat2std_xfm')]),
+            (concat_std2anat_buffer, select_infant_mni, [('out', 'std2anat_xfm')]),
+            (select_infant_mni, concat_reg_wf, [
+                ('template', 'inputnode.intermediate'),
+                ('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')
+            ]),
+            (concat_reg_wf, ds_concat_reg_wf, [
+                ('outputnode.template', 'inputnode.template'),
+                ('outputnode.anat2std_xfm', 'inputnode.anat2std_xfm'),
+                ('outputnode.std2anat_xfm', 'inputnode.std2anat_xfm'),
+            ]),
+            (concat_reg_wf, template_buffer, [('outputnode.template', 'in3')]),
+            (concat_reg_wf, anat2std_buffer, [('outputnode.anat2std_xfm', 'in3')]),
+            (concat_reg_wf, std2anat_buffer, [('outputnode.std2anat_xfm', 'in3')]),
+        ])  # fmt:skip
+
     if found_xfms:
-        LOGGER.info(f'ANAT Stage 5: Found pre-computed registrations for {found_xfms}')
+        LOGGER.info(f'ANAT Stage 5c: Found pre-computed registrations for {found_xfms}')
 
     # Only refine mask if necessary
     if anat_mask or recon_method is None or not refine_mask:
@@ -1394,9 +1467,9 @@ def init_infant_single_anat_fit_wf(
         name='seg_buffer',
     )
     # Stage 4 - collated template names, forward and reverse transforms
-    template_buffer = pe.Node(niu.Merge(2), name='template_buffer')
-    anat2std_buffer = pe.Node(niu.Merge(2), name='anat2std_buffer')
-    std2anat_buffer = pe.Node(niu.Merge(2), name='std2anat_buffer')
+    template_buffer = pe.Node(niu.Merge(3), name='template_buffer')
+    anat2std_buffer = pe.Node(niu.Merge(3), name='anat2std_buffer')
+    std2anat_buffer = pe.Node(niu.Merge(3), name='std2anat_buffer')
 
     # Stage 5 results: Refined stage 2 results; may be direct copy if no refinement
     refined_buffer = pe.Node(
@@ -1710,15 +1783,41 @@ def init_infant_single_anat_fit_wf(
         seg_buffer.inputs.anat_tpms = anat_tpms
 
     # Stage 4: Normalization
+    # If selected adult templates are requested (MNI152 6th Gen or 2009)
+    # opt to concatenate transforms first from native -> infant template (MNIInfant),
+    # and then use a previously computed MNIInfant<cohort> -> MNI transform
+    # this minimizes the chance of a bad registration.
+
     templates = []
+    concat_xfms = []
     found_xfms = {}
+    intermediate = None  # The intermediate space when concatenating xfms - includes cohort
+    intermediate_targets = {
+        'MNI152NLin6Asym',
+    }  # TODO: 'MNI152NLin2009cAsym'
+
     for template in spaces.get_spaces(nonstandard=False, dim=(3,)):
+        # resolution / spec will not differentiate here
+        if template.startswith('MNIInfant'):
+            intermediate = template
         xfms = precomputed.get('transforms', {}).get(template, {})
         if set(xfms) != {'forward', 'reverse'}:
-            templates.append(template)
+            if template in intermediate_targets:
+                concat_xfms.append(template)
+            else:
+                templates.append(template)
         else:
             found_xfms[template] = xfms
 
+    # Create another set of buffers to handle the case where we aggregate found and generated
+    # xfms to be concatenated
+    concat_template_buffer = pe.Node(niu.Merge(2), name='concat_template_buffer')
+    concat_template_buffer.inputs.in1 = list(found_xfms)
+    concat_anat2std_buffer = pe.Node(niu.Merge(2), name='concat_anat2std_buffer')
+    concat_anat2std_buffer.inputs.in1 = [xfm['forward'] for xfm in found_xfms.values()]
+    concat_std2anat_buffer = pe.Node(niu.Merge(2), name='concat_std2anat_buffer')
+    concat_std2anat_buffer.inputs.in1 = [xfm['reverse'] for xfm in found_xfms.values()]
+
     template_buffer.inputs.in1 = list(found_xfms)
     anat2std_buffer.inputs.in1 = [xfm['forward'] for xfm in found_xfms.values()]
     std2anat_buffer.inputs.in1 = [xfm['reverse'] for xfm in found_xfms.values()]
@@ -1748,9 +1847,54 @@ def init_infant_single_anat_fit_wf(
                 ('outputnode.std2anat_xfm', 'inputnode.std2anat_xfm'),
             ]),
             (register_template_wf, template_buffer, [('outputnode.template', 'in2')]),
+            (register_template_wf, concat_template_buffer, [('outputnode.template', 'in2')]),
             (register_template_wf, std2anat_buffer, [('outputnode.std2anat_xfm', 'in2')]),
+            (register_template_wf, concat_std2anat_buffer, [('outputnode.std2anat_xfm', 'in2')]),
             (register_template_wf, anat2std_buffer, [('outputnode.anat2std_xfm', 'in2')]),
+            (register_template_wf, concat_anat2std_buffer, [('outputnode.anat2std_xfm', 'in2')]),
+        ])  # fmt:skip
+
+    if concat_xfms:
+        LOGGER.info(f'ANAT Stage 5b: Concatenating normalization for {concat_xfms}')
+        # 1. Select intermediate's transforms
+        select_infant_mni = pe.Node(
+            KeySelect(fields=['template', 'anat2std_xfm', 'std2anat_xfm']),
+            name='select_infant_mni',
+            run_without_submitting=True,
+        )
+        select_infant_mni.inputs.key = intermediate
+
+        concat_reg_wf = init_concat_registrations_wf(templates=concat_xfms)
+        ds_concat_reg_wf = init_ds_template_registration_wf(
+            output_dir=str(output_dir),
+            image_type=reference_anat,
+            name='ds_concat_registration_wf',
+        )
+
+        workflow.connect([
+            (concat_template_buffer, select_infant_mni, [('out', 'keys')]),
+            (concat_template_buffer, select_infant_mni, [('out', 'template')]),
+            (concat_anat2std_buffer, select_infant_mni, [('out', 'anat2std_xfm')]),
+            (concat_std2anat_buffer, select_infant_mni, [('out', 'std2anat_xfm')]),
+            (select_infant_mni, concat_reg_wf, [
+                ('template', 'inputnode.intermediate'),
+                ('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')
+            ]),
+            (concat_reg_wf, ds_concat_reg_wf, [
+                ('outputnode.template', 'inputnode.template'),
+                ('outputnode.anat2std_xfm', 'inputnode.anat2std_xfm'),
+                ('outputnode.std2anat_xfm', 'inputnode.std2anat_xfm'),
+            ]),
+            (concat_reg_wf, template_buffer, [('outputnode.template', 'in3')]),
+            (concat_reg_wf, anat2std_buffer, [('outputnode.anat2std_xfm', 'in3')]),
+            (concat_reg_wf, std2anat_buffer, [('outputnode.std2anat_xfm', 'in3')]),
         ])  # fmt:skip
+
     if found_xfms:
         LOGGER.info(f'ANAT Stage 4: Found pre-computed registrations for {found_xfms}')