%**************************************************************************
% Demo file for the proposed color demosaicking algorithm JM-CDM
% Author: Kan Chang, Guangxi University
% Date: 4/5/2015
%**************************************************************************

clc;
clear all;

addpath(genpath('./Test_image')); 
addpath(genpath('./Initial denosing and demosaicking algorithms'));
addpath(genpath('./JM_CDM'));              

%% Read the test image
Images{1} = 'img1';    
Images{2} = 'img2';
Images{3} = 'img3';
Images{4} = 'img4';
Images{5} = 'img5';
Images{6} = 'img6';
Images{7} = 'img7';
Images{8} = 'img8';
Images{9} = 'img9';
Images{10} = 'img10';
Images{11} = 'img11';
Images{12} = 'img12';

Images{13} = '1';    
Images{14} = '2';
Images{15} = '3';
Images{16} = '4';
Images{17} = '5';
Images{18} = '6';
Images{19} = '7';
Images{20} = '8';
Images{21} = '9'; 
Images{22} = '10'; 
Images{23} = '11';
Images{24} = '12';
Images{25} = '13'; 
Images{26} = '14';
Images{27} = '15'; 
Images{28} = '16';
Images{29} = '17';
Images{30} = '18'; 
 
test_set = [29,4];
sigma_set = [0,5,10];
    
num_Kodak_images = length(find(test_set< 13));
num_IMAX_images = length(find(test_set> 12));
num_sigma = length(sigma_set);

Kodak_SNR = zeros(num_sigma,4,num_Kodak_images);
IMAX_SNR = zeros(num_sigma,4,num_IMAX_images); 

Kodak_Time = zeros(num_sigma,num_Kodak_images);
IMAX_Time = zeros(num_sigma,num_IMAX_images); 

Kodak_num = 0;
IMAX_num = 0;

ave_result_file_name = sprintf('Ave_result.txt');
fp_ave = fopen(ave_result_file_name,'a'); 

