Add synthstrip for ASLRef (#605)

Co-authored-by: Taylor Salo <[email protected]>

Author: mattcieslak

aslprep/interfaces/freesurfer.py

@@ -0,0 +1,156 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+"""
+SynthStrip interfaces
+~~~~~~~~~~~~~~~~~~~~~
+
+"""
+
+import os.path as op
+
+import nibabel as nb
+from nipype.interfaces.afni import Zeropad
+from nipype.interfaces.base import (
+    BaseInterfaceInputSpec,
+    File,
+    SimpleInterface,
+    TraitedSpec,
+    traits,
+)
+from nipype.interfaces.freesurfer.base import FSCommandOpenMP, FSTraitedSpec
+from nipype.utils.filemanip import fname_presuffix
+from niworkflows.utils.images import _copyxform
+
+
+class FSTraitedSpecOpenMP(FSTraitedSpec):
+    num_threads = traits.Int(desc='allows for specifying more threads', nohash=True)
+
+
+class _PrepareSynthStripGridInputSpec(BaseInterfaceInputSpec):
+    input_image = File(exists=True, mandatory=True)
+
+
+class _PrepareSynthStripGridOutputSpec(TraitedSpec):
+    prepared_image = File(exists=True)
+
+
+class PrepareSynthStripGrid(SimpleInterface):
+    input_spec = _PrepareSynthStripGridInputSpec
+    output_spec = _PrepareSynthStripGridOutputSpec
+
+    def _run_interface(self, runtime):
+        out_fname = fname_presuffix(
+            self.inputs.input_image,
+            newpath=runtime.cwd,
+            suffix='_SynthStripGrid.nii',
+            use_ext=False,
+        )
+        self._results['prepared_image'] = out_fname
+
+        # possibly downsample the image for sloppy mode. Always ensure float32
+        img = nb.load(self.inputs.input_image)
+        if not img.ndim == 3:
+            raise Exception('3D inputs are required for Synthstrip')
+        xvoxels, yvoxels, zvoxels = img.shape
+
+        def get_padding(nvoxels):
+            extra_slices = nvoxels % 64
+            if extra_slices == 0:
+                return 0
+            complete_64s = nvoxels // 64
+            return 64 * (complete_64s + 1) - nvoxels
+
+        def split_padding(padding):
+            halfpad = padding // 2
+            return halfpad, halfpad + halfpad % 2
+
+        spad = get_padding(zvoxels)
+        rpad, lpad = split_padding(get_padding(xvoxels))
+        apad, ppad = split_padding(get_padding(yvoxels))
+
+        zeropad = Zeropad(
+            S=spad,
+            R=rpad,
+            L=lpad,
+            A=apad,
+            P=ppad,
+            in_files=self.inputs.input_image,
+            out_file=out_fname,
+        )
+
+        _ = zeropad.run()
+        if not op.exists(out_fname):
+            raise Exception('zeropad failed! You may need to increase the memory limit.')
+        return runtime
+
+
+class _SynthStripInputSpec(FSTraitedSpecOpenMP):
+    input_image = File(argstr='-i %s', exists=True, mandatory=True)
+    no_csf = traits.Bool(argstr='--no-csf', desc='Exclude CSF from brain border.')
+    border = traits.Int(argstr='-b %d', desc='Mask border threshold in mm. Default is 1.')
+    gpu = traits.Bool(argstr='-g')
+    out_brain = File(
+        argstr='-o %s',
+        name_template='%s_brain.nii.gz',
+        name_source=['input_image'],
+        keep_extension=False,
+        desc='skull stripped image with corrupt sform',
+    )
+    out_brain_mask = File(
+        argstr='-m %s',
+        name_template='%s_mask.nii.gz',
+        name_source=['input_image'],
+        keep_extension=False,
+        desc='mask image with corrupt sform',
+    )
+
+
+class _SynthStripOutputSpec(TraitedSpec):
+    out_brain = File(exists=True)
+    out_brain_mask = File(exists=True)
+
+
+class SynthStrip(FSCommandOpenMP):
+    input_spec = _SynthStripInputSpec
+    output_spec = _SynthStripOutputSpec
+    _cmd = 'mri_synthstrip'
+
+    def _num_threads_update(self):
+        if self.inputs.num_threads:
+            self.inputs.environ.update({'OMP_NUM_THREADS': '1'})
+
+
+class FixHeaderSynthStrip(SynthStrip):
+    def _run_interface(self, runtime, correct_return_codes=(0,)):
+        # Run normally
+        runtime = super()._run_interface(runtime, correct_return_codes)
+
+        outputs = self._list_outputs()
+        if not op.exists(outputs['out_brain']):
+            raise Exception('mri_synthstrip failed!')
+
+        if outputs.get('out_brain_mask'):
+            _copyxform(self.inputs.input_image, outputs['out_brain_mask'])
+
+        _copyxform(self.inputs.input_image, outputs['out_brain'])
+
+        return runtime
+
+
+class MockSynthStrip(SimpleInterface):
+    input_spec = _SynthStripInputSpec
+    output_spec = _SynthStripOutputSpec
+
+    def _run_interface(self, runtime):
+        from nipype.interfaces.fsl import BET
+
+        this_bet = BET(
+            mask=True,
+            in_file=self.inputs.input_image,
+            output_type='NIFTI_GZ',
+        )
+        result = this_bet.run()
+        self._results['out_brain'] = result.outputs.out_file
+        self._results['out_brain_mask'] = result.outputs.mask_file
+
+        return runtime

aslprep/workflows/asl/cbf.py

@@ -33,7 +33,7 @@
 )
 from aslprep.utils.atlas import get_atlas_nifti
 from aslprep.utils.bids import find_atlas_entities
