ENH: Generate a simple mask after correction

Author: oesteban

sdcflows/fieldmaps.py

@@ -5,7 +5,7 @@
 import re
 import attr
 from json import loads
-from bids.layout import BIDSFile, parse_file_entities
+from bids.layout import parse_file_entities
 from bids.utils import listify
 from niworkflows.utils.bids import relative_to_root
 from .utils.bimap import EstimatorRegistry

sdcflows/workflows/apply/correction.py

@@ -6,7 +6,7 @@
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
 
 
-def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf"):
+def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf", mask=True):
     """
     Set up a workflow that unwarps the input :abbr:`EPI (echo-planar imaging)` dataset.
 
@@ -24,6 +24,10 @@ def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf"):
         Maximum number of threads an individual process may use.
     name : :obj:`str`
         Unique name of this workflow.
+    debug : :obj:`bool`
+        Whether to run in *sloppy* mode.
+    mask : obj:`bool`
+        Generate a quick mask.
 
     Inputs
     ------
@@ -41,7 +45,7 @@ def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf"):
 
     """
     from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
-    from ...interfaces.epi import GetReadoutTime
+    from ...interfaces.epi import GetReadoutTime, EPIMask
     from ...interfaces.bspline import Coefficients2Warp
 
     workflow = Workflow(name=name)
@@ -50,7 +54,8 @@ def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf"):
         name="inputnode",
     )
     outputnode = pe.Node(
-        niu.IdentityInterface(fields=["corrected", "fieldmap"]), name="outputnode"
+        niu.IdentityInterface(fields=["corrected", "fieldmap", "corrected_mask"]),
+        name="outputnode",
     )
 
     rotime = pe.Node(GetReadoutTime(), name="rotime")
@@ -75,4 +80,16 @@ def init_unwarp_wf(omp_nthreads=1, debug=False, name="unwarp_wf"):
         (unwarp, outputnode, [("output_image", "corrected")]),
     ])
     # fmt:on
+
+    if not mask:
+        return workflow
+
+    epi_mask = pe.Node(EPIMask(), name="epi_mask")
+
+    # fmt:off
+    workflow.connect([
+        (unwarp, epi_mask, [("output_image", "in_file")]),
+        (epi_mask, outputnode, [("out_file", "corrected_mask")]),
+    ])
+    # fmt:on
     return workflow

sdcflows/workflows/apply/tests/test_correct.py

@@ -96,7 +96,8 @@ def test_unwarp_wf(tmpdir, datadir, workdir, outdir):
         # fmt: off
         workflow.connect([
             (epi_ref_wf, report, [("outputnode.fmap_ref", "before")]),
-            (unwarp_wf, report, [("outputnode.corrected", "after")]),
+            (unwarp_wf, report, [("outputnode.corrected", "after"),
+                                 ("outputnode.corrected_mask", "wm_seg")]),
             (report, ds_report, [("out_report", "in_file")]),
             (epi_ref_wf, rep, [("outputnode.fmap_ref", "reference"),
                                ("outputnode.fmap_mask", "mask")]),