ENH: Add n4_only workflow -- to skip brain extraction

niworkflows/anat/ants.py

@@ -28,6 +28,7 @@
     FixHeaderApplyTransforms as ApplyTransforms,
 )
 from ..interfaces.utils import CopyXForm
+from ..interfaces.nibabel import Binarize
 
 
 ATROPOS_MODELS = {
@@ -310,7 +311,7 @@ def init_brain_extraction_wf(name='brain_extraction_wf',
             ('reverse_transforms', 'transforms'),
             ('reverse_invert_flags', 'invert_transform_flags')]),
         (map_brainmask, thr_brainmask, [('output_image', 'input_image')]),
-        (thr_brainmask, dil_brainmask, [('output_image', 'op1')]),
+        (thr_brainmask, dil_brainmask, [('out_mask', 'op1')]),
         (dil_brainmask, get_brainmask, [('output_image', 'op1')]),
         (inu_n4_final, apply_mask, [('output_image', 'in_file')]),
         (get_brainmask, apply_mask, [('output_image', 'mask_file')]),
@@ -613,6 +614,164 @@ def init_atropos_wf(name='atropos_wf',
     return wf
 
 
+def init_n4_only_wf(atropos_model=None,
+                    atropos_refine=True,
+                    atropos_use_random_seed=True,
+                    bids_suffix='T1w',
+                    mem_gb=3.0,
+                    name='n4_only_wf',
+                    omp_nthreads=None
+                    ):
+    """
+    Build a workflow to sidetrack brain extraction on skull-stripped datasets.
+
+    An alternative workflow to "init_brain_extraction_wf", for anatomical
+    images which have already been brain extracted.
+
+      1. Creates brain mask assuming all zero voxels are outside the brain
+      2. Applies N4 bias field correction
+      3. (Optional) apply ATROPOS and massage its outputs
+      4. Use results from 3 to refine N4 bias field correction
+
+    Workflow Graph
+        .. workflow::
+            :graph2use: orig
+            :simple_form: yes
+
+            from niworkflows.anat.ants import init_n4_only_wf
+            wf = init_n4_only_wf()
+
+    Parameters
+    ----------
+    omp_nthreads : int
+        Maximum number of threads an individual process may use
+    mem_gb : float
+        Estimated peak memory consumption of the most hungry nodes
+    bids_suffix : str
+        Sequence type of the first input image. For a list of acceptable values see
+        https://bids-specification.readthedocs.io/en/latest/04-modality-specific-files/01-magnetic-resonance-imaging-data.html#anatomy-imaging-data
+    atropos_refine : bool
+        Enables or disables the whole ATROPOS sub-workflow
+    atropos_use_random_seed : bool
+        Whether ATROPOS should generate a random seed based on the
+        system's clock
+    atropos_model : tuple or None
+        Allows to specify a particular segmentation model, overwriting
+        the defaults based on ``bids_suffix``
+    name : str, optional
+        Workflow name (default: ``'n4_only_wf'``).
+
+    Inputs
+    ------
+    in_files
+        List of input anatomical images to be bias corrected,
+        typically T1-weighted.
+        If a list of anatomical images is provided, subsequently
+        specified images are used during the segmentation process.
+        However, only the first image is used in the registration
+        of priors.
+        Our suggestion would be to specify the T1w as the first image.
+
+    Outputs
+    -------
+    out_file
+        :abbr:`INU (intensity non-uniformity)`-corrected ``in_files``
+    out_mask
+        Calculated brain mask
+    bias_corrected
+        Same as "out_file", provided for consistency with brain extraction
+    bias_image
+        The :abbr:`INU (intensity non-uniformity)` field estimated for each
+        input in ``in_files``
+    out_segm
+        Output segmentation by ATROPOS
+    out_tpms
+        Output :abbr:`TPMs (tissue probability maps)` by ATROPOS
+
+    """
+    wf = pe.Workflow(name)
+
+    inputnode = pe.Node(niu.IdentityInterface(fields=['in_files', 'in_mask']),
+                        name='inputnode')
+
+    outputnode = pe.Node(niu.IdentityInterface(
+        fields=['out_file', 'out_mask', 'bias_corrected', 'bias_image',
+                'out_segm', 'out_tpms']),
+        name='outputnode')
+
+    # Create brain mask
+    thr_brainmask = pe.Node(Binarize(thresh_low=2), name='binarize')
+
+    # INU correction
+    inu_n4_final = pe.MapNode(
+        N4BiasFieldCorrection(
+            dimension=3, save_bias=True, copy_header=True,
+            n_iterations=[50] * 5, convergence_threshold=1e-7, shrink_factor=4,
+            bspline_fitting_distance=200),
+        n_procs=omp_nthreads, name='inu_n4_final', iterfield=['input_image'])
+
+    # Check ANTs version
+    try:
+        inu_n4_final.inputs.rescale_intensities = True
+    except ValueError:
+        warn("The installed ANTs version too old. Please consider upgrading to "
+             "2.1.0 or greater.", DeprecationWarning)
+
+    wf.connect([
+        (inputnode, inu_n4_final, [('in_files', 'input_image')]),
+        (inputnode, thr_brainmask, [(('in_files', _pop), 'in_file')]),
+        (thr_brainmask, outputnode, [('out_mask', 'out_mask')]),
+        (inu_n4_final, outputnode, [('output_image', 'out_file')]),
+        (inu_n4_final, outputnode, [('output_image', 'bias_corrected')]),
+        (inu_n4_final, outputnode, [('bias_image', 'bias_image')])
+    ])
+
+    # If atropos refine, do in4 twice
+    if atropos_refine:
+        # Morphological dilation, radius=2
+        dil_brainmask = pe.Node(ImageMath(operation='MD', op2='2'),
+                                name='dil_brainmask')
+        # Get largest connected component
+        get_brainmask = pe.Node(ImageMath(operation='GetLargestComponent'),
+                                name='get_brainmask')
+        atropos_model = atropos_model or list(
+            ATROPOS_MODELS[bids_suffix].values())
+        atropos_wf = init_atropos_wf(
+            use_random_seed=atropos_use_random_seed,
+            omp_nthreads=omp_nthreads,
+            mem_gb=mem_gb,
+            in_segmentation_model=atropos_model,
+        )
+        sel_wm = pe.Node(niu.Select(index=atropos_model[-1] - 1), name='sel_wm',
+                         run_without_submitting=True)
+
+        inu_n4 = pe.MapNode(
+            N4BiasFieldCorrection(
+                dimension=3, save_bias=False, copy_header=True,
+                n_iterations=[50] * 4, convergence_threshold=1e-7,
+                shrink_factor=4, bspline_fitting_distance=200),
+            n_procs=omp_nthreads, name='inu_n4', iterfield=['input_image'])
+
+        wf.connect([
+            (inputnode, inu_n4, [('in_files', 'input_image')]),
+            (inu_n4, atropos_wf, [
+                ('output_image', 'inputnode.in_files')]),
+            (thr_brainmask, atropos_wf, [
+                ('out_mask', 'inputnode.in_mask')]),
+            (thr_brainmask, dil_brainmask, [('output_image', 'op1')]),
+            (dil_brainmask, get_brainmask, [('output_image', 'op1')]),
+            (get_brainmask, atropos_wf, [
+                ('output_image', 'inputnode.in_mask_dilated')]),
+            (atropos_wf, sel_wm, [('outputnode.out_tpms', 'inlist')]),
+            (sel_wm, inu_n4_final, [('out', 'weight_image')]),
+            (atropos_wf, outputnode, [
+                ('outputnode.out_segm', 'out_segm'),
+                ('outputnode.out_tpms', 'out_tpms')]),
+        ])
+
+    return wf
+
+
 def _pop(in_files):
     if isinstance(in_files, (list, tuple)):
         return in_files[0]