Merge pull request #1227 from nipreps/enh/iqms-piesno

ENH: Integrate PIESNO noise mask and sigma estimation

mriqc/interfaces/diffusion.py

@@ -56,6 +56,7 @@
     'ExtractOrientations',
     'FilterShells',
     'NumberOfShells',
+    'PIESNO',
     'ReadDWIMetadata',
     'SpikingVoxelsMask',
     'SplitShells',
@@ -129,6 +130,7 @@ class _DiffusionQCInputSpec(_BaseInterfaceInputSpec):
         minlen=1,
         desc='q-space nearest neighbor pairs',
     )
+    piesno_sigma = traits.Float(-1.0, usedefault=True, desc='noise sigma calculated with PIESNO')
 
 
 class _DiffusionQCOutputSpec(TraitedSpec):
@@ -139,6 +141,8 @@ class _DiffusionQCOutputSpec(TraitedSpec):
     fber = traits.Dict
     fd = traits.Dict
     ndc = traits.Float
+    sigma_cc = traits.Float
+    sigma_piesno = traits.Float
     spikes_ppm = traits.Dict
     # gsr = traits.Dict
     # tsnr = traits.Float
@@ -214,13 +218,12 @@ def _run_interface(self, runtime):
         rois = {
             'fg': mskdata,
             'bg': 1.0 - mskdata,
-            'cc': ccdata,
             'wm': wmdata,
         }
         stats = aqc.summary_stats(b0data, rois)
         self._results['summary'] = stats
 
-        self._results['cc_snr'] = dqc.cc_snr(
+        self._results['cc_snr'], cc_sigma = dqc.cc_snr(
             in_b0=b0data,
             dwi_shells=shelldata,
             cc_mask=ccdata,
@@ -275,6 +278,10 @@ def _run_interface(self, runtime):
             )
         )
 
+        # PIESNO
+        self._results['sigma_piesno'] = round(self.inputs.piesno_sigma, 4)
+        self._results['sigma_cc'] = round(float(cc_sigma), 4)
+
         self._results['out_qc'] = _flatten_dict(self._results)
 
         return runtime
@@ -1027,6 +1034,51 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class _PIESNOInputSpec(_BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc='a DWI 4D file')
+    n_channels = traits.Int(4, usedefault=True, min=1, desc='number of channels')
+
+
+class _PIESNOOutputSpec(_TraitedSpec):
+    sigma = traits.Float(desc='noise sigma calculated with PIESNO')
+    out_mask = File(exists=True, desc='a 4D binary mask of spiking voxels')
+
+
+class PIESNO(SimpleInterface):
+    """Computes :abbr:`QC (Quality Control)` measures on the input DWI EPI scan."""
+
+    input_spec = _PIESNOInputSpec
+    output_spec = _PIESNOOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results['out_mask'] = fname_presuffix(
+            self.inputs.in_file,
+            suffix='piesno',
+            newpath=runtime.cwd,
+        )
+
+        in_nii = nb.load(self.inputs.in_file)
+        data = np.round(in_nii.get_fdata(), 4).astype('float32')
+
+        sigma, maskdata = noise_piesno(data)
+
+        header = in_nii.header.copy()
+        header.set_data_dtype(np.uint8)
+        header.set_xyzt_units('mm')
+        header.set_intent('estimate', name='PIESNO noise voxels mask')
+        header['cal_max'] = 1
+        header['cal_min'] = 0
+
+        nb.Nifti1Image(
+            maskdata.astype(np.uint8),
+            in_nii.affine,
+            header,
+        ).to_filename(self._results['out_mask'])
+
+        self._results['sigma'] = round(float(np.median(sigma)), 5)
+        return runtime
+
+
 def _rms(estimator, X):
     """
     Callable to pass to GridSearchCV that will calculate a distance score.
@@ -1252,3 +1304,36 @@ def _find_qspace_neighbors(
         dwi_neighbors.append((dwi_index, neighbor_index))
 
     return dwi_neighbors
+
+
+def noise_piesno(data: np.ndarray, n_channels: int = 4) -> (np.ndarray, np.ndarray):
+    """
+    Estimates noise in raw diffusion MRI (dMRI) data using the PIESNO algorithm.
+
+    This function implements the PIESNO (Probabilistic Identification and Estimation
+    of Noise) algorithm [Koay2009]_ to estimate the standard deviation (sigma) of the
+    noise in each voxel of a 4D dMRI data array. The PIESNO algorithm assumes Rician
+    distributed signal and exploits the statistical properties of the noise to
+    separate it from the underlying signal.
+
+    Parameters
+    ----------
+    data : :obj:`~numpy.ndarray`
+        The 4D raw dMRI data array.
+    n_channels : :obj:`int`, optional (default=4)
+        The number of diffusion-encoding channels in the data. This value is used
+        internally by the PIESNO algorithm.
+
+    Returns
+    -------
+    sigma : :obj:`~numpy.ndarray`
+        The estimated noise standard deviation for each voxel in the data array.
+    mask : :obj:`~numpy.ndarray`
+        A brain mask estimated by PIESNO. This mask identifies voxels containing
+        mostly noise and can be used for further processing.
+
+    """
+    from dipy.denoise.noise_estimate import piesno
+
+    sigma, mask = piesno(data, N=n_channels, return_mask=True)
+    return sigma, mask

mriqc/qc/diffusion.py

@@ -58,7 +58,7 @@
 
 Noise in raw dMRI estimated with PIESNO (``piesno_sigma``)
     Employs PIESNO (Probabilistic Identification and Estimation
-    of Noise) algorithm [1]_ to estimate the standard deviation (sigma) of the
+    of Noise) algorithm [Koay2009]_ to estimate the standard deviation (sigma) of the
     noise in each voxel of a 4D dMRI data array.
 
 .. _iqms_cc_snr:
@@ -88,6 +88,12 @@
     The spikes mask is calculated by identifying voxels with signal intensities
     exceeding a threshold based on standard deviations above the mean.
 
+References
+----------
+.. [Koay2009] Koay C.G., E. Ozarslan, C. Pierpaoli. Probabilistic Identification
+       and Estimation of Noise (PIESNO): A self-consistent approach and
+       its applications in MRI. JMR, 199(1):94-103, 2009.
+
 """
 from __future__ import annotations
 
