%%     demo for MLRsub algorithm
%--------------Brief description-------------------------------------------
%
% 
% This demo implements the  spectral-spatial hyperspectral image segmentation 
% using subspace multinomial logistic regression and Markov random fields
% descibed in [1].
%
%   
% Summary of the segmentation algorithm:
%
%    1- Learn the logistic regressors (see Algorithm 1 in [1])
%
%    2- Based on a multi-level logistic (MLL) prior, to model the
%       contextual spatial information of the hyperspectral images, 
%       compute the MAP segmentation  via the \alpha-Expanasion 
%       Graph-cut algorithm.
%
% More details in:
%
% [1] J. Li, J. Bioucas-Dias and A. Plaza. Spectral-Spatial Hyperspectral
% Image Segmentation Using Subspace Multinomial Logistic Regression and
% Markov Random Fields. IEEE Transactions on Geoscience and Remote Sensing,
% in press, 2011
%
%-------------------------- Input parameters-----------------------------
%
%-------------------------Algorithm parameters----------------------------
%
% beta   -  quadratic prior regularization parameter to cope with bad 
%           or ill conditioning 
%           of the underlying inverse problem: beta = [1e-15 1e-7]
%
% tau    -  threshold parameter controlling the amount of spectral information 
% no_r   -  alternative parameter of tau controlling the amount of spectral
%           information
%
% mu      - the spatial prior regularization  parameter, controlling 
%           the degree of spatial smoothness. Tune this parameter to 
%           obtain a good segmentation results.
%
%--------------------hyperpsectral data parameters------------------------
%
% no_classes        - number of classes 
% no_bands          - dimentionality of the hyperspectral data
% no_rows           - number of rows in the image
% no_lines  - number of lines in the image
% no_samples        - number of pixels in the image: n = no_rows*no_lines
%
% ---------------------------- demo  parameters ---------------------------
%
% GC_toolbox        - this variable takes values in  {Bagon, Bioucas}
%                     default = Bagon.
%                     Graph cut toolbox to implement the alpha-Expansion
%                     algorithm. Both toolboxes are compiled for 
%                     Matlab R2009.
%
%                     If you face problems, compile the Bagon toolbox in 
%                     you computer by running compile_gc.m. See details in
%                     GCmex1.2
%
%                       
% no_train          - size of labeled samples
%
% MCiter            - number of Monte Carlo runs to compute the overall
%                     accuracy
%
%
%--------------------------- output parameters ---------------------------
%
% class_results      - classification accuracies with the MLRsub algorithm 
%                      (Algorithm 1) in [2]
% seg_results        - segmentation accuracies with the MLRsubMLL algorithm
%                       (Algorithm 2) in [1]
%
%  
%
%% Author: Jun Li
%         Jose M. Bioucas-Dias
%         Antonio Plaza,  Dec. 2011


close all
clear all
clc

beta             = 1e-8;
no_r             = 190;
mu               = 1.0;
no_classes       = 16;
no_train         = 1036;

% load the ground truth and the hyperspectral image
load AVIRIS_Indiana_16class
trainall = trainall'; 
load imgreal;

% size of image
[no_lines, no_rows, no_bands] = size(img);

img = ToVector(img);
img = img';
% remove the noise
img([104:108 150:163 220],:) =[];

no_samples = no_lines*no_rows;

%  randomly select training samples from the ground truth image
 indexes=train_test_random_new(trainall(2,:),fix(no_train/no_classes),no_train);
 
 train_samples = trainall(:,indexes);
 train_vectors = img(:,train_samples(1,:));
 
 % the remaining used for test
 test_samples = trainall;
 test_samples(:,indexes) = [];  

 

%%  subspace classifier

% first step: build the class independent subspace
%  class independent subspace for the training set
 sub_training = [];
%  class independent subspace for the whole image
 sub_image= [];
  
 for k_iter = 1:no_classes
     index_k = train_samples(2,:) == k_iter;          
     train_k = train_vectors(:,index_k);     
     n_k   = size(train_k,1);
     [v,d]   = eig(train_k*train_k'/n_k);

     P = v(:,1:no_r)*v(:,1:no_r)'*train_vectors;
    
     sub_train = zeros(1,no_train);
     for num_k = 1:no_train
         sub_train(num_k) = sqrt(P(:,num_k)'*P(:,num_k));
     end
     sub_training = [sub_training; sub_train];       
     
     P_all = v(:,1:no_r)*v(:,1:no_r)'*img;
    
     sub_all=zeros(1,no_samples);
     for num_all = 1:no_samples
         sub_all(num_all) = sqrt(P_all(:,num_all)'*P_all(:,num_all));
     end
     sub_image = [sub_image;sub_all];
      
 end
        
        
 % second step: learn the MLRsub logistic regressors
 w = subspace_classifier(-sub_training,train_samples(2,:),beta);
 
 % compute the probablity
 p= subspace_mlogistic(w,-sub_image);
 
 [maxp,class_results.map] = max(p);
 [class_results.OA,class_results.kappa,class_results.AA,class_results.CA] = calcError( test_samples(2,:)-1, class_results.map(test_samples(1,:))-1, 1: no_classes);

 
 %% segmentation

 Dc = reshape((log(p+eps))',[no_lines, no_rows, no_classes]);
 Sc = ones(no_classes) - eye(no_classes);        

 % Expansion Algorithm
 gch = GraphCut('open', -Dc, mu*Sc);
 [gch seg_results.map] = GraphCut('expand',gch);
 gch = GraphCut('close', gch);

 [seg_results.OA,seg_results.kappa,seg_results.AA,seg_results.CA] = calcError( test_samples(2,:)-1, seg_results.map(test_samples(1,:)), 1: no_classes);
 
% display the output results
fprintf(' \n');
fprintf('Classification results:   \n');
disp(class_results);
fprintf(' \n');
fprintf('Segmentation results:   \n');
disp(seg_results);

% plot the classification and segmentation maps
mycolormap = [  1    1    1;
                1    1  0.5;
                0    0    1;
                1    0    0;
                0    1 0.25;
                1    0    1;
              0.5    0    1;
                0  0.5    1;
                0    1    0;
              0.5  0.5 0.25;
              0.5    0  0.5;
                0  0.5    1;
                0 0.25  0.5;
                0  0.5 0.25;
              0.5 0.25    0;
                0    1  0.5;
                1    1    0;   ];
            
            

figure
groundtruth = zeros(no_lines,no_rows);
groundtruth(trainall(1,:)) = trainall(2,:);
imagesc(groundtruth')
colormap(mycolormap)
title('Ground Truth Image')

figure
imagesc(reshape(class_results.map,no_lines,no_rows)')
colormap(mycolormap(2:end,:))
title(['Classification OA:  ',num2str(class_results.OA)])

figure
imagesc((seg_results.map+1)')
colormap(mycolormap(2:end,:))
title(['Segmentation OA:  ',num2str(seg_results.OA)])
