ENH: b0 reference and skullstrip workflow

dmriprep/conftest.py

@@ -5,6 +5,7 @@
 import numpy as np
 import nibabel as nb
 import pytest
+from bids.layout import BIDSLayout
 from dipy.data.fetcher import _make_fetcher, UW_RW_URL
 
 _dipy_datadir_root = os.getenv('DMRIPREP_TESTS_DATA') or Path.home()
@@ -28,6 +29,20 @@
     'bvals': np.loadtxt(dipy_datadir / "HARDI193.bval"),
 }
 
+test_data_env = os.getenv('TEST_DATA_HOME', '/tmp/data')
+test_output_dir = os.getenv('TEST_OUTPUT_DIR')
+test_workdir = os.getenv('TEST_WORK_DIR')
+
+layouts = {p.name: BIDSLayout(str(p), validate=False, derivatives=True)
+           for p in Path(test_data_env).glob('*') if p.is_dir()}
+
+
+def pytest_report_header(config):
+    msg = "Datasets found: %s" % ', '.join([v.root for v in layouts.values()])
+    if test_output_dir is not None:
+        msg += '\nOutput folder: %s' % Path(test_output_dir).resolve()
+    return msg
+
 
 @pytest.fixture(autouse=True)
 def doctest_autoimport(doctest_namespace):
@@ -48,3 +63,18 @@ def doctest_autoimport(doctest_namespace):
 def dipy_test_data(scope='session'):
     """Create a temporal directory shared across tests to pull data in."""
     return _sherbrooke_data
+
+
+@pytest.fixture
+def workdir():
+    return None if test_workdir is None else Path(test_workdir)
+
+
+@pytest.fixture
+def output_path():
+    return None if test_output_dir is None else Path(test_output_dir)
+
+
+@pytest.fixture
+def bids_layouts():
+    return layouts

dmriprep/interfaces/images.py

@@ -0,0 +1,131 @@
+"""Image tools interfaces."""
+
+from nipype import logging
+from nipype.interfaces.base import (
+    traits, TraitedSpec, BaseInterfaceInputSpec, SimpleInterface, File
+)
+
+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')
+
+
+class ExtractB0OutputSpec(TraitedSpec):
+    out_file = File(exists=True, desc='b0 file')
+
+
+class ExtractB0(SimpleInterface):
+    """
+    Extract all b=0 volumes from a dwi series.
+
+    Example
+    -------
+
+    >>> os.chdir(tmpdir)
+    >>> extract_b0 = ExtractB0()
+    >>> extract_b0.inputs.in_file = str(data_dir / 'dwi.nii.gz')
+    >>> extract_b0.inputs.b0_ixs = [0, 1, 2]
+    >>> res = extract_b0.run()  # doctest: +SKIP
+
+    """
+    input_spec = ExtractB0InputSpec
+    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."""
+    import numpy as np
+    import nibabel as nib
+    from nipype.utils.filemanip import fname_presuffix
+
+    out_file = fname_presuffix(
+        in_file, suffix='_b0', newpath=newpath)
+
+    img = nib.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')
+    hdr.set_data_dtype(np.float32)
+    nib.Nifti1Image(b0, img.affine, hdr).to_filename(out_file)
+    return out_file
+
+
+class RescaleB0InputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc='b0s file')
+    mask_file = File(exists=True, mandatory=True, desc='mask file')
+
+
+class RescaleB0OutputSpec(TraitedSpec):
+    out_file = File(exists=True, desc='b0 file')
+
+
+class RescaleB0(SimpleInterface):
+    """
+    Rescale the b0 volumes to deal with average signal decay over time
+    and output the median image.
+
+    Example
+    -------
+
+    >>> os.chdir(tmpdir)
+    >>> rescale_b0 = RescaleB0()
+    >>> rescale_b0.inputs.in_file = str(data_dir / 'dwi.nii.gz')
+    >>> rescale_b0.inputs.mask_file = str(data_dir / 'dwi_mask.nii.gz')
+    >>> res = rescale_b0.run()  # doctest: +SKIP
+
+    """
+
+    input_spec = RescaleB0InputSpec
+    output_spec = RescaleB0OutputSpec
+
+    def _run_interface(self, runtime):
+        self._results['out_file'] = rescale_b0(
+            self.inputs.in_file,
+            self.inputs.mask_file,
+            newpath=runtime.cwd)
+        return runtime
+
+
+def rescale_b0(in_file, mask_file, newpath=None):
+    """
+    Rescale the *b0* volumes from a DWI dataset."""
+    import numpy as np
+    import nibabel as nib
+    from nipype.utils.filemanip import fname_presuffix
+
+    out_file = fname_presuffix(
+        in_file, suffix='_median_b0', newpath=newpath)
+
+    img = nib.load(in_file)
+    data = img.get_fdata()
+
+    mask_img = nib.load(mask_file)
+    mask_data = mask_img.get_fdata()
+
+    mean_b0_signals = data[mask_data > 0, ...].mean(axis=0)
+
+    rescale_b0 = 1000 * data / mean_b0_signals
+
+    median_b0 = np.median(rescale_b0, axis=-1)
+
+    hdr = img.header.copy()
+    hdr.set_data_shape(median_b0.shape)
+    hdr.set_xyzt_units('mm')
+    hdr.set_data_dtype(np.float32)
+    nib.Nifti1Image(median_b0, img.affine, hdr).to_filename(out_file)
+    return out_file