@@ -147,45 +153,6 @@ def noise_b0(
     return noise_estimates
 
 
-def noise_piesno(data: np.ndarray, n_channels: int = 4) -> (np.ndarray, np.ndarray):
-    """
-    Estimates noise in raw diffusion MRI (dMRI) data using the PIESNO algorithm.
-
-    This function implements the PIESNO (Probabilistic Identification and Estimation
-    of Noise) algorithm [1]_ to estimate the standard deviation (sigma) of the
-    noise in each voxel of a 4D dMRI data array. The PIESNO algorithm assumes Rician
-    distributed signal and exploits the statistical properties of the noise to
-    separate it from the underlying signal.
-
-    Parameters
-    ----------
-    data : :obj:`~numpy.ndarray`
-        The 4D raw dMRI data array.
-    n_channels : :obj:`int`, optional (default=4)
-        The number of diffusion-encoding channels in the data. This value is used
-        internally by the PIESNO algorithm.
-
-    Returns
-    -------
-    sigma : :obj:`~numpy.ndarray`
-        The estimated noise standard deviation for each voxel in the data array.
-    mask : :obj:`~numpy.ndarray`
-        A brain mask estimated by PIESNO. This mask identifies voxels containing
-        mostly noise and can be used for further processing.
-
-    Notes
-    -----
-
-    .. [1] Koay C.G., E. Ozarslan, C. Pierpaoli. Probabilistic Identification
-           and Estimation of Noise (PIESNO): A self-consistent approach and
-           its applications in MRI. JMR, 199(1):94-103, 2009.
-    """
-    from dipy.denoise.noise_estimate import piesno
-
-    sigma, mask = piesno(data, N=n_channels, return_mask=True)
-    return sigma, mask
-
-
 def cc_snr(
     in_b0: np.ndarray,
     dwi_shells: list[np.ndarray],
@@ -268,7 +235,7 @@ def cc_snr(
             np.mean(mean_signal_best / std_signal),
         )
 
-    return cc_snr_estimates
+    return cc_snr_estimates, std_signal
 
 
 def spike_ppm(

mriqc/workflows/diffusion/base.py

@@ -71,6 +71,7 @@ def dmri_qc_workflow(name='dwiMRIQC'):
     from niworkflows.interfaces.images import RobustAverage
 
     from mriqc.interfaces.diffusion import (
+        PIESNO,
         CCSegmentation,
         CorrectSignalDrift,
         DipyDTI,
@@ -195,6 +196,9 @@ def dmri_qc_workflow(name='dwiMRIQC'):
     # Calculate CC mask
     cc_mask = pe.Node(CCSegmentation(), name='cc_mask')
 
+    # Run PIESNO noise estimation
+    piesno = pe.Node(PIESNO(), name='piesno')
+
     # EPI to MNI registration
     spatial_norm = epi_mni_align()
 
@@ -215,6 +219,7 @@ def dmri_qc_workflow(name='dwiMRIQC'):
         (datalad_get, sanitize, [('in_file', 'in_file')]),
         (sanitize, dwi_ref, [('out_file', 'in_file')]),
         (sanitize, sp_mask, [('out_file', 'in_file')]),
+        (sanitize, piesno, [('out_file', 'in_file')]),
         (shells, dwi_ref, [(('b_masks', _first), 't_mask')]),
         (shells, sp_mask, [('b_masks', 'b_masks')]),
         (meta, shells, [('out_bval_file', 'in_bvals')]),
@@ -252,6 +257,7 @@ def dmri_qc_workflow(name='dwiMRIQC'):
         (meta, iqms_wf, [('qspace_neighbors', 'inputnode.qspace_neighbors')]),
         (averages, iqms_wf, [(('out_file', _first), 'inputnode.in_b0')]),
         (sp_mask, iqms_wf, [('out_mask', 'inputnode.spikes_mask')]),
+        (piesno, iqms_wf, [('sigma', 'inputnode.piesno_sigma')]),
         (hmcwf, iqms_wf, [('outputnode.out_fd', 'inputnode.framewise_displacement')]),
         (dti, iqms_wf, [('out_fa', 'inputnode.in_fa'),
                         ('out_cfa', 'inputnode.in_cfa'),
@@ -324,6 +330,7 @@ def compute_iqms(name='ComputeIQMs'):
                 'spikes_mask',
                 'framewise_displacement',
                 'qspace_neighbors',
+                'piesno_sigma',
             ]
         ),
         name='inputnode',
@@ -380,7 +387,8 @@ def compute_iqms(name='ComputeIQMs'):
                                ('in_fa_nans', 'in_fa_nans'),
                                ('in_fa_degenerate', 'in_fa_degenerate'),
                                ('framewise_displacement', 'in_fd'),
-                               ('qspace_neighbors', 'qspace_neighbors')]),
+                               ('qspace_neighbors', 'qspace_neighbors'),
+                               ('piesno_sigma', 'piesno_sigma')]),
         (inputnode, addprov, [('in_file', 'in_file')]),
         (addprov, datasink, [('out_prov', 'provenance')]),
         (meta, datasink, [('subject', 'subject_id'),