2023_Kalantari_AIDAqc/code/adding_noising_example_for_voting.py

import os
import glob
import nibabel as nib
import numpy as np
from skimage.util import random_noise
from tqdm import tqdm
import random
import skimage.exposure as exposure

init_path = r"C:\Users\aswen\Desktop\TestingData\Aswendt_qc_rsfmri_plot"

# Define the search paths for noisy files and all files
searchpath_delete = os.path.join(init_path,"**" ,"dwi", "brkraw", "*artifact*.nii.gz")
searchpath = os.path.join(init_path, "**","dwi", "brkraw", "*.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
all_files = glob.glob(searchpath, recursive=True)
# Randomly choose 5 files
randomly_chosen = random.sample(all_files, 2)
# Process each file
for file in tqdm(randomly_chosen, desc="Processing", unit="file"):
    print(file)
    Im = nib.load(file)
    data = Im.get_fdata()  # Getting image data from the nibabel object
    
    # Normalize the data to [0, 1] range
    data_min = np.min(data)
    data_max = np.max(data)
    max_abs = np.max(np.abs(data))
    
    # Normalize the data to [0, 1] range
    normalized_data = data/max_abs
    
    # Add weak artifact: Gaussian noise only
    #Im_weak_artifact = random_noise(normalized_data, mode='gaussian', mean=0, var=0.001)  # Adjust the variance for noise
    motion_factor = 0.7
    # Add strong artifact: Gaussian noise + sudden shifts along the appropriate axis
    if normalized_data.ndim == 2:  # Check if the data is 4D
        # Generate motion artifacts along the third and fourth dimensions for 4D images
     
        translation_third_dim = int(normalized_data.shape[2] * motion_factor)
        translation_fourth_dim = int(normalized_data.shape[3] * motion_factor)
        motion_artifact = np.roll(normalized_data, translation_third_dim, axis=1)
        #motion_artifact = np.roll(motion_artifact, translation_fourth_dim, axis=3)
        motion_artifact = (2 * motion_artifact + 3 * normalized_data) / 5  # Weighted combination
        normalized_data_with_added_motion =  normalized_data
        normalized_data_with_added_motion[:, :, :, normalized_data.shape[3] // 2:normalized_data.shape[3] // 2 + 40:2] = motion_artifact[:, :, :, normalized_data.shape[3] // 2:normalized_data.shape[3] // 2 + 40:2]

        
        Im_strong_artifact = random_noise(normalized_data_with_added_motion, mode='speckle', var=0.2)  # Adjust the variance for noise
        gamma = 0.6  # Adjust the gamma value to control the darkness 
        Im_strong_artifact = exposure.adjust_gamma(Im_strong_artifact, gamma)
    
    else:
 
        Im_strong_artifact = random_noise(normalized_data, mode='speckle', var=0.2)  # Adjust the variance for noise
        # Darken the intensity using gamma correction
        gamma = 0.8  # Adjust the gamma value to control the darkness 
        Im_strong_artifact = exposure.adjust_gamma(Im_strong_artifact, gamma)

    # Combine weak and strong artifacts by averaging with weights
    weight_weak = 0.5  # Weight for weak artifact
    weight_strong = 0.2  # Weight for strong artifact
    #Im_combined_artifact = (weight_weak * Im_weak_artifact + weight_strong * Im_strong_artifact) / (weight_weak + weight_strong)
    
    # Creating new file names for weak and strong artifact images
    weak_artifact_filename = os.path.splitext(file)[0] + "_weak_artifact.nii.gz"
    strong_artifact_filename = os.path.splitext(file)[0] + "_strong_artifact.nii.gz"
    
    # Save the artifact images
    #nib.save(nib.Nifti1Image(Im_weak_artifact*max_abs, Im.affine), weak_artifact_filename)
    nib.save(nib.Nifti1Image(Im_strong_artifact*max_abs, Im.affine), strong_artifact_filename)