% demo4_SegSALSA_LORSAL_seg

% Brief description
%
% This demo illustrates the use of the SegSALSA algorithm applied to data
% classified with the  multinomial  logistic regression (MLR). The classes
% correspond to spectral linear mixtures immersed in additive Gaussian
% noise
%
%
%
% The approach is supervised; based on a training set containing  spectral
% vectors and the respective labels, learn the regressors of a multinomial
% logistic regression (MLR) using LORSAL (see appendix of [1])
% and then apply spatial regularization [4]
%
%
% Relevant papers:
%
%  [1] J. Li, J. Bioucas-Dias, and A. Plaza, "Hyperspectral image segmentation
%  using a new Bayesian approach with active learning", IEEE Transactions
%  on Geoscience and Remote Sensing, 2011. doi:10.1109/TGRS.2011.28330
%
%  [2] J. Bioucas-Dias, Filipe Condessa, and Jelena Kovacevic,
%  "Alternating direction optimization for image segmentation using hidden
%  Markov measure field models", SPIE Electronic Imaging, 2014.
%
%%
%
% Summary:
%
%      The demo implements the following steps:
%
%      1) generates an image Y of labels by sampling  an MRF
%         Multilogistic level prior
%
%      2) generates feature vectors according to the  mixture model
%
%                 X(:,i) = M A(:,i) + N
%
%         where M is a matrix holding spectral signatures of different
%         materials, A(:,i)  depends on Y(i), the class of pixel i, and
%         sets  the weights of for the mixture of signatures in pixel i,
%         and N is Gaussian additive noise
%
%
%      3) Learn the multinomial logistic regressors  (MLR) based on a
%         training set (X_t, Y_t)
%
%
%      4) Compute the posteriori probabilities p(y_i = k | X(:,i))
%
%
%      5) Segment the image of features using SegSALSA
%
%
% -----------------------------------------------------------------------
%
%
% ---Algorithm parameters
%
% lambda - the Laplace parameter controlling the degree of sparsity of the
%          regressor components.
%
% beta   -  LORSAL parameter setting the augmented Lagrange weight  and
%           algorithm convergency speed (see appendix of [1]; reasonable
%           values: beta < lambda)
%
% tau   -  SegSALSA regularization parameter
%
% Author: Jose Bioucas Dias, 2014 

close all
clear all

rand('seed',314);
randn('seed',314);

%--------------------------------------------------------------------------
% generate an image of labels
%--------------------------------------------------------------------------
% % image dimension 
% dim=256;
% n = dim^2;
% % number of classes 
% no_classes = 3;
% % generate an MLL image
% beta=1.5;
% neigh = 2; % second order neighborhood
% generate classes
% Y = mll(dim,beta,no_classes,neigh);



%%
% -------------------------------------------------------------------------
% for faster execution, load precomputed image of classes
% -------------------------------------------------------------------------

no_classes = 3;
beta=1.5;
neigh = 2;
load classes_K3_n256_b1_5_neig2 Y

% ------------------------------------------------------------------------
% no_classes = 3;
% beta=3;
% neigh = 2; 
% load classes_K3_n256_b3_neig2 Y

% ------------------------------------------------------------------------
% no_classes = 10;
% beta=1.5;
% neigh = 2;
% load classes_K10_n256_b1_5_neig2 Y

% ------------------------------------------------------------------------
no_classes = 10;
beta=3;
neigh = 2; 
load classes_K10_n256_b3_neig2 Y
% ------------------------------------------------------------------------

[nr,nc] = size(Y);
n=nr*nc;

% confirm no_classes
no_classes = length(unique(Y(:)));

figure(10); figure(gcf)
imagesc(Y)
title('image of labels')



%% 
% -----------------------------------------------------------------------
%        generate feature  vectors 
% -----------------------------------------------------------------------

% size libray elements from USGS
dim_M = no_classes;
% define the minimum angle between any pair of spectral signatures 
% thus controlling the difficulty of the classification problem
min_angle = 10;  % in degress 
% standard deviation of the additive noise (component-wise)
sigma =  0.1;

% number of training set samples per class
samples_per_class = 100;


% sparsity parameter (LORSAL)
lambda = 0.001;
% LORSAL AL parameter
theta = lambda/10;
% kernel  relative width RBFs
kern_sigma = 1;


