Denoise nl_means script

scilpy/stats/stats.py

@@ -1,12 +1,60 @@
 # -*- coding: utf-8 -*-
 
-import logging
-
 from itertools import combinations
+import logging
 
+from dipy.denoise.noise_estimate import estimate_sigma, piesno
+import numpy as np
 import scipy.stats
 
 
+def estimate_piesno_sigma(data, number_coils=0):
+    """
+    Here are Dipy's note on this method:
+     > It is expected that
+     >   1. The data has a noisy, non-masked background and
+     >   2. The data is a repetition of the same measurements along the last
+     >      axis, i.e. dMRI or fMRI data, not structural data like T1/T2."
+
+    Parameters
+    ----------
+    data: np.ndarray
+        The 4D volume.
+    number_coils: int
+        The number of coils in the scanner.
+    """
+    assert len(data.shape) == 4
+
+    sigma = np.zeros(data.shape[:3], dtype=np.float32)
+    mask_noise = np.zeros(data.shape[:3], dtype=np.int16)
+
+    for idx in range(data.shape[-2]):
+        logging.info('Now processing slice {} / {}'
+                     .format(idx + 1, data.shape[-2]))
+        sigma[..., idx], mask_noise[..., idx] = \
+            piesno(data[..., idx, :], N=number_coils, return_mask=True)
+
+    # If the noise mask has few voxels, the detected noise standard
+    # deviation can be very low and maybe something went wrong. We
+    # check here that at least 1% of noisy voxels were found and warn
+    # the user otherwise.
+    frac_noisy_voxels = np.sum(mask_noise) / np.size(mask_noise) * 100
+
+    if frac_noisy_voxels < 1.:
+        logging.warning(
+            'PIESNO was used with N={}, but it found only {:.3f}% of voxels '
+            'as pure noise with a mean standard deviation of {:.5f}. This is '
+            'suspicious, so please check the resulting sigma volume if '
+            'something went wrong. In cases where PIESNO is not working, '
+            'you might want to try basic sigma estimation.'
+            .format(number_coils, frac_noisy_voxels, np.mean(sigma)))
+    else:
+        logging.info('The noise standard deviation from piesno is %s',
+                     np.array_str(sigma[0, 0, :]))
+
+    return sigma, mask_noise
+
+
 def verify_normality(data, alpha=0.05):
     """
     Parameters

scripts/scil_denoising_nlmeans.py

@@ -2,11 +2,52 @@
 # -*- coding: utf-8 -*-
 
 """
-Script to denoise a dataset with the Non Local Means algorithm.
+Script to denoise a dataset with the Non-Local Means algorithm (from Dipy's
+package).
+
+This method requires an estimate of the noise (sigma). There are three methods
+available:
+    1. Provide your own sigma value, if you already know it.
+    2. Estimate it through Dipy's "noise_estimate" method.
+    Here are Dipy's note on this method:
+     > This function is the same as manually taking the standard deviation of
+     > the background and gives one value for each 3D volume. It also includes
+     > the coil-dependent correction factor of Koay 2006 (see [1], equation 18)
+     > with theta = 0. Since this function was introduced in [2] for T1
+     > imaging, it is expected to perform ok on diffusion MRI data, but might
+     > oversmooth some regions and leave others un-denoised for spatially
+     > varying noise profiles. Consider using --piesno to estimate sigma
+     > instead if visual inaccuracies are apparent in the denoised result.
+    (Note. We then use the median of the noise for all volumes)
+    3. Estimate it through using Dipy's "Piesno" method, as described in [3].
+    Here are Dipy's note on this method:
+     > It is expected that
+     >   1. The data has a noisy, non-masked background and
+     >   2. The data is a repetition of the same measurements along the last
+     >      axis, i.e. dMRI or fMRI data, not structural data like T1 or T2
+     >      images.
+
+The two last methods require the number of coils in the data to be known.
+Typically, if you know that the noise follows a Gaussian distribution, you
+may use --number_coils=0.
+
+References:
+[1] Koay, C. G., & Basser, P. J. (2006). Analytically exact correction scheme
+for signal extraction from noisy magnitude MR signals. Journal of Magnetic
+Resonance, 179(2), 317-22.
+[2] Coupe, P., Yger, P., Prima, S., Hellier, P., Kervrann, C., Barillot,
+C., (2008). An optimized blockwise nonlocal means denoising filter for 3-D
+magnetic resonance images, IEEE Trans. Med. Imaging 27, 425-41.
+[3] St-Jean, S., Coupé, P., & Descoteaux, M. (2016). Non Local Spatial and
+Angular Matching: Enabling higher spatial resolution diffusion MRI datasets
+through adaptive denoising. Medical image analysis, 32, 115-130.
+[4] Koay CG, Ozarslan E and Pierpaoli C. "Probabilistic Identification and
+Estimation of Noise (PIESNO): A self-consistent approach and its applications
+in MRI." Journal of Magnetic Resonance 2009; 199: 94-103.
 
 Formerly: scil_run_nlmeans.py
