Merge pull request #459 from mgxd/fix/preproc-hdr

ENH: 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,29 @@ 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')]),
+                (conform_t1w_mask_wf, apply_t1w_mask, [('outputnode.out_file', 'in_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
+            apply_t1w_mask.inputs.in_mask = t1w_mask
 
     # T2w masking logic:
     #
@@ -701,21 +714,30 @@ 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')]),
+                (conform_t2w_mask_wf, apply_t2w_mask, [('outputnode.out_file', 'in_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
+            apply_t2w_mask.inputs.in_mask = t2w_mask
 
     # Stage 3: Coregistration
     t1w2t2w_xfm = precomputed.get('t1w2t2w_xfm')
@@ -819,7 +841,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 +1748,47 @@ 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')]),
+                (conform_anat_mask_wf, apply_mask, [('outputnode.out_file', 'in_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
+            apply_mask.inputs.in_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, [

nibabies/workflows/anatomical/tests/test_preproc.py

@@ -1,55 +1,93 @@
-import typing as ty
 from pathlib import Path
 
 import nibabel as nb
 import numpy as np
 import pytest
+from nibabel.orientations import aff2axcodes
 
-from nibabies.workflows.anatomical.preproc import _normalize_roi, init_csf_norm_wf
+from nibabies.workflows.anatomical.preproc import (
+    _normalize_roi,
+    init_conform_derivative_wf,
+    init_csf_norm_wf,
+)
 
-EXPECTED_CSF_NORM = np.array([[[10, 73], [73, 29]], [[77, 80], [6, 16]]], dtype='uint8')
+EXPECTED_CSF_NORM = np.array([[[49, 75], [23, 75]], [[77, 80], [33, 3]]], dtype='uint8')
 
 
 @pytest.fixture
-def csf_norm_data(tmp_path) -> ty.Generator[tuple[Path, list[Path]], None, None]:
-    np.random.seed(10)
-
-    in_file = tmp_path / 'input.nii.gz'
-    data = np.random.randint(1, 101, size=(2, 2, 2), dtype='uint8')
+def anat_file(tmp_path):
+    data = np.array([[[49, 73], [23, 73]], [[77, 80], [33, 3]]], dtype='uint8')
     img = nb.Nifti1Image(data, np.eye(4))
-    img.to_filename(in_file)
+    out = tmp_path / 'input.nii.gz'
+    img.to_filename(out)
+    return out
+
 
-    masks = []
-    for tpm in ('gm', 'wm', 'csf'):
+def test_csf_norm_wf(tmp_path, anat_file):
+    tpms = []
+    for tpm, thresh in (('gm', 25), ('wm', 75), ('csf', 50)):
         name = tmp_path / f'{tpm}.nii.gz'
-        binmask = data > np.random.randint(10, 90)
+        anat_img = nb.load(anat_file)
+        anat_data = np.asanyarray(nb.load(anat_file).dataobj)
+
+        binmask = anat_data > thresh
         masked = (binmask * 1).astype('uint8')
-        mask = nb.Nifti1Image(masked, img.affine)
+        mask = nb.Nifti1Image(masked, anat_img.affine)
         mask.to_filename(name)
-        masks.append(name)
-
-    yield in_file, masks
+        tpms.append(name)
 
-    in_file.unlink()
-    for m in masks:
-        m.unlink()
-
-
-def test_csf_norm_wf(tmp_path, csf_norm_data):
-    anat, tpms = csf_norm_data
     wf = init_csf_norm_wf()
     wf.base_dir = tmp_path
-
-    wf.inputs.inputnode.anat_preproc = anat
+    wf.inputs.inputnode.anat_preproc = anat_file
     wf.inputs.inputnode.anat_tpms = tpms
 
     # verify workflow runs
     wf.run()
 
     # verify function works as expected
-    outfile = _normalize_roi(anat, tpms[2])
+    outfile = _normalize_roi(anat_file, tpms[2])
     assert np.array_equal(
         np.asanyarray(nb.load(outfile).dataobj),
         EXPECTED_CSF_NORM,
     )
     Path(outfile).unlink()
+
+
+@pytest.mark.parametrize(
+    ('affine_mismatch', 'ornt_mismatch'),
+    [
+        (False, False),
+        (True, False),
+        (False, True),
+        (True, True),
+    ],
+)
+def test_conform_derivative_wf(tmp_path, anat_file, affine_mismatch, ornt_mismatch):
+    deriv = tmp_path / 'mask.nii.gz'
+    ref_img = nb.load(anat_file)
+    aff = ref_img.affine.copy()
+    if affine_mismatch:
+        # Alter affine slightly
+        aff[:3, :3] += 0.01
+        assert not np.array_equal(aff, ref_img.affine)
+
+    img = ref_img.__class__(ref_img.dataobj, affine=aff)
+    if ornt_mismatch:
+        from niworkflows.interfaces.nibabel import reorient_image
+
+        img = reorient_image(img, target_ornt='LPI')
+        assert aff2axcodes(img.affine) != aff2axcodes(ref_img.affine)
+
+    img.to_filename(deriv)
+    wf = init_conform_derivative_wf(in_file=deriv)
+    wf.base_dir = tmp_path
+    wf.inputs.inputnode.ref_file = anat_file
+
+    wf.run()
+
+    output = list((tmp_path / 'conform_derivative_wf' / 'match_header').glob('*.nii.gz'))
+    assert output
+    out_file = output[0]
+    out_img = nb.load(out_file)
+    assert np.array_equal(out_img.affine, ref_img.affine)
+    assert aff2axcodes(out_img.affine) == aff2axcodes(ref_img.affine)