%Multiview Image Compression method
%EUSIPCO 2012
%Imperial College London 
%Andriy Gelman, Jon Onativia and Pier Luigi Dragotti 

%Description - Function implements compression of multiview images 

%data_set = dataset name
%bpp = target bit rate

function [psnr, output_rate, output] = run_method(data_set, bpp)
clear global all;

global information_data index_distortion rate distortion passes...
    layer_header size_of_largest_block number_of_blocks...
    distortion_scaled;

addpath('../Data');

initialize_folders;

%switch off warning
warning('off', 'MATLAB:inpolygon:ModelingWorldLower');

size_of_largest_block = [64 64];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%Compression%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

[compressed_output, layer_header, r,c,k, r_modified, c_modified, output_location, total_segmentation_cost, disparity] = compress_data(data_set, bpp);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%Rate Allocation%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

[decoded_data, block_rate] = rate_allocation(compressed_output, bpp, layer_header,  r, c, k, total_segmentation_cost, rate, passes, distortion, distortion_scaled, information_data, index_distortion);
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%% Decoding %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
[output] = decode_data(decoded_data,layer_header, r, r_modified, c, c_modified, k, number_of_blocks, disparity, output_location);
    
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%Evaluate PSNR and output bit rate%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%evaluate the largest bit plane in each block
largest_bit_plane = zeros(1, length(decoded_data));
for l = 1:length(information_data)
    buffer = information_data{l};
    largest_bit_plane(l) = buffer(9);    
end
    
%evaluate the additional overhead    
[sig_overhead, block_rate_overhead, bit_plane_overhead] = calculate_overhead(block_rate, largest_bit_plane);

%load original data
load(data_set);
input=  in;
input = squeeze(input);
output = double(output);
mean_error = (sum(sum(sum(((output - input).^2)))))/(r*c*k);
total_error = (sum(sum(sum(((output - input).^2)))));
psnr = 10*log10((255^2)/mean_error);
psnr

number_of_layers = length(layer_header(:,1))-1;
%evaluate total rate
output_rate = sum(block_rate)*8 + total_segmentation_cost  + sig_overhead + block_rate_overhead + bit_plane_overhead + number_of_layers*(16);
output_rate = output_rate/(r*c*k);

end


function output = extract_layer(input, label, layer)

if isempty(input)
    load label.mat
    input = label;
end
%label(find(label == 0))= number_of_layers; %Assuming final layer is the background

%Extract Signal
label_layer = label; clear label;
label_layer(label_layer ~= layer) = 0;
label_layer =  label_layer/layer;
output = input.*label_layer;

end 

function [compressed_output, layer_header, r,c,k, r_modified, c_modified, output_location, total_segmentation_cost, disparity] = compress_data(data_set, bpp)

global MSE information_data passes index_distortion rate distortion size_of_largest_block number_of_blocks...
    distortion_scaled layer_block_location temporal_subband_location;


stack_phi_location = 1;

path = data_set;
load(path);

input = in;
clear in;

input = squeeze(input);
label = squeeze(label);

[r,c,k] = size(input);
save input_data.mat input;

%Zero pad input and disparity
original_label = label;

%load lambda layer stack and D
[label, D] = remove_unused_layer(label, D);

%allocate unused pixels to closest layers
for l = 1:length(label(1,1,:))
    label(:,:,l) = allocate_background(label(:,:,l));
end

%generate layer stack 
layer_stack = generate_layer_stack(label, D, 100, 'ceil', stack_phi_location);

max_disparity = ceil(max(D)) + 1;

disparity = ceil(([0:1: k-1] - (stack_phi_location - 1))'*D');

%if strcmp(string2, 'ceil')
%    disparity = ceil(([0:1: k-1] - (stack_phi_location - 1))'*D');
%else
%    disparity = round(([0:1: k-1] - (stack_phi_location - 1))'*D');
%end

save disparity_data_set.mat disparity;
save D.mat D


for l = 1:length(layer_stack(1,1,:))
    individual_labels{l} = layer_stack(:,:,l);
end



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%Encode the layer contours%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%set the segmentation target rate to be 10% of the total bit budget
segmentation_target_rate = bpp*r*c*k*0.1;
%segmentation_target_rate = 1e10;

