%% DICTIONARY LEARNING FOR IMAGE DENOISING
%   This file contains an example of how SMALLbox can be used to test different
%   dictionary learning techniques in Image Denoising problem.
%   It calls generateImageDenoiseProblem that will let you to choose image,
%   add noise and use noisy image to generate training set for dictionary
%   learning.
%   Three dictionary learning techniques were compared:
%   -   KSVD - M. Elad, R. Rubinstein, and M. Zibulevsky, "Efficient
%              Implementation of the K-SVD Algorithm using Batch Orthogonal
%              Matching Pursuit", Technical Report - CS, Technion, April 2008.
%   -   KSVDS - R. Rubinstein, M. Zibulevsky, and M. Elad, "Learning Sparse
%               Dictionaries for Sparse Signal Approximation", Technical
%               Report - CS, Technion, June 2009.
%   -   SPAMS - J. Mairal, F. Bach, J. Ponce and G. Sapiro. Online
%               Dictionary Learning for Sparse Coding. International
%               Conference on Machine Learning,Montreal, Canada, 2009
%
%
% Ivan Damnjanovic 2010
%%



%   If you want to load the image outside of generateImageDenoiseProblem
%   function uncomment following lines. This can be useful if you want to
%   denoise more then one image for example.
clear;
TMPpath=pwd;
FS=filesep;
[pathstr1, name, ext, versn] = fileparts(which('SMALLboxSetup.m'));
cd([pathstr1,FS,'data',FS,'images']);
load('test_image.mat');
cd(TMPpath);
% [filename,pathname] = uigetfile({'*.png;'},'Select a file containin pre-calculated notes');
% [pathstr, name, ext, versn] = fileparts(filename);
% test_image = imread(filename);
% test_image = double(test_image);
% cd(TMPpath);
% SMALL.Problem.name=name;

noise_level=[10 20 25 50 100];
% Defining Image Denoising Problem as Dictionary Learning
% Problem. As an input we set the number of training patches.
for noise_ind=1:1
for im_num=2:2
SMALL.Problem = generateImageDenoiseProblem(test_image(im_num).i, 40000, '',256, noise_level(noise_ind));
SMALL.Problem.name=int2str(im_num);

results(noise_ind,im_num).noisy_psnr=SMALL.Problem.noisy_psnr;

%%
%   Use KSVD Dictionary Learning Algorithm to Learn overcomplete dictionary

%   Initialising Dictionary structure
%   Setting Dictionary structure fields (toolbox, name, param, D and time)
%   to zero values

SMALL.DL(1)=SMALL_init_DL();

% Defining the parameters needed for dictionary learning

SMALL.DL(1).toolbox = 'KSVD';
SMALL.DL(1).name = 'ksvd';

%   Defining the parameters for KSVD
%   In this example we are learning 256 atoms in 20 iterations, so that
%   every patch in the training set can be represented with target error in
%   L2-norm (EData)
%   Type help ksvd in MATLAB prompt for more options.

Edata=sqrt(prod(SMALL.Problem.blocksize)) * SMALL.Problem.sigma * SMALL.Problem.gain;
maxatoms = floor(prod(SMALL.Problem.blocksize)/2);
SMALL.DL(1).param=struct(...
    'Edata', Edata,...
    'initdict', SMALL.Problem.initdict,...
    'dictsize', SMALL.Problem.p,...
    'exact', 1, ...
    'iternum', 20,...
    'memusage', 'high');

%   Learn the dictionary

SMALL.DL(1) = SMALL_learn(SMALL.Problem, SMALL.DL(1));

%   Set SMALL.Problem.A dictionary
%   (backward compatiblity with SPARCO: solver structure communicate
%   only with Problem structure, ie no direct communication between DL and
%   solver structures)

SMALL.Problem.A = SMALL.DL(1).D;
SMALL.Problem.reconstruct = @(x) ImgDenoise_reconstruct(x, SMALL.Problem);

%%
%   Initialising solver structure
%   Setting solver structure fields (toolbox, name, param, solution,
%   reconstructed and time) to zero values

SMALL.solver(1)=SMALL_init_solver;

% Defining the parameters needed for image denoising

SMALL.solver(1).toolbox='ompbox';
SMALL.solver(1).name='omp2';
SMALL.solver(1).param=struct(...
    'epsilon',Edata,...
    'maxatoms', maxatoms); 

%   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
%   but backward compatible with KSVD definition of denoising

SMALL.solver(1)=SMALL_solve(SMALL.Problem, SMALL.solver(1));
SMALL.solver(1).reconstructed.psnr
%%
% Use KSVDS Dictionary Learning Algorithm to denoise image

