RF: Converge anatomical template workflow to mimic smriprep's

mgxd · mgxd · commit 792a293762d6 · 2023-05-08T22:14:03.000-04:00
diff --git a/nibabies/workflows/anatomical/preproc.py b/nibabies/workflows/anatomical/preproc.py
@@ -1,164 +1,146 @@
 """Prepare anatomical images for processing."""
 from nipype.interfaces import utility as niu
 from nipype.pipeline import engine as pe
+from niworkflows.engine.workflows import LiterateWorkflow
 
 
-def init_anat_average_wf(
+def init_anat_template_wf(
     *,
-    bspline_fitting_distance=200,
-    longitudinal=False,
-    name="anat_average_wf",
-    num_maps=1,
-    omp_nthreads=None,
-    sloppy=False,
-):
+    contrast: str,
+    num_files: int,
+    omp_nthreads: int,
+    longitudinal: bool = False,
+    bspline_fitting_distance: int = 200,
+    sloppy: bool = False,
+    name: str = "anat_template_wf",
+) -> LiterateWorkflow:
     """
-    Create an average from several images of the same modality.
-
-    Each image undergoes a clipping step, removing background noise and
-    high-intensity outliers, which is required by INU correction with the
-    N4 algorithm.
-    Then INU correction is performed for each of the inputs and the range
-    of the image clipped again to fit within uint8.
-    Finally, each image is reoriented to have RAS+ data matrix and, if
-    more than one inputs, aligned and averaged with FreeSurfer's
-    ``mri_robust_template``.
-
+    Generate a canonically-oriented, structural average from all input images.
+    Workflow Graph
+        .. workflow::
+            :graph2use: orig
+            :simple_form: yes
+            from nibabies.workflows.anatomical.template import init_anat_template_wf
+            wf = init_anat_template_wf(
+                longitudinal=False, omp_nthreads=1, num_files=1, contrast="T1w"
+            )
     Parameters
     ----------
-    bspline_fitting_distance : :obj:`float`
-        Distance in mm between B-Spline control points for N4 INU estimation.
+    contrast : :obj:`str`
+        Name of contrast, for reporting purposes, e.g., T1w, T2w, PDw
+    num_files : :obj:`int`
+        Number of images
     longitudinal : :obj:`bool`
-        Whether an unbiased middle point should be calculated.
-    name : :obj:`str`
-        This particular workflow's unique name (Nipype requirement).
-    num_maps : :obj:`int`
-        Then number of input 3D volumes to be averaged.
+        Create unbiased structural average, regardless of number of inputs
+        (may increase runtime)
     omp_nthreads : :obj:`int`
-        The number of threads for individual processes in this workflow.
+        Maximum number of threads an individual process may use
+    bspline_fitting_distance : :obj:`float`
+        Distance in mm between B-Spline control points for N4 INU estimation.
     sloppy : :obj:`bool`
         Run in *sloppy* mode.
-
+    name : :obj:`str`, optional
+        Workflow name (default: anat_template_wf)
     Inputs
     ------
-    in_files : :obj:`list`
-        A list of one or more input files. They can be 3D or 4D.
-
+    anat_files
+        List of structural images
     Outputs
     -------
-    out_file : :obj:`str`
-        The output averaged reference file.
-    valid_list : :obj:`list`
-        A list of accepted/discarded volumes from the input list.
-    realign_xfms : :obj:`list`
-        List of rigid-body transformation matrices that bring every volume
-        into alignment with the average reference.
-    out_report : :obj:`str`
-        Path to a reportlet summarizing what happened in this workflow.
-
+    anat_ref
+        Structural reference averaging input images
+    anat_valid_list
+        List of structural images accepted for combination
+    anat_realign_xfm
+        List of affine transforms to realign input images to final reference
+    out_report
+        Conformation report
     """
     from nipype.interfaces.ants import N4BiasFieldCorrection
     from nipype.interfaces.image import Reorient
-    from niworkflows.engine.workflows import LiterateWorkflow as Workflow
     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, SplitSeries
+    from niworkflows.interfaces.nibabel import IntensityClip
     from niworkflows.interfaces.nitransforms import ConcatenateXFMs
     from niworkflows.utils.misc import add_suffix
-    from pkg_resources import resource_filename as pkgr
 
-    wf = Workflow(name=name)
+    from nibabies.utils.misc import get_file
+
+    wf = LiterateWorkflow(name=name)
+    if num_files > 1:
+        import nipype.interfaces.freesurfer as fs
 