w2 = cd;
cd('lossless_segmentation_encoding');
%do not encode the background
[clear_var, total_segmentation_cost] = lossless_encode_segmentation(individual_labels(1:end-1));
cd(w2);


if segmentation_target_rate < 0.85*total_segmentation_cost
    
    layer_stack = padarray(layer_stack, [1,1], 'both');
    cd('lossy_segmentation_coding\lossy_segmentation');
        [layer_stack, total_segmentation_cost] = lossy_progressive_segmentation_encoding(layer_stack, segmentation_target_rate);
    cd(w2);
    layer_stack = layer_stack(2:end-1, 2:end-1, :);
    
else
    layer_stack = layer_stack;
end
 

label = generate_3D_label(layer_stack, [r,c], D, k, 'ceil', stack_phi_location);
%cd('Level_set_label_computation');
[label, D] = remove_unused_layer(label, D);
%cd('..');


max_disparity = ceil(max(D)) + 1;

disparity = ceil(([0:1: k-1] - (stack_phi_location - 1))'*D');

%if strcmp(string2, 'ceil')
%    disparity = ceil(([0:1: k-1] - (stack_phi_location - 1))'*D');
%else
%    disparity = round(([0:1: k-1] - (stack_phi_location - 1))'*D');
%end

save disparity_data_set.mat disparity;
save D.mat D

for l = 1:length(label(1,1,:))
    label(:,:,l) = allocate_background(label(:,:,l));
end


input = uint8(input);
label = uint8(label);

input = padarray(input, [0, floor(k)*(max_disparity), 0], 'both');   
label = padarray(label, [0, floor(k)*(max_disparity), 0], 'both');

[r_modified, c_modified] = size(input(:,:,1));

label= round(label);
save label.mat label;

number_of_layers = double(max(max(max(label))));

layer_block_location = zeros(1, number_of_layers+1);
temporal_subband_location = zeros(1, 1 + number_of_layers*k);
ind_loc = 1;

[r_temp, c_temp, number_of_images] = size(input);

combined_output = zeros(r_temp, c_temp, number_of_images);
combined_output_label = zeros(r_temp, c_temp, number_of_images);

for i = 1:number_of_layers                                                                                                                                                    
    
        
    original_label = extract_layer(label, label, i)/i;
    
    out = extract_layer(input,label,i);  
  
    load disparity_data_set.mat; 
    
    disparity_vec = disparity(:,i);
    save disparity_vec disparity_vec;
      
    input_label = original_label;
    
    input_label = double(input_label);
    string = ['operation/extracted_layer_label/original_label', num2str(i)];
    save(string, 'input_label')
     
    out = double(out);
   
    cd('Inter view lifting DWT');
        [temporal_out, dwt_label, intra_pixels, uncovered_pixels, motion_Vect, output_location] = lifting_haar_transform(double(out), double(input_label), D(i), 'ceil', []);
    cd('..');
     
    
    temporal_out = change_view_location(temporal_out, disparity(:,i), output_location, [1:number_of_images]);
    temporal_label = dwt_label;
    temporal_label = change_view_location(temporal_label, disparity(:,i), output_location, [1:number_of_images]);
    
    %save the data
    w = cd;
    cd('original_label');
    string = ['label', num2str(i)];
    save(string, 'input_label'); 
    string = ['lifting_params', num2str(i)];
    save(string, 'dwt_label', 'dwt_label', 'intra_pixels', 'uncovered_pixels', 'motion_Vect');
    cd(w);
    
    number_of_recombined_images = number_of_images;
    recombined_label = zeros(r_temp, c_temp, number_of_images);
    
    recombined_label(temporal_label(:,:,1:number_of_recombined_images) == 1) = 1;
    recombined_label(combined_output_label > 0) = 0;
    
    combined_output_label = combined_output_label + recombined_label*i;
    
    recombined_label_ind = find(recombined_label == 1);
    
    recombined = zeros(r_temp, c_temp, number_of_images);
   
    recombined(recombined_label_ind) = temporal_out(recombined_label_ind);
    %low_pass_main_label(low_pass_main_ind) = 1;
    
    residual = temporal_out;
    residual_label = temporal_label;
    
    residual(recombined_label_ind) = 0;
    residual_label(recombined_label_ind) = 0; 
    
    combined_output = combined_output + recombined;
    
    w = cd;
    cd('operation/recombined');
    string = ['recombined', num2str(i)];
    save(string, 'recombined');
    
    string_label = ['recombined_label', num2str(i)];
    save(string_label, 'recombined_label');
    
    clear recombined;
    clear recombined_label;
    clear recombined_label_ind;
    cd(w);
    
   
    w = cd;
    cd('operation/residual');
    string = ['residual', num2str(i)];
    save(string, 'residual');
    
    string_label = ['residual_label', num2str(i)];
    save(string_label, 'residual_label');
    clear residual;
    clear residual_label;
    cd(w);
    