%   Initialising solver structure
%   Setting solver structure fields (toolbox, name, param, solution,
%   reconstructed and time) to zero values
% 
% SMALL.DL(2)=SMALL_init_DL();
% 
% % Defining the parameters needed for dictionary learning
% 
% SMALL.DL(2).toolbox = 'KSVDS';
% SMALL.DL(2).name = 'ksvds';
% 
% %   Defining the parameters for KSVDS
% %   In this example we are learning 256 atoms in 20 iterations, so that
% %   every patch in the training set can be represented with target error in
% %   L2-norm (EDataS). We also impose "double sparsity" - dictionary itself
% %   has to be sparse in the given base dictionary (Tdict - number of
% %   nonzero elements per atom).
% %   Type help ksvds in MATLAB prompt for more options.
% 
% 
% SMALL.DL(2).param=struct(...
%     'Edata', Edata, ...
%     'Tdict', 6,...
%     'stepsize', 1,...
%     'dictsize', SMALL.Problem.p,...
%     'iternum', 20,...
%     'memusage', 'high');
% SMALL.DL(2).param.initA = speye(SMALL.Problem.p);
% SMALL.DL(2).param.basedict{1} = odctdict(8,16);
% SMALL.DL(2).param.basedict{2} = odctdict(8,16);
% 
% % Learn the dictionary
% 
% SMALL.DL(2) = SMALL_learn(SMALL.Problem, SMALL.DL(2));

%   Set SMALL.Problem.A dictionary and SMALL.Problem.basedictionary
%   (backward compatiblity with SPARCO: solver structure communicate
%   only with Problem structure, ie no direct communication between DL and
%   solver structures)

SMALL.Problem.A = SMALL.Problem.initdict;
% SMALL.Problem.basedict{1} = SMALL.DL(2).param.basedict{1};
% SMALL.Problem.basedict{2} = SMALL.DL(2).param.basedict{2};
SMALL.DL(2).D=SMALL.Problem.initdict;
SparseDict=0;
SMALL.Problem.reconstruct = @(x) ImgDenoise_reconstruct(x, SMALL.Problem, SparseDict);

%%
%   Initialising solver structure
%   Setting solver structure fields (toolbox, name, param, solution,
%   reconstructed and time) to zero values

SMALL.solver(2)=SMALL_init_solver;

% Defining the parameters needed for image denoising

SMALL.solver(2).toolbox='ompbox';
SMALL.solver(2).name='omp2';
SMALL.solver(2).param=struct(...
    'epsilon',Edata,...
    'maxatoms', maxatoms); 

%   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
%   but backward compatible with KSVD definition of denoising
%   Pay attention that since implicit base dictionary is used, denoising
%   can be much faster then using explicit dictionary in KSVD example.

SMALL.solver(2)=SMALL_solve(SMALL.Problem, SMALL.solver(2));
%%

for i =1:1

X=SMALL.Problem.b1;
X_norm=sqrt(sum(X.^2, 1));
[X_norm_sort, p]=sort(X_norm);
p1=p(X_norm_sort>Edata);
if size(p1,2)>140000
    p2 = randperm(size(p1,2));
    p2=sort(p2(1:40000));
    size(p2,2)
    SMALL.Problem.b=X(:,p1(p2));
else 
    size(p1,2)
    SMALL.Problem.b=X(:,p1);

end

lambda=0.9998

%   Use Recursive Least Squares
%   to Learn overcomplete dictionary 

%   Initialising Dictionary structure
%   Setting Dictionary structure fields (toolbox, name, param, D and time)
%   to zero values

SMALL.DL(3)=SMALL_init_DL();

%   Defining fields needed for dictionary learning

SMALL.DL(3).toolbox = 'SMALL';
SMALL.DL(3).name = 'SMALL_rlsdla';
SMALL.DL(3).param=struct(...
    'Edata', Edata,...
    'initdict', SMALL.Problem.initdict,...
    'dictsize', SMALL.Problem.p,...
    'forgettingMode', 'FIX',...
    'forgettingFactor', lambda);

% %   Type 'help mexTrainDL in MATLAB prompt for explanation of parameters.
% 
% SMALL.DL(3).param=struct(...
%     'D', SMALL.Problem.initdict,...
%     'K', SMALL.Problem.p,...
%     'lambda', 2,...
%     'iter', 200,...
%     'mode', 3, ...
%     'modeD', 0);

%   Learn the dictionary

SMALL.DL(3) = SMALL_learn(SMALL.Problem, SMALL.DL(3));
%SMALL.DL(3).D(:,1)=SMALL.DL(1).D(:,1);
% 
% %   Set SMALL.Problem.A dictionary
% %   (backward compatiblity with SPARCO: solver structure communicate
% %   only with Problem structure, ie no direct communication between DL and
% %   solver structures)
% 
% 
% 
% %%
% %   Initialising solver structure
% %   Setting solver structure fields (toolbox, name, param, solution,
% %   reconstructed and time) to zero values
% SMALL.Problem.A = SMALL.DL(1).D;
% SMALL.Problem.reconstruct = @(x) ImgDenoise_reconstruct(x, SMALL.Problem);
% maxatoms=5;
% SMALL.solver(3)=SMALL_init_solver;
% 
% % Defining the parameters needed for denoising
% 
% % SMALL.solver(3).toolbox='SPAMS';
% % SMALL.solver(3).name='mexLasso';
% % SMALL.solver(3).param=struct(...
% %     'mode', 1, ...
% %     'lambda',Edata*Edata,...
% %     'L', maxatoms); 
% % %   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
% % %   but backward compatible with KSVD definition of denoising
% % 
% % SMALL.solver(3)=SMALL_solve(SMALL.Problem, SMALL.solver(3));
% SMALL.solver(3).toolbox='SMALL';
% SMALL.solver(3).name='SMALL_cgp';
% SMALL.solver(3).param=sprintf('%d, %.2f', maxatoms, sqrt(Edata));
% %   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
% %   but backward compatible with KSVD definition of denoising
% 
% SMALL.solver(3)=SMALL_solve(SMALL.Problem, SMALL.solver(3));