-    inputnode = pe.Node(niu.IdentityInterface(fields=["in_files"]), name="inputnode")
+        fs_ver = fs.Info().looseversion() or "<ver>"
+        wf.__desc__ = f"""\
+An anatomical {contrast}-reference map was computed after registration of
+{num_files} {contrast} images (after INU-correction) using
+`mri_robust_template` [FreeSurfer {fs_ver}, @fs_template].
+"""
+
+    inputnode = pe.Node(niu.IdentityInterface(fields=["anat_files"]), name="inputnode")
     outputnode = pe.Node(
         niu.IdentityInterface(fields=["out_file", "valid_list", "realign_xfms", "out_report"]),
         name="outputnode",
     )
 
-    # 1. Validate each of the input images
-    validate = pe.MapNode(
-        ValidateImage(),
-        iterfield="in_file",
-        name="validate",
-        run_without_submitting=True,
-    )
-
-    # 2. Ensure we don't have two timepoints and implicitly squeeze image
-    split = pe.MapNode(SplitSeries(), iterfield="in_file", name="split")
+    # 0. Reorient anatomical image(s) to RAS and resample to common voxel space
+    anat_ref_dimensions = pe.Node(TemplateDimensions(), name="anat_ref_dimensions")
+    anat_conform = pe.MapNode(Conform(), iterfield="in_file", name="anat_conform")
 
-    # 3. INU correction of all independent volumes
-    clip_preinu = pe.MapNode(IntensityClip(p_min=50), iterfield="in_file", name="clip_preinu")
-    correct_inu = pe.MapNode(
-        N4BiasFieldCorrection(
-            dimension=3,
-            save_bias=False,
-            copy_header=True,
-            n_iterations=[50] * (5 - 2 * sloppy),
-            convergence_threshold=1e-7,
-            shrink_factor=4,
-            bspline_fitting_distance=bspline_fitting_distance,
-        ),
-        iterfield="input_image",
-        n_procs=omp_nthreads,
-        name="correct_inu",
-    )
-    clip_postinu = pe.MapNode(
-        IntensityClip(p_min=10.0, p_max=99.5), iterfield="in_file", name="clip_postinu"
-    )
-
-    # 4. Reorient T2w image(s) to RAS and resample to common voxel space
-    ref_dimensions = pe.Node(TemplateDimensions(), name="ref_dimensions")
-    conform = pe.MapNode(Conform(), iterfield="in_file", name="conform")
     # fmt:off
     wf.connect([
-        (inputnode, ref_dimensions, [("in_files", "t1w_list")]),
-        (inputnode, validate, [("in_files", "in_file")]),
-        (validate, split, [("out_file", "in_file")]),
-        (split, clip_preinu, [(("out_files", _flatten), "in_file")]),
-        (clip_preinu, correct_inu, [("out_file", "input_image")]),
-        (correct_inu, clip_postinu, [("output_image", "in_file")]),
-        (ref_dimensions, conform, [
-            ("t1w_valid_list", "in_file"),
-            ("target_zooms", "target_zooms"),
-            ("target_shape", "target_shape")]),
-        (ref_dimensions, outputnode, [("out_report", "out_report"),
-                                      ("t1w_valid_list", "valid_list")]),
+        (inputnode, anat_ref_dimensions, [('anat_files', 't1w_list')]),
+        (anat_ref_dimensions, anat_conform, [
+            ('t1w_valid_list', 'in_file'),
+            ('target_zooms', 'target_zooms'),
+            ('target_shape', 'target_shape')]),
+        (anat_ref_dimensions, outputnode, [('out_report', 'out_report'),
+                                           ('t1w_valid_list', 'anat_valid_list')]),
     ])
     # fmt:on
 
-    # 5. Reorient template to RAS, if needed (mri_robust_template may set to LIA)
-    ensure_ras = pe.Node(Reorient(), name="ensure_ras")
-
-    if num_maps == 1:
+    if num_files == 1:
         get1st = pe.Node(niu.Select(index=[0]), name="get1st")
-        outputnode.inputs.realign_xfms = [pkgr("smriprep", "data/itkIdentityTransform.txt")]
+        outputnode.inputs.anat_realign_xfm = [
+            get_file("smriprep", "data/itkIdentityTransform.txt")
+        ]
+
         # fmt:off
         wf.connect([
-            (conform, get1st, [("out_file", "inlist")]),
-            (get1st, ensure_ras, [("out", "in_file")]),
-            (ensure_ras, outputnode, [("out_file", "out_file")]),
+            (anat_conform, get1st, [('out_file', 'inlist')]),
+            (get1st, outputnode, [('out', 'anat_ref')]),
         ])
         # fmt:on
         return wf
 
