% Experiment 4: Spatial domain filtering of image
% a. Averaging (mean) filter  b. Median filter
clear; close all; clc;

I = imread('cameraman.tif');
Ins = imnoise(I, 'salt & pepper');      % add salt & pepper noise (density 0.05)
% window size m x n of the filter (the handout reads these with input())
m = 3;
n = 3;

%% Part A: Using inbuilt MATLAB functions (lab handout code)
% Mean (average) filter
H = fspecial('average', [m n]);
Im = imfilter(Ins, H, 'replicate');     % carry out filtering on the noisy image

figure('Name', 'A_mean');
subplot(1, 3, 1); imshow(I);   title('Original Image');
subplot(1, 3, 2); imshow(Ins); title('Noisy Image');
subplot(1, 3, 3); imshow(Im);  title('Mean Filtered Image');

% Median filter
Ifilt = medfilt2(Ins, [m n]);

figure('Name', 'A_median');
subplot(1, 3, 1); imshow(I);     title('Original Image');
subplot(1, 3, 2); imshow(Ins);   title('Noisy Image');
subplot(1, 3, 3); imshow(Ifilt); title('Median Filtered Image');

%% Part B: Custom implementation
myMean3 = myMeanFilter(Ins, 3, 3);
myMed3  = myMedianFilter(Ins, 3, 3);
myMean5 = myMeanFilter(Ins, 5, 5);
myMed5  = myMedianFilter(Ins, 5, 5);

% Same filters on Gaussian noise for comparison
Ign = imnoise(I, 'gaussian', 0, 0.01);
gnMean = myMeanFilter(Ign, 3, 3);
gnMed  = myMedianFilter(Ign, 3, 3);

fprintf('Max |custom - inbuilt| : mean %d, median %d\n', ...
    max(abs(double(myMean3(:)) - double(Im(:)))), max(abs(double(myMed3(:)) - double(Ifilt(:)))));
fprintf('\nPSNR (dB) against the original image\n');
fprintf('%-28s %8s %8s\n', 'Case', 'Mean', 'Median');
fprintf('%-28s %8.2f %8.2f\n', 'Salt & pepper, noisy', myPsnr(Ins, I), myPsnr(Ins, I));
fprintf('%-28s %8.2f %8.2f\n', 'Salt & pepper, 3x3', myPsnr(myMean3, I), myPsnr(myMed3, I));
fprintf('%-28s %8.2f %8.2f\n', 'Salt & pepper, 5x5', myPsnr(myMean5, I), myPsnr(myMed5, I));
fprintf('%-28s %8.2f %8.2f\n', 'Gaussian, noisy', myPsnr(Ign, I), myPsnr(Ign, I));
fprintf('%-28s %8.2f %8.2f\n', 'Gaussian, 3x3', myPsnr(gnMean, I), myPsnr(gnMed, I));

figure('Name', 'B_custom');
subplot(2, 4, 1); imshow(Ins);     title('Salt & pepper noise');
subplot(2, 4, 2); imshow(myMean3); title('Custom mean 3x3');
subplot(2, 4, 3); imshow(myMean5); title('Custom mean 5x5');
subplot(2, 4, 4); imshow(myMed3);  title('Custom median 3x3');
subplot(2, 4, 5); imshow(myMed5);  title('Custom median 5x5');
subplot(2, 4, 6); imshow(Ign);     title('Gaussian noise');
subplot(2, 4, 7); imshow(gnMean);  title('Custom mean 3x3');
subplot(2, 4, 8); imshow(gnMed);   title('Custom median 3x3');

%% Local functions
function out = myMeanFilter(img, m, n)
    % average of the m x n neighbourhood (replicate padding)
    P = padarray(double(img), [floor(m/2) floor(n/2)], 'replicate');
    out = zeros(size(img));
    for x = 1:size(img, 1)
        for y = 1:size(img, 2)
            win = P(x:x+m-1, y:y+n-1);
            out(x, y) = sum(win(:)) / (m * n);
        end
    end
    out = uint8(round(out));
end

function out = myMedianFilter(img, m, n)
    % middle value of the sorted m x n neighbourhood (zero padding, like medfilt2)
    P = padarray(double(img), [floor(m/2) floor(n/2)], 0);
    out = zeros(size(img));
    mid = ceil(m * n / 2);
    for x = 1:size(img, 1)
        for y = 1:size(img, 2)
            win = sort(reshape(P(x:x+m-1, y:y+n-1), [], 1));
            out(x, y) = win(mid);
        end
    end
    out = uint8(out);
end

function p = myPsnr(a, ref)
    mse = mean((double(a(:)) - double(ref(:))).^2);
    p = 10 * log10(255^2 / mse);
end