Merge pull request #102 from mgxd/enh/cifti-rework

ENH: Incorporate subcortical CIFTI alignment to functional processing

Author: mgxd

nibabies/data/FreeSurferSubcorticalLabelTableLut.txt

@@ -0,0 +1,38 @@
+ACCUMBENS_LEFT
+26 255 165 0 255
+ACCUMBENS_RIGHT
+58 255 165 0 255
+AMYGDALA_LEFT
+18 103 255 255 255
+AMYGDALA_RIGHT
+54 103 255 255 255
+BRAIN_STEM
+16 119 159 176 255
+CAUDATE_LEFT
+11 122 186 220 255
+CAUDATE_RIGHT
+50 122 186 220 255
+CEREBELLUM_LEFT
+8 230 148 34 255
+CEREBELLUM_RIGHT
+47 230 148 34 255
+DIENCEPHALON_VENTRAL_LEFT
+28 165 42 42 255
+DIENCEPHALON_VENTRAL_RIGHT
+60 165 42 42 255
+HIPPOCAMPUS_LEFT
+17 220 216 20 255
+HIPPOCAMPUS_RIGHT
+53 220 216 20 255
+PALLIDUM_LEFT
+13 12 48 255 255
+PALLIDUM_RIGHT
+52 13 48 255 255
+PUTAMEN_LEFT
+12 236 13 176 255
+PUTAMEN_RIGHT
+51 236 13 176 255
+THALAMUS_LEFT
+10 0 118 14 255
+THALAMUS_RIGHT
+49 0 118 14 255

nibabies/data/MNIInfant_to_MNI1526NLinAsym.mat

@@ -0,0 +1,4 @@
+1.330660649  -0.004427136451  0.01025602509  -41.20370054
+-0.003753457765  1.339140797  0.153721562  -68.37553187
+-0.01546533827  -0.05324436952  1.426601839  -22.76900701
+0  0  0  1

nibabies/data/tpl-MNI152NLin6Asym_res-01_desc-avgwmparc_dseg.nii.gz

@@ -1377,8 +1377,6 @@ class VolumeLabelImportInputSpec(CommandLineInputSpec):
         desc="set any voxels with values not mentioned in the label list to the ??? label",
     )
     unlabeled_values = traits.Int(
-        0,
-        usedefault=True,
         argstr="-unlabeled-value %d",
         desc="the value that will be interpreted as unlabeled",
     )
@@ -1433,7 +1431,7 @@ class VolumeLabelImport(WBCommand):
     >>> label_import.inputs.label_list_file = data_dir / 'label_list.txt'
     >>> label_import.cmdline  #doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
     'wb_command -volume-label-import .../atlas.nii .../label_list.txt \
-    atlas_labels.nii.gz -unlabeled-value 0'
+    atlas_labels.nii.gz'
     """
 
     input_spec = VolumeLabelImportInputSpec

nibabies/interfaces/workbench.py

@@ -2,7 +2,86 @@
 Subcortical alignment into MNI space
 """
 
