% learn hierarchical sparse representation for object recognition
% written by Jimei Yang
% updated on 4/4/2010
clear all;
clc;
%% set path
addpath(genpath('sparse_coding'));
addpath(genpath('liblinear-1.51'));

%% parameters setting
img_dir = 'image'; % folder for dataset images
data_dir = 'data'; % folder to save the features of the chosen dataset
dataSet = 'Caltech2';

param.gridSpacing = 1; % image patch sampling grid in S1 layer
param.patchSize = 8; % the size of image patch in S1 layer
param.maxImgSize = 151; % image size after preprocessing
param.localNormRange = 9; % local normalization range in preprocessing step

param.nBasesS1 = 64; % number of bases in S1 dictionary
param.poolingS1 = 4; % grid for max pooling in C1 layer 
param.nSmpForS1Dict = 200000; % number of image patch samples for training S1 dictionary
param.biasC1 = 1; % spatal grouping bias in C1 layer
param.cellSizeC1 = 4; % the size of spatial grouping
param.pca = 96; % PCA dimensions in C1 layer

param.nBasesS2 = 2048; % number of bases in S1 dictionary
param.nSmpForS2Dict = 200000; % number of C1 samples for training S2 dictionary
param.poolingS2 = [1,2,4]; % spatial block number on each level of the pyramid in C2 layer
param.weightsC2 = [1,1,1]; % weights of features in different scales used in multi-scale max pooling

param.beta = 1e-3; % a small regularization for stablizing sparse coding
param.maxIter = 20; % maximum iterations for sparse coding
param.gamma = 0.15; % a small regulariztion for sparsity in sparse coding

param.tr_num = 30; % number of training samples
param.nRounds = 5; % number of trials
%% build database
rt_img_dir = fullfile(img_dir, dataSet);
imgPath = [rt_img_dir,'\database.mat'];
try
    load(imgPath);
catch
    database = build_database(rt_img_dir,'*.jpg');
    save(imgPath,'database');
end

%% preprocess the images
fprintf('preprocess the images...\n');
rt_data_dir = ['data\pp_', dataSet];
imgPath = [rt_data_dir,'\database.mat'];
try
    load(imgPath);
catch
    imgDatabase = preprocess(database,param);
    save(imgPath,'imgDatabase');
end
fprintf('done!\n');
%% learn s1 dictionary
fprintf('learn s1 dictionary...\n');
s1DictPath = ['dictionary\s1_dict_' dataSet '_' num2str(param.nBasesS1) '.mat'];
try 
    load(s1DictPath);
catch
    [ s1Dictionary ] = learn_s1_dictionary(imgDatabase,param);
    save( s1DictPath, 's1Dictionary' );
    clear s1DictPath;
end
fprintf('done!\n');
%% build s1 layer -- s1 feature extraction by sparse coding
fprintf('build s1 layer...\n');
rt_data_dir = ['data\s1_pp_', dataSet];
s1Path = [rt_data_dir,'\s1Database.mat'];
try 
    load(s1Path);
catch
    s1Database = build_s1_layer(imgDatabase,s1Dictionary,param);
    save(s1Path,'s1Database');
    clear s1Path;
end
fprintf('done!\n');
%% build c1 layer -- spatial max pooling
fprintf('build c1 layer...\n');
rt_data_dir = ['data\c1_s1_pp_', dataSet];
c1Path = [rt_data_dir,'\c1Database.mat'];
try 
    load(c1Path);
catch
    [c1Database] = build_c1_layer(s1Database,param);
    save(c1Path,'c1Database');
    clear c1Path;
end
fprintf('done!\n');
%% learn s2 dictionary
fprintf('learn s2 dictionary...\n');
s2DictPath = ['dictionary\s2_dict_' dataSet '_' num2str(param.nBasesS2) '_pca', num2str(param.pca),'.mat'];
try
    load(s2DictPath);
catch
    [s2Eigvector,s2Mean,s2Dictionary] = learn_s2_dictionary(c1Database,param);
    save(s2DictPath,'s2Dictionary','s2Eigvector','s2Mean');
    clear s2DictPath;
end
fprintf('done!\n');
%% build s2 layer -- s2 feature extraction by sparse coding
fprintf('build s2 layer...\n');
rt_data_dir = ['data\s2_c1_s1_pp_', dataSet];%, '_pca', num2str(param.pca)];
s2Path = [rt_data_dir,'\s2Database.mat'];
try 
    load(s2Path);
catch
    s2Database = build_s2_layer(c1Database,s2Eigvector,s2Mean,s2Dictionary,param);
    save(s2Path,'s2Database');
    clear s2Patch;
end
fprintf('done!\n');
%% learn saliency 
fprintf('learn saiency...\n');
rt_data_dir = ['data\sm_', dataSet];
smPath = [rt_data_dir,'\smDatabase.mat'];
try 
    load(smPath);
catch
    smDatabase = learn_saliency(database,s2Database);
    save(smPath,'smDatabase');
    clear smPatch;
end
fprintf('done!\n');

%% build c2 layer
fprintf('build c2 layer...\n');
c2FeaPath = ['data\c2_fea_', dataSet, '.mat'];
try 
    load(c2FeaPath);
catch
    [c2FeaTotal,c2Label] = build_c2_layer(s2Database,smDatabase,param);
    save(c2FeaPath,'c2FeaTotal','c2Label','-v7.3');

end
fprintf('done!\n');

%% classification
accuracy = linear_svm(c2FeaTotal,c2Label,param);
save accuracy accuracy;
fprintf('Mean accuracy: %f\n', mean(accuracy));
fprintf('Standard deviation: %f\n', std(accuracy));
