Saturday, October 1, 2016

Example code: 4x4 Haar transform on ETC1 selectors

Someone asked me for the code that implements the selector frequency domain filtering experiment I did on one of my previous posts. Here you go. The linear algebra and utility code is all very similar to crunch's. The Haar stuff was copied straight from wikipedia.

If you understand how coefficient quantization works in JPEG, then this example should be pretty easy to follow. Let me know if you have any questions or want to see the 8x8 version, but it's basically the same thing with larger matrices (and the selectors within 2x2 ETC1 blocks).

Notes:
gpu_image: A compressed GPU texture (in this case, ETC1). image_u8: A 32bpp ARGB raster image. etc_block: A simple helper class to get/set ETC1/2 selectors, subblock colors, codework table indices, etc.
static void selector_haar_transform_test(const gpu_image &g)
{
image_u8 g_unpacked;
g.unpack_image(g_unpacked);

image_utils::write_to_file("g_orig.png", g_unpacked);

gpu_image g_temp(g);

#define sqrt2 (1.4142135623730951f)

matrix44F haar_matrix(
.5f * 1, .5f * 1, .5f * 1, .5f * 1,
.5f * 1, .5f * 1, .5f * -1, .5f * -1,
.5f * sqrt2, .5f * -sqrt2, 0, 0,
0, 0, .5f * sqrt2, .5f * -sqrt2);

matrix44F inv_haar_matrix(haar_matrix.get_transposed());

for (uint block_y = 0; block_y < g_temp.get_blocks_y(); block_y++)
{
for (uint block_x = 0; block_x < g_temp.get_blocks_x(); block_x++)
{
etc_block &blk = g_temp.get_element_of_type<etc_block>(block_x, block_y);

matrix44F s;
for (uint y = 0; y < 4; y++)
for (uint x = 0; x < 4; x++)
s(x, y) = (float)blk.get_selector(x, y) + .5f;

matrix44F h(haar_matrix * s * inv_haar_matrix);

for (uint i = 0; i < 4; i++)
{
for (uint j = 0; j < 4; j++)
{
float q = h(i, j);

const float scale = 4.0f;
int iq = math::float_to_int_nearest(q * scale);

int quant = (i+j+1);

if (quant < 1)
quant = 1;

iq /= quant;

iq *= quant;

h(i, j) = iq / scale;
}
}

matrix44F ih(inv_haar_matrix * h * haar_matrix);

for (uint y = 0; y < 4; y++)
for (uint x = 0; x < 4; x++)
//blk.set_selector(x, y, math::clamp<int>((int)ih(x, y) + .5f, 0, 3));
blk.set_selector(x, y, math::clamp<int>((int)ih(x, y), 0, 3));
}
}

write_etc1_vis_images(g_temp, "g_temp_");

g_temp.unpack_image(g_unpacked);

image_utils::write_to_file("g.png", g_unpacked);
exit(0);
}