path: root/lib/libsquish/extra/squishpng.cpp

/* -----------------------------------------------------------------------------

	Copyright (c) 2006 Simon Brown                          si@sjbrown.co.uk

	Permission is hereby granted, free of charge, to any person obtaining
	a copy of this software and associated documentation files (the 
	"Software"), to	deal in the Software without restriction, including
	without limitation the rights to use, copy, modify, merge, publish,
	distribute, sublicense, and/or sell copies of the Software, and to 
	permit persons to whom the Software is furnished to do so, subject to 
	the following conditions:

	The above copyright notice and this permission notice shall be included
	in all copies or substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
	OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 
	MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
	IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 
	CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, 
	TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 
	SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
	
   -------------------------------------------------------------------------- */
   
/*! @file

	@brief	Test program that compresses images loaded using the PNG format.
	
	This program requires libpng for PNG input and output, and is designed to
	test the RMS error for DXT compression for a set of test images.

	This program uses the high-level image compression and decompression
	functions that process an entire image at a time.
*/

#include <iostream>
#include <string>
#include <sstream>
#include <ctime>
#include <cmath>
#include <squish.h>
#include <png.h>

#ifdef _MSC_VER
#pragma warning( disable: 4511 4512 )
#endif // def _MSC_VER

using namespace squish;

//! Simple exception class.
class Error : public std::exception
{
public:
	Error( std::string const& excuse ) : m_excuse( excuse ) {}
	~Error() throw() {}
	
	virtual char const* what() const throw() { return m_excuse.c_str(); }
	
private:
	std::string m_excuse;
};

//! Base class to make derived classes non-copyable
class NonCopyable
{
public:
	NonCopyable() {}
	
private:
	NonCopyable( NonCopyable const& );
	NonCopyable& operator=( NonCopyable const& );
};

//! Memory object.
class Mem : NonCopyable
{
public:
	Mem() : m_p( 0 ) {}
	explicit Mem( int size ) : m_p( new u8[size] ) {}
	~Mem() { delete[] m_p; }

	void Reset( int size )
	{
		u8 *p = new u8[size];
		delete m_p;
		m_p = p;
	}
	
	u8* Get() const { return m_p; }
	
private:
	u8* m_p;
};

//! File object.
class File : NonCopyable
{
public:
	explicit File( FILE* fp ) : m_fp( fp ) {}
	~File() { if( m_fp ) fclose( m_fp ); }
	
	bool IsValid() const { return m_fp != 0; }
	FILE* Get() const { return m_fp; }

private:
	FILE* m_fp;
};

//! PNG read object.
class PngReadStruct : NonCopyable
{
public:
	PngReadStruct()
	  : m_png( 0 ), 
		m_info( 0 ), 
		m_end( 0 )
	{
		m_png = png_create_read_struct( PNG_LIBPNG_VER_STRING, 0, 0, 0 );
		if( !m_png )
			throw Error( "failed to create png read struct" );	
			
		m_info = png_create_info_struct( m_png );
		m_end = png_create_info_struct( m_png );
		if( !m_info || !m_end )
		{
			png_infopp info = m_info ? &m_info : 0;
			png_infopp end = m_end ? &m_end : 0;
			png_destroy_read_struct( &m_png, info, end );
			throw Error( "failed to create png info structs" );
		}
	}
	
	~PngReadStruct() 
	{ 
		png_destroy_read_struct( &m_png, &m_info, &m_end );
	}

	png_structp GetPng() const { return m_png; }
	png_infop GetInfo() const { return m_info; }

private:
	png_structp m_png;
	png_infop m_info, m_end;
};

//! PNG write object.
class PngWriteStruct : NonCopyable
{
public:
	PngWriteStruct()
	  : m_png( 0 ), 
		m_info( 0 )
	{
		m_png = png_create_write_struct( PNG_LIBPNG_VER_STRING, 0, 0, 0 );
		if( !m_png )
			throw Error( "failed to create png read struct" );	
			
		m_info = png_create_info_struct( m_png );
		if( !m_info )
		{
			png_infopp info = m_info ? &m_info : 0;
			png_destroy_write_struct( &m_png, info );
			throw Error( "failed to create png info structs" );
		}
	}
	
	~PngWriteStruct()
	{
		png_destroy_write_struct( &m_png, &m_info );
	}
	
	png_structp GetPng() const { return m_png; }
	png_infop GetInfo() const { return m_info; }

private:
	png_structp m_png;
	png_infop m_info;
};

//! PNG rows object.
class PngRows : NonCopyable
{
public:
	PngRows( int pitch, int height ) : m_height( height )
	{
		m_rows = new png_bytep[m_height];
		for( int i = 0; i < m_height; ++i )
			m_rows[i] = new png_byte[pitch];
	}
	
	~PngRows() 
	{
		for( int i = 0; i < m_height; ++i )
			delete[] m_rows[i];
		delete[] m_rows;
	}
	
	png_bytep* Get() const { return m_rows; }

	png_bytep operator[](int y) const { return m_rows[y]; }
	
private:
	png_bytep* m_rows;
	int m_height;
};

