Remap labels from FreeSurfer aseg to MCRIB (#270)

* use MapLabel interface to remap aseg labels to mcrib-compatible

* STY: black

* Update nibabies/workflows/anatomical/segmentation.py

Co-authored-by: Mathias Goncalves <[email protected]>

* Update nibabies/workflows/anatomical/segmentation.py

Co-authored-by: Mathias Goncalves <[email protected]>

* Update segmentation.py

add suggestions from PR review

* Update segmentation.py

formatting

* Update segmentation.py

define aseg2mcrib dict before using

Co-authored-by: Mathias Goncalves <[email protected]>

Author: madisoth

nibabies/workflows/anatomical/base.py

@@ -147,6 +147,7 @@ def init_infant_anat_wf(
                 "surfaces",
                 "morphometrics",
                 "anat_aseg",
+                "anat_mcrib",
                 "anat_aparc",
                 "template",
             ]

nibabies/workflows/anatomical/segmentation.py

@@ -6,6 +6,7 @@
 from nipype.pipeline import engine as pe
 from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
 from niworkflows.interfaces.fixes import FixHeaderRegistration as Registration
+from niworkflows.interfaces.nibabel import MapLabels
 from pkg_resources import resource_filename as pkgr_fn
 from smriprep.utils.misc import apply_lut as _apply_bids_lut
 from smriprep.workflows.anatomical import (
@@ -51,7 +52,7 @@ def init_anat_segmentations_wf(
         name="inputnode",
     )
     outputnode = pe.Node(
-        niu.IdentityInterface(fields=["anat_aseg", "anat_dseg", "anat_tpms"]),
+        niu.IdentityInterface(fields=["anat_aseg", "anat_mcrib_aseg", "anat_dseg", "anat_tpms"]),
         name="outputnode",
     )
 
@@ -175,15 +176,61 @@ def init_anat_segmentations_wf(
     lut_anat_dseg.inputs.lut = _aseg_to_three()
     split_seg = pe.Node(niu.Function(function=_split_segments), name="split_seg")
     out_aseg = validate_aseg if precomp_aseg else jf_label
+
+    wf.__desc__ += """Segmentation volume label indices were remapped
+                     from FreeSurfer to M-CRIB-compatible format. """
+
+    # dictionary to map labels from aseg to mcrib
+    aseg2mcrib = {
+        2: 51,
+        3: 21,
+        4: 49,
+        5: 0,
+        7: 17,
+        8: 17,
+        10: 43,
+        11: 41,
+        12: 47,
+        13: 47,
+        14: 0,
+        15: 0,
+        16: 19,
+        17: 1,
+        18: 3,
+        26: 41,
+        28: 45,
+        31: 49,
+        41: 52,
+        42: 20,
+        43: 50,
+        44: 0,
+        46: 18,
+        47: 18,
+        49: 42,
+        50: 40,
+        51: 46,
+        52: 46,
+        53: 2,
+        54: 4,
+        58: 40,
+        60: 44,
+        63: 50,
+        253: 48,
+    }
+    map_labels = pe.Node(MapLabels(mappings=aseg2mcrib), name="map_labels")
+
     # fmt: off
     wf.connect([
         (out_aseg, outputnode, [('out_file', 'anat_aseg')]),
         (out_aseg, lut_anat_dseg, [('out_file', 'in_dseg')]),
         (lut_anat_dseg, outputnode, [('out', 'anat_dseg')]),
         (lut_anat_dseg, split_seg, [('out', 'in_file')]),
         (split_seg, outputnode, [('out', 'anat_tpms')]),
+        (out_aseg, map_labels, [("out_file", "in_file")]),
+        (map_labels, outputnode, [("out_file", "anat_mcrib_aseg")]),
     ])
     # fmt: on
+
     return wf