%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%% Demo code %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%% Spectral Total-Variation Local Scale Signatures %%%%%%%%
%%%%%%%% for Image Manipulation and Fusion %%%%%%%%%%%%%%%%%%%%%%
%%%%%%%% Ester Hait and Guy Gilboa %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%% EE department, Technion - Israel Institue of Technology%

% If you use this code, please cite our paper:
% Hait, Ester, and Guy Gilboa. "Spectral Total-Variation 
% Local Scale Signatures for Image Manipulation and Fusion." 
% IEEE Transactions on Image Processing 28.2 (2019): 880 - 895.
% https://ieeexplore.ieee.org/document/8476587

% Saliency extaction and image fusion 
% of a thermal image with its corresponding optical image %

clc,clear,close all

%%%%%%%%%%%%% Input image %%%%%%%%%%%%%%%%%%%%%%%%5
thermal_img=rgb2gray(imread('img_00011.bmp'));
thermal_img=double(thermal_img)/255;
figure;imshow(thermal_img);title('Thermal input image');

%%%%%%%%%%%%% Specral TV decomposition, time=scale %%%%%%%%%%%%%
% output = comprehensive scale and location information, Phi
addpath('specTV')
dt=0.06; % time step
dtf=100; % number of iterations
Max_time=dt*dtf; % total time
tic;[S,T,Phi,f_r,u2] = specTV_evolve_sym(thermal_img, Max_time, dt);toc;
figure, imshow(f_r,[]); title('Residual image')
figure; plot((1:dtf)*dt, S,'LineWidth',4); xlabel('\bf t'); title('\bf Global spectrum','FontSize',11)

%%%%%%%%%%%%% K-means clustering %%%%%%%%%%%%%
% load and display the clustering result
K=20; % # of clusters
load('lawn_06_no_den_20_1.mat','clust')
figure;imagesc(clust);colormap('jet');axis off; title('Clustering result');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% % perform clustering yourself; due to its randomized initialization,
% % K-means clustering may produce different results.

% [M,N,z]=size(Phi);
% X=reshape(Phi,M*N,z);
% 
% % Preprocessing: signature enhancement
% p=1;
% for i=1:M*N
%     n(i) = norm(X(i,:),p);
% end
% % L-1 norm map
% norm_map=reshape(n,M,N);
% figure;imagesc(norm_map);colormap('jet');title('L-1 norm map')
% 
% K=20; % # of clusters
% [idx,C] = kmeans(X.*n',K,'replicates',10);
% clust=reshape(idx,M,N);
% figure;imagesc(clust);colormap('jet');axis off; title('Clustering result');

%%%%%%%%%%%%% Saliency extraction %%%%%%%%%%%%%
% Prepare for postprocessing: image matting (Levin et al., 2008)
% http://webee.technion.ac.il/people/anat.levin/
addpath('matting')
levels_num=3;
active_levels_num=1;
thr_alpha=[];
epsilon=[];
win_size=[];

for i=1:K
    vec_mean_clust(i)=mean(thermal_img(find(clust==i)));
end
[~,ind_clust]=sort(vec_mean_clust,'descend');
K1=14; K2=6;

% Foreground map = K1 clusters of strongest mean
fg_ind=ind_clust(1:K1);
fg_map=zeros(size(thermal_img));
for i=fg_ind
    fg_map=fg_map+(clust==i);
end
figure;imshow(fg_map); title('Foreground map for reduced saliency map');

% Background map = K2 clusters of weakest mean
bg_ind=ind_clust(end-K2+1:end);
bg_map=zeros(size(thermal_img));
for i=bg_ind
    bg_map=bg_map+(clust==i);
end
figure;imshow(bg_map); title('Background map for reduced saliency map');

% Option 1: reduced saliency map
% Image matting 
alpha_thermal=solveAlphaC2F(thermal_img,or(fg_map,bg_map),fg_map,levels_num,active_levels_num,thr_alpha,epsilon,win_size);
drawnow;
[F_thermal,B_thermal]=solveFB(thermal_img,alpha_thermal);
saliency_thermal_reduced=F_thermal.*alpha_thermal;
figure;imshow(saliency_thermal_reduced); title('Reduced thermal saliency map');

% Option 2: detailed saliency map by incorporating residual
fg_map=or(fg_map,f_r>=0.5);
figure;imshow(fg_map); title('Foreground map for detailed saliency map');
bg_map=and(bg_map,f_r<0.5);
figure;imshow(bg_map); title('Background map for detailed saliency map');

% Image matting
alpha_thermal=solveAlphaC2F(thermal_img,or(fg_map,bg_map),fg_map,levels_num,active_levels_num,thr_alpha,epsilon,win_size);
drawnow;
[F_thermal,B_thermal]=solveFB(thermal_img,alpha_thermal);
saliency_thermal_detailed=F_thermal.*alpha_thermal;
figure;imshow(saliency_thermal_detailed); title('Detailed thermal saliency map');

%%%%%%%%%%%%%% Image fusion visualization %%%%%%%%%%%%%
%%%%%%%%%%%%%% Temprature Gradient Coloring %%%%%%%%%%%

% Read optical input image
opt_img=imread('img_00010.bmp');
opt_img=double(opt_img)/255;
figure; imshow(opt_img);title('Corresponding optical input image');

% Option 1: reduced saliency map
% Choose T_ = enhancement function, e.g.:
% T_=(1+saliency_thermal_reduced); % weaker
T_=exp(saliency_thermal_reduced); % stronger

opt_img_R=opt_img(:,:,1).*T_;
opt_img_G=opt_img(:,:,2);
opt_img_B=opt_img(:,:,3);
max_c=max(max(max(opt_img_R,opt_img_G),opt_img_B),1);
fusion_reduced=cat(3,opt_img_R./max_c,opt_img_G./max_c,opt_img_B./max_c);
figure;imshow(fusion_reduced);title('Image fusion result - reduced version');

% Option 2: detailed saliency map
% Choose T_ = enhancement function, e.g.:
% T_=(1+saliency_thermal_detailed); % weaker
T_=exp(saliency_thermal_detailed); % stronger

opt_img_R=opt_img(:,:,1).*T_;
opt_img_G=opt_img(:,:,2);
opt_img_B=opt_img(:,:,3);
max_c=max(max(max(opt_img_R,opt_img_G),opt_img_B),1);
fusion_detailed=cat(3,opt_img_R./max_c,opt_img_G./max_c,opt_img_B./max_c);
figure;imshow(fusion_detailed);title('Image fusion result - detailed version');
