Skip to content

Commit

Permalink
Merge pull request #81 from josephmje/enh/skullstrip
Browse files Browse the repository at this point in the history
ENH: Update image utility output path behaviour
  • Loading branch information
oesteban authored Mar 20, 2020
2 parents 199a496 + a0e205c commit 23d8db4
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 73 deletions.
112 changes: 39 additions & 73 deletions dmriprep/interfaces/images.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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):
Expand Down Expand Up @@ -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
70 changes: 70 additions & 0 deletions dmriprep/utils/images.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 23d8db4

Please sign in to comment.