//! Represents a DXT compressed image in memory.
struct DxtData
{
	int width;
	int height;
	int format;		//!< Either kDxt1, kDxt3 or kDxt5.
	Mem data;
	bool isColour;
	bool isAlpha;
};

//! Represents an uncompressed RGBA image in memory.
class Image
{
public:
	Image();

	void LoadPng( std::string const& fileName );
	void SavePng( std::string const& fileName ) const;

	void Decompress( DxtData const& dxt );
	void Compress( DxtData& dxt, int flags ) const;

	double GetRmsError( Image const& image ) const;

private:
	int m_width;
	int m_height;
	bool m_isColour;	//!< Either colour or luminance.
	bool m_isAlpha;		//!< Either alpha or not.
	Mem m_pixels;
};

Image::Image() 
  : m_width( 0 ), 
  	m_height( 0 ),
	m_isColour( false ),
	m_isAlpha( false )
{
}

void Image::LoadPng( std::string const& fileName )
{
	// open the source file
	File file( fopen( fileName.c_str(), "rb" ) );
	if( !file.IsValid() )
	{
		std::ostringstream oss;
		oss << "failed to open \"" << fileName << "\" for reading";
		throw Error( oss.str() );
	}
	
	// check the signature bytes
	png_byte header[8];
	size_t check = fread( header, 1, 8, file.Get() );
	if( check != 8 )
		throw Error( "file read error" );
	if( png_sig_cmp( header, 0, 8 ) )
	{
		std::ostringstream oss;
		oss << "\"" << fileName << "\" does not look like a png file";
		throw Error( oss.str() );
	}
	
	// read the image into memory
	PngReadStruct png;
	png_init_io( png.GetPng(), file.Get() );
	png_set_sig_bytes( png.GetPng(), 8 );
	png_read_png( png.GetPng(), png.GetInfo(), PNG_TRANSFORM_EXPAND, 0 );

	// get the image info
	png_uint_32 width;
	png_uint_32 height;
	int bitDepth;
	int colourType;
	png_get_IHDR( png.GetPng(), png.GetInfo(), &width, &height, &bitDepth, &colourType, 0, 0, 0 );
	
	// check the image is 8 bit
	if( bitDepth != 8 )
	{
		std::ostringstream oss;
		oss << "cannot process " << bitDepth << "-bit image (bit depth must be 8)";
		throw Error( oss.str() );
	}

	// copy the data into a contiguous array
	m_width = width;
	m_height = height;
	m_isColour = ( ( colourType & PNG_COLOR_MASK_COLOR ) != 0 );
	m_isAlpha = ( ( colourType & PNG_COLOR_MASK_ALPHA ) != 0 );
	m_pixels.Reset(4*width*height);

	// get the image rows
	png_bytep const *rows = png_get_rows( png.GetPng(), png.GetInfo() );
	if( !rows )
		throw Error( "failed to get image rows" );

	// copy the pixels into the storage
	u8 *dest = m_pixels.Get();
	for( int y = 0; y < m_height; ++y )
	{
		u8 const *src = rows[y];
		for( int x = 0; x < m_width; ++x )
		{
			if( m_isColour )
			{
				dest[0] = src[0];
				dest[1] = src[1];
				dest[2] = src[2];
				src += 3;
			}
			else
			{
				u8 lum = *src++;
				dest[0] = lum;
				dest[1] = lum;
				dest[2] = lum;
			}
			
			if( m_isAlpha )
				dest[3] = *src++;
			else
				dest[3] = 255;

			dest += 4;
		}
	}
}

void Image::SavePng( std::string const& fileName ) const
{
	// create the target rows
	int const pixelSize = ( m_isColour ? 3 : 1 ) + ( m_isAlpha ? 1 : 0 );
	PngRows rows( m_width*pixelSize, m_height );

	// fill the rows with pixel data
	u8 const *src = m_pixels.Get();
	for( int y = 0; y < m_height; ++y )
	{
		u8 *dest = rows[y];
		for( int x = 0; x < m_width; ++x )
		{
			if( m_isColour )
			{
				dest[0] = src[0];
				dest[1] = src[1];
				dest[2] = src[2];
				dest += 3;
			}
			else
				*dest++ = src[1];
			
			if( m_isAlpha )
				*dest++ = src[3];

			src += 4;
		}
	}

	// set up the image
	PngWriteStruct png;
	png_set_IHDR(
		png.GetPng(), png.GetInfo(), m_width, m_height,
		8, ( m_isColour ? PNG_COLOR_MASK_COLOR : 0) | ( m_isAlpha ? PNG_COLOR_MASK_ALPHA : 0 ), 
		PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT 
	);
	   
	// open the target file
	File file( fopen( fileName.c_str(), "wb" ) );
	if( !file.IsValid() )
	{
		std::ostringstream oss;
		oss << "failed to open \"" << fileName << "\" for writing";
		throw Error( oss.str() );
	}
	
	// write the image
	png_set_rows( png.GetPng(), png.GetInfo(), rows.Get() );
	png_init_io( png.GetPng(), file.Get() );
	png_write_png( png.GetPng(), png.GetInfo(), PNG_TRANSFORM_IDENTITY, 0 );
}