%% Loop over different test images
for img_num = 1:length(test_set)   
    
    if test_set(img_num) < 13
        Last_Name = 'bmp';
        Dataset = 'Kodak';
    else
        Last_Name = 'tif'; 
        Dataset = 'IMAX';
    end
    
    fprintf(1,'\n*********Current Test Image:%s in %s*********\n', Images{test_set(img_num)}, Dataset);
  
    Img_Name = Images{test_set(img_num)};  
    imgfile = sprintf('%s.%s', Img_Name, Last_Name);

    Ori = imread(imgfile);
    Ori = double(Ori);
    [n,m,c] = size(Ori);   
    
    %% Loop over noise std, 0 means noiseless
    for noise_num = 1:length(sigma_set);       
    
        sigma = sigma_set(noise_num);
        randn('seed', 0);
        Noise_Ori = Ori + sigma*randn(n,m,c);   
    
        result_file_name = sprintf('%s_sigma%d.txt', Img_Name, sigma);
        fp=fopen(result_file_name,'a'); 

        %% Generate CFA data and bayer filter 
        [CFA_image, Bayer_Mask] = mosaic_bayer(Noise_Ori, 'grbg');
        
        if sigma == 0  % Noiseless case
            tic;
            %% CDM by MLRI method (Kiku, EI 2014) for initial demosaiced value
            Rec = demosaick(CFA_image, 'grbg', 1.4, Bayer_Mask);
            MLRI_Time = toc;
            [snr_MLRI,cpsnr_MLRI] = csnr(Rec(21:n-20,21:m-20,:),Ori(21:n-20,21:m-20,:));
            fprintf(1, 'PSNR for MLRI, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\n', snr_MLRI(1), snr_MLRI(2), snr_MLRI(3), cpsnr_MLRI);    
            fprintf(fp, 'PSNR for MLRI, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\r\n', snr_MLRI(1), snr_MLRI(2), snr_MLRI(3), cpsnr_MLRI); 
        else  % Noisy case
            tic;
            A = zeros(n,m);
            A(1:2:n,1:2:m)= Noise_Ori(1:2:n,1:2:m,2);    % G at up left
            A(2:2:n,2:2:m)= Noise_Ori(2:2:n,2:2:m,2);    % G at down right
            A(1:2:n,2:2:m)= Noise_Ori(1:2:n,2:2:m,1);    % R
            A(2:2:n,1:2:m)= Noise_Ori(2:2:n,1:2:m,3);    % B         
        
          %% Denoising the CFA image via PCA based algorithm (Lei Zhang, TIP 2009)
            A = PCA_CFA_Denosing(A, sigma, sigma, sigma);
            CFA_image(1:2:n,1:2:m,2) = A(1:2:n,1:2:m);
            CFA_image(2:2:n,2:2:m,2) = A(2:2:n,2:2:m);
            CFA_image(1:2:n,2:2:m,1) = A(1:2:n,2:2:m);
            CFA_image(2:2:n,1:2:m,3) = A(2:2:n,1:2:m);
        
          %% CDM by MLRI method (Kiku, EI 2014) for initial demosaicked value
            Rec = demosaick(CFA_image, 'grbg', 1.4, Bayer_Mask);
            PCA_MLRI_Time = toc;
            [snr_PCA_MLRI, csnr_PCA_MLRI] = csnr(Rec(21:n-20,21:m-20,:),Ori(21:n-20,21:m-20,:));
            fprintf(1, 'PSNR for PCA_MLRI, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\n', snr_PCA_MLRI(1), snr_PCA_MLRI(2), snr_PCA_MLRI(3), csnr_PCA_MLRI);  
            fprintf(fp, 'PSNR for PCA_MLRI, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\r\n', snr_PCA_MLRI(1), snr_PCA_MLRI(2), snr_PCA_MLRI(3), csnr_PCA_MLRI);
        end 
        
        % Measurements for the proposed JM-CDM method
        Y = zeros(size(Noise_Ori));      
        Y(:,:,1) = Noise_Ori(:,:,1) .* Bayer_Mask(:,:,1);
        Y(:,:,2) = Noise_Ori(:,:,2) .* Bayer_Mask(:,:,2);
        Y(:,:,3) = Noise_Ori(:,:,3) .* Bayer_Mask(:,:,3);     
    
       %% Parameters setting for the proposed JM-CDM
        Par = [];
        Par.Group_size = 40;
        Par.Patch_size = 6;
        Par.Window_size = 20;
        Par.Sliding_step = 4;
                  
        Par.sigma = sigma;
        Par.RL_iter = 20; 
        Par.RL_tol = 1e-4;
        Par.SB_iter = 10; 
        Par.SB_tol = 1e-4; 
        if sigma == 0
            Par.Grouping_iter = 1;
        else
            Par.Grouping_iter = 2;
        end
        
        % Note that alpha and beta are automatically set inside the algorithm.
        % We don't recommand changing the value of mu and lambda here.
        Par.mu = 1/10.0;                              
        Par.lambda = 50;
        
        Rec_Result = zeros(size(Rec));
        Partition_y_size = n/2;
        Partition_x_size = m/2;
        num_Par = 0;
    
        tic;  
        % Note: To reduce artifacts, the best way is to set overlaps for adjacent image blocks. However, for simplicity, overlaps were not considered in this experiment. 
        for k = 1:Partition_y_size:n
            for l = 1:Partition_x_size:m
                num_Par = num_Par + 1;
                fprintf(1,'current processing Partition %d...\n', num_Par);
                Par.Mask = Bayer_Mask(k:k+Partition_y_size-1, l:l+Partition_x_size-1,:);
                [Rec_Result(k:k+Partition_y_size-1, l:l+Partition_x_size-1,:),Par] = JM_CDM(Y(k:k+Partition_y_size-1, l:l+Partition_x_size-1,:), ...
                                            Rec(k:k+Partition_y_size-1, l:l+Partition_x_size-1,:), Par);   
            end
        end
        JM_CDM_time = toc;
        [snr_JM_CDM,csnr_JM_CDM] = csnr(Rec_Result(21:n-20,21:m-20,:),Ori(21:n-20,21:m-20,:));
        fprintf(1, 'PSNR for Proposed, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\n', snr_JM_CDM(1), snr_JM_CDM(2), snr_JM_CDM(3), csnr_JM_CDM); 
        fprintf(fp, 'PSNR for Proposed, R: %.3fdB, G: %.3fdB, B: %.3fdB, CPSNR: %.3fdB\r\n', snr_JM_CDM(1), snr_JM_CDM(2), snr_JM_CDM(3), csnr_JM_CDM); 
 
        fprintf(1, 'Time: %.3f \n', JM_CDM_time);
        fprintf(fp, 'Time: %.3f\r\n', JM_CDM_time);                
        fclose(fp);    
        
        if test_set(img_num) < 13
            if noise_num == 1
                Kodak_num = Kodak_num + 1;
            end
            Kodak_SNR(noise_num,1,Kodak_num) = snr_JM_CDM(1);
            Kodak_SNR(noise_num,2,Kodak_num) = snr_JM_CDM(2);
            Kodak_SNR(noise_num,3,Kodak_num) = snr_JM_CDM(3);
            Kodak_SNR(noise_num,4,Kodak_num) = csnr_JM_CDM;
            Kodak_Time(noise_num,Kodak_num) = JM_CDM_time;
        else
            if noise_num == 1
                IMAX_num = IMAX_num + 1;
            end
            IMAX_SNR(noise_num,1,IMAX_num) = snr_JM_CDM(1);
            IMAX_SNR(noise_num,2,IMAX_num) = snr_JM_CDM(2);
            IMAX_SNR(noise_num,3,IMAX_num) = snr_JM_CDM(3);
            IMAX_SNR(noise_num,4,IMAX_num) = csnr_JM_CDM;
            IMAX_Time(noise_num,IMAX_num) = JM_CDM_time;
        end
   
    end  % End for noise_num = 1:length(sigma_set);
          
end


for sig_num = 1: length(sigma_set)
    fprintf(fp_ave, 'Kodak Average PSNR and CPSNR, sigma = %d: %.3f, %.3f, %.3f, %.3f \r\n', sigma_set(sig_num), mean(Kodak_SNR(sig_num,1,:)),...
        mean(Kodak_SNR(sig_num,2,:)), mean(Kodak_SNR(sig_num,3,:)), mean(Kodak_SNR(sig_num,4,:)));
    fprintf(fp_ave, 'Kodak Average Time, sigma = %d: %.2f\r\n', sigma_set(sig_num), mean(Kodak_Time(sig_num,:)));
end

for sig_num = 1: length(sigma_set)
    fprintf(fp_ave, 'IMAX Average PSNR and CPSNR, sigma = %d: %.3f, %.3f, %.3f, %.3f \r\n', sigma_set(sig_num), mean(IMAX_SNR(sig_num,1,:)),...
        mean(IMAX_SNR(sig_num,2,:)), mean(IMAX_SNR(sig_num,3,:)), mean(IMAX_SNR(sig_num,4,:)));
    fprintf(fp_ave, 'IMAX Average Time, sigma = %d: %.2f\r\n', sigma_set(sig_num), mean(IMAX_Time(sig_num,:)));
end

fclose(fp_ave);
 
