Merge pull request #81 from josephmje/enh/skullstrip

ENH: Update image utility output path behaviour

Author: oesteban

dmriprep/interfaces/images.py

@@ -1,23 +1,27 @@
 """Image tools interfaces."""
-import numpy as np
-import nibabel as nb
-from nipype.utils.filemanip import fname_presuffix
+from pathlib import Path
+
 from nipype import logging
 from nipype.interfaces.base import (
-    traits, TraitedSpec, BaseInterfaceInputSpec, SimpleInterface, File
+    BaseInterfaceInputSpec,
+    File,
+    SimpleInterface,
+    TraitedSpec,
+    traits
 )
 
+from dmriprep.utils.images import extract_b0, median, rescale_b0
+
 LOGGER = logging.getLogger('nipype.interface')
 
 
 class _ExtractB0InputSpec(BaseInterfaceInputSpec):
     in_file = File(exists=True, mandatory=True, desc='dwi file')
-    b0_ixs = traits.List(traits.Int, mandatory=True,
-                         desc='Index of b0s')
+    b0_ixs = traits.List(traits.Int, mandatory=True, desc='Index of b0s')
 
 
 class _ExtractB0OutputSpec(TraitedSpec):
-    out_file = File(exists=True, desc='b0 file')
+    out_file = File(exists=True, desc='output b0 file')
 
 
 class ExtractB0(SimpleInterface):
@@ -38,29 +42,19 @@ class ExtractB0(SimpleInterface):
     output_spec = _ExtractB0OutputSpec
 
     def _run_interface(self, runtime):
-        self._results['out_file'] = extract_b0(
-            self.inputs.in_file,
-            self.inputs.b0_ixs,
-            newpath=runtime.cwd)
-        return runtime
-
-
-def extract_b0(in_file, b0_ixs, newpath=None):
-    """Extract the *b0* volumes from a DWI dataset."""
-    out_file = fname_presuffix(
-        in_file, suffix='_b0', newpath=newpath)
-
-    img = nb.load(in_file)
-    data = img.get_fdata(dtype='float32')
+        from nipype.utils.filemanip import fname_presuffix
 
-    b0 = data[..., b0_ixs]
+        out_file = fname_presuffix(
+            self.inputs.in_file,
+            suffix='_b0',
+            use_ext=True,
+            newpath=str(Path(runtime.cwd).absolute()),
+        )
 
-    hdr = img.header.copy()
-    hdr.set_data_shape(b0.shape)
-    hdr.set_xyzt_units('mm')
-    hdr.set_data_dtype(np.float32)
-    nb.Nifti1Image(b0, img.affine, hdr).to_filename(out_file)
-    return out_file
+        self._results['out_file'] = extract_b0(
+            self.inputs.in_file, self.inputs.b0_ixs, out_file
+        )
+        return runtime
 
 
 class _RescaleB0InputSpec(BaseInterfaceInputSpec):
@@ -91,55 +85,27 @@ class RescaleB0(SimpleInterface):
     output_spec = _RescaleB0OutputSpec
 
     def _run_interface(self, runtime):
+        from nipype.utils.filemanip import fname_presuffix
+
+        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.inputs.in_file,
-            self.inputs.mask_file,
-            newpath=runtime.cwd
+            self.inputs.mask_file, out_b0s
         )
         self._results['out_ref'] = median(
             self._results['out_b0s'],
-            newpath=runtime.cwd
+            out_path=out_ref
         )
         return runtime
-
-
-def rescale_b0(in_file, mask_file, newpath=None):
-    """Rescale the input volumes using the median signal intensity."""
-    out_file = fname_presuffix(
-        in_file, suffix='_rescaled_b0', newpath=newpath)
-
-    img = nb.load(in_file)
-    if img.dataobj.ndim == 3:
-        return in_file
-
-    data = img.get_fdata(dtype='float32')
-    mask_img = nb.load(mask_file)
-    mask_data = mask_img.get_fdata(dtype='float32')
-
-    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(img.get_fdata(dtype='float32'),
-                            axis=-1)
-
-    hdr = img.header.copy()
-    hdr.set_xyzt_units('mm')
-    hdr.set_data_dtype(np.float32)
-    nb.Nifti1Image(median_data, img.affine, hdr).to_filename(out_file)
-    return out_file

dmriprep/utils/images.py

@@ -0,0 +1,70 @@
+import nibabel as nb
+import numpy as np
+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)
+
+    img = nb.load(in_file)
+    data = img.get_fdata()
+
+    b0 = data[..., b0_ixs]
+
+    hdr = img.header.copy()
+    hdr.set_data_shape(b0.shape)
+    hdr.set_xyzt_units('mm')
+    nb.Nifti1Image(b0.astype(hdr.get_data_dtype()), img.affine, hdr).to_filename(out_path)
+    return out_path
+
+
+def rescale_b0(in_file, mask_file, out_path=None):
+    """Rescale the input volumes using the median signal intensity."""
+    if out_path is None:
+        out_path = fname_presuffix(
+            in_file, suffix='_rescaled', use_ext=True)
+
+    img = nb.load(in_file)
+    if img.dataobj.ndim == 3:
+        return in_file
+
+    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
+
+
+def median(in_file, dtype=None, out_path=None):
+    """Average a 4D dataset across the last dimension using median."""
+    if out_path is None:
+        out_path = fname_presuffix(
+            in_file, suffix='_b0ref', use_ext=True)
+
+    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_path)
+        return out_path
+
+    median_data = np.median(img.get_fdata(dtype=dtype),
+                            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)
+    return out_path