void Image::Decompress( DxtData const& dxt )
{
	// allocate storage
	m_width = dxt.width;
	m_height = dxt.height;
	m_isColour = dxt.isColour;
	m_isAlpha = dxt.isAlpha;
	m_pixels.Reset( 4*m_width*m_height );

	// use the whole image decompression function to do the work
	DecompressImage( m_pixels.Get(), m_width, m_height, dxt.data.Get(), dxt.format );
}

void Image::Compress( DxtData& dxt, int flags ) const
{
	// work out how much memory we need
	int storageSize = GetStorageRequirements( m_width, m_height, flags );

	// set the structure fields and allocate it
	dxt.width = m_width;
	dxt.height = m_height;
	dxt.format = flags & ( kDxt1 | kDxt3 | kDxt5 );
	dxt.isColour = m_isColour;
	dxt.isAlpha = m_isAlpha;
	dxt.data.Reset( storageSize );

	// use the whole image compression function to do the work
	CompressImage( m_pixels.Get(), m_width, m_height, dxt.data.Get(), flags );
}

double Image::GetRmsError( Image const& image ) const
{
	if( m_width != image.m_width || m_height != image.m_height )
		throw Error( "image dimensions mismatch when computing RMS error" );

	// accumulate colour error
	double difference = 0;
	u8 const *a = m_pixels.Get();
	u8 const *b = image.m_pixels.Get();
	for( int y = 0; y < m_height; ++y )
	{
		for( int x = 0; x < m_width; ++x )
		{
			int d0 = ( int )a[0] - ( int )b[0];
			int d1 = ( int )a[1] - ( int )b[1];
			int d2 = ( int )a[2] - ( int )b[2];
			difference += ( double )( d0*d0 + d1*d1 + d2*d2 ); 
			a += 4;
			b += 4;
		}
	}
	return std::sqrt( difference/( double )( m_width*m_height ) );
}

int main( int argc, char* argv[] )
{
	try
	{
		// parse the command-line
		std::string sourceFileName;
		std::string targetFileName;
		int format = kDxt1;
		int fit = kColourClusterFit;
		int extra = 0;
		bool help = false;
		bool arguments = true;
		bool error = false;
		for( int i = 1; i < argc; ++i )
		{
			// check for options
			char const* word = argv[i];
			if( arguments && word[0] == '-' )
			{
				for( int j = 1; word[j] != '\0'; ++j )
				{
					switch( word[j] )
					{
					case 'h': help = true; break;
					case '1': format = kDxt1; break;
					case '3': format = kDxt3; break;
					case '5': format = kDxt5; break;
					case 'r': fit = kColourRangeFit; break;
					case 'i': fit = kColourIterativeClusterFit; break;
					case 'w': extra = kWeightColourByAlpha; break;
					case '-': arguments = false; break;
					default:
						std::cerr << "squishpng error: unknown option '" << word[j] << "'" << std::endl;
						error = true;
					}
				}
			}
			else
			{
				if( sourceFileName.empty() )
					sourceFileName.assign( word );
				else if( targetFileName.empty() )
					targetFileName.assign( word );
				else
				{
					std::cerr << "squishpng error: unexpected argument \"" << word << "\"" << std::endl;
					error = true;
				}
			}
		}
		
		// check arguments
		if( sourceFileName.empty() )
		{
			std::cerr << "squishpng error: no source file given" << std::endl;
			error = true;
		}
		if( help || error )
		{
			std::cout 
				<< "SYNTAX" << std::endl
				<< "\tsquishpng [-135riw] <source> [<target>]" << std::endl 
				<< "OPTIONS" << std::endl
				<< "\t-h\tPrint this help message" << std::endl
				<< "\t-135\tSpecifies whether to use DXT1 (default), DXT3 or DXT5 compression" << std::endl
				<< "\t-r\tUse the fast but inferior range-based colour compressor" << std::endl
				<< "\t-i\tUse the very slow but slightly better iterative colour compressor" << std::endl
				<< "\t-w\tWeight colour values by alpha in the cluster colour compressor" << std::endl
				;
			
			return error ? -1 : 0;
		}

		// load the source image
		Image sourceImage;
		sourceImage.LoadPng( sourceFileName );

		// compress to DXT
		DxtData dxt;
		sourceImage.Compress( dxt, format | fit | extra );

		// decompress back
		Image targetImage;
		targetImage.Decompress( dxt );

		// compare the images
		double rmsError = sourceImage.GetRmsError( targetImage );
		std::cout << sourceFileName << " " << rmsError << std::endl;

		// save the target image if necessary
		if( !targetFileName.empty() )
			targetImage.SavePng( targetFileName );
	}
	catch( std::exception& excuse )
	{
		// complain
		std::cerr << "squishpng error: " << excuse.what() << std::endl;
		return -1;
	}
	
	// done
	return 0;
}