%% 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.
%   We tested time and psnr for two dictionary learning techniques. This
%   example does not represnt any extensive testing. The aim of this
%   example is just to show how SMALL structure can be used for testing.
%
%   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.
%   -   SPAMS - J. Mairal, F. Bach, J. Ponce and G. Sapiro. Online
%               Dictionary Learning for Sparse Coding. International
%               Conference on Machine Learning,Montreal, Canada, 2009

%
%   Centre for Digital Music, Queen Mary, University of London.
%   This file copyright 2010 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.%%
%%

clear all;

%% Load an image
TMPpath=pwd;
FS=filesep;
[pathstr1, name, ext, versn] = fileparts(which('SMALLboxSetup.m'));
cd([pathstr1,FS,'data',FS,'images']);
[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);

%   number of different values we want to test
n =5;
step = floor((size(test_image,1)-8+1)*(size(test_image,2)-8+1)/n);
Training_size=zeros(1,n);
time = zeros(2,n);
psnr = zeros(2,n);
for i=1:n
    
    %   Here we want to test time spent and quality of denoising for
    %   different sizes of training sample.
    Training_size(i)=i*step;
    
    SMALL.Problem = generateImageDenoiseProblem(test_image,Training_size(i));
    SMALL.Problem.name=name;
    %%
    %   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,...
        '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 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

    
    %%
    %   Use SPAMS Online Dictionary Learning Algorithm
    %   to Learn overcomplete dictionary (Julien Mairal 2009)
    %   (If you have not installed SPAMS please comment the following two cells)
    
    %   Initialising Dictionary structure
    %   Setting Dictionary structure fields (toolbox, name, param, D and time)
    %   to zero values
    
    SMALL.DL(2)=SMALL_init_DL();
    
    %   Defining fields needed for dictionary learning
    
    SMALL.DL(2).toolbox = 'SPAMS';
    SMALL.DL(2).name = 'mexTrainDL';
    
    %   Type 'help mexTrainDL in MATLAB prompt for explanation of parameters.
    
    SMALL.DL(2).param=struct(...
        'D', SMALL.Problem.initdict,...
        'K', SMALL.Problem.p,...
        'lambda', 2,...
        'iter', 300,...
        'mode', 3,...
        'modeD', 0 );
    
    %   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) ImgDenoise_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 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 results %%
    %   This will show denoised images and dictionaries for all training sets.
    %   If you are not interested to see them and do not want clutter your
    %   screen comment following line
    
    SMALL_ImgDeNoiseResult(SMALL);
    
    time(1,i) = SMALL.DL(1).time;
    psnr(1,i) = SMALL.solver(1).reconstructed.psnr;
    
    time(2,i) = SMALL.DL(2).time;
    psnr(2,i) = SMALL.solver(2).reconstructed.psnr;
    
    clear SMALL
end

%% show time and psnr %%
figure('Name', 'KSVD vs SPAMS');

subplot(1,2,1); plot(Training_size, time(1,:), 'ro-', Training_size, time(2,:), 'b*-');
legend('KSVD','SPAMS',0);
title('Time vs Training size');
xlabel('Training Size (Num. of patches)');
ylabel('Time(s)');
subplot(1,2,2); plot(Training_size, psnr(1,:), 'ro-', Training_size, psnr(2,:), 'b*-');
legend('KSVD','SPAMS',0);
title('PSNR vs Training size');
xlabel('Training Size (Num. of patches)');
ylabel('PSNR(dB)');