fix: various nodes failing

Author: oesteban

dmriprep/interfaces/images.py

@@ -1,5 +1,7 @@
 """Image tools interfaces."""
-
+import numpy as np
+import nibabel as nb
+from nipype.utils.filemanip import fname_presuffix
 from nipype import logging
 from nipype.interfaces.base import (
     traits, TraitedSpec, BaseInterfaceInputSpec, SimpleInterface, File
@@ -45,10 +47,6 @@ def _run_interface(self, runtime):
 
 def extract_b0(in_file, b0_ixs, newpath=None):
     """Extract the *b0* volumes from a DWI dataset."""
-    import numpy as np
-    import nibabel as nb
-    from nipype.utils.filemanip import fname_presuffix
-
     out_file = fname_presuffix(
         in_file, suffix='_b0', newpath=newpath)
 
@@ -71,13 +69,13 @@ class _RescaleB0InputSpec(BaseInterfaceInputSpec):
 
 
 class _RescaleB0OutputSpec(TraitedSpec):
-    out_file = File(exists=True, desc='b0 file')
+    out_ref = File(exists=True, desc='One average b0 file')
+    out_b0s = File(exists=True, desc='series of rescaled b0 volumes')
 
 
 class RescaleB0(SimpleInterface):
     """
-    Rescale the b0 volumes to deal with average signal decay over time
-    and output the median image.
+    Rescale the b0 volumes to deal with average signal decay over time.
 
     Example
     -------
@@ -93,44 +91,54 @@ class RescaleB0(SimpleInterface):
     output_spec = _RescaleB0OutputSpec
 
     def _run_interface(self, runtime):
-        self._results['out_file'] = rescale_b0(
+        self._results['out_b0s'] = rescale_b0(
             self.inputs.in_file,
             self.inputs.mask_file,
-            newpath=runtime.cwd)
+            newpath=runtime.cwd
+        )
+        self._results['out_ref'] = median(
+            self._results['out_b0s'],
+            newpath=runtime.cwd
+        )
         return runtime
 
 
 def rescale_b0(in_file, mask_file, newpath=None):
-    """
-    Rescale the input volumes using the median signal intensity
-    and output a median image.
-    """
-    import numpy as np
-    import nibabel as nb
-    from nipype.utils.filemanip import fname_presuffix
-
+    """Rescale the input volumes using the median signal intensity."""
     out_file = fname_presuffix(
-        in_file, suffix='_median_b0', newpath=newpath)
+        in_file, suffix='_rescaled_b0', newpath=newpath)
 
     img = nb.load(in_file)
     if img.dataobj.ndim == 3:
         return in_file
-    if img.shape[-1] == 1:
-        nb.squeeze_image(img).to_filename(out_file)
-        return out_file
 
     data = img.get_fdata(dtype='float32')
     mask_img = nb.load(mask_file)
     mask_data = mask_img.get_fdata(dtype='float32')
 
-    median_signal = data[mask_data > 0, ...].median(axis=0)
-
+    median_signal = np.median(data[mask_data > 0, ...], axis=0)
     rescaled_data = 1000 * data / median_signal
+    hdr = img.header.copy()
+    nb.Nifti1Image(rescaled_data, img.affine, hdr).to_filename(out_file)
+    return out_file
+
+
+def median(in_file, newpath=None):
+    """Average a 4D dataset across the last dimension using median."""
+    out_file = fname_presuffix(
+        in_file, suffix='_b0ref', newpath=newpath)
+
+    img = nb.load(in_file)
+    if img.dataobj.ndim == 3:
+        return in_file
+    if img.shape[-1] == 1:
+        nb.squeeze_image(img).to_filename(out_file)
+        return out_file
 
-    median_data = np.median(rescaled_data, axis=-1)
+    median_data = np.median(img.get_fdata(dtype='float32'),
+                            axis=-1)
 
     hdr = img.header.copy()
-    hdr.set_data_shape(median_data.shape)
     hdr.set_xyzt_units('mm')
     hdr.set_data_dtype(np.float32)
     nb.Nifti1Image(median_data, img.affine, hdr).to_filename(out_file)

dmriprep/workflows/dwi/base.py

@@ -152,6 +152,7 @@ def init_dwi_preproc_wf(
 
     summary = pe.Node(
         DiffusionSummary(
+            distortion_correction='Not implemented',
             pe_direction=metadata.get("PhaseEncodingDirection")),
         name='summary', mem_gb=DEFAULT_MEMORY_MIN_GB, run_without_submitting=True)
 
@@ -190,6 +191,7 @@ def init_dwi_preproc_wf(
         mem_gb=DEFAULT_MEMORY_MIN_GB)
 
     workflow.connect([
+        (inputnode, ds_report_summary, [('dwi_file', 'in_file')]),
         (summary, ds_report_summary, [('out_report', 'in_file')]),
         (dwi_reference_wf, ds_report_validation, [
             ('outputnode.validation_report', 'in_file')]),

dmriprep/workflows/dwi/util.py

@@ -102,7 +102,7 @@ def init_dwi_reference_wf(omp_nthreads, name='dwi_reference_wf', gen_report=Fals
         (reg_b0, pre_mask, [('out_file', 'in_file')]),
         (reg_b0, rescale_b0, [('out_file', 'in_file')]),
         (pre_mask, rescale_b0, [('out_file', 'mask_file')]),
-        (rescale_b0, enhance_and_skullstrip_dwi_wf, [('out_file', 'inputnode.in_file')]),
+        (rescale_b0, enhance_and_skullstrip_dwi_wf, [('out_ref', 'inputnode.in_file')]),
         (pre_mask, enhance_and_skullstrip_dwi_wf, [('out_file', 'inputnode.pre_mask')]),
         (validate, outputnode, [('out_file', 'dwi_file'),
                                 ('out_report', 'validation_report')]),