ENH: Reflect downstream changes in anatomical preprocessing workflow

Author: mgxd

nibabies/workflows/anatomical/base.py

@@ -1,4 +1,6 @@
 """Base anatomical preprocessing."""
+import warnings
+
 from nipype.interfaces import utility as niu
 from nipype.pipeline import engine as pe
 
@@ -93,6 +95,7 @@ def init_infant_anat_wf(
         init_anat_reports_wf,
         init_coreg_report_wf,
     )
+    from .preproc import init_anat_preproc_wf
     from .registration import init_coregistration_wf
     from .segmentation import init_anat_segmentations_wf
     from .surfaces import init_anat_ribbon_wf
@@ -107,13 +110,23 @@ def init_infant_anat_wf(
 
     # verify derivatives are relatively similar to T1w
     if precomp_mask or precomp_aseg:
-        from ...utils.validation import validate_t1w_derivatives
-
-        validated_derivatives = validate_t1w_derivatives(  # compare derivatives to the first T1w
-            t1w[0], anat_mask=precomp_mask, anat_aseg=precomp_aseg
-        )
-        precomp_mask = validated_derivatives.get("anat_mask")
-        precomp_aseg = validated_derivatives.get("anat_aseg")
+        if num_t1w > 1:
+            precomp_mask = None
+            precomp_aseg = None
+            warnings.warn(
+                "Multiple T1w files were found; precomputed derivatives will not be used."
+            )
+
+        else:
+            from ...utils.validation import validate_t1w_derivatives
+
+            validated_derivatives = (
+                validate_t1w_derivatives(  # compare derivatives to the first T1w
+                    t1w[0], anat_mask=precomp_mask, anat_aseg=precomp_aseg
+                )
+            )
+            precomp_mask = validated_derivatives.get("anat_mask")
+            precomp_aseg = validated_derivatives.get("anat_aseg")
 
     wf = Workflow(name=name)
     desc = f"""\n
@@ -159,10 +172,10 @@ def init_infant_anat_wf(
 
     desc += (
         """\
-All of the T1-weighted images were corrected for intensity non-uniformity (INU)"""
+All of the T1-weighted images were denoised and corrected for intensity non-uniformity (INU)"""
         if num_t1w > 1
         else """\
-The T1-weighted (T1w) image was corrected for intensity non-uniformity (INU)"""
+The T1-weighted (T1w) image was denoised and corrected for intensity non-uniformity (INU)"""
     )
 
     desc += """\
@@ -200,13 +213,15 @@ def init_infant_anat_wf(
         cifti_output=cifti_output,
     )
 
-    # Multiple T1w files -> generate average reference
+    # Multiple anatomical files -> generate average reference
     t1w_template_wf = init_anat_template_wf(
         contrast="T1w",
         num_files=num_t1w,
         longitudinal=longitudinal,
         omp_nthreads=omp_nthreads,
         sloppy=sloppy,
+        precomputed_mask=bool(precomp_mask),
+        precomputed_aseg=bool(precomp_aseg),
         name="t1w_template_wf",
     )
 
@@ -219,38 +234,27 @@ def init_infant_anat_wf(
         name="t2w_template_wf",
     )
 
-    # INU + Brain Extraction
+    # Clean up each anatomical template
+    # Denoise, INU, + Clipping
+    t1w_preproc_wf = init_anat_preproc_wf(name="t1w_preproc_wf")
+    t2w_preproc_wf = init_anat_preproc_wf(name="t2w_preproc_wf")
+
     if skull_strip_mode != "force":
         raise NotImplementedError("Skull stripping is currently required.")
 
-    if precomp_mask:
-        precomp_mask_wf = init_precomputed_mask_wf(omp_nthreads=omp_nthreads)
-        precomp_mask_wf.inputs.inputnode.t1w_mask = precomp_mask
-        sdc_brain_extraction_wf = init_sdc_brain_extraction_wf(
-            name="sdc_brain_extraction_wf",
-        )
-    brain_extraction_wf = init_infant_brain_extraction_wf(
-        age_months=age_months,
-        ants_affine_init=ants_affine_init,
-        skull_strip_template=skull_strip_template.space,
-        template_specs=skull_strip_template.spec,
-        omp_nthreads=omp_nthreads,
-        sloppy=sloppy,
-        debug="registration" in config.execution.debug,
-    )
     coregistration_wf = init_coregistration_wf(
         omp_nthreads=omp_nthreads,
         sloppy=sloppy,
         debug="registration" in config.execution.debug,
+        precomputed_mask=bool(precomp_mask),
     )
     coreg_report_wf = init_coreg_report_wf(
         output_dir=output_dir,
     )
-    t1w_preproc_wf = precomp_mask_wf if precomp_mask else coregistration_wf
 
     # Segmentation - initial implementation should be simple: JLF
     anat_seg_wf = init_anat_segmentations_wf(
-        anat_modality=anat_modality.capitalize(),
+        anat_modality=anat_modality.capitalize(),  # TODO: Revisit this option
         template_dir=segmentation_atlases,
         sloppy=sloppy,
         omp_nthreads=omp_nthreads,
@@ -264,143 +268,117 @@ def init_infant_anat_wf(
         templates=spaces.get_spaces(nonstandard=False, dim=(3,)),
     )
 
-    # Anatomical ribbon file using HCP signed-distance volume method
-    # if config.workflow.project_goodvoxels:
-    anat_ribbon_wf = init_anat_ribbon_wf()
-
     # fmt:off
     wf.connect([
         (inputnode, t1w_template_wf, [("t1w", "inputnode.in_files")]),
+        (inputnode, t2w_template_wf, [("t2w", "inputnode.in_files")]),
+        (inputnode, anat_reports_wf, [("t1w", "inputnode.source_file")]),
+        (inputnode, coreg_report_wf, [("t1w", "inputnode.source_file")]),
+        (inputnode, anat_norm_wf, [(("t1w", fix_multi_source_name), "inputnode.orig_t1w")]),
+
         (t1w_template_wf, outputnode, [
-            ("outputnode.realign_xfms", "anat_ref_xfms"),
-        ]),
+            ("outputnode.realign_xfms", "anat_ref_xfms")]),
+        (t1w_template_wf, t1w_preproc_wf, [("outputnode.out_file", "inputnode.in_anat")]),
+        (t2w_template_wf, t2w_preproc_wf, [("outputnode.out_file", "inputnode.in_anat")]),
+        (t1w_template_wf, anat_derivatives_wf, [
+            ("outputnode.valid_list", "inputnode.source_files"),
+            ("outputnode.realign_xfms", "inputnode.t1w_ref_xfms")]),
+        (t2w_template_wf, anat_derivatives_wf, [
+            ("outputnode.valid_list", "inputnode.t2w_source_files")]),
+
+        (t1w_preproc_wf, coregistration_wf, [("outputnode.anat_preproc", "inputnode.in_t1w")]),
+        (t1w_preproc_wf, coreg_report_wf, [("outputnode.anat_preproc", "inputnode.t1w_preproc")]),
+        (t1w_preproc_wf, anat_norm_wf, [
+            ("outputnode.t1w_preproc", "inputnode.moving_image"),
+            ("outputnode.t1w_mask", "inputnode.moving_mask")]),
+
+        (coregistration_wf, coreg_report_wf, [
+            ("outputnode.t1w_mask", "inputnode.in_mask"),
+            ("outputnode.t2w_preproc", "inputnode.t2w_preproc")]),
+
         (anat_seg_wf, outputnode, [
             ("outputnode.anat_dseg", "anat_dseg"),
-            ("outputnode.anat_tpms", "anat_tpms"),
+            ("outputnode.anat_tpms", "anat_tpms")]),
+        (anat_seg_wf, anat_derivatives_wf, [
+            ("outputnode.anat_dseg", "inputnode.t1w_dseg"),
+            ("outputnode.anat_tpms", "inputnode.t1w_tpms"),
         ]),
         (anat_seg_wf, anat_norm_wf, [
             ("outputnode.anat_dseg", "inputnode.moving_segmentation"),
-            ("outputnode.anat_tpms", "inputnode.moving_tpms"),
-        ]),
+            ("outputnode.anat_tpms", "inputnode.moving_tpms")]),
+
+        (anat_norm_wf, anat_reports_wf, [("poutputnode.template", "inputnode.template")]),
         (anat_norm_wf, outputnode, [
             ("poutputnode.standardized", "std_preproc"),
             ("poutputnode.std_mask", "std_mask"),
             ("poutputnode.std_dseg", "std_dseg"),
             ("poutputnode.std_tpms", "std_tpms"),
             ("outputnode.template", "template"),
             ("outputnode.anat2std_xfm", "anat2std_xfm"),
-            ("outputnode.std2anat_xfm", "std2anat_xfm"),
-        ]),
-        (inputnode, anat_norm_wf, [
-            (("t1w", fix_multi_source_name), "inputnode.orig_t1w"),  # anat_validate? not used...
-        ]),
-        (t1w_preproc_wf, anat_norm_wf, [
-            ("outputnode.t1w_preproc", "inputnode.moving_image"),
-            ("outputnode.t1w_mask", "inputnode.moving_mask"),
-        ]),
-        (t1w_preproc_wf, anat_derivatives_wf, [
+            ("outputnode.std2anat_xfm", "std2anat_xfm")]),
+        (anat_norm_wf, anat_derivatives_wf, [
+            ("outputnode.template", "inputnode.template"),
+            ("outputnode.anat2std_xfm", "inputnode.anat2std_xfm"),
+            ("outputnode.std2anat_xfm", "inputnode.std2anat_xfm")]),
+
+        (coregistration_wf, anat_seg_wf, [("outputnode.t1w_brain", "inputnode.anat_brain")]),
+        (coregistration_wf, anat_derivatives_wf, [
             ("outputnode.t1w_mask", "inputnode.t1w_mask"),
             ("outputnode.t1w_preproc", "inputnode.t1w_preproc"),
-        ]),
-        (coregistration_wf, anat_derivatives_wf, [
-            ("outputnode.t2w_preproc", "inputnode.t2w_preproc")
+            ("outputnode.t2w_preproc", "inputnode.t2w_preproc"),
          ]),