-from aslprep.workflows.asl.reference import init_enhance_and_skullstrip_asl_wf
+from aslprep.workflows.asl.reference import init_synthstrip_aslref_wf
 
 
 def init_cbf_wf(
@@ -395,9 +395,8 @@ def _getfiledir(file):
             (mean_m0, extract_deltam, [('out_file', 'm0scan')]),
         ])  # fmt:skip
 
-        enhance_and_skullstrip_m0scan_wf = init_enhance_and_skullstrip_asl_wf(
+        enhance_and_skullstrip_m0scan_wf = init_synthstrip_aslref_wf(
             disable_n4=config.workflow.disable_n4,
-            omp_nthreads=1,
             name='enhance_and_skullstrip_m0scan_wf',
         )
         workflow.connect([

aslprep/workflows/asl/fit.py

@@ -53,7 +53,7 @@
 from aslprep.utils.asl import select_processing_target
 from aslprep.workflows.asl.hmc import init_asl_hmc_wf
 from aslprep.workflows.asl.outputs import init_asl_fit_reports_wf, init_ds_aslref_wf
-from aslprep.workflows.asl.reference import init_enhance_and_skullstrip_asl_wf, init_raw_aslref_wf
+from aslprep.workflows.asl.reference import init_raw_aslref_wf, init_synthstrip_aslref_wf
 
 
 def get_sbrefs(
@@ -621,9 +621,8 @@ def init_asl_fit_wf(
             )
             workflow.connect(raw_sbref_wf, 'outputnode.aslref', fmapref_buffer, 'sbref_files')
 
-        enhance_aslref_wf = init_enhance_and_skullstrip_asl_wf(
+        enhance_aslref_wf = init_synthstrip_aslref_wf(
             disable_n4=config.workflow.disable_n4,
-            omp_nthreads=omp_nthreads,
         )
 
         ds_coreg_aslref_wf = init_ds_aslref_wf(

aslprep/workflows/asl/reference.py

@@ -422,3 +422,103 @@ def init_enhance_and_skullstrip_asl_wf(
     ])  # fmt: skip
 
     return workflow
+
+
+def init_synthstrip_aslref_wf(
+    disable_n4=False,
+    name='synthstrip_aslref_wf',
+):
+    """Enhance and run brain extraction on an ASL image.
+
+    This workflow uses SynthStrip to extract the brain from an ASLRef image.
+
+    Parameters
+    ----------
+    disable_n4 : :obj:`bool`
+        If True, N4 bias field correction will be disabled.
+    name : str
+        Name of workflow (default: ``synthstrip_aslref_wf``)
+
+    Inputs
+    ------
+    in_file : str
+        ASLRef image
+
+    Outputs
+    -------
+    bias_corrected_file : str
+        the ``in_file`` after N4 bias field correction.
+        If ``disable_n4`` is True, this will be the same as ``in_file``.
+    skull_stripped_file : str
+        the ``bias_corrected_file`` after skull-stripping
+    mask_file : str
+        mask of the skull-stripped input file
+    """
+    from nipype.interfaces import afni
+    from nipype.interfaces import utility as niu
+    from nipype.pipeline import engine as pe
+    from niworkflows.engine.workflows import LiterateWorkflow as Workflow
+    from niworkflows.interfaces.fixes import FixN4BiasFieldCorrection as N4BiasFieldCorrection
+    from niworkflows.interfaces.header import CopyXForm
+    from niworkflows.interfaces.nibabel import ApplyMask
+
+    from aslprep.workflows.segmentation import init_synthstrip_wf
+
+    workflow = Workflow(name=name)
+    inputnode = pe.Node(niu.IdentityInterface(fields=['in_file', 'pre_mask']), name='inputnode')
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=['mask_file', 'skull_stripped_file', 'bias_corrected_file']),
+        name='outputnode',
+    )
+
+    n4_buffer = pe.Node(niu.IdentityInterface(fields=['bias_corrected_file']), name='n4_buffer')
+
+    if not disable_n4:
+        # Run N4 normally, force num_threads=1 for stability (images are small, no need for >1)
+        n4_correct = pe.Node(
+            N4BiasFieldCorrection(dimension=3, copy_header=True, bspline_fitting_distance=200),
+            shrink_factor=2,
+            name='n4_correct',
+            n_procs=1,
+        )
+        n4_correct.inputs.rescale_intensities = True
+
+        workflow.connect([
+            (inputnode, n4_correct, [('in_file', 'input_image')]),
+            (n4_correct, n4_buffer, [('output_image', 'bias_corrected_file')]),
+        ])  # fmt:skip
+    else:
+        workflow.connect([(inputnode, n4_buffer, [('in_file', 'bias_corrected_file')])])
+
+    # Use AFNI's unifize for T2 contrast & fix header
+    unifize = pe.Node(
+        afni.Unifize(
+            t2=True,
+            outputtype='NIFTI_GZ',
+            # Default -clfrac is 0.1, 0.4 was too conservative
+            # -rbt because I'm a Jedi AFNI Master (see 3dUnifize's documentation)
+            args='-clfrac 0.2 -rbt 18.3 65.0 90.0',
+            out_file='uni.nii.gz',
+        ),
+        name='unifize',
+    )
+    fixhdr_unifize = pe.Node(CopyXForm(), name='fixhdr_unifize', mem_gb=0.1)
+
+    synthstrip_wf = init_synthstrip_wf(name='synthstrip_wf')
+
+    # Compute masked brain
+    apply_mask = pe.Node(ApplyMask(), name='apply_mask')
+
+    workflow.connect([
+        (inputnode, fixhdr_unifize, [('in_file', 'hdr_file')]),
+        (inputnode, synthstrip_wf, [('in_file', 'inputnode.original_image')]),
+        (inputnode, unifize, [('in_file', 'in_file')]),
+        (unifize, fixhdr_unifize, [('out_file', 'in_file')]),
+        (synthstrip_wf, apply_mask, [('outputnode.brain_mask', 'in_mask')]),
+        (fixhdr_unifize, apply_mask, [('out_file', 'in_file')]),
+        (apply_mask, outputnode, [('out_file', 'skull_stripped_file')]),
+        (n4_buffer, outputnode, [('bias_corrected_file', 'bias_corrected_file')]),
+        (synthstrip_wf, outputnode, [('outputnode.brain_mask', 'mask_file')]),
+    ])  # fmt: skip
+
+    return workflow