% 
% Depth from Defocus via Diffusion
%
% Copyright 2006 Paolo Favaro (p.favaro@hw.ac.uk)
% 
% School of Engineering and Physical Sciences
% Heriot-Watt University, Edinburgh, UK
% 
% Last revision: August 2006
%
% This program can be used only for research purposes.
% This program is distributed WITHOUT ANY WARRANTY; 
% without even the implied warranty of MERCHANTABILITY 
% or FITNESS FOR A PARTICULAR PURPOSE.

clear all
close all

No = 0;
Yes = 1;

% set the maximum number of iterations
MaxIterations = 100;

fprintf('Shape estimation via relative diffusion\n\n');
% shape type (wave, slope, box, sin, plane, realdata)
notValid = Yes;
while notValid
    shape = input(['What test do you want to perform?'...
        '\nChoose wave, slope, box, sin, plane, realdata: '],'s'); 
    notValid = ~strcmp(shape,'wave')&...
        ~strcmp(shape,'slope')&...
        ~strcmp(shape,'box')&...
        ~strcmp(shape,'sin')&...
        ~strcmp(shape,'plane')&...
        ~strcmp(shape,'realdata');
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Dataset initialization
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% choose between loading or generating defocused images
LoadImage = strcmp(shape,'realdata');
if LoadImage == No
    % generate synthetic image
    fprintf('Generating synthetic dataset...');
    % set image size
    M = 101; % use odd numbers
    N = 101; % use odd numbers
    % lens parameters
    F = 12e-3; % focal length
    Fnumb = 2; % F number
    D = F/Fnumb; % aperture (effective lens diameter)
    gamma = 2e4; % calibration parameter
    p = 1./(1/F-1./[.52 .85]); % distance CCD-lens
    K = length(p); % number of focus settings
    % generate depth map
    [X,Y] = meshgrid([-(N-1)/2:(N-1)/2],[-(M-1)/2:(M-1)/2]);
    switch shape
        case 'wave'
            DepthMap = (1+cos(-(X.^2+Y.^2)/200))./...
                sqrt(X.^2/8600+Y.^2/8600+1)/2*.33+.52;
        case 'slope'
            % along X
            DepthMap = (X-min(X(:)))/(max(X(:))-min(X(:)))*.33+.52;
        case 'box'
            block1 = floor(M/5:M/5*3);
            block2 = floor(M/5*2:M/5*4);
            block3 = floor(N/10:N/2);
            block4 = floor(N/10*4:N/10*9);
            DepthMap = ones(M,N)*.52;
            DepthMap(block1,block3) = ...
                DepthMap(block1,block3)+.2;
            DepthMap(block2,block4) = ...
                DepthMap(block2,block4)+.13;
        case 'sin'
            DepthMap = (sin(X/3)+1)/2*.33+.52;
        otherwise % plane
            DepthMap = .05+ones(M,N)*sum(p)*F/(sum(p)-2*F);
    end
    % compute true domain
    OmegaTruth = double(DepthMap<sum(p)*F/(sum(p)-2*F));
    % generate random texture (test algorithm on different intensities)
    a1 = 50; a2 = 120; a3 = 255;
    Radiance = [a1*rand(floor(M/3),N);...
        a2*rand(floor(M/3),N);...
        a3*rand(M-2*floor(M/3),N)]+1; % skip zero
    % add space-varying blurring to test performance of algorithm with
    % different textures (high-freq to low-freq)
    TextureDepth = p(1)*F/(p(1)-F)*ones(M,N)+...
        [zeros(M,floor(N/3)) ...
        .15*ones(M,floor(N/3)) ...
        .3*ones(M,N-2*floor(N/3))];
    % set camera parameters
    parms.m = M;
    parms.n = N;
    parms.F = F;
    parms.Fnumb = Fnumb;
    parms.D = D;
    parms.gamma = gamma;
    % same settings for second image
    parms(2) = parms(1);
    % plane in focus distance (from the lens)
    parms(1).p = p(1);
    parms(2).p = p(2);
    % blur the radiance in three columns
    % compute the relative diffusion map
    Sigma = (gamma*D/2*abs(1-p(1).*(1/F-1./TextureDepth))).^2;
    Radiance = synthesize(Radiance,Sigma);
    % synthesizing blurred images
    Sigma1 = (gamma*D/2*abs(1-p(1).*(1/F-1./DepthMap))).^2;
    Sigma2 = (gamma*D/2*abs(1-p(2).*(1/F-1./DepthMap))).^2;
    I(:,:,1) = synthesize(Radiance,Sigma1);
    I(:,:,2) = synthesize(Radiance,Sigma2);
    % select the domain where the data term is considered
    ww = 3;
    mask = ones(M,N);
    mask(1:ww,:) = 0;
    mask(end-ww+1:end,:) = 0;
    mask(:,1:ww) = 0;
    mask(:,end-ww+1:end) = 0;
    dt_data = 2e-1;% data term step
    dt_reg  = 2e-1;% regularization term step
    maxDelta = 1; % max amount of relative blur
    win = 2; % regularization of initial depth map
    fprintf('done\n');