-"""
 
+"""
 import argparse
 import logging
 import warnings
@@ -22,98 +63,166 @@
                              add_verbose_arg,
                              assert_inputs_exist,
                              assert_outputs_exist,
-                             assert_headers_compatible)
+                             assert_headers_compatible, ranged_type)
+from scilpy.stats.stats import estimate_piesno_sigma
 
 
 def _build_arg_parser():
     p = argparse.ArgumentParser(
         description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
 
     p.add_argument('in_image',
-                   help='Path of the image file to denoise.')
+                   help='Path of the image file to denoise (3D or 4D data)')
     p.add_argument('out_image',
                    help='Path to save the denoised image file.')
-    p.add_argument('number_coils', type=int,
-                   help='Number of receiver coils of the scanner.\nUse '
-                        'number_coils=1 in the case of a SENSE (GE, Philips) '
-                        'reconstruction and \nnumber_coils >= 1 for GRAPPA '
-                        'reconstruction (Siemens). number_coils=4 works well '
-                        'for the 1.5T\n in Sherbrooke. Use number_coils=0 if '
-                        'the noise is considered Gaussian distributed.')
-
-    p.add_argument('--mask', metavar='',
-                   help='Path to a binary mask. Only the data inside the mask'
-                        ' will be used for computations')
-    p.add_argument('--sigma', metavar='float', type=float,
-                   help='The standard deviation of the noise to use instead '
-                        'of computing  it automatically.')
-    p.add_argument('--log', dest="logfile",
-                   help='If supplied, name of the text file to store '
-                        'the logs.')
+    p.add_argument('--mask_denoise',
+                   help="Path to a binary mask. Only the data inside the mask "
+                        "will be denoised. If not provided, only non-zero "
+                        "voxels will be denoised.")
+    p.add_argument('--gaussian', action='store_true',
+                   help="If you know that your data contains gaussian noise, "
+                        "use this option. Otherwise, Rician is assumed.")
+
+    g = p.add_argument_group("Noise estimation methods")
+    g = g.add_mutually_exclusive_group(required=True)
+    g.add_argument(
+        '--sigma', metavar='float', type=float,
+        help='Provide your own standard deviation of the noise.')
+    g.add_argument(
+        '--basic_sigma', action='store_true',
+        help="Use dipy's basic estimation of sigma.")
+    g.add_argument(
+        '--piesno', action='store_true',
+        help="Estimate sigma using Piesno's method.")
+
+    g = p.add_argument_group("Noise estimation options: piesno and basic")
+    g.add_argument(
+        '--number_coils', type=ranged_type(int, min_value=0),
+        help="Option for Dipy's noise estimation. Here is their description:\n"
+             ">Number of coils of the receiver array. Use N = 1 in case of a\n"
+             ">SENSE reconstruction (Philips scanners) or the number of \n"
+             ">coils for a GRAPPA reconstruction (Siemens and GE). Use 0 to \n"
+             ">disable the correction factor, as for example if the noise is\n"
+             ">Gaussian distributed. See [1] for more information.\n"
+             "Note. If you don't know the number of coils, 0 will probably "
+             "work.")
+
+    g = p.add_argument_group("Noise estimation options: basic")
+    gg = g.add_mutually_exclusive_group()
+    gg.add_argument(
+        '--mask_sigma',
+        help='Path to a binary mask for --basic_sigma estimation. Only the '
+             'data inside the mask will be used to estimate sigma. If not '
+             'provided, only non-zero voxels will be used.')
+    gg.add_argument(
+        '--sigma_from_all_voxels', action='store_true',
+        help="If set, all voxels are used for the --basic_sigma estimation, "
+             "even zeros.")
+
+    g = p.add_argument_group("Noise estimation options: piesno")
+    g.add_argument('--save_piesno_mask', metavar='filepath',
+                   help="If set, save piesno mask.")
 
     add_processes_arg(p)
     add_verbose_arg(p)
     add_overwrite_arg(p)
     return p
 
 
