function MuscleDenoise

% ----------------- INPUTS -----------------
% Original dti file: original diffusion-weighted NIFTI image, should be in working directory (EX: 'dti.nii.gz')
% Mask for dti image: arm mask for NIFTI image, should be in working directory (EX: 'dti_arm_mask.nii.gz')

% ----------------- OUTPUT ----------------
% Denoised NIFTI file of diffusion-weighted iamges: file name is 'dti_dn.nii.gz'

% ----------------- NOTES -----------------
%   - Download and use MRI_lmmse function: SANTIAGO AJA-FERNANDEZ (2023). LMMSE filter for Rician MRI data (https://www.mathworks.com/matlabcentral/fileexchange/36741-lmmse-filter-for-rician-mri-data), MATLAB Central File Exchange. Retrieved May 9, 2023.
%   - Download and use MATLAB NIFTI toolbox: Jimmy Shen (2023). Tools for NIfTI and ANALYZE image (https://www.mathworks.com/matlabcentral/fileexchange/8797-tools-for-nifti-and-analyze-image), MATLAB Central File Exchange. Retrieved May 9, 2023.
%   - Lines 28-29: replace with desired sigma and Ws values

%% Load Original Diffusion-Weighted Image and Mask
disp('Loading original data');

% Load DTI data from current working directory
DWI_orig = load_untouch_nii('dti.nii.gz');              % Load the diffusion-weighted NIFTI image as a structure
DWI_orig_img = double(DWI_orig.img);                    % Get the image from the structure
DWI_mask = load_untouch_nii('dti_arm_mask.nii.gz');     % Load the mask NIFTI image as a structure
DWI_mask_img = double(DWI_mask.img);                    % Load the mask NIFTI image as a structure

%% Initialize Containers and Values

% Initialize containers
DWI_dim = size(DWI_orig_img);   % Find size of original diffusion-weighted image
DWI_dn_img = zeros(DWI_dim);    % Create an array of this size

% Define denoising parameters
sigma = 7;  % Standard deviation estimated in diffusion-weighted images
Ws = [3,3]; % Square window used for estimation

%% Denoise Image
disp('Denoising data');

% Double loop to apply LMMSE filter on each slice
for i=1:DWI_dim(4)                                              % Loop through each DWI volume (4th dimension)
    for j=1:DWI_dim(3)                                          % Loop through each z-slice (3rd dimension)
        DWI_orig_2d = squeeze(DWI_orig_img(:,:,j,i));           % Make into a 2D array
        dn=MRI_lmmse(DWI_orig_2d,Ws,'sigma',sigma);             % Apply MRI_lmmse filter
        DWI_dn_img(:,:,j,i) = dn.*squeeze(DWI_mask_img(:,:,j)); % Mask the denoised image
    end
    fprintf('Volume %s denoised\n',num2str(i));
end

%% Save Denoised Image
disp('Saving NIFTI structures');

DWI_dn_nii = make_nii(DWI_dn_img);      % Make NIFTI structure and save it to current working directory
DWI_dn_nii.hdr = DWI_orig.hdr;          % Save original header to new NIFTI structure
save_nii(DWI_dn_nii,'dti_dn.nii.gz');   % Save denoised NIFTI file

disp('Denoising done!');

exit;

end