pyrex/rximg
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Zigzagging works, but uncommetning this quantization logic breaks things

Browse Source
This commit is contained in:
Pyrex 2024-04-07 18:42:52 -07:00
parent 7de68f03bf
commit f61ec3dddc
8 changed files with 347 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
Cargo.lock generated
View File

@ -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"

1
Cargo.toml
View File

@ -8,3 +8,4 @@ edition = "2021"
[dependencies]
png = "0.17.13"
quickcheck = "1.0.3"
thiserror = "1.0.58"

56
src/compression.rs
View File

@ -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<W: Write>(
width: u32,
height: u32,
// n_components: u16,
layers: &[&[u8]],
) -> Vec<u8> {
let mut output = vec![];
writer: &mut protocol::ProtocolWriter<W>
) -> 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<u8>) {
output.write(&self.coefs[0].to_le_bytes()).unwrap();
}
fn write_rest(&self, output: &mut Vec<u8>) {
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<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
writer.write_i32_packed(self.quants[0])?;
Ok(())
}
fn write_rest<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
for i in 1..self.quants.len() {
writer.write_i32_packed(self.quants[i])?;
}
Ok(())
}
}

7
src/constants.rs
View File

@ -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;

89
src/decompression.rs
View File

@ -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<Vec<u8>>) {
let mut reader = BufReader::new(input);
#[derive(Clone)]
struct QuantTile {
quants: [i32; TILE_SZ * TILE_SZ]
}
pub fn decompress<R: Read>(
reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> {
// 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();
@ -109,3 +93,28 @@ 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<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
self.quants[0] = reader.read_i32_packed()?;
Ok(())
}
fn load_rest<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
for i in 1..self.quants.len() {
self.quants[i] = reader.read_i32_packed()?;
}
Ok(())
}
}

19
src/main.rs
View File

@ -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 mut reader = ProtocolReader::new(Cursor::new(compressed));
let (width2, height2, decompressed) =
decompression::decompress(&compressed);
decompression::decompress(&mut reader).unwrap();
assert_eq!(3, decompressed.len());
save_image(width2 as usize, height2 as usize, decompressed);

138
src/protocol.rs Normal file
View File

@ -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<T> = Result<T, ProtocolWriterError>;
pub struct ProtocolWriter<W: Write> {
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<T> = Result<T, ProtocolReaderError>;
pub struct ProtocolReader<R: Read> {
reader: R
}
impl<W: Write> ProtocolWriter<W> {
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<R: Read> ProtocolReader<R> {
pub fn new(reader: R) -> Self {
Self { reader }
}
pub fn read_header(&mut self) -> ProtocolReaderResult<Header> {
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<u32> {
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<i32> {
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));
}
}
}
}

47
src/quantization.rs Normal file
View File

@ -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
}