From f61ec3dddc58ec18e859c1aee833c6e38749bf9e Mon Sep 17 00:00:00 2001 From: Nyeogmi Date: Sun, 7 Apr 2024 18:42:52 -0700 Subject: [PATCH] Zigzagging works, but uncommetning this quantization logic breaks things --- Cargo.lock | 56 ++++++++++++++++++ Cargo.toml | 1 + src/compression.rs | 58 +++++++++++------- src/constants.rs | 7 ++- src/decompression.rs | 89 +++++++++++++++------------- src/main.rs | 21 ++++--- src/protocol.rs | 138 +++++++++++++++++++++++++++++++++++++++++++ src/quantization.rs | 47 +++++++++++++++ 8 files changed, 347 insertions(+), 70 deletions(-) create mode 100644 src/protocol.rs create mode 100644 src/quantization.rs diff --git a/Cargo.lock b/Cargo.lock index d5c2054..7e8a792 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -119,6 +119,15 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "proc-macro2" +version = "1.0.79" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e" +dependencies = [ + "unicode-ident", +] + [[package]] name = "quickcheck" version = "1.0.3" @@ -130,6 +139,15 @@ dependencies = [ "rand", ] +[[package]] +name = "quote" +version = "1.0.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef" +dependencies = [ + "proc-macro2", +] + [[package]] name = "rand" version = "0.8.5" @@ -183,6 +201,7 @@ version = "0.1.0" dependencies = [ "png", "quickcheck", + "thiserror", ] [[package]] @@ -191,6 +210,43 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +[[package]] +name = "syn" +version = "2.0.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44cfb93f38070beee36b3fef7d4f5a16f27751d94b187b666a5cc5e9b0d30687" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index f8d7e39..85415ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,4 @@ edition = "2021" [dependencies] png = "0.17.13" quickcheck = "1.0.3" +thiserror = "1.0.58" diff --git a/src/compression.rs b/src/compression.rs index 056b7f8..31f3263 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -1,6 +1,6 @@ use std::io::Write; -use crate::{constants::{MAGIC, TILE_SZ}, transform}; +use crate::{constants::{MAGIC, TILE_SZ}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform}; struct PixelTile { @@ -11,14 +11,18 @@ struct CoefTile { coefs: [i32; TILE_SZ * TILE_SZ] } +struct QuantTile { + quants: [i32; TILE_SZ * TILE_SZ] +} -pub fn compress( - width: u16, - height: u16, + +pub fn compress( + width: u32, + height: u32, // n_components: u16, layers: &[&[u8]], -) -> Vec { - let mut output = vec![]; + writer: &mut protocol::ProtocolWriter +) -> ProtocolWriterResult<()> { // validation for l in 0..layers.len() { @@ -26,11 +30,10 @@ pub fn compress( } // write header - output.write(&MAGIC).unwrap(); - output.write(&(width as u16).to_le_bytes()).unwrap(); - output.write(&(height as u16).to_le_bytes()).unwrap(); - // output.write(&n_components.to_le_bytes()).unwrap(); - output.write(&(layers.len() as u16).to_le_bytes()).unwrap(); + writer.write_header(MAGIC)?; + writer.write_u32_wide(width)?; + writer.write_u32_wide(height)?; + writer.write_u32_wide(layers.len() as u32)?; let mut tiles = vec![]; for layer in layers.iter() { @@ -41,19 +44,20 @@ pub fn compress( width as usize, height as usize ); let coef_tile = CoefTile::from_pixel_tile(&pixel_tile); - tiles.push(coef_tile); + let quant_tile = QuantTile::from_coef_tile(&coef_tile); + tiles.push(quant_tile); } } } for t in tiles.iter() { - t.write_zero(&mut output); + t.write_zero(writer)?; } for t in tiles.iter() { - t.write_rest(&mut output); + t.write_rest(writer)?; } - output + Ok(()) } impl PixelTile { @@ -98,14 +102,24 @@ impl CoefTile { return CoefTile { coefs } } +} - fn write_zero(&self, output: &mut Vec) { - output.write(&self.coefs[0].to_le_bytes()).unwrap(); - } - - fn write_rest(&self, output: &mut Vec) { - for i in 1..self.coefs.len() { - output.write(&self.coefs[i].to_le_bytes()).unwrap(); +impl QuantTile { + fn from_coef_tile(pt: &CoefTile) -> QuantTile { + QuantTile { + quants: quantization::to_quantized(pt.coefs) } } + + fn write_zero(&self, writer: &mut ProtocolWriter) -> ProtocolWriterResult<()> { + writer.write_i32_packed(self.quants[0])?; + Ok(()) + } + + fn write_rest(&self, writer: &mut ProtocolWriter) -> ProtocolWriterResult<()> { + for i in 1..self.quants.len() { + writer.write_i32_packed(self.quants[i])?; + } + Ok(()) + } } \ No newline at end of file diff --git a/src/constants.rs b/src/constants.rs index 06366a0..f53238b 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -1,2 +1,7 @@ -pub const MAGIC: [u8; 4] = *b"rxi\0"; +use crate::protocol::Header; + +pub const MAGIC: Header = Header { + bytes: *b"rxi", + version: 0 +}; pub const TILE_SZ: usize = 8; \ No newline at end of file diff --git a/src/decompression.rs b/src/decompression.rs index c9323d4..ac37012 100644 --- a/src/decompression.rs +++ b/src/decompression.rs @@ -1,6 +1,6 @@ -use std::io::{BufReader, Read}; +use std::io::Read; -use crate::{constants::{MAGIC, TILE_SZ}, transform}; +use crate::{constants::{MAGIC, TILE_SZ}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform}; struct PixelTile { // i32: representation that supports Walsh-Hadamard @@ -12,35 +12,37 @@ struct CoefTile { coefs: [i32; TILE_SZ * TILE_SZ] } -pub fn decompress( - input: &[u8] -) -> (u16, u16, Vec>) { - let mut reader = BufReader::new(input); +#[derive(Clone)] +struct QuantTile { + quants: [i32; TILE_SZ * TILE_SZ] +} +pub fn decompress( + reader: &mut ProtocolReader +) -> ProtocolReaderResult<(u32, u32, Vec>)> { // read header - let mut magic_buffer = [0; MAGIC.len()]; - reader.read_exact(&mut magic_buffer).unwrap(); + let header = reader.read_header()?; + if header.bytes != MAGIC.bytes { + return Err(ProtocolReaderError::WrongHeader); + } + if header.version != MAGIC.version { + return Err(ProtocolReaderError::WrongVersion); + } - assert_eq!(MAGIC, magic_buffer); - - let mut u16_buffer = [0; 2]; - reader.read_exact(&mut u16_buffer).unwrap(); - let width = u16::from_le_bytes(u16_buffer); - reader.read_exact(&mut u16_buffer).unwrap(); - let height = u16::from_le_bytes(u16_buffer); - reader.read_exact(&mut u16_buffer).unwrap(); - let n_layers = u16::from_le_bytes(u16_buffer); + let width = reader.read_u32_wide()?; + let height = reader.read_u32_wide()?; + let n_layers = reader.read_u32_wide()?; let width_in_tiles = (width as usize + TILE_SZ - 1) / TILE_SZ; let height_in_tiles = (height as usize + TILE_SZ - 1) / TILE_SZ; let n_tiles = width_in_tiles * height_in_tiles * (n_layers as usize); - let mut tiles = vec![CoefTile::new(); n_tiles]; + let mut tiles = vec![QuantTile::new(); n_tiles]; for i in 0..n_tiles { - tiles[i].load_zero(&mut reader); + tiles[i].load_zero(reader)?; } for i in 0..n_tiles { - tiles[i].load_rest(&mut reader); + tiles[i].load_rest(reader)?; } let mut tile_i = 0; @@ -48,7 +50,7 @@ pub fn decompress( for layer in 0..n_layers { for x0 in (0..width).step_by(TILE_SZ) { for y0 in (0..height).step_by(TILE_SZ) { - let pixel_tile = tiles[tile_i].to_pixel_tile(); + let pixel_tile = tiles[tile_i].to_coef_tile().to_pixel_tile(); tile_i += 1; pixel_tile.to_layer( x0 as usize, y0 as usize, &mut layers[layer as usize], @@ -57,28 +59,10 @@ pub fn decompress( } } } - return (width, height, layers) + return Ok((width, height, layers)) } impl CoefTile { - fn new() -> CoefTile { - CoefTile { coefs: [0; TILE_SZ * TILE_SZ] } - } - - fn load_zero(&mut self, reader: &mut BufReader<&[u8]>) { - let mut coef_zero_buffer = [0; 4]; - reader.read_exact(&mut coef_zero_buffer).unwrap(); - self.coefs[0] = i32::from_le_bytes(coef_zero_buffer); - } - - fn load_rest(&mut self, reader: &mut BufReader<&[u8]>) { - let mut coef_buffer = [0; 4]; - for i in 1..self.coefs.len() { - reader.read_exact(&mut coef_buffer).unwrap(); - self.coefs[i] = i32::from_le_bytes(coef_buffer); - } - } - fn to_pixel_tile(&self) -> PixelTile { let mut pixels = self.coefs.clone(); @@ -108,4 +92,29 @@ impl PixelTile { } } } +} + +impl QuantTile { + fn new() -> QuantTile { + QuantTile { quants: [0; TILE_SZ * TILE_SZ] } + } + + fn to_coef_tile(&self) -> CoefTile { + CoefTile { + coefs: quantization::from_quantized(self.quants) + } + } + + fn load_zero(&mut self, reader: &mut ProtocolReader) -> ProtocolReaderResult<()> { + self.quants[0] = reader.read_i32_packed()?; + Ok(()) + } + + fn load_rest(&mut self, reader: &mut ProtocolReader) -> ProtocolReaderResult<()> { + for i in 1..self.quants.len() { + self.quants[i] = reader.read_i32_packed()?; + } + Ok(()) + } + } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 01027aa..6272f8d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,14 @@ -use std::fs::File; +use std::{fs::File, io::{Cursor, Write}}; use png::{BitDepth, ColorType}; +use crate::protocol::{ProtocolReader, ProtocolWriter}; + mod compression; mod constants; mod decompression; +mod protocol; +mod quantization; mod transform; @@ -25,16 +29,19 @@ fn main() { } fn hard_main() { let (width, height, r, g, b) = load_image(); - let compressed = - compression::compress(width as u16, height as u16, &[&r, &g, &b]); + let mut writer = ProtocolWriter::new(vec![]); + compression::compress( + width as u32, height as u32, &[&r, &g, &b], + &mut writer + ).unwrap(); + let compressed = writer.destroy(); - /* let mut output_file = File::create("outputs/avatar2.rxi").unwrap(); output_file.write_all(&compressed).unwrap(); - */ - let (width2, height2, decompressed) = - decompression::decompress(&compressed); + let mut reader = ProtocolReader::new(Cursor::new(compressed)); + let (width2, height2, decompressed) = + decompression::decompress(&mut reader).unwrap(); assert_eq!(3, decompressed.len()); save_image(width2 as usize, height2 as usize, decompressed); diff --git a/src/protocol.rs b/src/protocol.rs new file mode 100644 index 0000000..726a95d --- /dev/null +++ b/src/protocol.rs @@ -0,0 +1,138 @@ +use std::io::{self, ErrorKind, Read, Write}; + +use thiserror::Error; + +#[derive(Clone, Copy)] +pub struct Header { + pub bytes: [u8; 3], + pub version: u8 +} + + +#[derive(Error, Debug)] +pub enum ProtocolWriterError { + #[error("error using underlying writer")] + Io(#[from] io::Error) +} + +pub type ProtocolWriterResult = Result; + +pub struct ProtocolWriter { + writer: W +} + +#[derive(Error, Debug)] +pub enum ProtocolReaderError { + // this is explicitly supported: rxi images are progressive + #[error("EOF before end of image")] + EarlyEof, + + #[error("wrong header")] + WrongHeader, + + #[error("header has unsupported version")] + WrongVersion, + + #[error("error using underlying reader")] + Io(#[from] io::Error), +} + +pub type ProtocolReaderResult = Result; + +pub struct ProtocolReader { + reader: R +} + +impl ProtocolWriter { + pub fn new(writer: W) -> Self { + Self { writer } + } + + pub fn destroy(self) -> W { + self.writer + } + + pub fn write_header(&mut self, value: Header) -> ProtocolWriterResult<()> { + self.writer.write_all(&value.bytes)?; + self.writer.write_all(&[value.version])?; + Ok(()) + } + pub fn write_u32_wide(&mut self, value: u32) -> ProtocolWriterResult<()> { + self.writer.write_all(&value.to_le_bytes())?; + Ok(()) + } + pub fn write_i32_packed(&mut self, value: i32) -> ProtocolWriterResult<()> { + // We reserve i8::MAX and i16::MAX as signal values that `value` is too big + // for each respective type + if (i8::MIN as i32..i8::MAX as i32).contains(&value) { + self.writer.write_all(&(value as i8).to_le_bytes())?; + return Ok(()) + } + self.writer.write_all(&(i8::MAX).to_le_bytes())?; + + if (i16::MIN as i32..i16::MAX as i32).contains(&value) { + self.writer.write_all(&(value as i16).to_le_bytes())?; + return Ok(()) + } + self.writer.write_all(&(i16::MAX).to_le_bytes())?; + self.writer.write_all(&value.to_le_bytes())?; + Ok(()) + } +} + +impl ProtocolReader { + pub fn new(reader: R) -> Self { + Self { reader } + } + + pub fn read_header(&mut self) -> ProtocolReaderResult
{ + let mut bytes_buf = [0; 3]; + self.read_exact(&mut bytes_buf)?; + + let mut version_buf = [0; 1]; + self.read_exact(&mut version_buf)?; + + Ok(Header { + bytes: bytes_buf, + version: version_buf[0], + }) + } + + pub fn read_u32_wide(&mut self) -> ProtocolReaderResult { + let mut u32_buf = [0; 4]; + self.read_exact(&mut u32_buf)?; + + Ok(u32::from_le_bytes(u32_buf)) + } + + pub fn read_i32_packed(&mut self) -> ProtocolReaderResult { + let mut i8_buf = [0; 1]; + self.read_exact(&mut i8_buf)?; + let i8_value = i8::from_le_bytes(i8_buf); + if i8_value != i8::MAX { return Ok(i8_value as i32) } + + let mut i16_buf = [0; 2]; + self.read_exact(&mut i16_buf)?; + let i16_value = i16::from_le_bytes(i16_buf); + if i16_value != i16::MAX { return Ok(i16_value as i32) } + + let mut i32_buf = [0; 4]; + self.read_exact(&mut i32_buf)?; + let i32_value = i32::from_le_bytes(i32_buf); + Ok(i32_value) + } + + // wrap UnexpectedEof, since that's an error for other Io clients, + // but it's specifically not an error for us + fn read_exact(&mut self, buf: &mut [u8]) -> ProtocolReaderResult<()> { + match self.reader.read_exact(buf) { + Ok(_) => Ok(()), + Err(e) => { + if e.kind() == ErrorKind::UnexpectedEof { + return Err(ProtocolReaderError::EarlyEof) + } + return Err(ProtocolReaderError::Io(e)); + } + } + } +} \ No newline at end of file diff --git a/src/quantization.rs b/src/quantization.rs new file mode 100644 index 0000000..4f516cc --- /dev/null +++ b/src/quantization.rs @@ -0,0 +1,47 @@ +use crate::constants::TILE_SZ; + +const ZIGZAG: [u8; TILE_SZ * TILE_SZ] = [ + 0 , 1 , 5 , 6 , 14, 15, 27, 28, + 2 , 4 , 7 , 13, 16, 26, 29, 42, + 3 , 8 , 12, 17, 25, 30, 41, 43, + 9 , 11, 18, 24, 31, 40, 44, 53, + 10, 19, 23, 32, 39, 45, 52, 54, + 20, 22, 33, 38, 46, 51, 55, 60, + 21, 34, 37, 47, 50, 56, 59, 61, + 35, 36, 48, 49, 57, 58, 62, 63, +]; +const DIVISORS: [u8; TILE_SZ * TILE_SZ] = [ + // source: https://en.wikipedia.org/wiki/Quantization_(image_processing) + 16, 11, 10, 16, 24, 40, 51, 61, + 12, 12, 14, 19, 26, 58, 60, 55, + 14, 13, 16, 24, 40, 57, 69, 56, + 14, 17, 22, 29, 51, 87, 80, 62, + 18, 22, 37, 56, 68, 109, 103, 77, + 24, 35, 55, 64, 81, 104, 113, 92, + 49, 64, 78, 87, 103, 121, 120, 101, + 72, 92, 95, 98, 112, 100, 103, 99, +]; + +pub fn to_quantized( + coefs: [i32; TILE_SZ * TILE_SZ] +) -> [i32; TILE_SZ * TILE_SZ] { + let mut quant: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ]; + + for cf_ix in 0..TILE_SZ * TILE_SZ { + quant[ZIGZAG[cf_ix] as usize] = + coefs[cf_ix] + // (coefs[cf_ix] + DIVISORS[cf_ix] as i32 / 2) / + // (DIVISORS[cf_ix] as i32); + } + quant +} + +pub fn from_quantized( + quant: [i32; TILE_SZ* TILE_SZ] +) -> [i32; TILE_SZ * TILE_SZ] { + let mut coefs: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ]; + for cf_ix in 0..TILE_SZ * TILE_SZ { + coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize]; // * DIVISORS[cf_ix] as i32; + } + coefs +} \ No newline at end of file