%% =========================================================================
% Demo codes for Patch-Inpainting using bi-directional similarity 
% Code Developed by Pulak Purkait : Date oct 10, 2012
% ----------------------------------------------------
% ECSU, Indian Statistical Institute Kolkata
% For any questions/bug related issues, send email to: pulak_r@isical.ac.in
% Referece 
%     P. Purkait and B chanda, Digital Restoration of Damaged Mural images,
%     ICVGIP 2012, Decembar 2012.
% =========================================================================
%%
image = 'inpainting/HillWeb';
% image = 'inpainting/Road_Trip';
% image = 'inpainting/man';


Aorg = imread([image '.png']);
% Aorg = imread([image '.bmp']);
Morg = imread([image '-mask.png']);

close all

cores = 4;    % Use more cores for more speed
algo = 'cputiled';

nn_iters = 5; %defaut = 5
rs_iters = 1; % default = 1;
verbose = true;

rs_factor = 1.5;
width = 29;

% Also accept RGB masks, but only use first channel
if size(Morg,3) > 1
    Morg = Morg(:,:,1);
end

% Determinte starting scale
[m n rgb] = size(Aorg);
startscale = -ceil(log2(min(m,n))/log2(rs_factor)) + 13;
scale = rs_factor^(startscale);

% Resize image to starting scale
A = imresize(Aorg,scale);
M = imresize(Morg,scale);
M(M>0)=1;
M3 = repmat(M,[1 1 3])==1;

% Random starting guess for inpainted image

width2 = (width-1)/2;

B = A;
R = rand(size(A)) > 0.5;
Rnd = uint8(255*R);
B(M3)=Rnd(M3);
R = rand(size(A)) > 0.5;
Rnd = uint8(255*R);

M4 = M3;
se = strel('disk', 1);
ittt = 2;

%% Initial Approximation in lower scale
figure,
while sum(M4(:))   
    M5 = zeros(size(M3));
    M5(1:end-width2, 1:end-width2, :) = double(M3(width2+1:end, width2+1:end, :));
    if verbose
%         fprintf('  Iteration %2d\n',iterations);
        imshow(A);
        pause(0.01)
    end
    M4 = imerode(M4, se);
    for inner = 1:ittt
        A(~M3) = B(~M3);
        A(M4)=Rnd(M4);
        PM_ann = nnmex(A, B, algo, width, [], [], [], [], [], cores, M5);
        A = votemex(B, PM_ann,  [], algo, width, [], []);
    end
end

%% Go through all scales
ittt = 18;

for logscale = startscale+1:0

    scale = rs_factor^(logscale);
    Adata = imresize(Aorg,scale);
    [m n rgb] = size(Adata);
    A = imresize(A,[m n]);
    M = imresize(Morg,[m n]);
    M(M>0)=1;
    M3 = repmat(M,[1 1 3])==1;
    A(~M3) = Adata(~M3);
    
    iterations = ittt;
    M4 = zeros(size(M3));
    M4(1:end-width2, 1:end-width2, :) = double(M3(width2+1:end, width2+1:end, :));
    se = strel('square', width+logscale);
    M4 = imdilate(M4, se);
    B = A;
    for iter = 1:iterations       
        if verbose
            fprintf('  Iteration %2d/%2d\n',iter,iterations);
            imshow(A);
            pause(0.001)
        end
        
        % Compute NN field
        PM_ann = nnmex(A, B, algo, width, [], [], [], [], [], cores, M4);
        PM_bnn = nnmex(B, A, algo, width, [], [], [], [], [], cores, M4);
        A = votemex(B, PM_ann,  [], algo, width);
%         A = votemex(B, PM_ann,  PM_bnn, algo, width, M4);
        if logscale < 0
            A(~M3) = B(~M3);
        end
    end
end


imshow(Aorg);
figure, imshow(A);
imwrite(A, [image '-out.png']);
