%%  Dictionary Learning for Image Denoising - KSVD vs KSVDS vs SPAMS
%
%   *WARNING!* You should have SPAMS in your search path in order for this
%   script to work.Due to licensing issues SPAMS can not be automatically 
%   provided in SMALLbox (http://www.di.ens.fr/willow/SPAMS/downloads.html).
%
%   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.
%   -   KSVDS - R. Rubinstein, M. Zibulevsky, and M. Elad, "Learning Sparse
%               Dictionaries for Sparse Signal Approximation", Technical
%               Report - CS, Technion, June 2009.
%

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

%   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.

% TMPpath=pwd;
% FS=filesep;
% [pathstr1, name, ext] = fileparts(which('SMALLboxSetup.m'));
% cd([pathstr1,FS,'data',FS,'images']);
% [filename,pathname] = uigetfile({'*.png;'},'Select a file containin pre-calculated notes');
% [pathstr, name, ext] = fileparts(filename);
% test_image = imread(filename);
% test_image = double(test_image);
% cd(TMPpath);
% SMALL.Problem.name=name;


% Defining Image Denoising Problem as Dictionary Learning
% Problem. As an input we set the number of training patches.

SMALL.Problem = generateImageDenoiseProblem('', 40000);


%%
%   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) ImageDenoise_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 - 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 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.

EdataS=sqrt(prod(SMALL.Problem.blocksize)) * SMALL.Problem.sigma * SMALL.Problem.gain;
SMALL.DL(2).param=struct(...
    'Edata', EdataS, ...
    '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.DL(2).D;
SMALL.Problem.basedict{1} = SMALL.DL(2).param.basedict{1};
SMALL.Problem.basedict{2} = SMALL.DL(2).param.basedict{2};

%   Setting up reconstruction function

SparseDict=1;
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(2)=SMALL_init_solver;

% Defining the parameters needed for image denoising

SMALL.solver(2).toolbox='ompsbox';
SMALL.solver(2).name='omps2';
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));


%%
% Plot results and save midi files

% show results %

SMALL_ImgDeNoiseResult(SMALL);
