initial b0 registration adapted from niworkflows

Author: josephmje

dmriprep/interfaces/registration.py

@@ -0,0 +1,72 @@
+"""Registration tools interfaces."""
+
+import os
+import nibabel as nb
+import numpy as np
+from nipype.interfaces.base import (
+    traits, TraitedSpec, BaseInterfaceInputSpec, File, SimpleInterface)
+from nipype.interfaces import fsl, afni
+
+
+class EstimateReferenceImageInputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc="b0 file")
+    pre_mask = File(exists=True, mandatory=True, desc="rough mask of skull-stripped b0")
+    mc_method = traits.Enum(
+        "AFNI", "FSL", usedefault=True,
+        desc="Which software to use to perform motion correction")
+
+
+class EstimateReferenceImageOutputSpec(TraitedSpec):
+    ref_image = File(exists=True, desc="3D reference image")
+
+
+class EstimateReferenceImage(SimpleInterface):
+    """
+    Given an 4D EPI file estimate an optimal reference image that could be later
+    used for motion estimation and coregistration purposes. If detected uses
+    T1 saturated volumes (non-steady state) otherwise a median of
+    of a subset of motion corrected volumes is used.
+    """
+    input_spec = EstimateReferenceImageInputSpec
+    output_spec = EstimateReferenceImageOutputSpec
+
+    def _run_interface(self, runtime):
+        ref_name = self.inputs.in_file
+        pre_mask_name = self.inputs.pre_mask
+
+        ref_nii = nb.load(ref_name)
+        pre_mask_nii = nb.load(pre_mask_name)
+
+        out_ref_fname = os.path.join(runtime.cwd, "ref_dwi.nii.gz")
+
+        # If reference is only 1 volume, return it directly
+        if len(ref_nii.shape) == 3:
+            ref_nii.header.extensions.clear()
+            ref_nii.to_filename(out_ref_fname)
+            self._results['ref_image'] = out_ref_fname
+            return runtime
+        else:
+            ref_nii.header.extensions.clear()
+
+            if self.inputs.mc_method == "AFNI":
+                res = afni.Volreg(in_file=ref_name, args='-Fourier -twopass',
+                                  zpad=4, outputtype='NIFTI_GZ').run()
+            elif self.inputs.mc_method == "FSL":
+                res = fsl.MCFLIRT(in_file=ref_name,
+                                  ref_vol=0, interpolation='sinc').run()
+            mc_ref_nii = nb.load(res.outputs.out_file)
+
+            pre_mask_data = pre_mask_nii.get_fdata()
+
+            image_data = mc_ref_nii.get_fdata()
+
+            normalized_image_data = 1000 * image_data / image_data[pre_mask_data, ...].mean(axis=-1)
+
+            median_image_data = np.median(normalized_image_data, axis=-1)
+
+            nb.Nifti1Image(median_image_data, ref_nii.affine,
+                           ref_nii.header).to_filename(out_ref_fname)
+
+            self._results["ref_image"] = out_ref_fname
+
+            return runtime

dmriprep/workflows/dwi/util.py

@@ -1,6 +1,3 @@
-# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
-# vi: set ft=python sts=4 ts=4 sw=4 et:
-
 """
 Utility workflows
 ^^^^^^^^^^^^^^^^^
@@ -19,7 +16,8 @@
 from niworkflows.interfaces.masks import SimpleShowMaskRPT
 from niworkflows.interfaces.utils import CopyXForm
 
-from ...interfaces.images import MeanB0
+from ...interfaces.images import ExtractB0
+from ...interfaces.registration import EstimateReferenceImage
 
 DEFAULT_MEMORY_MIN_GB = 0.01
 
@@ -57,7 +55,7 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
         dwi_file
             Validated dwi NIfTI file
         raw_ref_image
-            Reference image to which dwi series is motion corrected
+            Reference image
         ref_image
             Contrast-enhanced reference image
         ref_image_brain
@@ -92,19 +90,21 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
 
     validate = pe.Node(ValidateImage(), name='validate', mem_gb=DEFAULT_MEMORY_MIN_GB)
 
-    gen_ref = pe.Node(MeanB0(), name="gen_ref")
+    extract_b0 = pe.Node(ExtractB0(), name="extract_b0")
 
     pre_mask = pe.Node(afni.Automask(outputtype="NIFTI_GZ"), name="pre_mask")
 
+    gen_ref = pe.Node(EstimateReferenceImage(), name='gen_ref')
+
     enhance_and_skullstrip_dwi_wf = init_enhance_and_skullstrip_dwi_wf(
         omp_nthreads=omp_nthreads)
 
     workflow.connect([
         (inputnode, validate, [('dwi_file', 'in_file')]),
-        (validate, gen_ref, [('out_file', 'in_file')]),
-        (inputnode, gen_ref, [('bvec_file', 'in_bvec'),
-                              ('bval_file', 'in_bval')]),
-        (gen_ref, pre_mask, [('out_file', 'in_file')]),
+        (validate, extract_b0, [('out_file', 'in_file')]),
+        (extract_b0, pre_mask, [('out_file', 'in_file')]),
+        (extract_b0, gen_ref, [('out_file', 'in_file')]),
+        (pre_mask, gen_ref, [('out_file', 'pre_mask')]),
         (gen_ref, enhance_and_skullstrip_dwi_wf, [('out_file', 'inputnode.in_file')]),
         (pre_mask, enhance_and_skullstrip_dwi_wf, [('out_file', 'inputnode.pre_mask')]),
         (validate, outputnode, [('out_file', 'dwi_file'),
@@ -132,7 +132,7 @@ def init_enhance_and_skullstrip_dwi_wf(
         name='enhance_and_skullstrip_dwi_wf',
         omp_nthreads=1):
     """
-    This workflow takes in a dwi reference iamge and sharpens the histogram
+    This workflow takes in a dwi reference image and sharpens the histogram
     with the application of the N4 algorithm for removing the
     :abbr:`INU (intensity non-uniformity)` bias field and calculates a signal
     mask.