end

layer_header =[];
maximum_block_size = size_of_largest_block;

spatial_decompositions_high = 0;

[layer_header, min_row, max_row, min_col, max_col] =  compute_header_file(layer_header, combined_output, 6, spatial_decompositions_high, number_of_layers+1);
combined_output = combined_output(min_row:max_row, min_col:max_col,:);
combined_output_label = combined_output_label(min_row:max_row, min_col:max_col,:);

spatial_decompositions_low = layer_header(end,5);

block_data = [];
block_label_data = [];
wavelet_state = [];

%Apply DWT on main data and separate into blocks
number_of_blocks = zeros(number_of_layers + 1, k);

string_dat = 'operation/low_pass_main/combined_output.mat';
string_label = 'operation/low_pass_main/combined_output_label.mat';

save(string_dat, 'combined_output');
save(string_label, 'combined_output_label');

ind_loc = 1;

for i = 1:number_of_images
   
    buffer_combined_label = combined_output_label(:,:,i);
    buffer_combined_label(buffer_combined_label >0)=1;
    [buffer_output, buffer_label, wavelet_state_buffer] = spatial_shape_adaptive_DWT(combined_output(:,:,i), 255*buffer_combined_label, spatial_decompositions_low, maximum_block_size);    
    %[buffer_output, buffer_label] = shape_adapative_residual_coding_combined_output(combined_output, combined_output_label, spatial_decompositions_low, maximum_block_size);

    [block_data, block_label_data, wavelet_state, number_of_blocks]=concantenate_blocks_of_main_data(block_data, block_label_data, buffer_output, buffer_label, wavelet_state, wavelet_state_buffer, number_of_blocks, (number_of_layers+1), i, k);       

    layer_block_location(ind_loc) = length(block_data);
    temporal_subband_location(ind_loc) = length(block_data);
    ind_loc = ind_loc+1;
   
end

for i = 1:number_of_layers
    
    w = cd;
    cd('operation/residual');
    
    spatial_decompositions_high = 0;
    
    string = ['residual', num2str(i)];
    string_label = ['residual_label', num2str(i)];
    
    load(string); 
    load(string_label);
    
    [layer_header, min_row, max_row, min_col, max_col, spatial_decompositions_low] =  compute_header_file(layer_header, residual_label, [], spatial_decompositions_high, i);
    
    layer_header(i, end) = layer_header(i, end-2);
    
    if ~isempty(min_row)
        residual = residual(min_row:max_row, min_col:max_col,:);
        residual_label = residual_label(min_row:max_row, min_col:max_col,:);
    else        
        residual = [];
        residual_label = [];
    end     
    cd(w);  
    

    for l = 1:k       
        if ~isempty(residual)
            
            if (l == 1)
                spatial_decompositions = layer_header(i,end-2);
                [buffer_output, buffer_label, wavelet_state_buffer] = spatial_shape_adaptive_DWT(residual(:,:,l), 255*residual_label(:,:,l), spatial_decompositions, maximum_block_size);    
                [block_data, block_label_data, wavelet_state, number_of_blocks] = concantenate_blocks_of_main_data(block_data, block_label_data, buffer_output, buffer_label, wavelet_state, wavelet_state_buffer, number_of_blocks, i, l, k);
                temporal_subband_location(ind_loc) = length(block_data);
                ind_loc = ind_loc+1;
            else
                spatial_decompositions = layer_header(i, end-1);
                [buffer_output, buffer_label, wavelet_state_buffer] = spatial_shape_adaptive_DWT(residual(:,:,l), 255*residual_label(:,:,l), spatial_decompositions, maximum_block_size);    
                [block_data, block_label_data, wavelet_state, number_of_blocks] = concantenate_blocks_of_main_data(block_data, block_label_data, buffer_output, buffer_label, wavelet_state, wavelet_state_buffer, number_of_blocks, i, l, k);
                temporal_subband_location(ind_loc) = length(block_data);
                ind_loc = ind_loc+1;
            end
               
        end
    end
    layer_block_location(i+1) = length(block_data);
                       