-from nibabies.interfaces.nibabel import MergeROIs
+from nipype.interfaces import utility as niu, fsl
+from nipype.pipeline import engine as pe
+from niworkflows.engine.workflows import LiterateWorkflow as Workflow
+
+from ...interfaces.workbench import VolumeLabelImport
+
+from pkg_resources import resource_filename
+
+
+def init_subcortical_rois_wf(*, name="subcortical_rois_wf"):
+    """
+    Refine segmentations into volumes of expected CIFTI subcortical structures.
+
+
+    Parameters
+    ----------
+    name : :obj:`str`
+        Name of the workflow
+
+    Inputs
+    ------
+    MNIInfant_aseg : :obj:`str`
+        FreeSurfer's aseg in MNIInfant space
+
+    Outputs
+    -------
+    MNIInfant_rois : :obj:`str`
+        Subcortical ROIs in MNIInfant space
+    MNI152_rois : :obj:`str`
+        Subcortical ROIs in `MNI152NLin6Asym` space
+    """
+    from templateflow.api import get as get_template
+
+    # TODO: Implement BOLD refinement once InfantFS outputs subj/mri/wmparc.mgz
+    # The code is found at
+    # https://github.com/DCAN-Labs/dcan-infant-pipeline/blob/
+    # 0e9c2fe32fb4a5032d0a2a3e0905ad97fa52b398/PostFreeSurfer/scripts/
+    # FreeSurfer2CaretConvertAndRegisterNonlinear.sh
+    # Lines 70-78 & 116-127
+    # #
+    # For now, just use the aseg
+
+    workflow = Workflow(name=name)
+    inputnode = pe.Node(niu.IdentityInterface(fields=["MNIInfant_aseg"]), name='inputnode')
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=["MNIInfant_rois", "MNI152_rois"]),
+        name='outputnode',
+    )
+    # Fetch the HCP volumetric template
+    tpl_rois = get_template(
+        'MNI152NLin6Asym', resolution=2, atlas="HCP", suffix="dseg", raise_empty=True
+    )
+    outputnode.inputs.MNI152_rois = tpl_rois
+
+    # This will only used for the wmparc in subject space
+    # For now, define it and don't run it
+    # TODO: Move to TemplateFlow
+
+    # tpl_avgwmparc = resource_filename(
+    #     'nibabies', 'data/tpl-MNI152NLin6Asym_res-01_desc-avgwmparc_dseg.nii.gz'
+    # )
+    # applywarp_tpl = pe.Node(
+    #     fsl.ApplyWarp(in_file=tpl_avgwmparc, ref_file=tpl_rois, interp="nn"),
+    #     name="applywarp_std"
+    # )
+
+    subcortical_labels = resource_filename(
+        'nibabies', 'data/FreeSurferSubcorticalLabelTableLut.txt'
+    )
+    refine_bold_rois = pe.Node(
+        VolumeLabelImport(label_list_file=subcortical_labels, discard_others=True),
+        name="refine_bold_rois",
+    )
+
+    workflow.connect([
+        (inputnode, refine_bold_rois, [("MNIInfant_aseg", "in_file")]),
+        # (applywarp_tpl, refine_std_rois, [("out_file", "in_file")]),
+        (refine_bold_rois, outputnode, [("out_file", "MNIInfant_rois")]),
+    ])
+    return workflow
 
 
 def init_subcortical_mni_alignment_wf(*, vol_sigma=0.8, name='subcortical_mni_alignment_wf'):
