%%  Dictionary Learning on mimages - KSVD vs MOD vs large step gradient descent
%
%   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.
%   Two 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.



%   Centre for Digital Music, Queen Mary, University of London.
%   This file copyright 2011 Ivan Damnjanovic.
%
%   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 2 of the
%   License, or (at your option) any later version.  See the file
%   COPYING included with this distribution for more information.
%
%%

% Add SMALLbox to the MATLAB path and create the SMALL_path variable to the
% root if it is not already defined
if ~exist('SMALL_path', 'var')
    SMALLboxInit;
end

% clear variables used
clear SMALL results;

%%   Load test_image.mat.
% The test_image structure contains 5 images : lena, barbara, boat, house and peppers.

FS=filesep;
load([SMALL_path FS 'data' FS 'images' FS 'test_image.mat']);

%% Set the parameters

% Choose the image to test
im_num = 1;

% Choose the noise level
noise_level= 50;

%% Create the Problem structure.

SMALL.Problem = generateImageDenoiseProblem(test_image(im_num).i, 40000,...
    '', 256, noise_level);
SMALL.Problem.name=int2str(im_num);

% Define the error threshold to reach with OMP
Edata=sqrt(prod(SMALL.Problem.blocksize)) * ...
    SMALL.Problem.sigma * SMALL.Problem.gain;
maxatoms = floor(prod(SMALL.Problem.blocksize)/2);

%  structure is to store all results
results.noisy_psnr=SMALL.Problem.noisy_psnr;

%% Denoising with a redundant DCT dictionary

%   Set SMALL.Problem.A dictionary to the overcomplete DCT dictionary used
%   for as an initializer for dictionary learning
SMALL.DL(1)=SMALL_init_DL();
SMALL.Problem.A = SMALL.Problem.initdict;
SMALL.DL(1).D=SMALL.Problem.initdict;

%   Setting up reconstruction function
SparseDict=0;
SMALL.Problem.reconstruct = @(x) ImageDenoise_reconstruct(x, SMALL.Problem, SparseDict);

%   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 - find the sparse solution in the learned
%   dictionary for all patches in the image and the end it uses
%   reconstruction function to reconstruct the patches and put them into a
%   denoised image
SMALL.solver(1)=SMALL_solve(SMALL.Problem, SMALL.solver(1));

%   Show PSNR after reconstruction
SMALL.solver(1).reconstructed.psnr

%% Denoising with a dictionary learnt using with KSVDbox

%   Initialising Dictionary structure
%   Setting Dictionary structure fields (toolbox, name, param, D and time)
%   to zero values
SMALL.DL(2)=SMALL_init_DL();

% Defining the parameters needed for dictionary learning
SMALL.DL(2).toolbox = 'KSVD';
SMALL.DL(2).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.

SMALL.DL(2).param=struct(...
    'Edata', Edata,...
    'initdict', SMALL.Problem.initdict,...
    'dictsize', SMALL.Problem.p,...
    'exact', 1, ...
    'iternum', 20,...
    'memusage', 'high');

%   Learn the dictionary

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

%   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(2).D;
SMALL.Problem.reconstruct = @(x) ImageDenoise_reconstruct(x, SMALL.Problem);

%   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 - find the sparse solution in the learned
%   dictionary for all patches in the image and the end it uses
%   reconstruction function to reconstruct the patches and put them into a
%   denoised image
SMALL.solver(2)=SMALL_solve(SMALL.Problem, SMALL.solver(2));

%   Show PSNR after reconstruction
SMALL.solver(2).reconstructed.psnr

%% Denoising with a dictionary learnt using LGD
%   For comparison purposes we will denoise image with overcomplete DCT
%   here
%   Set SMALL.Problem.A dictionary to be oDCT (i.e. Problem.initdict -
%   since initial dictionaruy is already set to be oDCT when generating the
%   denoising problem


%   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);

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

SMALL.DL(3)=SMALL_init_DL('TwoStepDL', 'LGD');


%   Defining the parameters for MOD
%   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.


SMALL.DL(3).param=struct(...
    'solver', SMALL.solver(3),...
    'initdict', SMALL.Problem.initdict,...
    'dictsize', SMALL.Problem.p,...
    'iternum', 20,...
    'show_dict', 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)
SMALL.Problem.A = SMALL.DL(3).D;
SMALL.Problem.reconstruct = @(x) ImageDenoise_reconstruct(x, SMALL.Problem);

%   Denoising the image - find the sparse solution in the learned
%   dictionary for all patches in the image and the end it uses
%   reconstruction function to reconstruct the patches and put them into a
%   denoised image
SMALL.solver(3)=SMALL_solve(SMALL.Problem, SMALL.solver(3));

%   Show PSNR after reconstruction
SMALL.solver(3).reconstructed.psnr

% show results %

SMALL_ImgDeNoiseResult(SMALL);

results(im_num).psnr.ksvd=SMALL.solver(1).reconstructed.psnr;
results(im_num).psnr.mod=SMALL.solver(2).reconstructed.psnr;
results(im_num).psnr.lgd=SMALL.solver(3).reconstructed.psnr;
results(im_num).vmrse.ksvd=SMALL.solver(1).reconstructed.vmrse;
results(im_num).vmrse.mod=SMALL.solver(2).reconstructed.vmrse;
results(im_num).vmrse.lgd=SMALL.solver(3).reconstructed.vmrse;
results(im_num).ssim.ksvd=SMALL.solver(1).reconstructed.ssim;
results(im_num).ssim.mod=SMALL.solver(2).reconstructed.ssim;
results(im_num).ssim.lgd=SMALL.solver(3).reconstructed.ssim;

results(im_num).time.ksvd=SMALL.solver(1).time+SMALL.DL(1).time;
results(im_num).time.rlsdla.time=SMALL.solver(3).time+SMALL.DL(3).time;
%clear SMALL;

% save results.mat results
