Merge pull request #101 from oesteban/fix/90-bad-b0-reference

FIX: Minor refactor of DWI/brainmasks utilities and dtypes

Author: josephmje

dmriprep/interfaces/images.py

@@ -47,12 +47,12 @@ def _run_interface(self, runtime):
         out_file = fname_presuffix(
             self.inputs.in_file,
             suffix='_b0',
-            use_ext=True,
             newpath=str(Path(runtime.cwd).absolute()),
         )
 
         self._results['out_file'] = extract_b0(
-            self.inputs.in_file, self.inputs.b0_ixs, out_file
+            self.inputs.in_file, self.inputs.b0_ixs,
+            out_path=out_file
         )
         return runtime
 
@@ -65,6 +65,7 @@ class _RescaleB0InputSpec(BaseInterfaceInputSpec):
 class _RescaleB0OutputSpec(TraitedSpec):
     out_ref = File(exists=True, desc='One average b0 file')
     out_b0s = File(exists=True, desc='series of rescaled b0 volumes')
+    signal_drift = traits.List(traits.Float, desc='estimated signal drift factors')
 
 
 class RescaleB0(SimpleInterface):
@@ -90,17 +91,15 @@ def _run_interface(self, runtime):
         out_b0s = fname_presuffix(
             self.inputs.in_file,
             suffix='_rescaled',
-            use_ext=True,
             newpath=str(Path(runtime.cwd).absolute())
         )
         out_ref = fname_presuffix(
             self.inputs.in_file,
             suffix='_ref',
-            use_ext=True,
             newpath=str(Path(runtime.cwd).absolute())
         )
 
-        self._results['out_b0s'] = rescale_b0(
+        self._results['out_b0s'], self._results['signal_drift'] = rescale_b0(
             self.inputs.in_file,
             self.inputs.mask_file, out_b0s
         )

dmriprep/utils/images.py

@@ -1,23 +1,21 @@
-import nibabel as nb
+"""Utilities to handle images."""
 import numpy as np
+import nibabel as nb
 from nipype.utils.filemanip import fname_presuffix
 
 
 def extract_b0(in_file, b0_ixs, out_path=None):
     """Extract the *b0* volumes from a DWI dataset."""
     if out_path is None:
-        out_path = fname_presuffix(
-            in_file, suffix='_b0', use_ext=True)
+        out_path = fname_presuffix(in_file, suffix='_b0')
 
     img = nb.load(in_file)
-    data = img.get_fdata()
-
-    b0 = data[..., b0_ixs]
+    bzeros = np.squeeze(np.asanyarray(img.dataobj)[..., b0_ixs])
 
     hdr = img.header.copy()
-    hdr.set_data_shape(b0.shape)
+    hdr.set_data_shape(bzeros.shape)
     hdr.set_xyzt_units('mm')
-    nb.Nifti1Image(b0.astype(hdr.get_data_dtype()), img.affine, hdr).to_filename(out_path)
+    nb.Nifti1Image(bzeros, img.affine, hdr).to_filename(out_path)
     return out_path
 
 
@@ -27,24 +25,25 @@ def rescale_b0(in_file, mask_file, out_path=None):
         out_path = fname_presuffix(
             in_file, suffix='_rescaled', use_ext=True)
 
-    img = nb.load(in_file)
+    img = nb.squeeze_image(nb.load(in_file))
     if img.dataobj.ndim == 3:
-        return in_file
+        return in_file, [1.0]
 
+    mask_data = nb.load(mask_file).get_fdata() > 0
+
+    dtype = img.get_data_dtype()
     data = img.get_fdata()
-    mask_img = nb.load(mask_file)
-    mask_data = mask_img.get_fdata()
 
-    median_signal = np.median(data[mask_data > 0, ...], axis=0)
-    rescaled_data = 1000 * data / median_signal
-    hdr = img.header.copy()
-    nb.Nifti1Image(
-        rescaled_data.astype(hdr.get_data_dtype()), img.affine, hdr
-    ).to_filename(out_path)
-    return out_path
+    median_signal = np.median(data[mask_data, ...], axis=0)
+    # Normalize to the first volume
+    signal_drift = median_signal[0] / median_signal
+    data /= signal_drift
 
+    nb.Nifti1Image(data.astype(dtype), img.affine, img.header).to_filename(out_path)
+    return out_path, signal_drift.tolist()
 
-def median(in_file, dtype=None, out_path=None):
+
+def median(in_file, out_path=None):
     """Average a 4D dataset across the last dimension using median."""
     if out_path is None:
         out_path = fname_presuffix(
@@ -57,14 +56,8 @@ def median(in_file, dtype=None, out_path=None):
         nb.squeeze_image(img).to_filename(out_path)
         return out_path
 
-    median_data = np.median(img.get_fdata(dtype=dtype),
-                            axis=-1)
+    dtype = img.get_data_dtype()
+    median_data = np.median(img.get_fdata(), axis=-1)
 
-    hdr = img.header.copy()
-    hdr.set_xyzt_units('mm')
-    if dtype is not None:
-        hdr.set_data_dtype(dtype)
-    else:
-        dtype = hdr.get_data_dtype()
-    nb.Nifti1Image(median_data.astype(dtype), img.affine, hdr).to_filename(out_path)
+    nb.Nifti1Image(median_data.astype(dtype), img.affine, img.header).to_filename(out_path)
     return out_path