% %%
% %   Use RLS-DLA 
% 
% %   Initialising Dictionary structure
% %   Setting Dictionary structure fields (toolbox, name, param, D and time)
% %   to zero values
% 
% SMALL.DL(3)=SMALL_init_DL();
% 
% %   Defining fields needed for dictionary learning
% 
% SMALL.DL(3).toolbox = 'mpv2';
% SMALL.DL(3).name = 'rlsdla';
% 
% %   Type 'help mexTrainDL in MATLAB prompt for explanation of parameters.
% 
% SMALL.DL(3).param=struct(...
%     'D', SMALL.Problem.initdict,...
%     'K', SMALL.Problem.p,...
%     'abs', Edata*Edata,...
%     'lambda', 0.995,...
%     'iternum',1);
% 
% %   Learn the dictionary
% 
% SMALL.DL(3) = SMALL_learn(SMALL.Problem, SMALL.DL(3));
% 
% %   Set SMALL.Problem.A dictionary
% %   (backward compatiblity with SPARCO: solver structure communicate
% %   only with Problem structure, ie no direct communication between DL and
% %   solver structures)
% 
% 

%%
%   Initialising solver structure
%   Setting solver structure fields (toolbox, name, param, solution,
%   reconstructed and time) to zero values
%SMALL.DL(3).D(:,225:256)=0;
SMALL.Problem.A = SMALL.DL(3).D;
SMALL.Problem.reconstruct = @(x) ImgDenoise_reconstruct(x, SMALL.Problem);
%maxatoms=32;
SMALL.solver(3)=SMALL_init_solver;

% Defining the parameters needed for denoising

% SMALL.solver(3).toolbox='SPAMS';
% SMALL.solver(3).name='mexLasso';
% SMALL.solver(3).param=struct(...
%     'mode', 1, ...
%     'lambda',Edata*Edata,...
%     'L', maxatoms); 
% %   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
% %   but backward compatible with KSVD definition of denoising
% 
% SMALL.solver(3)=SMALL_solve(SMALL.Problem, SMALL.solver(3));

%   Initialising solver structure
%   Setting solver structure fields (toolbox, name, param, solution,
%   reconstructed and time) to zero values

SMALL.solver(3)=SMALL_init_solver;

% Defining the parameters needed for image denoising

SMALL.solver(3).toolbox='ompbox';
SMALL.solver(3).name='omp2';
SMALL.solver(3).param=struct(...
    'epsilon',Edata,...
    'maxatoms', maxatoms); 
% SMALL.solver(3).toolbox='SPAMS';
% SMALL.solver(3).name='mexLasso';
% SMALL.solver(3).param=struct(...
%     'mode', 2, ...
%     'lambda',40,...
%     'L', maxatoms); 

%   Denoising the image - SMALL_denoise function is similar to SMALL_solve,
%   but backward compatible with KSVD definition of denoising

SMALL.solver(3)=SMALL_solve(SMALL.Problem, SMALL.solver(3));
% Plot results and save midi files
SMALL.solver(3).reconstructed.psnr
% show results %

SMALL_ImgDeNoiseResult(SMALL);
end
results(noise_ind,im_num).psnr.ksvd=SMALL.solver(1).reconstructed.psnr;
results(noise_ind,im_num).psnr.odct=SMALL.solver(2).reconstructed.psnr;
results(noise_ind,im_num).psnr.rlsdla=SMALL.solver(3).reconstructed.psnr;
results(noise_ind,im_num).vmrse.ksvd=SMALL.solver(1).reconstructed.vmrse;
results(noise_ind,im_num).vmrse.odct=SMALL.solver(2).reconstructed.vmrse;
results(noise_ind,im_num).vmrse.rlsdla=SMALL.solver(3).reconstructed.vmrse;
results(noise_ind,im_num).ssim.ksvd=SMALL.solver(1).reconstructed.ssim;
results(noise_ind,im_num).ssim.odct=SMALL.solver(2).reconstructed.ssim;
results(noise_ind,im_num).ssim.rlsdla=SMALL.solver(3).reconstructed.ssim;

results(noise_ind,im_num).time.ksvd=SMALL.solver(1).time+SMALL.DL(1).time;
results(noise_ind,im_num).time.rlsdla.time=SMALL.solver(3).time+SMALL.DL(3).time;
%clear SMALL;
end
end
save results.mat results