-        (t1w_preproc_wf, outputnode, [
+        (coregistration_wf, outputnode, [
             ("outputnode.t1w_preproc", "anat_preproc"),
             ("outputnode.t1w_brain", "anat_brain"),
             ("outputnode.t1w_mask", "anat_mask"),
         ]),
-    ])
 
-    if not precomp_aseg:
-        wf.connect([
-            (t1w_preproc_wf, anat_seg_wf, [("outputnode.t1w_brain", "inputnode.anat_brain")]),
-        ])
-    wf.connect([
-        (inputnode, t2w_template_wf, [("t2w", "inputnode.in_files")]),
+        (t1w_template_wf, anat_reports_wf, [
+            ("outputnode.out_report", "inputnode.t1w_conform_report")]),
+        (outputnode, anat_reports_wf, [
+            ("anat_preproc", "inputnode.t1w_preproc"),
+            ("anat_mask", "inputnode.t1w_mask"),
+            ("anat_dseg", "inputnode.t1w_dseg"),
+            ("std_preproc", "inputnode.std_t1w"),
+            ("std_mask", "inputnode.std_mask"),
+        ]),
     ])
+
     if precomp_mask:
+        # Ensure the mask is conformed along with the T1w
+        t1w_template_wf.inputs.inputnode.anat_mask = precomp_mask
         wf.connect([
-            (t1w_template_wf, precomp_mask_wf, [
-                ("outputnode.out_file", "inputnode.t1w"),
-            ]),
-            (t2w_template_wf, sdc_brain_extraction_wf, [
-                ("outputnode.out_file", "inputnode.in_file"),
-            ]),
-            (sdc_brain_extraction_wf, coregistration_wf, [
-                ("outputnode.out_file", "inputnode.in_t2w_preproc"),
-                ("outputnode.out_mask", "inputnode.in_mask"),
-                ("outputnode.out_probseg", "inputnode.in_probmap"),
-            ]),
+            (t1w_template_wf, coregistration_wf, [("outputnode.anat_mask", "inputnode.in_mask")]),
         ])
     else:
