2023_Kalantari_AIDAqc/code/adding_noising_example.py

import os
import glob
import nibabel as nib
from skimage.restoration import denoise_tv_chambolle
from skimage.util import random_noise
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor

init_path = r"C:\Users\aswen\Desktop\TestingData\Aswendt_qc_data\proc_data"

# Define the search paths for noisy files and all files
searchpath_delete = os.path.join(init_path, "**", "anat", "*ois*.nii.gz")
searchpath = os.path.join(init_path, "**", "anat", "*.nii.gz")

# Delete all noisy files matching the pattern before proceeding
noisy_files = glob.glob(searchpath_delete, recursive=True)
for file in noisy_files:
    os.remove(file)

# Get the list of all files
T2files = glob.glob(searchpath, recursive=True)

def process_file(file):
    Im = nib.load(file)
    data = Im.get_fdata()  # Getting image data from the nibabel object
    
    # Perform denoising using total variation denoising from scikit-image
    Im_denoised = denoise_tv_chambolle(data, weight=0.1, multichannel=False)
    
    # Add random noise to the original image
    Im_random_noisy = random_noise(data, mode='gaussian', mean=0, var=0.01)  # Adjust the variance for noise
    
    # Add salt and pepper noise to the original image
    Im_salt_pepper_noisy = random_noise(data, mode='s&p', amount=0.05)  # Adjust the amount for noise density
    
    # Add speckle noise to the original image
    Im_speckle_noisy = random_noise(data, mode='speckle', mean=0, var=0.01)  # Adjust the variance for noise
    
    # Creating new file names for denoised and noisy images
    denoised_filename = os.path.splitext(file)[0] + "_denoised.nii"
    random_noisy_filename = os.path.splitext(file)[0] + "_randomNoisy.nii.gz"
    salt_pepper_noisy_filename = os.path.splitext(file)[0] + "_saltPepperNoisy.nii.gz"
    speckle_noisy_filename = os.path.splitext(file)[0] + "_speckleNoisy.nii.gz"
    
    # Save the denoised image
    nib.save(nib.Nifti1Image(Im_denoised, Im.affine), denoised_filename)
    
    # Save the randomly noisy image
    nib.save(nib.Nifti1Image(Im_random_noisy, Im.affine), random_noisy_filename)
    
    # Save the salt and pepper noisy image
    nib.save(nib.Nifti1Image(Im_salt_pepper_noisy, Im.affine), salt_pepper_noisy_filename)
    
    # Save the speckle noisy image
    nib.save(nib.Nifti1Image(Im_speckle_noisy, Im.affine), speckle_noisy_filename)

if __name__ == '__main__':
    # Use ProcessPoolExecutor for parallel processing with 10 cores
    with ProcessPoolExecutor(max_workers=10) as executor:
        list(tqdm(executor.map(process_file, T2files), total=len(T2files), desc="Processing", unit="file"))