@@ -23,22 +102,22 @@ def init_subcortical_mni_alignment_wf(*, vol_sigma=0.8, name='subcortical_mni_al
 
     Inputs
     ------
-    bold_file : :obj:`str`
-        BOLD file
-    bold_roi : :obj:`str`
-        File containing ROIs in BOLD space
-    atlas_roi : :obj:`str`
-        File containing ROIs in atlas space
-    std_xfm : :obj:`str`
-        File containing transform to the standard (MNI) space
+    MNIInfant_bold : :obj:`str`
+        BOLD file in MNI Infant space
+    MNIInfant_rois : :obj:`str`
+        File containing ROIs in MNI Infant space
+    MNI152_rois : :obj:`str`
+        File containing ROIs in MNI152NLin6Asym space
 
     Outputs
     -------
-    subcortical_file : :obj:`str`
+    subcortical_volume : :obj:`str`
         Volume file containing all ROIs individually aligned to standard
+    subcortical_labels : :obj:`str`
+        Volume file containing all labels
     """
-    from nipype.pipeline import engine as pe
-    from nipype.interfaces import utility as niu, fsl
+    from niworkflows.engine.workflows import LiterateWorkflow as Workflow
+    from ...interfaces.nibabel import MergeROIs
     from ...interfaces.workbench import (
         CiftiCreateDenseTimeseries,
         CiftiCreateLabel,
@@ -49,19 +128,21 @@ def init_subcortical_mni_alignment_wf(*, vol_sigma=0.8, name='subcortical_mni_al
         VolumeAffineResample,
         VolumeAllLabelsToROIs,
         VolumeLabelExportTable,
-        VolumeLabelImport,
     )
-    from niworkflows.engine.workflows import LiterateWorkflow as Workflow
-
+    # reuse saved atlas to atlas transform
+    atlas_xfm = resource_filename("nibabies", "data/MNIInfant_to_MNI1526NLinAsym.mat")
     inputnode = pe.Node(
-        niu.IdentityInterface(fields=["bold_file", "bold_roi", "atlas_roi", "atlas_xfm"]),
+        niu.IdentityInterface(fields=["MNIInfant_bold", "MNIInfant_rois", "MNI152_rois"]),
         name="inputnode",
     )
-    outputnode = pe.Node(niu.IdentityInterface(fields=["subcortical_file"]), name='outputnode')
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=["subcortical_volume", "subcortical_labels"]),
+        name='outputnode',
+    )
 
-    applyxfm_atlas = pe.Node(fsl.ApplyXFM(), name="applyxfm_atlas")
+    applyxfm_atlas = pe.Node(fsl.ApplyXFM(in_matrix_file=atlas_xfm), name="applyxfm_atlas")
     vol_resample = pe.Node(
-        VolumeAffineResample(method="ENCLOSING_VOXEL", flirt=True),
+        VolumeAffineResample(method="ENCLOSING_VOXEL", flirt=True, affine=atlas_xfm),
         name="vol_resample"
     )
     subj_rois = pe.Node(VolumeAllLabelsToROIs(label_map=1), name="subj_rois")
@@ -167,26 +248,24 @@ def init_subcortical_mni_alignment_wf(*, vol_sigma=0.8, name='subcortical_mni_al
     # fmt: off
     workflow.connect([
         (inputnode, applyxfm_atlas, [
-            ("bold_file", "in_file"),
-            ("atlas_roi", "reference"),
-            ("atlas_xfm", "in_matrix_file")]),
+            ("MNIInfant_bold", "in_file"),
+            ("MNI152_rois", "reference")]),
         (inputnode, vol_resample, [
-            ("bold_roi", "in_file"),
-            ("atlas_xfm", "affine"),
-            ("bold_roi", "flirt_source_volume")]),
+            ("MNIInfant_rois", "in_file"),
+            ("MNIInfant_rois", "flirt_source_volume")]),
         (applyxfm_atlas, vol_resample, [
             ("out_file", "volume_space"),
             ("out_file", "flirt_target_volume")]),
-        (inputnode, subj_rois, [("bold_roi", "in_file")]),
-        (inputnode, atlas_rois, [("atlas_roi", "in_file")]),
+        (inputnode, subj_rois, [("MNIInfant_rois", "in_file")]),
+        (inputnode, atlas_rois, [("MNI152_rois", "in_file")]),
         (subj_rois, split_rois, [("out_file", "in_file")]),
         (atlas_rois, split_atlas_rois, [("out_file", "in_file")]),
-        (inputnode, atlas_labels, [("atlas_roi", "in_file")]),
+        (inputnode, atlas_labels, [("MNI152_rois", "in_file")]),
         (atlas_labels, parse_labels, [("out_file", "label_file")]),
         # for loop across ROIs
         (split_rois, roi2atlas, [("out_files", "in_file")]),
         (split_atlas_rois, roi2atlas, [("out_files", "reference")]),
-        (inputnode, applyxfm_roi, [("bold_file", "in_file")]),
+        (inputnode, applyxfm_roi, [("MNIInfant_bold", "in_file")]),
         (split_atlas_rois, applyxfm_roi, [("out_files", "reference")]),
         (roi2atlas, applyxfm_roi, [("out_matrix_file", "in_matrix_file")]),
         (applyxfm_roi, bold_mask_roi, [("out_file", "in_file")]),
@@ -216,7 +295,8 @@ def init_subcortical_mni_alignment_wf(*, vol_sigma=0.8, name='subcortical_mni_al
             ("op_files", "operand_files"),
             ("op_string", "op_string")]),
         (separate, merge_rois, [("volume_all_file", "in_files")]),
-        (merge_rois, outputnode, [("out_file", "subcortical_file")]),
+        (merge_rois, outputnode, [("out_file", "subcortical_volume")]),
+        (inputnode, outputnode, [("MNI152_rois", "subcortical_labels")]),
     ])
     # fmt: on
     return workflow
@@ -267,3 +347,29 @@ def format_agg_rois(rois):
 
     """
     return rois[0], rois[1:], ("-add %s " * (len(rois) - 1)).strip()
+
+
+def drop_labels(in_file):
+    """Drop non-subcortical labels"""
+    from pathlib import Path
+    import nibabel as nb
+    import numpy as np
+    from niworkflows.interfaces.cifti import _reorient_image
+
+    # FreeSurfer LUT values
+    expected_labels = {
+        8, 10, 11, 12, 13, 16, 17, 18, 26, 28, 47, 49, 50, 51, 52, 53, 54, 58, 60,
+    }
+    img = _reorient_image(nb.load(in_file), orientation="LAS")
+    hdr = img.header
+    data = np.asanyarray(img.dataobj).astype("int16")
+    hdr.set_data_dtype("int16")
+    labels = np.unique(data)
+
+    for label in labels:
+        if label not in expected_labels:
+            data[data == label] = 0
+
+    out_file = str(Path("ROIs.nii.gz").absolute())
+    img.__class__(data, img.affine, header=hdr).to_filename(out_file)
+    return out_file

nibabies/workflows/bold/alignment.py

@@ -19,6 +19,7 @@
 
 from niworkflows.utils.connections import pop_file, listify
 from niworkflows.interfaces.header import ValidateImage
+from niworkflows.interfaces.utility import KeySelect
 from sdcflows.interfaces.brainmask import BrainExtraction
 
 from ...interfaces import DerivativesDataSink
@@ -765,22 +766,53 @@ def init_func_preproc_wf(bold_file, has_fieldmap=False):
 
         # CIFTI output
         if config.workflow.cifti_output:
+            from .alignment import init_subcortical_rois_wf, init_subcortical_mni_alignment_wf
             from .resampling import init_bold_grayords_wf
+
+            key = None
+            for space in spaces.get_spaces(nonstandard=False, dim=(3,)):
+                if "MNIInfant" in space:
+                    key = space.replace(":", "_")
+
+            assert key is not None, f"MNIInfant not found in SpatialReferences: {spaces}"
+
+            # BOLD/ROIs should be in MNIInfant space
+            cifti_select_std = pe.Node(
+                KeySelect(fields=['bold_std', 'bold_aseg_std'], key=key),
+                name='cifti_select_std',
+                run_without_submitting=True,
+            )
+
+            subcortical_rois_wf = init_subcortical_rois_wf()
+            subcortical_mni_alignment_wf = init_subcortical_mni_alignment_wf()
             bold_grayords_wf = init_bold_grayords_wf(
                 grayord_density=config.workflow.cifti_output,
                 mem_gb=mem_gb['resampled'],
-                repetition_time=metadata['RepetitionTime'])
+                repetition_time=metadata['RepetitionTime']
+            )
 
             workflow.connect([
-                (inputnode, bold_grayords_wf, [
-                    ('subjects_dir', 'inputnode.subjects_dir')]),
-                (bold_std_trans_wf, bold_grayords_wf, [
-                    ('outputnode.bold_std', 'inputnode.bold_std'),
-                    ('outputnode.spatial_reference', 'inputnode.spatial_reference')]),
+                (bold_std_trans_wf, cifti_select_std, [
+                    ("outputnode.bold_std", "bold_std"),
+                    ("outputnode.bold_aseg_std", "bold_aseg_std"),
+                    ("outputnode.spatial_reference", "keys")]),
+                (cifti_select_std, subcortical_rois_wf, [
+                    ("bold_aseg_std", "inputnode.MNIInfant_aseg")]),
+                (cifti_select_std, subcortical_mni_alignment_wf, [
+                    ("bold_std", "inputnode.MNIInfant_bold")]),
+                # (bold_t1_trans_wf, subcortical_rois_wf, [
+                #     ("outputnode.bold_aseg_t1", "inputnode.bold_aseg")]),
+                # (bold_bold_trans_wf, subcortical_mni_alignment_wf, [
+                #     ("outputnode.bold", "inputnode.bold_file")]),
+                (subcortical_rois_wf, subcortical_mni_alignment_wf, [
+                    ("outputnode.MNIInfant_rois", "inputnode.MNIInfant_rois"),
+                    ("outputnode.MNI152_rois", "inputnode.MNI152_rois")]),
+                (subcortical_mni_alignment_wf, bold_grayords_wf, [
+                    ("outputnode.subcortical_volume", "inputnode.subcortical_volume"),
+                    ("outputnode.subcortical_labels", "inputnode.subcortical_labels")]),
                 (bold_surf_wf, bold_grayords_wf, [
                     ('outputnode.surfaces', 'inputnode.surf_files'),
-                    ('outputnode.target', 'inputnode.surf_refs'),
-                ]),
+                    ('outputnode.target', 'inputnode.surf_refs')]),
                 (bold_grayords_wf, outputnode, [
                     ('outputnode.cifti_bold', 'bold_cifti'),
                     ('outputnode.cifti_variant', 'cifti_variant'),
@@ -867,7 +899,6 @@ def init_func_preproc_wf(bold_file, has_fieldmap=False):
     from niworkflows.interfaces.reportlets.registration import (
         SimpleBeforeAfterRPT as SimpleBeforeAfter,
     )
-    from niworkflows.interfaces.utility import KeySelect
     from sdcflows.workflows.apply.registration import init_coeff2epi_wf
     from sdcflows.workflows.apply.correction import init_unwarp_wf