-    from nipype.interfaces import freesurfer as fs
-
-    wf.__desc__ = f"""\
-An anatomical reference-map was computed after registration of
-{num_maps} images (after INU-correction) using
-`mri_robust_template` [FreeSurfer {fs.Info().looseversion() or "<ver>"}, @fs_template].
-"""
-
-    conform_xfm = pe.MapNode(
+    anat_conform_xfm = pe.MapNode(
         LTAConvert(in_lta="identity.nofile", out_lta=True),
         iterfield=["source_file", "target_file"],
-        name="conform_xfm",
+        name="anat_conform_xfm",
     )
-
-    # 6. StructuralReference is fs.RobustTemplate if > 1 volume, copying otherwise
-    merge = pe.Node(
+    # 1. Template (only if several anatomical images)
+    # 1a. Correct for bias field: the bias field is an additive factor
+    #     in log-transformed intensity units. Therefore, it is not a linear
+    #     combination of fields and N4 fails with merged images.
+    # 1b. Align and merge if several anatomical images are provided
+    clip_preinu = pe.MapNode(IntensityClip(p_min=50), iterfield="in_file", name="clip_preinu")
+    n4_correct = pe.MapNode(
+        N4BiasFieldCorrection(
+            dimension=3,
+            save_bias=False,
+            copy_header=True,
+            n_iterations=[50] * (5 - 2 * sloppy),
+            convergence_threshold=1e-7,
+            shrink_factor=4,
+            bspline_fitting_distance=bspline_fitting_distance,
+        ),
+        iterfield="input_image",
+        name="n4_correct",
+        n_procs=1,
+    )
+    # StructuralReference is fs.RobustTemplate if > 1 volume, copying otherwise
+    anat_merge = pe.Node(
         StructuralReference(
             auto_detect_sensitivity=True,
             initial_timepoint=1,  # For deterministic behavior
@@ -168,12 +150,12 @@ def init_anat_average_wf(
             no_iteration=not longitudinal,
             transform_outputs=True,
         ),
-        mem_gb=2 * num_maps - 1,
-        name="merge",
+        mem_gb=2 * num_files - 1,
+        name="anat_merge",
     )
 
-    # 7. Final intensity equalization/conformation
-    clip_final = pe.Node(IntensityClip(p_min=2.0, p_max=99.9), name="clip_final")
+    # 2. Reorient template to RAS, if needed (mri_robust_template may set to LIA)
+    anat_reorient = pe.Node(Reorient(), name="anat_reorient")
 
     merge_xfm = pe.MapNode(
         niu.Merge(2),
@@ -193,28 +175,22 @@ def _set_threads(in_list, maximum):
 
     # fmt:off
     wf.connect([
-        (ref_dimensions, conform_xfm, [("t1w_valid_list", "source_file")]),
-        (conform, conform_xfm, [("out_file", "target_file")]),
-        (conform, merge, [
-            ("out_file", "in_files"),
-            (("out_file", _set_threads, omp_nthreads), "num_threads"),
-            (("out_file", add_suffix, "_template"), "out_file")]),
-        (merge, ensure_ras, [("out_file", "in_file")]),
+        (anat_ref_dimensions, anat_conform_xfm, [('t1w_valid_list', 'source_file')]),
+        (anat_conform, anat_conform_xfm, [('out_file', 'target_file')]),
+        (anat_conform, clip_preinu, [('out_file', 'in_file')]),
+        (anat_conform, anat_merge, [
+            (('out_file', _set_threads, omp_nthreads), 'num_threads'),
+            (('out_file', add_suffix, '_template'), 'out_file')]),
+        (clip_preinu, n4_correct, [('out_file', 'input_image')]),
+        (n4_correct, anat_merge, [('output_image', 'in_files')]),
+        (anat_merge, anat_reorient, [('out_file', 'in_file')]),
         # Combine orientation and template transforms
-        (conform_xfm, merge_xfm, [("out_lta", "in1")]),
-        (merge, merge_xfm, [("transform_outputs", "in2")]),
-        (merge_xfm, concat_xfms, [("out", "in_xfms")]),
+        (anat_conform_xfm, merge_xfm, [('out_lta', 'in1')]),
+        (anat_merge, merge_xfm, [('transform_outputs', 'in2')]),
+        (merge_xfm, concat_xfms, [('out', 'in_xfms')]),
         # Output
-        (ensure_ras, clip_final, [("out_file", "in_file")]),
-        (clip_final, outputnode, [("out_file", "out_file")]),
-        (concat_xfms, outputnode, [("out_xfm", "realign_xfms")]),
+        (anat_reorient, outputnode, [('out_file', 'anat_ref')]),
+        (concat_xfms, outputnode, [('out_xfm', 'anat_realign_xfm')]),
     ])
     # fmt:on
-
     return wf
-
-
-def _flatten(inlist):
-    from bids.utils import listify
-
-    return [el for items in listify(inlist) for el in listify(items)]
