ENH: Add workflow to verify derivatives are compatible with anatomical reference

Author: mgxd

nibabies/workflows/anatomical/fit.py

@@ -38,7 +38,11 @@
 from nibabies.interfaces import DerivativesDataSink
 from nibabies.workflows.anatomical.brain_extraction import init_infant_brain_extraction_wf
 from nibabies.workflows.anatomical.outputs import init_anat_reports_wf, init_coreg_report_wf
-from nibabies.workflows.anatomical.preproc import init_anat_preproc_wf, init_csf_norm_wf
+from nibabies.workflows.anatomical.preproc import (
+    init_anat_preproc_wf,
+    init_conform_derivative_wf,
+    init_csf_norm_wf,
+)
 from nibabies.workflows.anatomical.registration import (
     init_concat_registrations_wf,
     init_coregistration_wf,
@@ -172,11 +176,11 @@ def init_infant_anat_fit_wf(
 
     # Stage 2 - Anatomicals
     t1w_buffer = pe.Node(
-        niu.IdentityInterface(fields=['t1w_preproc', 't1w_maskt1w_brain']),
+        niu.IdentityInterface(fields=['t1w_preproc', 't1w_mask', 't1w_brain']),
         name='t1w_buffer',
     )
     t2w_buffer = pe.Node(
-        niu.IdentityInterface(fields=['t2w_preproc', 't2w_maskt2w_brain', 't2w_probmap']),
+        niu.IdentityInterface(fields=['t2w_preproc', 't2w_mask', 't2w_brain', 't2w_probmap']),
         name='t2w_buffer',
     )
     anat_buffer = pe.Node(
@@ -323,6 +327,7 @@ def init_infant_anat_fit_wf(
 
     t1w_preproc = precomputed.get('t1w_preproc')
     t2w_preproc = precomputed.get('t2w_preproc')
+    anat_preproc = precomputed.get(f'{anat}_preproc')
 
     # Stage 1: Conform & valid T1w/T2w images
     # Note: Since stage 1 & 2 are tightly knit together, it may be more intuitive
@@ -575,21 +580,28 @@ def init_infant_anat_fit_wf(
                 'A pre-computed T1w brain mask was provided as input and used throughout the '
                 'workflow.'
             )
-        t1w_buffer.inputs.t1w_mask = t1w_mask
         apply_t1w_mask.inputs.in_mask = t1w_mask
         workflow.connect(apply_t1w_mask, 'out_file', t1w_buffer, 't1w_brain')
 
         if not t1w_preproc:
+            # Ensure compatibility with T1w template
+            conform_t1w_mask_wf = init_conform_derivative_wf(
+                in_file=t1w_mask, name='conform_t1w_mask_wf'
+            )
+
             LOGGER.info('ANAT Skipping skull-strip, INU-correction only')
             t1w_n4_wf = init_anat_preproc_wf(name='t1w_n4_wf')
             workflow.connect([
+                (t1w_validate, conform_t1w_mask_wf, [('out_file', 'inputnode.ref_file')]),
+                (conform_t1w_mask_wf, t1w_buffer, [('outputnode.out_file', 't1w_mask')]),
                 (t1w_validate, t1w_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (t1w_n4_wf, t1w_buffer, [('outputnode.anat_preproc', 't1w_preproc')]),
                 (t1w_n4_wf, apply_t1w_mask, [('outputnode.anat_preproc', 'in_file')]),
             ])  # fmt:skip
         else:
             LOGGER.info('ANAT Skipping T1w masking')
             workflow.connect(t1w_validate, 'out_file', apply_t1w_mask, 'in_file')
+            t1w_buffer.inputs.t1w_mask = t1w_mask
 
     # T2w masking logic:
     #
@@ -701,21 +713,29 @@ def init_infant_anat_fit_wf(
                 'A pre-computed T2w brain mask was provided as input and used throughout the '
                 'workflow.'
             )
-        t2w_buffer.inputs.t2w_mask = t2w_mask
         apply_t2w_mask.inputs.in_mask = t2w_mask
         workflow.connect(apply_t2w_mask, 'out_file', t2w_buffer, 't2w_brain')
 
         if not t2w_preproc:
+            # Ensure compatibility with T2w template
+            conform_t2w_mask_wf = init_conform_derivative_wf(
+                in_file=t2w_mask,
+                name='conform_t2w_mask_wf',
+            )
+
             LOGGER.info('ANAT Skipping skull-strip, INU-correction only')
             t2w_n4_wf = init_anat_preproc_wf(name='t2w_n4_wf')
             workflow.connect([
+                (t2w_validate, conform_t2w_mask_wf, [('out_file', 'inputnode.ref_file')]),
+                (conform_t2w_mask_wf, t2w_buffer, [('outputnode.out_file', 't2w_mask')]),
                 (t2w_validate, t2w_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (t2w_n4_wf, t2w_buffer, [('outputnode.anat_preproc', 't2w_preproc')]),
                 (t2w_n4_wf, apply_t2w_mask, [('outputnode.anat_preproc', 'in_file')]),
             ])  # fmt:skip
         else:
             LOGGER.info('ANAT Skipping T2w masking')
             workflow.connect(t2w_validate, 'out_file', apply_t2w_mask, 'in_file')
+            t2w_buffer.inputs.t2w_mask = t2w_mask
 
     # Stage 3: Coregistration
     t1w2t2w_xfm = precomputed.get('t1w2t2w_xfm')
@@ -819,7 +839,19 @@ def init_infant_anat_fit_wf(
 
     if anat_aseg:
         LOGGER.info('ANAT Found precomputed anatomical segmentation')
-        aseg_buffer.inputs.anat_aseg = anat_aseg
+        # Ensure compatibility with anatomical template
+        if not anat_preproc:
+            conform_aseg_wf = init_conform_derivative_wf(
+                in_file=anat_aseg,
+                name='conform_aseg_wf',
+            )
+
+            workflow.connect([
+                (anat_buffer, conform_aseg_wf, [('anat_preproc', 'inputnode.ref_file')]),
+                (conform_aseg_wf, aseg_buffer, [('outputnode.out_file', 'anat_aseg')]),
+            ])  # fmt:skip
+        else:
+            aseg_buffer.inputs.anat_aseg = anat_aseg
 
     if not (anat_dseg and anat_tpms):
         LOGGER.info('ANAT Stage 4: Tissue segmentation')
@@ -1714,27 +1746,46 @@ def init_infant_single_anat_fit_wf(
     else:
         LOGGER.info(f'ANAT Found {reference_anat} brain mask')
         desc += 'A pre-computed brain mask was provided as input and used throughout the workflow.'
-        anat_buffer.inputs.anat_mask = anat_mask
         apply_mask.inputs.in_mask = anat_mask
         workflow.connect(apply_mask, 'out_file', anat_buffer, 'anat_brain')
 
         if not anat_preproc:
+            conform_anat_mask_wf = init_conform_derivative_wf(
+                in_file=anat_mask,
+                name='conform_anat_mask_wf',
+            )
+
             LOGGER.info('ANAT Skipping skull-strip, INU-correction only')
             anat_n4_wf = init_anat_preproc_wf(name='anat_n4_wf')
             workflow.connect([
+                (anat_validate, conform_anat_mask_wf, [('out_file', 'inputnode.ref_file')]),
+                (conform_anat_mask_wf, anat_buffer, [('outputnode.out_file', 'anat_mask')]),
                 (anat_validate, anat_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (anat_n4_wf, anat_buffer, [('outputnode.anat_preproc', 'anat_preproc')]),
                 (anat_n4_wf, apply_mask, [('outputnode.anat_preproc', 'in_file')]),
             ])  # fmt:skip
         else:
             LOGGER.info(f'ANAT Skipping {reference_anat} masking')
             workflow.connect(anat_validate, 'out_file', apply_mask, 'in_file')
+            anat_buffer.inputs.anat_mask = anat_mask
 
     # Stage 3: Segmentation
     seg_method = 'jlf' if config.execution.segmentation_atlases_dir else 'fast'
     if anat_aseg:
         LOGGER.info('ANAT Found precomputed anatomical segmentation')
-        aseg_buffer.inputs.anat_aseg = anat_aseg
+        # Ensure compatibility with anatomical template
+        if not anat_preproc:
+            conform_aseg_wf = init_conform_derivative_wf(
+                in_file=anat_aseg,
+                name='conform_aseg_wf',
+            )
+
+            workflow.connect([
+                (anat_buffer, conform_aseg_wf, [('anat_preproc', 'inputnode.ref_file')]),
+                (conform_aseg_wf, aseg_buffer, [('outputnode.out_file', 'anat_aseg')]),
+            ])  # fmt:skip
+        else:
+            aseg_buffer.inputs.anat_aseg = anat_aseg
 
     if not (anat_dseg and anat_tpms):
         LOGGER.info('ANAT Stage 3: Tissue segmentation')

nibabies/workflows/anatomical/preproc.py

@@ -98,6 +98,41 @@ def init_csf_norm_wf(name: str = 'csf_norm_wf') -> Workflow:
     return workflow
 
 
+def init_conform_derivative_wf(
+    *, in_file: str = None, name: str = 'conform_derivative_wf'
+) -> pe.Workflow:
+    """
+    Ensure derivatives share the same space as anatomical references.
+
+    This workflow is used when a derivative is provided without a reference.
+    """
+    from niworkflows.interfaces.header import MatchHeader
+    from niworkflows.interfaces.images import Conform, TemplateDimensions
+
+    workflow = pe.Workflow(name=name)
+    inputnode = pe.Node(niu.IdentityInterface(fields=['in_file', 'ref_file']), name='inputnode')
+    inputnode.inputs.in_file = in_file
+    outputnode = pe.Node(niu.IdentityInterface(fields=['out_file']), name='outputnode')
+
+    ref_dims = pe.Node(TemplateDimensions(), name='ref_dims')
+    conform = pe.Node(Conform(), name='conform')
+    # Avoid mismatch tolerance from input
+    match_header = pe.Node(MatchHeader(), name='match_header')
+
+    workflow.connect([
+        (inputnode, ref_dims, [('ref_file', 'anat_list')]),
+        (ref_dims, conform, [
+            ('target_zooms', 'target_zooms'),
+            ('target_shape', 'target_shape'),
+        ]),
+        (inputnode, conform, [('in_file', 'in_file')]),
+        (conform, match_header, [('out_file', 'in_file')]),
+        (inputnode, match_header, [('ref_file', 'reference')]),
+        (match_header, outputnode, [('out_file', 'out_file')]),
+    ])  # fmt:skip
+    return workflow
+
+
 def _normalize_roi(in_file, mask_file, threshold=0.2, out_file=None):
     """Normalize low intensity voxels that fall within a given mask."""
     import nibabel as nb

nibabies/workflows/anatomical/surfaces.py

@@ -15,7 +15,6 @@
 from niworkflows.interfaces.freesurfer import (
     PatchedRobustRegister as RobustRegister,
 )
-from niworkflows.interfaces.header import MatchHeader
 from niworkflows.interfaces.morphology import BinaryDilation
 from niworkflows.interfaces.patches import FreeSurferSource
 from smriprep.interfaces.freesurfer import MakeMidthickness
@@ -128,9 +127,6 @@ def init_mcribs_surface_recon_wf(
     mask_dil = pe.Node(BinaryDilation(radius=3), name='mask_dil')
     mask_las = pe.Node(ReorientImage(target_orientation='LAS'), name='mask_las')
 
-    # N4 has low tolerance for mismatch between input / mask
-    match_header = pe.Node(MatchHeader(), name='match_header')
-
     # N4BiasCorrection occurs in MCRIBTissueSegMCRIBS (which is skipped)
     # Run it (with mask to rescale intensities) prior injection
     n4_mcribs = pe.Node(
@@ -182,9 +178,7 @@ def init_mcribs_surface_recon_wf(
             ('subjects_dir', 'subjects_dir'),
             ('subject_id', 'subject_id')]),
         (t2w_las, n4_mcribs, [('out_file', 'input_image')]),
-        (mask_las, match_header, [('out_file', 'in_file')]),
-        (t2w_las, match_header, [('out_file', 'reference')]),
-        (match_header, n4_mcribs, [('out_file', 'mask_image')]),
+        (mask_las, n4_mcribs, [('out_file', 'mask_image')]),
         (n4_mcribs, mcribs_recon, [('output_image', 't2w_file')]),
         (seg_las, mcribs_recon, [('out_file', 'segmentation_file')]),
         (inputnode, mcribs_postrecon, [