else
    % load images
    fprintf('Loading dataset...');
    load('DataSet.mat');
    % first resize
    scalingFactor = 1/2;
    I1 = imresize(I1,scalingFactor); % limit number of computations
    I2 = imresize(I2,scalingFactor); % limit number of computations
    % select region of interest
    I(:,:,2) = I1(1:end-10,60:end-50);
    I(:,:,1) = I2(1:end-10,60:end-50);
    [M,N,K] = size(I);
    % lens parameters
    F = 12e-3; % focal length
    Fnumb = 2; % F number
    D = F/Fnumb; % aperture (effective lens diameter)
    gamma = 3e4; % calibration parameter
    p = 1./(1/F-1./[.52 .85]); % distance CCD-lens)
    % set camera parameters
    parms.m = M;
    parms.n = N;
    parms.F = F;
    parms.Fnumb = Fnumb;
    parms.D = D;
    parms.gamma = gamma;
    % same settings for second image
    parms(2) = parms(1);
    % plane in focus distance (from the lens)
    parms(1).p = p(1);
    parms(2).p = p(2);
    % select the domain where the data term is considered
    ww = 3;
    mask = ones(M,N);
    mask(1:ww,:) = 0;
    mask(end-ww+1:end,:) = 0;
    mask(:,1:ww) = 0;
    mask(:,end-ww+1:end) = 0;
    dt_data = 2e-1;%6e-1; % data term step
    dt_reg  = 1e0;%2e0;  % regularization term step
    maxDelta = 1; % max amount of relative blur
    win = 5; % regularization of initial depth map
    K = length(p); % number of focus settings
    fprintf('done\n')
end

% preconditioning threshold
thresh = 6e-2;

% choose whether to use the fast initialization or not
% init type (Y,N)
notValid = Yes;
while notValid
    init = input(['Do you want to use '...
                  'the fast initialization (Y/N)?'],'s'); 
    notValid = ~strcmp(init,'Y')&...
        ~strcmp(init,'N');
end

if strcmp(init,'Y')
    % initialize depth map
    DepthMap_e = initialize(I,maxDelta,win,gamma,F,D,p,'l2');
else
    % initialize depth map with a plane
    DepthMap_e = sum(p)*F/(sum(p)-2*F)*ones(M,N);
end

% compute the relative diffusion map
Sigma1 = (gamma*D/2*abs(1-p(1).*(1/F-1./DepthMap_e))).^2;
Sigma2 = (gamma*D/2*abs(1-p(2).*(1/F-1./DepthMap_e))).^2;
Dsigma = Sigma2-Sigma1;

