Test and simulation for brain denoising

.gitignore

@@ -40,5 +40,10 @@ coverage.xml
 *.pyc
 phantoms/MR_XCAT_qMRI/*.json
 phantoms/MR_XCAT_qMRI/*.txt
+phantoms/brain/data/*.nii.gz
+phantoms/brain/ground_truth/*.nii.gz
 tests/IVIMmodels/unit_tests/models
 models
+temp
+.vscode
+.env

phantoms/brain/ground_truth/diffusive_relax_groundtruth.json

@@ -0,0 +1,8 @@
+{
+    "tissues": ["WM","GM","CSF"],
+    "D":[0.81e-3,0.86e-3,3e-3],
+    "f":[0.044,0.033,0],
+    "Dstar":[84e-3,76e-3,0],
+    "T1":[1.0, 1.5, 3.4],
+    "T2":[0.046,0.068,0.45]
+} 

phantoms/brain/sim_brain_phantom_preproc.py

@@ -0,0 +1,82 @@
+import os
+import shutil
+import json
+import numpy as np
+import nibabel as nib
+from scipy.ndimage import zoom
+from utilities.data_simulation.Download_data import download_data
+
+
+def simulate_brain_phantom(snr=100, TE=60e-3, TR=5, resolution=[3,3,3]):
+    '''
+    Simulation parameters can be set by changing the default values of the function arguments. The default values are chosen to be suitable for a diffusive regime phantom, but can be changed to simulate a ballistic regime phantom as well. The simulated image is saved in phantoms/brain/data with the name 'diffusive_sn{snr}_relax.nii.gz' or 'ballistic_sn{snr}_relax.nii.gz' depending on the regime. The corresponding bvals and cvals (if applicable) are also saved in the same folder.
+    snr: signal-to-noise ratio of the simulated image
+    TE: echo time in seconds
+    TR: repetition time in seconds
+    resolution: voxel size in mm
+    '''
+    download_data()
+
+    DIFFUSIVE_REGIME = 'diffusive'
+    BALLISTIC_REGIME = 'ballistic'
+
+    regime = DIFFUSIVE_REGIME # only adapted for diffusive regime for now, but can be easily changed to simulate ballistic regime as well by changing this variable and providing the corresponding ground truth parameters in the json file
+
+    folder = os.path.dirname(__file__)
+
+    # Ground truth
+    nii = nib.load(os.path.join(os.path.split(os.path.split(folder)[0])[0],'download','Phantoms','brain','ground_truth','hrgt_icbm_2009a_nls_3t.nii.gz'))
+    segmentation = np.squeeze(nii.get_fdata()[...,-1])
+
+    with open(os.path.join(folder,'ground_truth',regime+'_relax_groundtruth.json'), 'r') as f:
+        ivim_pars = json.load(f)
+    S0 = 1
+
+    # Sequence parameters
+    bval_file = os.path.join(folder,'ground_truth',regime+'.bval')
+    b = np.loadtxt(bval_file)
+    if regime == BALLISTIC_REGIME:
+        cval_file = bval_file.replace('bval','cval')
+        c = np.loadtxt(cval_file)
+
+    # Calculate signal
+    S = np.zeros(list(np.shape(segmentation))+[b.size])
+    print(ivim_pars)
+    if regime == BALLISTIC_REGIME:
+        Db = ivim_pars["Db"]
+        for i,(D,f,vd,T1,T2) in enumerate(zip(ivim_pars["D"],ivim_pars["f"],ivim_pars["vd"],ivim_pars["T1"],ivim_pars["T2"])):
+            S[segmentation==i+1,:] = S0*((1-f)*np.exp(-b*D)+f*np.exp(-b*Db-c**2*vd**2))*np.exp(-TE/T2)*(1-np.exp(-TR/T1))
+    else:
+        for i,(D,f,Dstar,T1,T2) in enumerate(zip(ivim_pars["D"],ivim_pars["f"],ivim_pars["Dstar"],ivim_pars["T1"],ivim_pars["T2"])):
+            S[segmentation==i+1,:] = S0*((1-f)*np.exp(-b*D)+f*np.exp(-b*Dstar))*np.exp(-TE/T2)*(1-np.exp(-TR/T1))
+
+    # Resample to suitable resolution
+    im = zoom(S,np.append(np.diag(nii.affine)[:3]/np.array(resolution),1),order=1)
+    sz = im.shape
+
+    # Save image without noise for reference
+    nii_out = nib.Nifti1Image(im,np.eye(4))
+    base_name = os.path.join(folder,'data','{}_reference_relax'.format(regime,snr))
+    nib.save(nii_out,base_name+'.nii.gz')
+    shutil.copyfile(bval_file,base_name+'.bval')
+
+    # Add Rician noise
+    im_noise = np.abs(im + S0/snr*(np.random.randn(sz[0],sz[1],sz[2],sz[3])+1j*np.random.randn(sz[0],sz[1],sz[2],sz[3])))
+
+    # Save as image and sequence parameters 
+    nii_out = nib.Nifti1Image(im_noise,np.eye(4))
+    base_name = os.path.join(folder,'data','{}_snr{}_relax'.format(regime,snr))
+    nib.save(nii_out,base_name+'.nii.gz')
+    shutil.copyfile(bval_file,base_name+'.bval')
+    if regime == BALLISTIC_REGIME:
+        shutil.copyfile(cval_file,base_name+'.cval')
+
+    # Resample and save segmentation
+    segmentation = zoom(segmentation,np.diag(nii.affine)[:3]/np.array(resolution),order=0)
+    nii_out = nib.Nifti1Image(segmentation,np.eye(4))
+    base_name = os.path.join(folder,'data','{}_snr{}_relax_mask'.format(regime,snr))
+    nib.save(nii_out,base_name+'.nii.gz')
+
+
+if __name__ == "__main__":
+    simulate_brain_phantom()

tests/IVIMpreproc/unit_tests/test_denoise.py

@@ -0,0 +1,63 @@
+'''
+1. Load brain phantom
+2. Add Rician noise
+3. Denoise using denoise algorithm
+4. Check that sum of squared residuals decreased after denoising
+'''
+import os
+
+import numpy as np
+from phantoms.brain.sim_brain_phantom_preproc import simulate_brain_phantom
+import nibabel as nib
+import matplotlib.pyplot as plt
+from src.original.EP_GU.brain_pipeline import denoise_wrap
+
+
+def test_denoise_wrap():
+    # Simulate phantom with noise
+    snr = 50
+    simulate_brain_phantom(snr = snr)
+
+    # Denoising
+    im_file = 'phantoms/brain/data/diffusive_snr{}_relax.nii.gz'.format(snr)
+    im_file_denoised = denoise_wrap(im_file)
+    S_denoised = nib.load(im_file_denoised).get_fdata()
+
+    # Compute sum of squared residuals before denoising
+    mask = nib.load(os.path.join('phantoms','brain','data','diffusive_snr{}_relax_mask.nii.gz'.format(snr))).get_fdata()
+    ref_file = 'phantoms/brain/data/diffusive_reference_relax.nii.gz'
+    S = nib.load(im_file).get_fdata()
+    S_ref = nib.load(ref_file).get_fdata()
+    ssr_before = np.sum((S[mask!=0,:]-S_ref[mask!=0,:])**2)
+    ssr_after = np.sum((S_denoised[mask!=0,:]-S_ref[mask!=0,:])**2)
+
+
+
+    # print('SSR before denoising: {}'.format(ssr_before))
+    # print('SSR after denoising: {}'.format(ssr_after))
+
+    # plt.subplot(1,3,1)
+    # plt.imshow(np.rot90(S_ref[:,:,30,10]),cmap='gray')
+    # plt.title('Reference')
+    # plt.xticks([])
+    # plt.yticks([])
+    # plt.subplot(1,3,2)
+    # plt.imshow(np.rot90(S[:,:,30,10]),cmap='gray')
+    # plt.title('Noisy')
+    # plt.xticks([])
+    # plt.yticks([])
+    # plt.subplot(1,3,3)
+    # plt.imshow(np.rot90(S_denoised[:,:,30,10]),cmap='gray')
+    # plt.title('Denoised')
+    # plt.xticks([])
+    # plt.yticks([])
+    # plt.show()
+
+    # Check that sum of squared residuals decreased after denoising
+    if ssr_after < ssr_before:
+        assert True
+    else:   
+        assert False
+
+
+