-def _get_basic_sigma(data):
-    # We force to zero as the 3T is either oversmoothed or still noisy, but
-    # we prefer the second option
-    logging.info("In basic noise estimation, number_coils=0 is enforced!")
-    sigma = estimate_sigma(data, N=0)
-
-    # Use a single value for all of the volumes.
-    # This is the same value for a given bval with this estimator
-    sigma = np.median(sigma)
-    logging.info('The noise standard deviation from the basic estimation '
-                 'is {}'.format(sigma))
-
-    # Broadcast the single value to a whole 3D volume for nlmeans
-    return np.ones(data.shape[:3]) * sigma
-
-
 def main():
     parser = _build_arg_parser()
     args = parser.parse_args()
     logging.getLogger().setLevel(logging.getLevelName(args.verbose))
 
-    assert_inputs_exist(parser, args.in_image, args.mask)
-    assert_outputs_exist(parser, args, args.out_image, args.logfile)
-    assert_headers_compatible(parser, args.in_image, args.mask)
+    # Verifications
+    if (args.number_coils is not None and args.number_coils == 0 and
+            not args.gaussian):
+        logging.warning("Usually, with --number_coils 0, the data ìs "
+                        "considered to have Gaussian noise, but you "
+                        "did not select --gaussian. Proceed with care.")
+
+    if (args.basic_sigma or args.piesno) and args.number_coils is None:
+        parser.error("Please provide the number of coils for basic_sigma "
+                     "and piesno options.")
+
+    if args.piesno or args.sigma:
+        if args.sigma_from_all_voxels:
+            parser.error("You selected --sigma_from_all_voxels, but this is "
+                         "only available for the --basic_sigma method.")
+        if args.mask_sigma:
+            parser.error("You selected --mask_sigma, but this is "
+                         "only available for the --basic_sigma method.")
+
+    if args.save_piesno_mask and not args.piesno:
+        parser.error("Option --save_piesno_mask cannot be used when --pieno "
+                     "is not selected.")
+
+    assert_inputs_exist(parser, args.in_image,
+                        [args.mask_denoise, args.mask_sigma])
+    assert_outputs_exist(parser, args, args.out_image, args.save_piesno_mask)
+    assert_headers_compatible(parser, args.in_image,
+                              [args.mask_denoise, args.mask_sigma])
+
+    # Loading
+    vol = nib.load(args.in_image)
+    vol_data = vol.get_fdata(dtype=np.float32)
 
-    if args.logfile is not None:
-        logging.getLogger().addHandler(logging.FileHandler(args.logfile,
-                                                           mode='w'))
+    if args.piesno and (len(vol_data.shape) != 4 or vol_data.shape[3] == 1):
+        parser.error("The piesno method requires 4D data.")
 
-    vol = nib.load(args.in_image)
-    data = vol.get_fdata(dtype=np.float32)
-    if args.mask is None:
-        mask = np.zeros(data.shape[0:3], dtype=bool)
-        if data.ndim == 4:
-            mask[np.sum(data, axis=-1) > 0] = 1
+    # Denoising mask
+    if args.mask_denoise is None:
+        mask_denoise = np.zeros(vol_data.shape[0:3], dtype=bool)
+        if vol_data.ndim == 4:
+            mask_denoise[np.sum(vol_data, axis=-1) > 0] = True
         else:
-            mask[data > 0] = 1
+            mask_denoise[vol_data > 0] = True
     else:
-        mask = get_data_as_mask(nib.load(args.mask), dtype=bool)
-
-    sigma = args.sigma
+        mask_denoise = get_data_as_mask(nib.load(args.mask_denoise),
+                                        dtype=bool)
+
+    # Processing
+    if args.sigma is not None:
+        logging.info('User supplied noise standard deviation is {}'
+                     .format(args.sigma))
+        sigma = args.sigma
+    elif args.basic_sigma:
+        if args.mask_sigma:
+            mask_sigma = get_data_as_mask(nib.load(args.mask_sigma))
+            tmp_vol_data = (vol_data * mask_sigma[:, :, :, None]
+                            ).astype(np.float32)
+        else:
+            tmp_vol_data = vol_data.astype(np.float32)
 