% start deblurring iterations
fprintf('Shape estimation (%3i iterations)\n     ',MaxIterations);
for i=1:MaxIterations
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % visualize estimated depth map
    figure(1)
    set(gcf,'name',sprintf('Iteration %3i out of %3i',i,MaxIterations));
    if LoadImage == No
        % data is synthetic; show ground truth
        subplot(221)
        image([I(:,:,1) 200*ones(size(I,1),10) I(:,:,2)])
        axis equal
        axis off
        set(gca,'XTick',[],'YTick',[]);
        set(gca,'Box','off');
        colormap gray(256)
        title('Input #1/#2')
        drawnow;
        subplot(222)
        imagesc([DepthMap_e>sum(p)*F/(sum(p)-2*F) ...
                 .8*ones(size(I,1),10) ...
                 DepthMap>sum(p)*F/(sum(p)-2*F)])
        axis equal
        axis off
        colormap gray(256)
        title('Estimated/True Segmentation')
        drawnow;
        step = ceil(min([size(DepthMap_e,1) ...
               size(DepthMap_e,2)])/50);
        pix2mm = 1/200;
        [X,Y] = meshgrid(pix2mm*[1:step:size(DepthMap_e,2)],...
                pix2mm*[1:step:size(DepthMap_e,1)]);
        [X2,Y2] = meshgrid(pix2mm*(N+[1:step:size(DepthMap_e,2)]),...
                pix2mm*([1:step:size(DepthMap_e,1)]));
        axisZ = 1./(1/F-1./p);
        subplot(223)
        h1 = surf(X2,-DepthMap(1:step:end,1:step:end),-Y2,DepthMap(1:step:end,1:step:end));
        shading interp
        hold on
        h = surf(X,-DepthMap_e(1:step:end,1:step:end),-Y,DepthMap_e(1:step:end,1:step:end));
        shading interp
        hold off
        axis([0 2*N -axisZ(2) -axisZ(1) -M 0])
        set(gca,'XTick',[],'YTick',[],'ZTick',[]);
        colormap gray(256)
        title('Estimated(red)/True(blue) Depth Map')
        %legend('True Depth','Estimated Depth');
        view([-1 -2 1]);
        axis equal
        axis off
        set(h1,'EdgeColor','b');
        set(h,'EdgeColor','r');
        drawnow;
        subplot(224)
        imagesc([DepthMap_e ...
                 max(DepthMap(:))*ones(size(I,1),10) ...
                 DepthMap]);
        colormap gray(256)
        %axis off
        axis equal
        set(gca,'XTick',[],'YTick',[]);
        set(gca,'Box','off');
        title('Estimated/True Depth Map')
        drawnow;
    else
        % data is real
        subplot(211)
        image([I(:,:,1) 200*ones(size(I,1),10) ...
               I(:,:,2) 200*ones(size(I,1),10) ...
               256*(DepthMap_e>sum(p)*F/(sum(p)-2*F))])
        axis equal
        axis off
        set(gca,'XTick',[],'YTick',[]);
        set(gca,'Box','off');
        colormap gray(256)
        title('Input #1/Input #2/Estimated Segmentation')
        drawnow;
        step = ceil(min([size(DepthMap_e,1) ...
               size(DepthMap_e,2)])/50);
        pix2mm = 1/600;
        [X,Y] = meshgrid(pix2mm*[1:step:size(DepthMap_e,2)],...
                pix2mm*[1:step:size(DepthMap_e,1)]);
        axisZ = 1./(1/F-1./p);
        subplot(223)
        h = surf(X,-DepthMap_e(1:step:end,1:step:end),-Y,DepthMap_e(1:step:end,1:step:end));
        shading interp
        hold off
        axis([0 N -axisZ(2) -axisZ(1) -M 0])
        set(gca,'XTick',[],'YTick',[],'ZTick',[]);
        colormap gray(256)
        title('Estimated Depth Map')
        %legend('True Depth','Estimated Depth');
        view([-1 -2 1]);
        axis equal
        axis off
        set(h,'EdgeColor','r');
        drawnow;
        subplot(224)
        imagesc(DepthMap_e);
        colormap gray(256)
        %axis off
        axis equal
        set(gca,'XTick',[],'YTick',[]);
        set(gca,'Box','off');
        title('Estimated Depth Map')
        drawnow;
    end
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % end visualization
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Optimization procedure starts here %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    modelnoise = 0;
    sensornoise = 0;
    Omega = double(Dsigma>0); % Heaviside function
    delta = zerocrossing(Dsigma); % Dirac delta
    % compute diffusion PDE solution by simulation
    [dummy,u1] = synthesize(I(:,:,1),Dsigma.*Omega,...
        modelnoise,sensornoise,5);
    [dummy,u2] = synthesize(I(:,:,2),-Dsigma.*(1-Omega),...
        modelnoise,sensornoise,5);
    % compute initial conditions for the adjoint PDE
    w01 = reshape(u1(:,:,end),M,N)-I(:,:,2);
    % compute adjoint diffusion PDE solution by simulation
    [dummy,w1] = synthesize(w01,Dsigma.*Omega,...
        modelnoise,sensornoise,5);
    w02 = reshape(u2(:,:,end),M,N)-I(:,:,1);
    % compute adjoint diffusion PDE solution by simulation
    [dummy,w2] = synthesize(w02,-Dsigma.*(1-Omega),...
        modelnoise,sensornoise,5);
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % cost functional gradients computation %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % time integration (from 0 to T) of grad v grad u
    [dvdu1] = timeintegrate(u1,w1);
    % time integration (from 0 to T) of grad v grad u
    [dvdu2] = timeintegrate(u2,w2);
    dEdc = -2*Omega.*dvdu1+2*(1-Omega).*dvdu2+...
            delta.*((u1(:,:,end)-I(:,:,2)).^2-...
            (u2(:,:,end)-I(:,:,1)).^2);
    % preconditioning
    [dummy,p1] = synthesize(reshape(I(:,:,2),M,N,1),...
                  Dsigma.*Omega,modelnoise,sensornoise,5);
    [dummy,p2] = synthesize(reshape(I(:,:,1),M,N,1),...
                 -Dsigma.*(1-Omega),modelnoise,sensornoise,5);
    % time integration (from 0 to T) of grad p grad u
    [dpdu1] = timeintegrate(u1,p1);
    [dpdu2] = timeintegrate(u2,p2);
    precond = 1+abs(-2*Omega.*dpdu1+2*(1-Omega).*dpdu2);
    dEdc = dEdc./precond;
    dEdc = (abs(dEdc) <  thresh).*dEdc+...
           (abs(dEdc) >= thresh).*thresh.*sign(dEdc);
    dDatads = dEdc.*mask;
    % make sure the number of iterations is an even number
    nIt = 2*(floor(dt_reg/1e-1)+1); 
    beta_data = dt_data/nIt;
    beta_regularization = dt_reg/nIt;
    for tt=1:nIt
        % compute gradient of the cost functional (data term +
        % regularization) wrt the depth map
        if tt-floor(tt/2)*2==0
            dsigma = beta_data*dDatads - beta_regularization.*...
                laplacian(Dsigma,'forward');
        else
            dsigma = beta_data*dDatads - beta_regularization.*...
                laplacian(Dsigma,'backward');
        end
        Dsigma = Dsigma-dsigma;
    end
    DepthMap_e = 1./(1/F-(p(2)-p(1)-sqrt((p(2)-p(1))^2+...
                 4*(p(2)^2-p(1)^2)*Dsigma/gamma^2/D^2))/...
                 (p(2)^2-p(1)^2));
end
fprintf('\ndone\n');

return