% ECE 420, Alan Brooks, 05-Mar-2003
% Computer Assignment 1 - Restoration Techniques
%   Starting with the 'cameraman.tif' image, simulate a blurred image
% using an 11x11 pill-box blur. Then, add white gaussian noise to
% obtain a Blurred SNR (BSNR) of (a) 30 dB and (b) 20 dB. BSNR is 
%
%   BSNR = 10*log10( sigmaY^2 / sigmaN^2 )
%
% where sigmaY^2=variance of blurred image, y, and sigmaN^2=variance
% of the noise.
%   Restore the 3 blurred images (no noise, (a), & (b)) with the
% following filters:
%
%   (1) Inverse filter
%   (2) Wiener filter
%           - using periodogram estimate of spectral density
%   (3) Constrained least-squares (CLS) filter 
%           - using a Laplacian
%           - regularization parameter values: 0.001, 0.01
%
%   For each restoration, compute the improvement in signal-to-noise
% ratio (ISNR) in dB, defined by
%
%   ISNR = 10*log10( sum([x(:)-y(:)].^2)/sum([x(:)-xr(:)].^2) )
%
% where x(i,j), y(i,j), and xr(i,j) represent respectively the original,
% blurred, and restored images.
%   Turn in all blurred & restored images (with ISNR). Turn in 
% difference image (orig-restored), linearly scaled to fit [0,255].
% Also, turn in 3D plots of image spectral densities used in Wiener
% filter and the operator C() used in the CLS filter.

function [] = ee420cmpAss1(plotsOn)

if ~exist('plotsOn','var'), plotsOn = 1; end;

% get original image, x
x = double(imread('cameraman.tif'))/255;
if plotsOn
    figure,imshow(x,[0 1]),title('Original Image')
end

% create pill-box (square) 11 x 11 kernel
h = ones(11,11)/(11^2); % NOT USED!

% circularly tile h to prepare for DFT
hc = zeros(256);
hc([1:6 252:256 1:6 252:256],[1:6 1:6 252:256 252:256]) = 1/(11^2);

% convolve kernel with image to get blurred image, y
%y = conv2(x,h,'same');
y = real(ifft2( fft2(x).*fft2(hc) ));
if plotsOn
    figure,imshow(y,[0 1]),title('Blurred Image, Original')
end

% add gaussian noise to get BSNR = 30, 20 dB
%    BSNR = 10*log10( sigmaY^2 / sigmaN^2 )
% -> sigmaN = sqrt( sigmaY^2 / 10^(BSNR/10) )
sigmaY = std2(y);
sigmaN30 = sqrt( sigmaY^2 / 10^(30/10) );
y30 = imnoise(y,'gaussian',0,(sigmaN30/1)^2);
if plotsOn
    figure,imshow(y30,[0 1]),title('Blurred Image, BSNR = 30dB')
end
sigmaN20 = sqrt( sigmaY^2 / 10^(20/10) );
y20 = imnoise(y,'gaussian',0,(sigmaN20/1)^2);
if plotsOn
    figure,imshow(y20,[0 1]),title('Blurred Image, BSNR = 20dB')
end

% ----------------------------------------
% restore using inverse filter
H = fft2(hc); % kernel
thresh = 0.000001;
bTh = abs(H)<thresh; % apply threshold
Y = fft2(y); % original blurred
XR = Y./H;
XR(bTh) = 0;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('Inv Restored Orig Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, Orig Inv')
end
thresh = 0.05;
bTh = abs(H)<thresh; % apply threshold
Y = fft2(y30); % 30dB blurred
XR = Y./H;
XR(bTh) = 0;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y30(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('Inv Restored 30dB Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 30dB Inv')
end
thresh = 0.11;
bTh = abs(H)<thresh; % apply threshold
Y = fft2(y20); % 20dB blurred
XR = Y./H;
XR(bTh) = 0;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y20(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('Inv Restored 20dB Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 20dB Inv')
end

% ----------------------------------------
% restore using Wiener filter
% H(k1,k2) = conj(B)/( abs(B)^2 + Snn(k1,k2)/Sxx(k1,k2) ); 
% note: can use original image, x, to find Sxx
% note: Snn = sigmaN^2 for white noise
B = fft2(hc); % blur kernel
Snn = 0;
Sxx = abs(fft2(x)).^2 ./ (256^2);
H = conj(B)./( abs(B).^2.+Snn./Sxx );
Y = fft2(y); % original
XR = Y.*H;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,mesh([1:256],[1:256],10*log10(Sxx))
    title('Spectral Density 10*log10(Sxx) of Orig Image for Wiener (dB)')
    figure,imshow(xr,[])
    title(sprintf('Wien Restored Orig Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, Orig Wien')
end
Snn = sigmaN30^2;
H = conj(B)./( abs(B).^2.+Snn./Sxx );
Y = fft2(y30); % 30dB blurred
XR = Y.*H;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y30(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('Wien Restored 30dB Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 30dB Wien')
end
Snn = sigmaN20^2;
H = conj(B)./( abs(B).^2.+Snn./Sxx );
Y = fft2(y20); % 20dB blurred
XR = Y.*H;
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y20(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('Wien Restored 20dB Image, ISNR = %.2f dB',ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 20dB Wien')
end

% ----------------------------------------
% restore using CLS filter
% XR(k1,k2) = Y * conj(H(k1,k2)) / (abs(H(k1,k2))^2 + a*abs(C(k1,k2))^2)
c = [0 .25 0; .25 -1 .25; 0 .25 0;]; % NOT USED!
cc = zeros(256);    % circularly tile to prep for DFT
cc(1,2) = .25; cc(2,1) = .25; cc(1,256) = .25; cc(256,1) = .25;
cc(1,1) = -1;
C = real(fft2(cc));
H = real(fft2(hc)); % kernel
Y = fft2(y); % orig blurred
a = 0.001; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,mesh([1:256],[1:256],C) %10*log10(Sxx))
    title('Operator C used in CLS filter')
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd Orig Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, Orig CLS')
end
a = 0.01; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd Orig Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, Orig CLS')
end
Y = fft2(y30); % 30dB blurred
a = 0.001; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y30(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd 30dB Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 30dB CLS')
end
a = 0.01; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y30(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd 30dB Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 30dB CLS')
end
Y = fft2(y20); % 20dB blurred
a = 0.001; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y20(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd 20dB Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 20dB CLS')
end
a = 0.01; % CLS parameter
XR = Y .* conj(H) ./ (abs(H).^2 + a.*abs(C).^2);
xr = real(ifft2(XR));
ISNR = 10*log10( sum([x(:)-y20(:)].^2)/sum([x(:)-xr(:)].^2) );
if plotsOn
    figure,imshow(xr,[])
    title(sprintf('CLS Rstrd 20dB Image, a = %.3f, ISNR = %.2f dB',a,ISNR))
    figure,imshow(x-xr,[])
    title('Difference Image, 20dB CLS')
end

%return