Merge pull request #402 from oesteban/enh/skullstrip-wf-no-fsl

ENH: Add a new ``Binarize`` interface using nibabel

Author: oesteban

niworkflows/anat/skullstrip.py

@@ -1,41 +1,39 @@
-# -*- coding: utf-8 -*-
 # emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
-from __future__ import absolute_import, division, print_function, unicode_literals
-from nipype.interfaces import ants, afni, fsl, utility as niu
+"""Brain extraction workflows."""
+from nipype.interfaces import afni, utility as niu
 from nipype.pipeline import engine as pe
+from ..interfaces.nibabel import Binarize
+from ..interfaces.fixes import FixN4BiasFieldCorrection as N4BiasFieldCorrection
 
 
 def afni_wf(name='AFNISkullStripWorkflow', unifize=False, n4_nthreads=1):
     """
-    Skull-stripping workflow
+    Create a skull-stripping workflow based on AFNI's tools.
 
-    Originally derived from the `codebase of the
-    QAP <https://github.com/preprocessed-connectomes-project/\
-quality-assessment-protocol/blob/master/qap/anatomical_preproc.py#L105>`_.
+    Originally derived from the `codebase of the QAP
+    <https://github.com/preprocessed-connectomes-project/quality-assessment-protocol/blob/master/qap/anatomical_preproc.py#L105>`_.
     Now, this workflow includes :abbr:`INU (intensity non-uniformity)` correction
     using the N4 algorithm and (optionally) intensity harmonization using
     ANFI's ``3dUnifize``.
 
-
     """
-
     workflow = pe.Workflow(name=name)
     inputnode = pe.Node(niu.IdentityInterface(fields=['in_file']),
                         name='inputnode')
     outputnode = pe.Node(niu.IdentityInterface(
         fields=['bias_corrected', 'out_file', 'out_mask', 'bias_image']), name='outputnode')
 
     inu_n4 = pe.Node(
-        ants.N4BiasFieldCorrection(dimension=3, save_bias=True, num_threads=n4_nthreads,
-                                   copy_header=True),
+        N4BiasFieldCorrection(dimension=3, save_bias=True, num_threads=n4_nthreads,
+                              rescale_intensities=True, copy_header=True),
         n_procs=n4_nthreads,
         name='inu_n4')
 
     sstrip = pe.Node(afni.SkullStrip(outputtype='NIFTI_GZ'), name='skullstrip')
     sstrip_orig_vol = pe.Node(afni.Calc(
         expr='a*step(b)', outputtype='NIFTI_GZ'), name='sstrip_orig_vol')
-    binarize = pe.Node(fsl.Threshold(args='-bin', thresh=1.e-3), name='binarize')
+    binarize = pe.Node(Binarize(thresh_low=0.0), name='binarize')
 
     if unifize:
         # Add two unifize steps, pre- and post- skullstripping.
@@ -64,7 +62,7 @@ def afni_wf(name='AFNISkullStripWorkflow', unifize=False, n4_nthreads=1):
         (sstrip, sstrip_orig_vol, [('out_file', 'in_file_b')]),
         (inputnode, inu_n4, [('in_file', 'input_image')]),
         (sstrip_orig_vol, binarize, [('out_file', 'in_file')]),
-        (binarize, outputnode, [('out_file', 'out_mask')]),
+        (binarize, outputnode, [('out_mask', 'out_mask')]),
         (inu_n4, outputnode, [('bias_image', 'bias_image')]),
     ])
     return workflow

niworkflows/interfaces/nibabel.py

@@ -0,0 +1,52 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+"""Nibabel-based interfaces."""
+
+import nibabel as nb
+from nipype import logging
+from nipype.utils.filemanip import fname_presuffix
+from nipype.interfaces.base import (
+    traits, TraitedSpec, BaseInterfaceInputSpec, File,
+    SimpleInterface
+)
+
+IFLOGGER = logging.getLogger('nipype.interface')
+
+
+class _BinarizeInputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc='input image')
+    thresh_low = traits.Float(mandatory=True,
+                              desc='non-inclusive lower threshold')
+
+
+class _BinarizeOutputSpec(TraitedSpec):
+    out_file = File(exists=True, desc='masked file')
+    out_mask = File(exists=True, desc='output mask')
+
+
+class Binarize(SimpleInterface):
+    """Binarizes the input image applying the given thresholds."""
+
+    input_spec = _BinarizeInputSpec
+    output_spec = _BinarizeOutputSpec
+
+    def _run_interface(self, runtime):
+        img = nb.load(self.inputs.in_file)
+
+        self._results['out_file'] = fname_presuffix(
+            self.inputs.in_file, suffix='_masked', newpath=runtime.cwd)
+        self._results['out_mask'] = fname_presuffix(
+            self.inputs.in_file, suffix='_mask', newpath=runtime.cwd)
+
+        data = img.get_fdata()
+        mask = data > self.inputs.thresh_low
+        data[~mask] = 0.0
+        masked = img.__class__(data, img.affine, img.header)
+        masked.to_filename(self._results['out_file'])
+
+        img.header.set_data_dtype('uint8')
+        maskimg = img.__class__(mask.astype('uint8'), img.affine,
+                                img.header)
+        maskimg.to_filename(self._results['out_mask'])
+
+        return runtime

niworkflows/interfaces/tests/test_nibabel.py

@@ -0,0 +1,24 @@
+"""test nibabel interfaces."""
+import os
+import numpy as np
+import nibabel as nb
+
+from ..nibabel import Binarize
+
+
+def test_Binarize(tmp_path):
+    """Test binarization interface."""
+    os.chdir(str(tmp_path))
+
+    mask = np.zeros((20, 20, 20), dtype=bool)
+    mask[5:15, 5:15, 5:15] = bool
+
+    data = np.zeros_like(mask, dtype='float32')
+    data[mask] = np.random.gamma(2, size=mask.sum())
+
+    in_file = tmp_path / 'input.nii.gz'
+    nb.Nifti1Image(data, np.eye(4), None).to_filename(str(in_file))
+
+    binif = Binarize(thresh_low=0.0, in_file=str(in_file)).run()
+    newmask = nb.load(binif.outputs.out_mask).get_fdata().astype(bool)
+    assert np.all(mask == newmask)