Patch2Self: Self-Supervised Denoising via Statistical Independence

Patch2Self [Fadnavis20] is a self-supervised learning method for denoising DWI data, which uses the entire volume to learn a full-rank locally linear denoiser for that volume. By taking advantage of the oversampled q-space of DWI data, Patch2Self can separate structure from noise without requiring an explicit model for either.

Classical denoising algorithms such as Local PCA [Manjon2013], [Veraa2016a], Non-local Means [Coupe08], Total Variation Norm [Knoll11], etc. assume certain properties on the signal structure. Patch2Self does not make any such assumption on the signal instead using the fact that the noise across different 3D volumes of the DWI signal originates from random fluctuations in the acquired signal.

Since Patch2Self only relies on the randomness of the noise, it can be applied at any step in the pre-processing pipeline. The design of Patch2Self is such that it can work on any type of diffusion data/ any body part without requiring a noise estimation or assumptions on the type of noise (such as its distribution).

The Patch2Self Framework:

https://github.com/dipy/dipy_data/blob/master/Patch2Self_Framework.PNG?raw=true

The above figure demonstrates the working of Patch2Self. The idea is to build a new regressor for denoising each 3D volume of the 4D diffusion data. This is done in the following 2 phases:

(A) Self-supervised training: First, we extract 3D Patches from all the ‘n’ volumes and hold out the target volume to denoise. Each patch from the rest of the ‘(n-1)’ volumes predicts the center voxel of the corresponding patch in the target volume.

This is done by using the self-supervised loss: \(\mathcal{L}\left(\Phi_{J}\right)=\mathbb{E}\left\|\Phi_{J}\left(Y_{*, *,-j}\right)-Y_{*, 0, j}\right\|^{2}\)

(B) Prediction: The same ‘n-1’ volumes which were used in the training are now fed into the regressor \(\Phi\) built in phase (A). The prediction is a denoised version of held-out volume.

Note: The volume to be denoised is merely used as the target in the training phase. But is not used in the training set for (A) nor is used to predict the denoised output in (B).

Let’s load the necessary modules:

import numpy as np
from dipy.data import get_fnames
from dipy.io.image import load_nifti, save_nifti
import matplotlib.pyplot as plt

from dipy.denoise.patch2self import patch2self

Now let’s load an example dataset and denoise it with Patch2Self. Patch2Self does not require noise estimation and should work with any kind of diffusion data.

hardi_fname, hardi_bval_fname, hardi_bvec_fname = get_fnames('stanford_hardi')
data, affine = load_nifti(hardi_fname)
bvals = np.loadtxt(hardi_bval_fname)
denoised_arr = patch2self(data, bvals)

The array denoised_arr contains the denoised output obtained from Patch2Self. Now lets visualize the output and the residuals obtained from the denoising.

# Gets the center slice and the middle volume of the 4D diffusion data.
sli = data.shape[2] // 2
gra = 60  # pick out a random volume for a particular gradient direction

orig = data[:, :, sli, gra]
den = denoised_arr[:, :, sli, gra]

# computes the residuals
rms_diff = np.sqrt((orig - den) ** 2)

fig1, ax = plt.subplots(1, 3, figsize=(12, 6),
                        subplot_kw={'xticks': [], 'yticks': []})

fig1.subplots_adjust(hspace=0.3, wspace=0.05)

ax.flat[0].imshow(orig.T, cmap='gray', interpolation='none',
                  origin='lower')
ax.flat[0].set_title('Original')
ax.flat[1].imshow(den.T, cmap='gray', interpolation='none',
                  origin='lower')
ax.flat[1].set_title('Denoised Output')
ax.flat[2].imshow(rms_diff.T, cmap='gray', interpolation='none',
                  origin='lower')
ax.flat[2].set_title('Residuals')

fig1.savefig('denoised_patch2self.png')

print("The result saved in denoised_patch2self.png")
../../_images/denoised_patch2self.png

Patch2Self preserved anatomical detail. This can be visually verified by inspecting the residuals obtained above. Since we do not see any structure in the difference residuals, it is clear that it preserved the underlying signal structure and got rid of the stochastic noise.

Below we show how the denoised data can be saved.

save_nifti('denoised_patch2self.nii.gz', denoised_arr, affine)

print("Entire denoised data saved in denoised_patch2self.nii.gz")

References

Fadnavis20

S. Fadnavis, J. Batson, E. Garyfallidis, Patch2Self: Denoising Diffusion MRI with Self-supervised Learning, Advances in Neural Information Processing Systems 33 (2020)

Manjon2013

Manjon JV, Coupe P, Concha L, Buades A, Collins DL “Diffusion Weighted Image Denoising Using Overcomplete Local PCA” (2013). PLoS ONE 8(9): e73021. doi:10.1371/journal.pone.0073021.

Veraa2016a

Veraart J, Fieremans E, Novikov DS. 2016. Diffusion MRI noise mapping using random matrix theory. Magnetic Resonance in Medicine. doi: 10.1002/mrm.26059.

Coupe08

P. Coupe, P. Yger, S. Prima, P. Hellier, C. Kervrann, C. Barillot, An Optimized Blockwise Non Local Means Denoising Filter for 3D Magnetic Resonance Images, IEEE Transactions on Medical Imaging, 27(4):425-441, 2008

Knoll11

F. Knoll, K. Bredies, T. Pock, R. Stollberger, Second order total generalized variation (TGV) for MRI. Magnetic resonance in medicine, 65(2), pp.480-491.

“””

Example source code

You can download the full source code of this example. This same script is also included in the dipy source distribution under the doc/examples/ directory.