fix masking

Author: josephmje

dmriprep/interfaces/images.py

@@ -25,9 +25,12 @@ class ExtractB0(SimpleInterface):
     -------
 
     >>> os.chdir(tmpdir)
-    >>> check = ExtractB0(
-    ...    in_file=str(data_dir / 'dwi.nii.gz'),
-    ...    b0_ixs=[0, 9, 18, 27, 36, 45, 54, 63, 72, 81, 90, 100]).run()
+    >>> extract_b0 = ExtractB0()
+    >>> extract_b0.inputs.in_file = str(data_dir / 'dwi.nii.gz')
+    >>> extract_b0.inputs.b0_ixs = [0, 9, 18, 27, 36, 45, 54, 63, 72, 81, 90, 100]
+    >>> res = extract_b0.run()
+    >>> import nibabel as nb
+    >>> nb.load(res.outputs.out_file).shape
 
     """
     input_spec = ExtractB0InputSpec
@@ -61,3 +64,58 @@ def extract_b0(in_file, b0_ixs, newpath=None):
     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):
+
+    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 to deal with average signal decay over time
+    and calculate the median.
+    """
+    from dipy.segment.mask import applymask
+    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

dmriprep/interfaces/registration.py

@@ -16,8 +16,7 @@
 from niworkflows.interfaces.masks import SimpleShowMaskRPT
 from niworkflows.interfaces.utils import CopyXForm
 
-from ...interfaces.images import ExtractB0
-from ...interfaces.registration import EstimateReferenceImage
+from ...interfaces.images import ExtractB0, RescaleB0
 
 DEFAULT_MEMORY_MIN_GB = 0.01
 
@@ -49,6 +48,8 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
 
         dwi_file
             dwi NIfTI file
+        b0_ixs : list
+            index of b0s in dwi NIfTI file
 
     **Outputs**
 
@@ -76,11 +77,11 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
 First, a reference volume and its skull-stripped version were generated
 using a custom methodology taken from *fMRIPrep*.
 """
-    inputnode = pe.Node(niu.IdentityInterface(fields=['dwi_file', 'bvec_file', 'bval_file']),
+    inputnode = pe.Node(niu.IdentityInterface(fields=['dwi_file', 'b0_ixs']),
                         name='inputnode')
     outputnode = pe.Node(
-        niu.IdentityInterface(fields=['dwi_file', 'raw_ref_image',
-                                      'ref_image', 'ref_image_brain',
+        niu.IdentityInterface(fields=['dwi_file', 'raw_ref_image', 'ref_image',
+                                      'ref_image_brain',
                                       'dwi_mask', 'validation_report']),
         name='outputnode')
 
@@ -90,26 +91,30 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
 
     validate = pe.Node(ValidateImage(), name='validate', mem_gb=DEFAULT_MEMORY_MIN_GB)
 
-    extract_b0 = pe.Node(ExtractB0(), name="extract_b0")
+    extract_b0 = pe.Node(ExtractB0(), name='extract_b0')
 
-    pre_mask = pe.Node(afni.Automask(outputtype="NIFTI_GZ"), name="pre_mask")
+    # b0_template_wf = init_b0_template_wf()
 
-    gen_ref = pe.Node(EstimateReferenceImage(), name='gen_ref')
+    pre_mask = pe.Node(afni.Automask(dilate=1, outputtype='NIFTI_GZ'),
+                       name='pre_mask')
+
+    rescale_b0 = pe.Node(RescaleB0(), name='rescale_b0')
 
     enhance_and_skullstrip_dwi_wf = init_enhance_and_skullstrip_dwi_wf(
         omp_nthreads=omp_nthreads)
 
     workflow.connect([
         (inputnode, validate, [('dwi_file', 'in_file')]),
         (validate, extract_b0, [('out_file', 'in_file')]),
+        (inputnode, extract_b0, [('b0_ixs', 'b0_ixs')]),
         (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')]),
+        (extract_b0, rescale_b0, [('out_file', 'in_file')]),
+        (pre_mask, rescale_b0, [('out_file', 'pre_mask')]),
+        (rescale_b0, 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'),
                                 ('out_report', 'validation_report')]),
-        (gen_ref, outputnode, [('out_file', 'raw_ref_image')]),
+        (rescale_b0, outputnode, [('ref_image', 'raw_ref_image')]),
         (enhance_and_skullstrip_dwi_wf, outputnode, [
             ('outputnode.bias_corrected_file', 'ref_image'),
             ('outputnode.mask_file', 'dwi_mask'),
@@ -128,6 +133,16 @@ def init_dwi_reference_wf(omp_nthreads, dwi_file=None,
     return workflow
 
 
+def init_b0_template_wf(
+    name='b0_template_wf'
+):
+   """[summary]
+
+   Keyword Arguments:
+       name {str} -- [description] (default: {'b0_template_wf'})
+   """
+
+
 def init_enhance_and_skullstrip_dwi_wf(
         name='enhance_and_skullstrip_dwi_wf',
         omp_nthreads=1):
@@ -204,16 +219,14 @@ def init_enhance_and_skullstrip_dwi_wf(
         internal_datatype='char'), name='skullstrip_first_dilate')
     bet_mask = pe.Node(fsl.ApplyMask(), name='skullstrip_first_mask')
 
-    # Use AFNI's unifize for T2 constrast & fix header
+    # 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')
+        out_file='uni.nii.gz'), name='unifize')
     fixhdr_unifize = pe.Node(CopyXForm(), name='fixhdr_unifize', mem_gb=0.1)
 
-    # Run ANFI's 3dAutomask to extract a refined brain mask
+    # Run AFNI's 3dAutomask to extract a refined brain mask
     skullstrip_second_pass = pe.Node(afni.Automask(dilate=1,
                                                    outputtype='NIFTI_GZ'),
                                      name='skullstrip_second_pass')