+        # Run brain extraction on the T2w
+        brain_extraction_wf = init_infant_brain_extraction_wf(
+            age_months=age_months,
+            ants_affine_init=ants_affine_init,
+            skull_strip_template=skull_strip_template.space,
+            template_specs=skull_strip_template.spec,
+            omp_nthreads=omp_nthreads,
+            sloppy=sloppy,
+            debug="registration" in config.execution.debug,
+        )
+
         wf.connect([
-            (t2w_template_wf, brain_extraction_wf, [
-                ("outputnode.out_file", "inputnode.in_t2w"),
-            ]),
+            (t1w_preproc_wf, brain_extraction_wf, [
+                ("outputnode.anat_preproc", "inputnode.in_t1w")]),
+            (t2w_preproc_wf, brain_extraction_wf, [
+                ("outputnode.anat_preproc", "inputnode.in_t2w")]),
             (brain_extraction_wf, coregistration_wf, [
-                ("outputnode.t2w_preproc", "inputnode.in_t2w_preproc"),
+                ("outputnode.t2w_preproc", "inputnode.in_t2w"),
                 ("outputnode.out_mask", "inputnode.in_mask"),
-                ("outputnode.out_probmap", "inputnode.in_probmap"),
-            ]),
+                ("outputnode.out_probmap", "inputnode.in_probmap")]),
         ])
