% Experiment 6: Frequency domain filtering of image
% a. Lowpass (smoothing) filter in frequency domain
clear; close all; clc;

%% Part A: Lab handout program (filtering using the 2D-DFT, steps 1 to 8)
% Step 1: Given an input image f(x,y) of size M x N, obtain the padding
% parameters P and Q. Typically, we select P = 2M and Q = 2N
Im = imread('cameraman.tif');
a = imnoise(Im, 'gaussian', 0, 0.04);   % adding high frequency noise
b = im2double(a);
[m, n] = size(b);
c = zeros(2*m, 2*n);
[p, q] = size(c);

% Step 2: form the padded image fp(x,y) of size P x Q by appending zeros
for i = 1:p
    for j = 1:q
        if i <= m && j <= n
            c(i, j) = b(i, j);
        else
            c(i, j) = 0;
        end
    end
end

% Step 3: multiply fp(x,y) by (-1)^(x+y) to centre the transform
d = zeros(p, q);
for i = 1:p
    for j = 1:q
        d(i, j) = c(i, j) .* (-1).^(i + j);
    end
end

% Step 4: compute the DFT F(u,v)
e = fft2(d);

% Step 5: real, symmetric ideal lowpass filter H(u,v) centred at (P/2, Q/2)
[x, y] = freqspace(p, 'meshgrid');
z = zeros(p, q);
for i = 1:p
    for j = 1:q
        z(i, j) = sqrt(x(i, j).^2 + y(i, j).^2);
    end
end
H = zeros(p, q);
for i = 1:p
    for j = 1:q
        if z(i, j) <= 0.4               % 0.4 is the cut-off frequency of the LPF
            H(i, j) = 1;
        else
            H(i, j) = 0;
        end
    end
end

% Step 6: G(u,v) = H(u,v) F(u,v)
h1 = e .* H;

% Step 7: gp(x,y) = real{inverse DFT[G(u,v)]} (-1)^(x+y)
h2 = ifft2(h1);
h3 = zeros(p, q);
for i = 1:p
    for j = 1:q
        h3(i, j) = real(h2(i, j)) .* ((-1).^(i + j));
    end
end

% Step 8: extract the M x N region from the top-left quadrant
out = zeros(m, n);
for i = 1:m
    for j = 1:n
        out(i, j) = h3(i, j);
    end
end

figure('Name', 'A_steps');
subplot(2, 4, 1); imshow(b);                    title('Input noisy image');
subplot(2, 4, 2); imshow(c);                    title('Padded image');
subplot(2, 4, 3); imshow(d);                    title('Pre-processed (-1)^{x+y}');
subplot(2, 4, 4); imshow(log(1 + abs(e)), []);  title('2D DFT (log magnitude)');
subplot(2, 4, 5); imshow(H);                    title('Low pass filter mask');
subplot(2, 4, 6); imshow(log(1 + abs(h1)), []); title('Low passed spectrum');
subplot(2, 4, 7); imshow(h3, []);               title('Post-processed image');
subplot(2, 4, 8); imshow(out);                  title('Output image');

figure('Name', 'A_result');
imshow([b out]); title('Input noisy image  |  Output image (ideal LPF, cut-off 0.4)');

%% Part B: Custom implementation (Ideal, Butterworth and Gaussian lowpass filters)
f = im2double(Im);
g = b;                                   % same noisy image as Part A
D0 = 50;                                 % cut-off distance in pixels of the P x Q grid
[outI, HI] = myFreqLPF(g, 'ideal', D0);
[outB, HB] = myFreqLPF(g, 'butterworth', D0, 2);
[outG, HG] = myFreqLPF(g, 'gaussian', D0);

figure('Name', 'B_filters');
subplot(2, 4, 1); imshow(g);   title('Noisy input');
subplot(2, 4, 2); imshow(outI); title(sprintf('Ideal LPF, D_0 = %d', D0));
subplot(2, 4, 3); imshow(outB); title(sprintf('Butterworth LPF (n=2), D_0 = %d', D0));
subplot(2, 4, 4); imshow(outG); title(sprintf('Gaussian LPF, D_0 = %d', D0));
subplot(2, 4, 6); imshow(HI);  title('H(u,v): ideal');
subplot(2, 4, 7); imshow(HB);  title('H(u,v): Butterworth');
subplot(2, 4, 8); imshow(HG);  title('H(u,v): Gaussian');
subplot(2, 4, 5);
mid = size(HI, 1) / 2 + 1;
plot(HI(mid, :), 'k', 'LineWidth', 1.2); hold on;
plot(HB(mid, :), 'b', 'LineWidth', 1.2); plot(HG(mid, :), 'r', 'LineWidth', 1.2); hold off;
xlim([mid - 150, mid + 150]); legend('Ideal', 'Butterworth', 'Gaussian'); title('Filter profiles');

fprintf('PSNR (dB) against the noise-free image\n');
fprintf('Noisy input                    : %.2f\n', myPsnr(g, f));
fprintf('Part A ideal LPF (cut-off 0.4) : %.2f\n', myPsnr(out, f));
cutoffs = [20 50 80];
fprintf('\n%-8s %10s %13s %10s\n', 'D0', 'Ideal', 'Butterworth', 'Gaussian');
gaussOut = cell(1, numel(cutoffs));
for k = 1:numel(cutoffs)
    oI = myFreqLPF(g, 'ideal', cutoffs(k));
    oB = myFreqLPF(g, 'butterworth', cutoffs(k), 2);
    gaussOut{k} = myFreqLPF(g, 'gaussian', cutoffs(k));
    fprintf('%-8d %10.2f %13.2f %10.2f\n', cutoffs(k), myPsnr(oI, f), myPsnr(oB, f), myPsnr(gaussOut{k}, f));
end

figure('Name', 'B_cutoff');
for k = 1:numel(cutoffs)
    subplot(1, numel(cutoffs), k); imshow(gaussOut{k});
    title(sprintf('Gaussian LPF, D_0 = %d', cutoffs(k)));
end

%% Local functions
function [out, H] = myFreqLPF(img, type, D0, order)
    % frequency domain lowpass filtering with zero padding to 2M x 2N
    [M, N] = size(img);
    P = 2 * M; Q = 2 * N;
    fp = zeros(P, Q);
    fp(1:M, 1:N) = img;
    [V, U] = meshgrid(0:Q-1, 0:P-1);
    fp = fp .* (-1).^(U + V);                    % centre the spectrum
    F = fft2(fp);
    D = sqrt((U - P/2).^2 + (V - Q/2).^2);       % distance from the centre
    switch type
        case 'ideal'
            H = double(D <= D0);
        case 'butterworth'
            H = 1 ./ (1 + (D / D0).^(2 * order));
        case 'gaussian'
            H = exp(-(D.^2) / (2 * D0^2));
    end
    gp = real(ifft2(H .* F)) .* (-1).^(U + V);
    out = min(max(gp(1:M, 1:N), 0), 1);
end

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