% define the mixtures for each  class  (mixtures of Dirichelet - see 
% function spectMixGen)
%
%
% EX: statistical mixture, with weights w1 and w2 of two pure 
%     components (m1 and m2)
%
%  souce_par(:,:,1) = [w1 1 0 0 0 0 0 0 0 0 0;
%                      w2 0 1 0 0 0 0 0 0 0 0 ];
% 
% EX: statistical mixture, with weights w1 and w2 of two 
%     linear mixtures (m1-m2 with Diriclet parameters 0.5-0.7
%     and  m3-m4 with Diriclet parameters 1-1) 
%
%  souce_par(:,:,1) = [w1 0.5 0.7 0 0 0 0 0 0 0 0;
%                      w2 0 0 1 1 0 0 0 0 0 0 ];
%

souce_par = zeros(2,dim_M+1,no_classes);

% linear mixtures of components 1 2 3
souce_par(:,:,1) = [1 1 1 1 0 0 0 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 0 0 ];       

% linear mixtures of components 2 3 4        
souce_par(:,:,2) = [1 0 1 1 1 0 0 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 1 1 ];  

% linear mixtures of components 3 4 5        
souce_par(:,:,3) = [1 0 0 1 1 1 0 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 
                
% linear mixtures of components 4 5 6        
souce_par(:,:,4) = [1 0 0 0 1 1 1 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 

% linear mixtures of components 5 6 7       
souce_par(:,:,5) = [1 0 0 0 0 1 1 1 0 0 0;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 
% linear mixtures of components 6 7  8       
souce_par(:,:,6) = [1 0 0 0 0 0 1 1 1 0 0;
                    0 0 0 0 0 0 0 0 0 0 1 ];       

% linear mixtures of components 7 8 9       
souce_par(:,:,7) = [1 0 0 0 0 0 0 1 1 1 0;
                    0 0 0 0 0 0 0 0 0 1 1 ];  

% linear mixtures of components 8 9 10       
souce_par(:,:,8) = [1 0 0 0 0 0 0 0 1 1 1;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 
                
% linear mixtures of components 9 10 1       
souce_par(:,:,9) = [1 1 0 0 0 0 0 0 0 1 1;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 

% linear mixtures of components 10 1 2       
souce_par(:,:,10) = [1 1 1 0 0 0 0 0 0 0 1;
                    0 0 0 0 0 0 0 0 0 1 1 ]; 
        
   
                
               
%%
%-------------------------------------------------------------------------
% generate the HS data sets
% Each class  is a mixture of spectral signatures  
%
% build the mixing mixture 
%-------------------------------------------------------------------------
load USGS_1995_Library.mat
%  order bands by increasing wavelength
[dummy index] = sort(datalib(:,1));
A =  datalib(index,4:end);

% prune the library 
% min angle (in degres) between any two signatures 

% the larger min_angle the easier the classification    
A = prune_library(A,min_angle);  
% select p endmembers  from A


indexperm = randperm(size(A,2));
supp = indexperm(1:dim_M);
M = A(:,supp);
[L,p] = size(M);  % L = number of bands; p = number of material


%--------------------------------------------------------------------------
% generate hyperspectral data set
%--------------------------------------------------------------------------
Y = Y(:)';
X = zeros(L,n);
for i=1:no_classes
    np_class_i = sum(Y == i);
    X(:,Y == i) = spectMixGen(M,np_class_i, 'Source_pdf', 'Diri_mix',...
                        'pdf_pars', souce_par(:,:,i), 'snr', inf);
end

Pan = sum(reshape(X',nr,nc,L),3)/L;
figure(20); figure(gcf)
imagesc(Pan)
title('Panchromatic image')

%--------------------------------------------------------------------------
% additive iid Gaussian noise
%--------------------------------------------------------------------------
X = X + sigma*randn(size(X));


%--------------------------------------------------------------------------
% normalize and remove the mean the vectors - optional
%--------------------------------------------------------------------------
X = X - repmat(mean(X,2),1,n);
X = X./repmat(sqrt(sum(X.^2)), L,1);


% show false color image
aux =  mean(X(1:60,:)); a = min(aux); b = max(aux); 
imag(:,:,1)  =  reshape( (aux - a)/(b-a) ,nr,nc);

aux =  mean(X(61:120,:)); a = min(aux); b = max(aux); 
imag(:,:,2)  =  reshape( (aux - a)/(b-a) ,nr,nc);

aux =  mean(X(121:180,:)); a = min(aux); b = max(aux); 
imag(:,:,3)  =  reshape( (aux - a)/(b-a) ,nr,nc);


figure(30); figure(gcf)
imagesc(imag)
title('False colour image of the features')



%%
%--------------------------------------------------------------------------
% CLASSIFICATION
%--------------------------------------------------------------------------

% select the training set randomly
indexes_t = [];
for i=1:no_classes
    indexes_class = find(Y == i);
    if length(indexes_class) >= samples_per_class;
        aux_index = randperm(length(indexes_class));
        indexes_class = indexes_class(aux_index(1:samples_per_class));
    end
    indexes_t = [indexes_t indexes_class(:)'];
end

Y_t = Y(indexes_t);  
X_t = [X(:,indexes_t)];
n_t = length(Y_t);

figure(40); figure(gcf);
plot(X_t)
title('training vectors')

%% MLR classification

% traning set dimensions 
[~,n_t] =size(X_t);
% build distance d_ij =||x_i-x_j|| matrix 
%  
px_t = sum(X_t.^2);
[PX,PY] = meshgrid(px_t);
dist=PY + PY -2*X_t'*X_t;

scale = mean(dist(:));
% build design matrix (kernel) 
K=exp(-dist/2/scale/kern_sigma^2);
% set first line to one 
K = [ones(1,n_t); K];

% learn the MLR regressors
[w,obj] = LORSAL(K,Y_t,lambda,theta,200);
figure(50); figure(gcf);
plot(obj)
title('Lorsal objective function')


% compute class probalilities
p = mlogistic(w,K);
% MAP classification
[maxp,Y_t_hat] = max(p);
% overall accuracy
OA_treino =  sum(Y_t_hat == Y_t)/n_t


%--------------------------------------------------------------------------
% compute the overall accuracy off the complete image in blocks of size 
% size_splits to be faster

size_splits = 5000;
n_splits = ceil(n/size_splits);

% norm of x_t
px_t = sum(X_t.^2);
p = [];
for i=1:n_splits  
    p_left = size_splits*(i-1)+1;
    if i ~= n_splits
        p_right = size_splits*i;
    else
        p_right = n;
    end
    n_aux = p_right - p_left +1;
    Xaux = X(:,p_left:p_right);   
    
    % norm of x
    px = sum(Xaux.^2);
    [PX,PY] = meshgrid(px,px_t);
    dist = PX + PY-2*X_t'*Xaux;
    
    clear PX PY
    
    % build kernel
    K=exp(-dist/2/scale/kern_sigma^2);
    
    %
    % set first line to one
    K = [ones(1,n_aux); K];
    p= [p mlogistic(w,K)];
    
    clear K
    
end

% 
% % build kernel K(x,z)
% % norm of x and z
% px = sum(X.^2);
% px_t = sum(X_t.^2);
% [PX,PY] = meshgrid(px,px_t);
% dist = PX + PY-2*X_t'*X;
% 
% % build kernel
% K=exp(-dist/2/scale/kern_sigma^2);
% 
% % 
% % set first line to one 
% K = [ones(1,n); K];
% p=mlogistic(w,K);


[maxp,Y_hat_MLR] = max(p);
OA_teste_MLR =  sum(Y_hat_MLR==Y)/n

figure(60); figure(gcf)
imagesc(reshape(Y_hat_MLR,nr,nc));
st = sprintf('LORSAL Classification: OA = %2.2f', OA_teste_MLR)
title(st)

%% 
% --------- segmentation with SegSALSA -------------------------------

tau = 2.5;
mu = 3;
iters = 100;
% segment the image
figure(1000)
P_seg= imseg_hidd_vf(reshape(p',nr,nc,no_classes), tau,mu, iters);
title('Latent probabilities for class 1 ')


[maxp,Y_hat_seg] = max(P_seg);
OA_teste_seg =  sum(Y_hat_seg==Y)/n


figure(70); figure(gcf)
imagesc(reshape(Y_hat_seg,nr,nc));
st = sprintf('SegSALSA segmantattion: OA = %2.2f', OA_teste_seg)
title(st)