-    wf.connect([
-        (t1w_template_wf, coregistration_wf, [
-            ("outputnode.out_file", "inputnode.in_t1w"),
-        ]),
 
-        (inputnode, coreg_report_wf, [
-            ("t1w", "inputnode.source_file"),
-        ]),
-        (t1w_preproc_wf, coreg_report_wf, [
-            ("outputnode.t1w_preproc", "inputnode.t1w_preproc"),
-            ("outputnode.t1w_mask", "inputnode.in_mask"),
-        ]),
-        (coregistration_wf, coreg_report_wf, [
-            ("outputnode.t2w_preproc", "inputnode.t2w_preproc")
-        ]),
-    ])
-
-    wf.connect([
-        # reports
-        (inputnode, anat_reports_wf, [
-            ("t1w", "inputnode.source_file"),
-        ]),
-        (outputnode, anat_reports_wf, [
-            ("anat_preproc", "inputnode.t1w_preproc"),
-            ("anat_mask", "inputnode.t1w_mask"),
-            ("anat_dseg", "inputnode.t1w_dseg"),
-            ("std_preproc", "inputnode.std_t1w"),
-            ("std_mask", "inputnode.std_mask"),
-        ]),
-        (t1w_template_wf, anat_reports_wf, [
-            ("outputnode.out_report", "inputnode.t1w_conform_report"),
-        ]),
-        (anat_norm_wf, anat_reports_wf, [
-            ("poutputnode.template", "inputnode.template"),
-        ]),
-        # derivatives
-        (t1w_template_wf, anat_derivatives_wf, [
-            ("outputnode.valid_list", "inputnode.source_files"),
-            ("outputnode.realign_xfms", "inputnode.t1w_ref_xfms"),
-        ]),
-        (t2w_template_wf, anat_derivatives_wf, [
-            ("outputnode.valid_list", "inputnode.t2w_source_files"),
-        ]),
-        (anat_norm_wf, anat_derivatives_wf, [
-            ("outputnode.template", "inputnode.template"),
-            ("outputnode.anat2std_xfm", "inputnode.anat2std_xfm"),
-            ("outputnode.std2anat_xfm", "inputnode.std2anat_xfm"),
-        ]),
-        (anat_ribbon_wf, anat_derivatives_wf, [
-            ("outputnode.anat_ribbon", "inputnode.anat_ribbon"),
-        ]),
-        (anat_seg_wf, anat_derivatives_wf, [
-            ("outputnode.anat_dseg", "inputnode.t1w_dseg"),
-            ("outputnode.anat_tpms", "inputnode.t1w_tpms"),
-        ]),
-    ])
-    # fmt:on
+    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, anat_seg_wf, [("outputnode.anat_aseg", "inputnode.anat_aseg")]),
+        ])
 
     if not freesurfer:
         return wf
@@ -419,20 +397,23 @@ def init_infant_anat_wf(
             use_aseg=use_aseg,
         )
 
+    # Anatomical ribbon file using HCP signed-distance volume method
+    anat_ribbon_wf = init_anat_ribbon_wf()
+
     # fmt:off
     wf.connect([
         (inputnode, surface_recon_wf, [
             ("subject_id", "inputnode.subject_id"),
-            ("subjects_dir", "inputnode.subjects_dir"),
-            ("t2w", "inputnode.t2w"),
-        ]),
+            ("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.out_file", "inputnode.t1w"),
         ]),
-        (t1w_preproc_wf, surface_recon_wf, [
+        (coregistration_wf, surface_recon_wf, [
             ("outputnode.t1w_brain", "inputnode.skullstripped_t1"),
             ("outputnode.t1w_preproc", "inputnode.corrected_t1"),
         ]),
@@ -446,7 +427,7 @@ def init_infant_anat_wf(
             ("outputnode.out_aparc", "anat_aparc"),
             ("outputnode.out_aseg", "anat_aseg"),
         ]),
-        (t1w_preproc_wf, anat_ribbon_wf, [
+        (coregistration_wf, anat_ribbon_wf, [
             ("outputnode.t1w_mask", "inputnode.t1w_mask"),
         ]),
         (surface_recon_wf, anat_ribbon_wf, [
@@ -455,6 +436,9 @@ def init_infant_anat_wf(
         (anat_ribbon_wf, outputnode, [
             ("outputnode.anat_ribbon", "anat_ribbon")
         ]),
+        (anat_ribbon_wf, anat_derivatives_wf, [
+            ("outputnode.anat_ribbon", "inputnode.anat_ribbon"),
+        ]),
         (surface_recon_wf, anat_reports_wf, [
             ("outputnode.subject_id", "inputnode.subject_id"),
             ("outputnode.subjects_dir", "inputnode.subjects_dir"),