function zDemo;
%% Readme
%   - This simulation evaluate four patterns:
%     Circular pattern, [Levin 2007], [Veeraraghavan 2008], [Zhou 2009]
%
%   - This is also a demo code of Weiner Deconvolution Algorithm using 1/f Law [Zhou, ICCP 2009]
%
%   - written by Changyin Zhou @ Columbia U, 2009'

close all;

%% Set Parameters
sigma = 0.005; %Standard deviation of the Gaussian white noise
blurSize = 7; %the blursize
temp = fspecial('disk', blurSize);
flow = max(temp(:));

disp(['Noise level = ', num2str(sigma), '; Blur size = ', num2str(blurSize)]);

% Set Environment Parameters
strFolder = '';
outFolder = '.\output\';
inFolder = '.\input\';

%% Read Image
f0 = im2double(imread([inFolder, 'CZP.png']));
[hei, wid, ch] = size(f0);
if ch==3
    f0 = rgb2gray(f0);
end;
figure; imshow(f0); title('The ground truth focused image');

%% Prepare the prior matrix: 1/f law
% For specific scenes, suggest to use related images to generate the prior matrix, A
AStar = eMakePrior(hei, wid)+0.00000001;
C = sigma.^2*hei*wid./AStar;

%% Circular Pattern
k1 = fspecial('disk', blurSize);
k1 = k1*(flow/max(k1(:)));

f1 = zDefocused(f0, k1, sigma);
f0_hat = zDeconvWNR(f1, k1, C*20); 
%The image boundary causes a severe problem here in the case of circular aperture. We have to use a
%very strong prior (e.g C*20) to suppress the ringing. But in most case, we can just use C.

figure; imshow([f1, zeros(hei, 10), f0_hat]); title('A Single Circular Aperture');

%% Levin's Pattern,  [Levin, Siggraph 2007]. 
disp('Note, this pattern was designed for DFD');

k1 = im2double(imread([inFolder, 'levin.bmp']));
k1 = imresize(k1, [2*blurSize+1, 2*blurSize+1]); 
%Note, if the scale is too small, imresize may not be able to keep the right shape of the pattern.
k1 = k1*(flow/max(k1(:)));

f1 = zDefocused(f0, k1, sigma);

f0_hat = zDeconvWNR(f1, k1, C);

figure; imshow([f1, zeros(hei, 10), f0_hat]); title('Levins Coded Aperture');

%% Veeraraghavan's Pattern [Veeraraghavan, Siggraph 2008]
k1 = im2double(imread([inFolder, 'raskar.bmp']));
k1 = imresize(k1, [2*blurSize+1, 2*blurSize+1]); 
%Note, if the scale is too small, imresize may not be able to keep the right shape of the pattern.
k1 = k1*(flow/max(k1(:)));

f1 = zDefocused(f0, k1, sigma);

f0_hat = zDeconvWNR(f1, k1, C);

figure; imshow([f1, zeros(hei, 10), f0_hat]); title('Raskars Coded Aperture');

%% Zhou's Patterns,  [Zhou, ICCP 2009]
disp('This is one of the optimized patterns in [Zhou, ICCP 2009]');
k1 = im2double(imread([inFolder, 'zhou.bmp']));
k1 = imresize(k1, [2*blurSize+1, 2*blurSize+1]);
%Note, if the scale is too small, imresize may not be able to keep the right shape of the pattern.
k1 = k1*(flow/max(k1(:)));

f1 = zDefocused(f0, k1, sigma);

f0_hat = zDeconvWNR(f1, k1, C);

figure; imshow([f1, zeros(hei, 10), f0_hat]); title('Zhous Coded Aperture');


function f0 = zDeconvWNR(f, k, C)
% This is the Weiner deconvolution algorithm using 1/f law
% f: defocused image
% k: defocus kernel
% C: sigma^2/A

[hei, wid] = size(f);

k = zPSFPad(k, hei, wid);
k = fft2(fftshift(k));

f0 = abs(ifft2((fft2(f).*conj(k))./(k.*conj(k) + C)));

function f = zDefocused(f0, k, sigma, flag)
% This function is to simulate the defocus
if nargin<4
    flag = 0;
end;

[hei, wid] = size(f0);

if flag==1 %Assume the defocus is circular convolution
    k = zPSFPad(k, hei, wid);
    k = fft2(fftshift(k));
    f = abs(ifft2(fft2(f0).*k)) + randn(hei, wid)*sigma;
else %Not make that assumption. But in this case, because of the boundary effects, the "effective" noise level is very high.
    k = fliplr(flipud(k));
    f = imfilter(f0, k) + randn(hei, wid)*sigma;
end;

function outK = zPSFPad(inK, hei, wid);
% This function is to zeropadding the psf
[shei, swid] = size(inK);
outK = zeros(hei, wid);
outK(floor(end/2-shei/2)+1:floor(end/2-shei/2)+shei, floor(end/2-swid/2)+1:floor(end/2-swid/2)+swid) = inK;
