%% Haar Image Compression
%
% *Wavelets Workshop*
%
% June 2-5, 2009
% 
% University of St. Thomas
%
% *Catherine Beneteau*, *Caroline Haddad*, *David Ruch*, *Patrick Van Fleet*

%% Objective
% In this m-file, we will learn how to perform naive image compression
% using the 2D Haar Wavelet Transform.

%% Conventions
% This m-file uses the DiscreteWavelets Toolbox.  Help is available on all
% functions in the toolbox by accessing the Help menu and clicking on the
% DiscreteWavelets Toolbox.  
%
% Demos are also available.  Click the Demos tab under Help to access all
% demos.

%% Available Images
% The DiscreteWavelets packages comes with 18 grayscale images.  You can
% see information about these images (name, size, etc.) by issuing the 
% command ImageList.

ImageList('ImageType','GrayScale');

%% 
% You can also get a look at these images by using the command
% ShowThumbnails.  

ShowThumbnails('ImageType','GrayScale');

%% 
% No matter where you installed the DiscreteWavelets package on your
% computer, you can retrieve the absolute path and file name for each 
% included image.  The command ImageNames produces a list of file names.  
% We will add a second directive to the call so that the routine will use
% smaller versions of the image.

gray = ImageNames('ImageType','GrayScale','ListThumbnails','True');
disp(gray{1});

%% Loading and Plotting Images
% Once you have the list of file names, it is very easy to load and plot
% images.  Let's load the first image in the list.

A = ImageRead(gray{1});
close;
ImagePlot(A);

%% Compute the Wavelet Transformation
% We first compute the 2D HWT.  Since we are interested in compression, we
% will use a modified version of the Haar transformation.  We will 
% multiply the Haar matrix by Sqrt[2] so that the output are integers.  
% Since we are not using the Haar filter per se, we will use the more 
% general DiscreteWavelets routine WT2D.  This routine requires three 
% arguments.  The first is the input matrix, the second is the filter used
% to construct the wavelet matrix, and the third is the number of 
% iterations.

its = 2;
h = sqrt(2)*Haar();
B = WT2D(A, h, its);
WaveletDensityPlot(B,its,'DivideLinesThickness',[2 2]);

%% Lossless Compression
% In this form of compression, we simply encode the transform.  The image
% can be recovered exactly in lossless compression.  We use the function 
% MakeHuffmanCodes from DiscreteWavelets to build the codes.  The routine 
% requires nonnegative integers.

[rows cols] = size(A);
[uniq, freq, codes, totalbits, encodedbits] = ...
    MakeHuffmanCodes(round(B - min(min(B))));
bpp = encodedbits/numel(A);

disp(sprintf('The number of bits in the original image is %i x %i x 8 = %i.',rows,cols,totalbits)); 

disp(sprintf('The bitstream length of the transformed data is %i or %f bpp.', encodedbits,bpp));

disp(sprintf('The new bitstream length is about %f%% of the original bitstream length.',round(10000*bpp/8)/100));

%% To Consider
% *What happens if we increase the number of iterations?
%
% *Try lossless compression with different images and different numbers of 
%  iterations.

%% Lossy Compression
% We now add a quantization step between transforming the data and encoding
% the transform.  The idea here is that we convert transform values that
% are "small" to zero and thus improve the performance by the Huffman 
% encoder.

%% Cumulative Energy
% We will use cumulative energy to perform quantization.  The command in 
% DiscreteWavelets to compute the cumulative energy is CE.  It takes either
% a vector or a matrix as input.
%
% Here is the cumulative energy vector for our original image.

ceA = CE(A);
close;
plot(ceA,'Color',[1 0 0]);

%%
% Here is the cumulative energy vector for our transformed image.   I've
% printed out an element from the vector to give you an idea how to read 
% the elements of the vector.

ceB = CE(B);
plot(ceB,'Color',[0 0 1]);
i = numel(A)/8;
pct = round(ceB(i)*10000)/100;

disp(sprintf('%f%% of the energy is stored in the largets (in absolute value) %i elements of B.',pct,i));

%% Quantizing with Cumulative Energy
% To quantize with cumulative energy, we will first pick an energy level 0
% <= p <= 1 and then determine the largest elements (in absolute value) in
% B that comprise p units of the energy.  There is a command in 
% DiscreteWavelets called nCE that will perform this task.    The function
% takes the cumulative energy vector and p and returns the number of 
% elements m from B that constitute p units of energy.

p = .998;
k = nCE(ceB, p);
disp(sprintf('The largest (in absolute value) %i elements of B constitute %f%% of the energy of B.',k,100*p));

%%
% We next convert all but the largest (in absolute value) k values of B to
% 0.  The module in DiscreteWavelets to perform this task is Comp.  Comp 
% takes a matrix or vector and the value k and first finds the kth largest 
% element (in absolute value) q in the input.    Comp then sets to 0 all 
% values in the input that are smaller (in absolute value) than q and 
% returns the result.
%
% Here is a simple example of Comp.

Comp([1  -2  3  4  -5], 3)

%%
% For our transform B, here is the result of Comp.  Try replacing k by 
% other values (smaller than rows*cols).  

