%Vector Quantization of an image using generalized Lloyd algorithm
%This version yields quantized images with various codebook size from 2 to 256, so that it can draw a RD curve.
%you can slightly modify the program to run with a fixed codebook size.

%	Author: Xiaofeng Xu, Polytechnic University  4/21/2002


%the end threshold of iteration 
e=0.01;

%Read gray image data from a file
img=double(imread('slena256.bmp'));
[H,W]=size(img);
M=H*W;

%Vector Block Size
BlockSize=4;

%convert the 2D image into the fixed-size vector
ImgVector=im2col(img,[BlockSize,BlockSize],'distinct');
ImgVector1=ImgVector;

[HeightV,WidthV]=size(ImgVector);

%quantization level
N=1;
while(1)
	%Set quantization level [2,4,8,16,32,64,128,256]
   N=N*2;
   
	MinIdx=ones(WidthV,1);
   % Previous distortion value
   PreDistortion=0.001;
   
   %Initiate codebooks
   %choose first N training vectors as codewords
	for i=1:N
   	CodeBook(1:HeightV,i)=ImgVector(:,i)
	end
   
   %50 is maximum no. of iterations
	for k=1:50
		TotalDistortion=0;
		TotalIdx=zeros(N,1);
      MeanVector=zeros(HeightV,N);
      
      %for each training vector compare with each codeword
      for i=1:WidthV
			minerr=inf;
         for j=1:N
            %Cacluate the MSE between current image block and codebook
            temp=sum([ImgVector(:,i)-CodeBook(:,j)].^2);            
            if temp<minerr
               %record the minimum MSE value and corresponding codebook index
               minerr=temp;
               %associate with i-th training vectors with j-th codeword
      		   MinIdx(i)=j;
		      end     
         end
         %caclulate new center vectors
         MeanVector(:,MinIdx(i))= MeanVector(:,MinIdx(i))+ImgVector(:,i);
         %accumulating the counter for training vectors quantized to the same codeword
         TotalIdx(MinIdx(i))=   TotalIdx(MinIdx(i))+1;   
         %caclulate total distortion
		   TotalDistortion=TotalDistortion+minerr;
		end
      TotalDistortion=TotalDistortion/WidthV;
      
      %Update the codebook.
      for i=1:N
		   CodeBook(1:HeightV,i)=MeanVector(:,i)/TotalIdx(i);
		end

		if (abs(TotalDistortion-PreDistortion)/PreDistortion)<e
		   break;
		end
      
      PreDistortion=TotalDistortion
      k
   end
   
   %replace a training vector with its corresponding codeword
	for i=1:WidthV
   	ImgVector1(:,i)=CodeBook(:,MinIdx(i));
	end

	Q=col2im(ImgVector1,[BlockSize,BlockSize],[H,W],'distinct');
   N
   figure;
   imshow(uint8(Q));
	if N==2,imwrite(uint8(Q),'c:\LBG002.tif','TIF'),LBGMSE02=TotalDistortion;end
	if N==4,imwrite(uint8(Q),'c:\LBG004.tif','TIF'),LBGMSE04=TotalDistortion;end
	if N==8,imwrite(uint8(Q),'c:\LBG008.tif','TIF'),LBGMSE08=TotalDistortion;end
	if N==16,imwrite(uint8(Q),'c:\LBG016.tif','TIF'),LBGMSE16=TotalDistortion;end
	if N==32,imwrite(uint8(Q),'c:\LBG032.tif','TIF'),LBGMSE32=TotalDistortion;end
	if N==64,imwrite(uint8(Q),'c:\LBG064.tif','TIF'),LBGMSE64=TotalDistortion;end
	if N==128,imwrite(uint8(Q),'c:\LBG128.tif','TIF'),LBGMSE128=TotalDistortion;end

	if N==256,break;end
end
imwrite(uint8(Q),'c:\LBG256.tif','TIF'),LBGMSE256=TotalDistortion;

%Draw Distortion Rate Curves
rate=log2([2 4 8 16 32 64 128 256]);
figure,plot(rate,[LBGMSE02,LBGMSE04,LBGMSE08,LBGMSE16,LBGMSE32,LBGMSE64,LBGMSE128,LBGMSE256],'bs')
title('rate-distortion curves'),xlabel('bit rate'),ylabel('mean-squared error')
legend('LBG vector quantizer')