-    if sigma is not None:
-        logging.info('User supplied noise standard deviation is '
-                     '{}'.format(sigma))
-        # Broadcast the single value to a whole 3D volume for nlmeans
-        sigma = np.ones(data.shape[:3]) * sigma
-    else:
         logging.info('Estimating noise')
-        sigma = _get_basic_sigma(vol.get_fdata(dtype=np.float32))
+        sigma = estimate_sigma(
+            tmp_vol_data,
+            disable_background_masking=args.sigma_from_all_voxels,
+            N=args.number_coils)
+        logging.info(
+            "The estimated noise for each volume is: {}".format(sigma))
+        sigma = np.median(sigma)  # Managing 4D data.
+        logging.info('The median noise is: {}'.format(sigma))
+    else:  # --piesno
+        sigma, mask_noise = estimate_piesno_sigma(vol_data, args.number_coils)
+
+        if args.save_piesno_mask:
+            logging.info("Saving resulting Piesno noise mask in {}"
+                         .format(args.save_piesno_mask))
+            nib.save(nib.Nifti1Image(mask_noise, vol.affine,
+                                     header=vol.header),
+                     args.save_piesno_mask)
+
+    # Broadcast the single value to a whole 3D volume for nlmeans
+    sigma = np.ones(vol_data.shape[:3]) * sigma
 
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", category=DeprecationWarning)
         data_denoised = nlmeans(
-            data, sigma, mask=mask, rician=args.number_coils > 0,
+            vol_data, sigma, mask=mask_denoise, rician=not args.gaussian,
             num_threads=args.nbr_processes)
 
+    # Saving
     nib.save(nib.Nifti1Image(data_denoised, vol.affine, header=vol.header),
              args.out_image)
 

scripts/tests/test_denoising_nlmeans.py

@@ -8,7 +8,7 @@
 from scilpy.io.fetcher import fetch_data, get_testing_files_dict
 
 # If they already exist, this only takes 5 seconds (check md5sum)
-fetch_data(get_testing_files_dict(), keys=['others.zip'])
+fetch_data(get_testing_files_dict(), keys=['others.zip', 'processing.zip'])
 tmp_dir = tempfile.TemporaryDirectory()
 
 
@@ -17,9 +17,31 @@ def test_help_option(script_runner):
     assert ret.success
 
 
-def test_execution_others(script_runner, monkeypatch):
+def test_execution_basic_3d(script_runner, monkeypatch):
     monkeypatch.chdir(os.path.expanduser(tmp_dir.name))
     in_img = os.path.join(SCILPY_HOME, 'others', 't1_resample.nii.gz')
     ret = script_runner.run('scil_denoising_nlmeans.py', in_img,
-                            't1_denoised.nii.gz',  '4', '--processes', '1')
+                            't1_denoised.nii.gz', '--processes', '1',
+                            '--basic_sigma', '--number_coils', 0,
+                            '--gaussian')
+    assert ret.success
+
+
+def test_execution_basic_4d_mask(script_runner, monkeypatch):
+    monkeypatch.chdir(os.path.expanduser(tmp_dir.name))
+    in_img = os.path.join(SCILPY_HOME, 'processing', 'dwi_crop_1000.nii.gz')
+    in_mask = os.path.join(SCILPY_HOME, 'processing', 'fa_thr.nii.gz')
+    ret = script_runner.run('scil_denoising_nlmeans.py', in_img,
+                            't1_denoised.nii.gz', '--processes', '1',
+                            '--basic_sigma', '--number_coils', 0,
+                            '--gaussian', '--mask_sigma', in_mask)
+    assert ret.success
+
+
+def test_execution_piesno(script_runner, monkeypatch):
+    monkeypatch.chdir(os.path.expanduser(tmp_dir.name))
+    in_img = os.path.join(SCILPY_HOME, 'processing', 'dwi.nii.gz')
+    ret = script_runner.run('scil_denoising_nlmeans.py', in_img,
+                            'dwi_denoised.nii.gz', '--processes', '1',
+                            '--piesno', '--number_coils', '4')
     assert ret.success