newB = Comp(B, 2500);
close;
WaveletDensityPlot(newB,its,'DivideLinesThickness',[2 2]);

%% Lossy Compression
% We are ready to put everything together and compress an image using 
% lossy compression.  Understand that we can never exactly recover the 
% original image using lossy compression.
%
% First we read an image and compute its 2D HWT.  Feel free to change the 
% image or the value for its.

A = ImageRead(gray{1});
close;
ImagePlot(A);
[rows  cols] = size(A);
disp(sprintf('The bit stream for this image has length %i.',numel(A)*8));

its = 3;
h = sqrt(2)*Haar();
B = WT2D(A, h, its);
figure;
WaveletDensityPlot(B,its,'DivideLinesThickness',[2 2 2]);

%%
% Next we compute the cumulative energy vector and compress it.  Feel free
% to change the value for p.  

ceB = CE(B);
close all;
plot(ceB);

p = .9995;
k = nCE(ceB, p);

disp(sprintf('The largest (in absolute value) %i elements of B constitute %f%% of the energy of B.',k,round(10000*p)/100)); 
disp(sprintf('%i of the %i elements of B are converted to 0.',numel(A)-k,numel(A)));

newB = Comp(B, k);

%%
% Now we Huffman encode the modified transform.

[uniq, freq, codes, bitstream, encodestream] = ...
    MakeHuffmanCodes(round(newB - min(min(newB))));
bpp = encodestream/numel(A);
disp(sprintf('The length of the encoded bit stream is %i or %f.',...
    encodestream,bpp));

%%
% If we compute the inverse transform, we can see the uncompressed image.
% We plot the original for comparative purposes.  Note that we have to 
% divide our filter now by sqrt(2).

h = Haar()/sqrt(2);
compressedA = IWT2D(newB, h, its);
close;
ImagePlot(compressedA);
title('Compressed Image');
figure;
ImagePlot(A);
title('Original Image');

%% To Consider
% 
% *What happens if you increase the number of iterations?  
%
% *What happens if you increase/decrease the energy level p?
%
% *Pick an image and then set its and p so that the encoded bit stream 
%  results in a compression rate of 1.5 bpp.  How does the uncompressed 
%  image look?

%% Color Image Compression
%
% It is easy to compress color images.    We load a color image and then
% convert the R,G,B channels to Y, Cb, Cr.  We perform grayscale lossy 
% compression on each of Y, Cb, and Cr to obtain modified Y, Cb, and Cr.  
% We then convert back to R,G,B.

%% Load an Image

color = ImageNames('ImageType','Color','ListThumbnails','True');
A = ImageRead(color{6});
ImagePlot(A)
[r,g,b]=Split3D(A);

%% Convert to YCbCr Space

B=RGBToYCbCr(A,'DisplayMode','True'); 
[y,cb,cr]=Split3D(B);

%% Compute the HWT of Each Channel
%
% We use a modified Haar transformation to map integers to integers.  In
% this case, we use the general |WT2D| function and pass it a filter |h|.

its=3;
h=sqrt(2)*Haar();

ywt=WT2D(y,h,its);
cbwt=WT2D(cb,h,its);
crwt=WT2D(cr,h,its);

WaveletDensityPlot(ywt,its);
figure;
WaveletDensityPlot(cbwt,its);
figure;
WaveletDensityPlot(crwt,its);

%% Use Cumulative Energy to Quantize the Tranforms
%

close all;

cey = CE(ywt);
cecb = CE(cbwt);
cecr = CE(crwt);

pct = .999;
ky = nCE(cey, pct);
kcb = nCE(cecb, pct);
kcr = nCE(cecr, pct);

newywt = Comp(ywt, ky);
newcbwt = Comp(cbwt, kcb);
newcrwt = Comp(crwt, kcr);

WaveletDensityPlot(newywt,its);
figure;
WaveletDensityPlot(newcbwt,its);
figure;
WaveletDensityPlot(newcrwt,its);

%% Encode the Quantized Transformations
%
close all;

[uniq, freq, codes, bitstream, encodestreamy] = MakeHuffmanCodes(round(newywt-min(min(newywt))));
[uniq, freq, codes, bitstream, encodestreamcb] = MakeHuffmanCodes(round(newcbwt-min(min(newcbwt))));
[uniq, freq, codes, bitstream, encodestreamcr] = MakeHuffmanCodes(round(newcrwt-min(min(newcrwt))));

bpp = (encodestreamy + encodestreamcb + encodestreamcr)/(3*rows*cols);
str=sprintf('The compression rate is %f bpp!',bpp);
disp(str);

%% Compute the Inverse Transformations

h = Haar()/sqrt(2);
newy = IWT2D(newywt, h, its);
newcb = IWT2D(newcbwt, h, its);
newcr = IWT2D(newcrwt, h, its);

%% Convert back to R, G, B

B=Make3D(newy,newcb,newcr);
newA = YCbCrToRGB(B,'DisplayMode','True');
ImagePlot(newA);
figure;
ImagePlot(A);
[newr, newg, newb]=Split3D(newA);
PSNR([r, g, b], [newr, newg, newb])

%%
close all;
displayEndOfDemoMessage(mfilename)