end

string = 'operation/layer_header.mat';
save(string, 'layer_header');

decoded_data = cell(1, length(block_data));
%Set decoded data

total_number_of_blocks = length(block_data);

MSE = zeros(1,total_number_of_blocks);
%Compute Initial Mean Square Error
for i = 1:(total_number_of_blocks)
    MSE(i) = sum(sum((block_data{i}.^2)));
end

for l = 1:length(block_data)
   decoded_data{l} = zeros(size(block_data{l})); 
end

index_distortion = find(MSE > 0);
passes = cell(1, total_number_of_blocks);
rate = cell(1, total_number_of_blocks);
distortion = cell(1, total_number_of_blocks);
distortion_scaled = cell(1, total_number_of_blocks);

compressed_output = cell(1, total_number_of_blocks);
%rate2 = cell(1, total_number_of_blocks);

clear temporal_out;
clear input;
save block_data.mat block_data;
save decoded_data.mat decoded_data; %This saves relevent blocks of output zeros in the data
save block_label_data.mat block_label_data;
%save dwt_output_label.mat 

clear decoded_data %Load this into the local memory later
%clear dwt_output_label;

%load statistics.mat;

w = cd;
cd('operation');
save layer_header layer_header
save r_modified r_modified
save c_modified c_modified
save D D
save label label
cd(w);

for i = 1:length(index_distortion)
    
    buffer = block_label_data{index_distortion(i)};
    buffer(buffer>0)= 1;
    block_label_data{index_distortion(i)} = buffer;
    
    [rate{index_distortion(i)}, distortion{index_distortion(i)}] = EBCOT_compress(block_data{index_distortion(i)}, block_label_data{index_distortion(i)});
    passes{index_distortion(i)} = [0:length(rate{index_distortion(i)})-1]';
    
    fid = fopen('output_data.txt', 'rb');
    compressed_output{index_distortion(i)} = fread(fid);
    fclose(fid);
    
    buffer = information_data{index_distortion(i)};
    low_pass_decompositions = buffer(5);
    high_pass_decompositions = buffer(6);
    
    %obtian the normalized l_2 norm of the distortion
    distortion_scaled{index_distortion(i)} = distortion{index_distortion(i)}* ((2)^(low_pass_decompositions)) * ((2)^high_pass_decompositions);
    
end

clear block_data
clear block_label_data
end



function [block_data, block_label, wavelet_state, number_of_blocks] = concantenate_blocks_of_main_data(block_data, block_label, input, input_label, wavelet_state, wavelet_state_buffer, number_of_blocks, i, ii, k)    

global information_data

number_of_blocks(i,ii) = length(input);
if ~isempty(input)
        
    index = length(information_data);

    low_decomp = zeros(1, k);
    high_decomp = zeros(1,k);
    l = k;
    end_high = k;

    for inc = 1:ceil(log2(k))
    
        new_l = 1:ceil(l/2);
        new_h = ceil(l/2)+1:end_high;
        low_decomp(new_l) = low_decomp(new_l)+1;
        high_decomp(new_h) = high_decomp(new_h)+1;  
        end_high = ceil(l/2);
        l = length(new_l);
        
    end

    if ii < 3
        temp_dim = ii;
    else
        temp_dim = ceil(log2(ii))+1;   
    end
    index = length(block_data)+1;
    for l = 1:length(input)
        %Layer    Temporal Subband  Location in image  Location in Layer
        %Number_of_lowpass_subband_decomp_temporal Number_of_highpass_subband_decomp_temporal
        
        [rows, cols] = size(input{l});
        
        max_val = max(max(abs(input{l})));
        largest_bit_plane = 0;
        
        while (max_val > 0)
            max_val = bitshift(max_val, -1);            
            largest_bit_plane = largest_bit_plane + 1; 
        end
        
        information_vector = [i temp_dim l ii low_decomp(ii) high_decomp(ii), rows, cols, largest_bit_plane];
        information_data{index} = information_vector; clear information_vector                                                               
        block_data{index} = round(input{l});
        block_label{index} = input_label{l};
        wavelet_state(index) = wavelet_state_buffer(l);
        
        index = index+1;

    end   
