RF: Move mask/aseg processing into preproc workflow

Author: mgxd

nibabies/workflows/anatomical/base.py

@@ -217,8 +217,6 @@ def init_infant_anat_wf(
         longitudinal=longitudinal,
         omp_nthreads=omp_nthreads,
         sloppy=sloppy,
-        precomputed_mask=bool(precomp_mask),
-        precomputed_aseg=bool(precomp_aseg),
         name="t1w_template_wf",
     )
 
@@ -233,7 +231,11 @@ def init_infant_anat_wf(
 
     # Clean up each anatomical template
     # Denoise, INU, + Clipping
-    t1w_preproc_wf = init_anat_preproc_wf(name="t1w_preproc_wf")
+    t1w_preproc_wf = init_anat_preproc_wf(
+        precomputed_mask=bool(precomp_mask),
+        precomputed_aseg=bool(precomp_aseg),
+        name="t1w_preproc_wf",
+    )
     t2w_preproc_wf = init_anat_preproc_wf(name="t2w_preproc_wf")
 
     if skull_strip_mode != "force":
@@ -343,10 +345,10 @@ def init_infant_anat_wf(
 
     if precomp_mask:
         # Ensure the mask is conformed along with the T1w
-        t1w_template_wf.inputs.inputnode.anat_mask = precomp_mask
+        t1w_preproc_wf.inputs.inputnode.in_mask = precomp_mask
         # fmt:off
         wf.connect([
-            (t1w_template_wf, coregistration_wf, [("outputnode.anat_mask", "inputnode.in_mask")]),
+            (t1w_preproc_wf, coregistration_wf, [("outputnode.anat_mask", "inputnode.in_mask")]),
             (t2w_preproc_wf, coregistration_wf, [("outputnode.anat_preproc", "inputnode.in_t2w")])
         ])
         # fmt:on
@@ -376,8 +378,8 @@ def init_infant_anat_wf(
 
     if precomp_aseg:
         # Ensure the segmentation is conformed along with the T1w
-        t1w_template_wf.inputs.inputnode.anat_aseg = precomp_aseg
-        wf.connect(t1w_template_wf, "outputnode.anat_aseg", anat_seg_wf, "inputnode.anat_aseg")
+        t1w_preproc_wf.inputs.inputnode.in_aseg = precomp_aseg
+        wf.connect(t1w_preproc_wf, "outputnode.anat_aseg", anat_seg_wf, "inputnode.anat_aseg")
 
     if not freesurfer:
         return wf
@@ -395,14 +397,46 @@ def init_infant_anat_wf(
             age_months=age_months,
             use_aseg=use_aseg,
         )
+
     elif config.workflow.surface_recon_method == 'mcribs':
+        from nipype.interfaces.ants import DenoiseImage
+
         from .surfaces import init_mcribs_surface_recon_wf
 
+        # Denoise raw T2w, since using the template / preproc resulted in intersection errors
+        denoise_raw_t2w = pe.Node(
+            DenoiseImage(dimension=3, noise_model="Rician"), name='denoise_raw_t2w'
+        )
+
         surface_recon_wf = init_mcribs_surface_recon_wf(
+            omp_nthreads=omp_nthreads,
             use_aseg=bool(precomp_aseg),
+            use_mask=bool(precomp_mask),
             mcribs_dir=str(config.execution.mcribs_dir),  # Needed to preserve runs
         )
 
+        # Transformed gives
+        if precomp_aseg:
+            surface_recon_wf.inputs.inputnode.ants_segs = precomp_aseg
+        if precomp_mask:
+            surface_recon_wf.inputs.inputnode.anat_mask = precomp_mask
+        # fmt:off
+        wf.connect([
+            (inputnode, denoise_raw_t2w, [('t2w', 'input_image')]),
+            (denoise_raw_t2w, surface_recon_wf, [('output_image', 'inputnode.t2w')])
+        ])
+        # fmt:on
+
+    if config.workflow.surface_recon_method in ('freesurfer', 'infantfs'):
+        # fmt:off
+        wf.connect([
+            (t2w_preproc_wf, surface_recon_wf, [
+                ("outputnode.anat_preproc", "inputnode.t2w")]),
+            (anat_seg_wf, surface_recon_wf, [
+                ("outputnode.anat_aseg", "inputnode.ants_segs")]),
+        ])
+        # fmt:on
+
     # Anatomical ribbon file using HCP signed-distance volume method
     anat_ribbon_wf = init_anat_ribbon_wf()
 
@@ -411,11 +445,6 @@ def init_infant_anat_wf(
         (inputnode, surface_recon_wf, [
             ("subject_id", "inputnode.subject_id"),
             ("subjects_dir", "inputnode.subjects_dir")]),
-        (t2w_preproc_wf, surface_recon_wf, [
-            ("outputnode.anat_preproc", "inputnode.t2w")]),
-        (anat_seg_wf, surface_recon_wf, [
-            ("outputnode.anat_aseg", "inputnode.ants_segs"),
-        ]),
         (t1w_template_wf, surface_recon_wf, [
             ("outputnode.anat_ref", "inputnode.t1w"),
         ]),

nibabies/workflows/anatomical/preproc.py

@@ -4,7 +4,11 @@
 
 
 def init_anat_preproc_wf(
-    *, bspline_fitting_distance: int = 200, name: str = "anat_preproc_wf"
+    *,
+    bspline_fitting_distance: int = 200,
+    precomputed_mask: bool = False,
+    precomputed_aseg: bool = False,
+    name: str = "anat_preproc_wf",
 ) -> LiterateWorkflow:
     """Polish up raw anatomical data.
 
@@ -27,12 +31,21 @@ def init_anat_preproc_wf(
         Preprocessed anatomical image (Denoising/INU/Clipping)
     """
     from nipype.interfaces.ants import DenoiseImage, N4BiasFieldCorrection
-    from niworkflows.interfaces.nibabel import IntensityClip
+    from niworkflows.interfaces.header import ValidateImage
+    from niworkflows.interfaces.nibabel import IntensityClip, RegridToZooms
 
     wf = LiterateWorkflow(name=name)
-    inputnode = pe.Node(niu.IdentityInterface(fields=["in_anat"]), name="inputnode")
-    outputnode = pe.Node(niu.IdentityInterface(fields=["anat_preproc"]), name="outputnode")
+    inputnode = pe.Node(
+        niu.IdentityInterface(fields=["in_anat", "in_mask", "in_aseg"]),
+        name="inputnode",
+    )
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=["anat_preproc", "anat_mask", "anat_aseg"]),
+        name="outputnode",
+    )
 
+    # validate image
+    validate = pe.Node(ValidateImage(), name="anat_validate", run_without_submitting=True)
     clip = pe.Node(IntensityClip(p_min=10.0, p_max=99.5), name="clip")
     denoise = pe.Node(DenoiseImage(dimension=3, noise_model="Rician"), name="denoise")
     n4_correct = pe.Node(
@@ -51,13 +64,43 @@ def init_anat_preproc_wf(
 
     final_clip = pe.Node(IntensityClip(p_min=5.0, p_max=99.5), name="final_clip")
 
+    if precomputed_mask:
+        validate_mask = pe.Node(ValidateImage(), name="validate_mask")
+        regrid_mask = pe.Node(RegridToZooms(), name="regrid_mask")
+        # fmt:off
+        wf.connect([
+            (inputnode, validate_mask, [("in_mask", "in_file")]),
+            (validate_mask, regrid_mask, [("out_file", "in_file")]),
+            (final_clip, regrid_mask, [(("out_file", _get_zooms), "zooms")]),
+            (regrid_mask, outputnode, [("out_file", "anat_mask")]),
+        ])
+        # fmt:on
+    if precomputed_aseg:
+        validate_aseg = pe.Node(ValidateImage(), name="validate_aseg")
+        regrid_aseg = pe.Node(RegridToZooms(), name="regrid_aseg")
+        # fmt:off
+        wf.connect([
+            (inputnode, validate_aseg, [("in_aseg", "in_file")]),
+            (validate_aseg, regrid_aseg, [("out_file", "in_file")]),
+            (final_clip, regrid_aseg, [(("out_file", _get_zooms), "zooms")]),
+            (regrid_aseg, outputnode, [("out_file", "anat_aseg")]),
+        ])
+        # fmt:on
+
     # fmt:off
     wf.connect([
-        (inputnode, clip, [("in_anat", "in_file")]),
+        (inputnode, validate, [("in_anat", "in_file")]),
+        (validate, clip, [("out_file", "in_file")]),
         (clip, denoise, [("out_file", "input_image")]),
         (denoise, n4_correct, [("output_image", "input_image")]),
         (n4_correct, final_clip, [("output_image", "in_file")]),
         (final_clip, outputnode, [("out_file", "anat_preproc")]),
     ])
     # fmt:on
     return wf
+
+
+def _get_zooms(in_file):
+    import nibabel as nb
+
+    return tuple(nb.load(in_file).header.get_zooms()[:3])

nibabies/workflows/anatomical/template.py

@@ -12,8 +12,6 @@ def init_anat_template_wf(
     longitudinal: bool = False,
     bspline_fitting_distance: int = 200,
     sloppy: bool = False,
-    precomputed_mask: bool = False,
-    precomputed_aseg: bool = False,
     name: str = "anat_template_wf",
 ) -> LiterateWorkflow:
     """
@@ -62,6 +60,7 @@ def init_anat_template_wf(
     from nipype.interfaces.image import Reorient
     from niworkflows.interfaces.freesurfer import PatchedLTAConvert as LTAConvert
     from niworkflows.interfaces.freesurfer import StructuralReference
+    from niworkflows.interfaces.header import ValidateImage
     from niworkflows.interfaces.images import Conform, TemplateDimensions
     from niworkflows.interfaces.nibabel import IntensityClip
     from niworkflows.interfaces.nitransforms import ConcatenateXFMs
@@ -90,9 +89,6 @@ def init_anat_template_wf(
                 "anat_valid_list",
                 "anat_realign_xfm",
                 "out_report",
-                # only if precomputed inputs
-                "anat_mask",
-                "anat_aseg",
             ],
         ),
         name="outputnode",
@@ -115,28 +111,6 @@ def init_anat_template_wf(
     # fmt:on
 
     if num_files == 1:
-        if precomputed_mask:
-            anat_mask_conform = pe.Node(Conform(), name='anat_mask_conform')
-            # fmt:off
-            wf.connect([
-                (inputnode, anat_mask_conform, [('anat_mask', 'in_file')]),
-                (anat_ref_dimensions, anat_mask_conform, [
-                    ('target_zooms', 'target_zooms'),
-                    ('target_shape', 'target_shape')]),
-                (anat_mask_conform, outputnode, [('out_file', 'anat_mask')]),
-            ])
-            # fmt:on
-        if precomputed_aseg:
-            anat_aseg_conform = pe.Node(Conform(), name='anat_aseg_conform')
-            # fmt:off
-            wf.connect([
-                (inputnode, anat_aseg_conform, [('anat_aseg', 'in_file')]),
-                (anat_ref_dimensions, anat_aseg_conform, [
-                    ('target_zooms', 'target_zooms'),
-                    ('target_shape', 'target_shape')]),
-                (anat_aseg_conform, outputnode, [('out_file', 'anat_aseg')])
-            ])
-            # fmt:on
         get1st = pe.Node(niu.Select(index=[0]), name="get1st")
         outputnode.inputs.anat_realign_xfm = [
             get_file("smriprep", "data/itkIdentityTransform.txt")