enh: minor refactor reorganizing base workflows, in prep for nipreps#97

Just to make the diff of nipreps#97 a bit cleaner.

dmriprep/workflows/base.py

@@ -12,9 +12,9 @@
     BIDSInfo, BIDSFreeSurferDir
 )
 from niworkflows.utils.misc import fix_multi_T1w_source_name
+from niworkflows.utils.spaces import Reference
 from smriprep.workflows.anatomical import init_anat_preproc_wf
 
-from niworkflows.utils.spaces import Reference
 from ..interfaces import DerivativesDataSink, BIDSDataGrabber
 from ..interfaces.reports import SubjectSummary, AboutSummary
 from ..utils.bids import collect_data
@@ -253,21 +253,28 @@ def init_single_subject_wf(subject_id):
     anat_preproc_wf.__postdesc__ = (anat_preproc_wf.__postdesc__ or "") + f"""
 Diffusion data preprocessing
 
-: For each of the {len(subject_data["dwi"])} dwi scans found per subject
- (across all sessions), the following preprocessing was performed."""
+: For each of the {len(subject_data["dwi"])} DWI scans found per subject
+ (across all sessions), the gradient table was vetted and converted into the *RASb*
+format (i.e., given in RAS+ scanner coordinates, normalized b-vectors and scaled b-values),
+and a *b=0* average for reference to the subsequent steps of preprocessing was calculated.
+"""
 
     layout = config.execution.layout
+    dwi_data = tuple([
+        (dwi, layout.get_metadata(dwi), layout.get_bvec(dwi), layout.get_bval(dwi))
+        for dwi in subject_data["dwi"]
+    ])
+
     inputnode = pe.Node(niu.IdentityInterface(fields=["dwi_data"]),
                         name="inputnode")
-    inputnode.iterables = [(
-        "dwi_data", tuple([
-            (dwi, layout.get_bvec(dwi), layout.get_bval(dwi),
-             layout.get_metadata(dwi)["PhaseEncodingDirection"])
-            for dwi in subject_data["dwi"]
-        ])
-    )]
+    inputnode.iterables = [("dwi_data", dwi_data)]
+
+    referencenode = pe.JoinNode(niu.IdentityInterface(
+        fields=["dwi_file", "metadata", "dwi_reference", "dwi_mask", "gradients_rasb"]),
+        name="referencenode", joinsource="inputnode", run_without_submitting=True)
+
     split_info = pe.Node(niu.Function(
-        function=_unpack, output_names=["dwi_file", "bvec", "bval", "pedir"]),
+        function=_unpack, output_names=["dwi_file", "metadata", "bvec", "bval"]),
         name="split_info", run_without_submitting=True)
 
     early_b0ref_wf = init_early_b0ref_wf()
@@ -276,6 +283,13 @@ def init_single_subject_wf(subject_id):
         (split_info, early_b0ref_wf, [("dwi_file", "inputnode.dwi_file"),
                                       ("bvec", "inputnode.in_bvec"),
                                       ("bval", "inputnode.in_bval")]),
+        (split_info, referencenode, [("dwi_file", "dwi_file"),
+                                     ("metadata", "metadata")]),
+        (early_b0ref_wf, referencenode, [
+            ("outputnode.dwi_reference", "dwi_reference"),
+            ("outputnode.dwi_mask", "dwi_mask"),
+            ("outputnode.gradients_rasb", "gradients_rasb"),
+        ]),
     ])
 
     return workflow

dmriprep/workflows/dwi/base.py

@@ -53,11 +53,6 @@ def init_early_b0ref_wf(
     """
     # Build workflow
     workflow = Workflow(name=name)
-    workflow.__postdesc__ = """
-For every run and subject, the gradient table was vetted and converted into the *RASb*
-format (i.e., given in RAS+ scanner coordinates, normalized b-vectors and scaled b-values),
-and a *b=0* average for reference to the subsequent steps of preprocessing was calculated.
-"""
 
     inputnode = pe.Node(niu.IdentityInterface(
         fields=['dwi_file', 'in_bvec', 'in_bval']),
@@ -84,7 +79,7 @@ def init_early_b0ref_wf(
         (dwi_reference_wf, outputnode, [
             ('outputnode.ref_image', 'dwi_reference'),
             ('outputnode.dwi_mask', 'dwi_mask')]),
-        (gradient_table, outputnode, [('out_rasb', 'out_rasb')])
+        (gradient_table, outputnode, [('out_rasb', 'gradients_rasb')])
     ])
 
     # REPORTING ############################################################

dmriprep/workflows/dwi/util.py

@@ -6,6 +6,7 @@
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
 from niworkflows.interfaces.images import ValidateImage
 from niworkflows.interfaces.fixes import FixN4BiasFieldCorrection as N4BiasFieldCorrection
+from niworkflows.interfaces.nibabel import ApplyMask
 from niworkflows.interfaces.utils import CopyXForm
 
 from ...interfaces.images import ExtractB0, RescaleB0
@@ -212,8 +213,10 @@ def init_enhance_and_skullstrip_dwi_wf(
     combine_masks = pe.Node(fsl.BinaryMaths(operation='mul'),
                             name='combine_masks')
 
+    normalize = pe.Node(niu.Function(function=_normalize), name="normalize")
+
     # Compute masked brain
-    apply_mask = pe.Node(fsl.ApplyMask(), name='apply_mask')
+    apply_mask = pe.Node(ApplyMask(), name='apply_mask')
 
     workflow.connect([
         (inputnode, n4_correct, [('in_file', 'input_image'),
@@ -230,11 +233,30 @@ def init_enhance_and_skullstrip_dwi_wf(
         (skullstrip_first_pass, combine_masks, [('mask_file', 'in_file')]),
         (skullstrip_second_pass, fixhdr_skullstrip2, [('out_file', 'in_file')]),
         (fixhdr_skullstrip2, combine_masks, [('out_file', 'operand_file')]),
-        (fixhdr_unifize, apply_mask, [('out_file', 'in_file')]),
-        (combine_masks, apply_mask, [('out_file', 'mask_file')]),
+        (combine_masks, apply_mask, [('out_file', 'in_mask')]),
         (combine_masks, outputnode, [('out_file', 'mask_file')]),
+        (n4_correct, normalize, [('output_image', 'in_file')]),
+        (normalize, apply_mask, [('out', 'in_file')]),
+        (normalize, outputnode, [('out', 'bias_corrected_file')]),
         (apply_mask, outputnode, [('out_file', 'skull_stripped_file')]),
-        (n4_correct, outputnode, [('output_image', 'bias_corrected_file')]),
     ])
 
     return workflow
+
+
+def _normalize(in_file, newmax=2000, perc=98.0):
+    from pathlib import Path
+    import numpy as np
+    import nibabel as nb
+
+    nii = nb.load(in_file)
+    data = nii.get_fdata()
+    data[data < 0] = 0
+    if data.max() >= 2**15 - 1:
+        data *= newmax / np.percentile(data.reshape(-1), perc)
+
+    out_file = str(Path("normalized.nii.gz").absolute())
+    hdr = nii.header.copy()
+    hdr.set_data_dtype('int16')
+    nii.__class__(data.astype('int16'), nii.affine, hdr).to_filename(out_file)
+    return out_file