
% -------------------------------------------------------------------------
%     This code implements the recoloring method described in the paper:
%     Musialski, P., Cui, M., Ye, J., Razdan, A., & Wonka, P. (2012). 
%     A Framework for Interactive Image Color Editing. 
%     The Visual Computer, 29(11), 1173–1186. doi:10.1007/s00371-012-0761-5
% 
%     If you use the code or parts of it in your academic work you need to
%     cite the paper: 
%
%     @article{Musialski2012,
%         author = {Musialski, Przemyslaw and Cui, Ming and Ye, Jieping and Razdan, Anshuman and Wonka, Peter},
%         doi = {10.1007/s00371-012-0761-5},
%         issn = {0178-2789},
%         journal = {The Visual Computer},
%         month = nov,
%         number = {11},
%         pages = {1173--1186},
%         title = {{A Framework for Interactive Image Color Editing}},
%         url = {http://link.springer.com/10.1007/s00371-012-0761-5},
%         volume = {29},
%         year = {2012}
%     }
%
% -------------------------------------------------------------------------
%
%     Copyright (C) 2011-2013, Przemyslaw Musialski, pm@cg.tuwien.ac.at
% 
%     This program is free software: you can redistribute it and/or modify
%     it under the terms of the GNU General Public License as published by
%     the Free Software Foundation, either version 3 of the License, or
%     (at your option) any later version.
% 
%     This program is distributed in the hope that it will be useful,
%     but WITHOUT ANY WARRANTY; without even the implied warranty of
%     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%     GNU General Public License for more details.
% 
%     You should have received a copy of the GNU General Public License
%     along with this program.  If not, see <http://www.gnu.org/licenses/>.
%
% -------------------------------------------------------------------------
%
%     The code has been tested with MATLAB 2012 and 2013
%
%     The included GetPointsFromBarycentric3D.mexw64 is currently compiled
%     with Microsoft Visual C++ 2010, Visual Studio 10.0
%     In order to compiled it call:
%     mex GetPointsFromBarycentric3D.cpp 
%      
% -------------------------------------------------------------------------
%
%     Test-images are contained in folder ./data/
%     banana1, man1, dance1, river1 are copyrighted by Tom Ang: http://tomang.com/
%     flowers4 is copyrithed by Norman Koren: http://www.normankoren.com/
%
% -------------------------------------------------------------------------
%% global constants

clc;
clear;

cformIn                 = makecform('srgb2lab','AdaptedWhitePoint', whitepoint('d50'));
cformOut                = makecform('lab2srgb','AdaptedWhitePoint', whitepoint('d50'));

% -------------------------------------------------------------------------
%% load the input image

% banana1 flowers4 butterfly1 man1 dance1 river1 flower2
dataset                 = 'flowers4';
basedir                 = '.\data\';
targetid                = '1';

fprintf(1,'Loading image \n'); 
t1=tic;
    targetfilename      = strcat(basedir,dataset,'\','target_',targetid,'.png');
    target              = im2double(imread(targetfilename)); 
    mask                = im2double(imread(strcat('.\data\',dataset,'\','mask.png')));  
    
    inputfilename       = strcat(basedir,dataset,'\','original.png');%.png')
    input               = im2double(imread(inputfilename)); 
    input               = input(:,:,1:3);

handles.basedir         = basedir;
handles.dataset         = dataset;
target                  = target.*repmat(mask,[1,1,3]);

[row, col, ~]           = size(target);
N                       = row*col;
srcRGB                  = WrapImage2Vector(input);
tarRGB                  = WrapImage2Vector(target);

srcLab                  = applycform(srcRGB,cformIn);
tarLab                  = applycform(tarRGB,cformIn);

% -- init handles
handles.mask            = mask;
handles.ourimg          = input;
handles.rgbimg          = target;
handles.srcRGB          = srcRGB;
handles.tarRGB          = tarRGB;
handles.srcLab          = srcLab;
handles.tarLab          = tarLab;
handles.imW             = col;
handles.imH             = row;
handles.targetid        = targetid;

% -------------------------------------------------------------------------
% plot images
figure(1); clf; 
subplot(2,2,1), imshow(input); title('Input Image');
subplot(2,2,2), imshow(target); title('Target Image');

fprintf(1,['...took ' num2str(toc(t1)) 's.\n']);
% -------------------------------------------------------------------------
%% init landmarks (need to be run only once for a given sampling)

ratio_t                 = 100;  % ratio for samples on target regions
ratio_n                 = 1000; % ratio for samples on non-target
handles.K               = 21;   % lle neighborss
handles.tol             = 1e-3;
handles.ratio_target    = 1/ratio_t;
handles.ratio_nontarget = 1/ratio_n;
handles.spatialscale    = 1.0;
handles.cformOut        = cformOut;

handles                 = InitLandmarks(handles);

gcf; subplot(2,2,4); 
imshow(handles.samplesimg);  title('Samples Image');      
imwrite(handles.samplesimg,strcat(handles.basedir,handles.dataset,'\','samples_',handles.targetid,'.png'),'png');

%% run optimization (run with new values of lambda)

handles.lambda          = 0.01; % optimization (pareto) lambda

% channels 3 for all 3 La*b* channels
% channels 2 for a*b* channels only
handles.channels        = 3;
outimg                  = RunOptimization(handles);

% plot image
gcf; subplot(2,2,3); 
imshow(outimg); title('Output Image');
imwrite(outimg, strcat('.\data\',dataset,'\','output_',targetid,'.png'), 'png');

% save current params to file
save(strcat('.\data\',dataset,'\','parameters_',targetid,'.mat'),...
    '-struct', 'handles', 'lambda','channels','K','ratio_target','ratio_nontarget');