end
    
end


function [layer_header, min_row, max_row, min_col, max_col, spatial_decompositions_low] =  compute_header_file(layer_header, input_label, spatial_decompositions_low, spatial_decompositions_high, layer)
    
%Compute header file
a = find(abs(input_label) > 0);
[row, col, temp] = ind2sub(size(input_label), a);
%Obtain minimum and maximum column location. Add one to ensure there
%are pixels around the border. 
    
min_col = min(col);
max_col = max(col);
min_row = min(row);
max_row = max(row);

if ~isempty(min_col)  
    if isempty(spatial_decompositions_low)
        spatial_decompositions_low = compute_number_of_decompositions([min_row, max_row], [min_col, max_col]);
    end    
    layer_header(layer,:) = [min_row, max_row, min_col, max_col, spatial_decompositions_low, spatial_decompositions_high, 0];   
else
    layer_header(layer,:) = [0, 0, 0, 0, 0, 0, 0];
end

end

function spatial_decompositions = compute_number_of_decompositions(rows, cols)

row_length = rows(2) - rows(1) + 1;
col_length = cols(2) - cols(1) + 1;
filter_length = 15;

minimum_length = min(row_length, col_length);
spatial_decompositions = 0;

while minimum_length > filter_length
    minimum_length = ceil(minimum_length/2);
    spatial_decompositions = spatial_decompositions + 1;
end
end


function layer_stack = smooth_layers(layer_stack)
%smooth layers

%apply smooting to the layer stack
alpha = 0.05;
for l = 1: length(layer_stack(1,1,:))
    dx = 1;
    dy = 1;
    [ny, nx] = size(layer_stack(:,:,l));
    n_reinit = 250;
    
    buffer = layer_stack(:,:,l);
    buffer = -buffer; buffer(buffer == 0) = 1; %create the level to be initialised
    buffer = reinit_sd2D(buffer, dx, dy, alpha, n_reinit, zeros(ny,nx), zeros(ny,nx), zeros(ny,nx), zeros(ny,nx));
    
    buffer(buffer >=0) = 0; buffer(buffer<0) = -1; buffer = -buffer;
    %place this back into the layer stack
    layer_stack(:,:,l) = buffer;        
    
end
end


function initialize_folders()

if ~exist('operation', 'dir')
    mkdir('low_pass_main');
end


cd('operation');
if ~exist('extracted_layer_label', 'dir')
    mkdir('extracted_layer_label');
end

if ~exist('low_pass_main', 'dir')
    mkdir('low_pass_main');
end

if ~exist('low_pass_main_reconstructed', 'dir')
    mkdir('low_pass_main_reconstructed');
end

if ~exist('recombined', 'dir')
    mkdir('recombined');
end

if ~exist('reconstructed_layers', 'dir')
    mkdir('reconstructed_layers');
end


if ~exist('reconstructed_layers', 'dir')
    mkdir('residual');
end


if ~exist('residual_reconstructed', 'dir')
    mkdir('residual_reconstructed');
end

cd('..');

if ~exist('original_label', 'dir')
    mkdir('original_label');
end

delete('original_label\*.mat');
delete('operation\extracted_layer_label\*.mat');
delete('operation\low_pass_main\*.mat');
delete('operation\low_pass_main_reconstructed\*.mat');
delete('operation\recombined\.*mat');
delete('operation\reconstructed_layers\.*mat');
delete('operation\residual\*.mat');
delete('operation\residual_reconstructed\.*mat');

end
