From f76be1b112a1d2d6df07da38f47dee9bbbc1a410 Mon Sep 17 00:00:00 2001 From: Nyeogmi Date: Sun, 7 Apr 2024 20:17:36 -0700 Subject: [PATCH] Fixes to bugs involving transform code --- src/compression.rs | 12 ++++----- src/constants.rs | 3 ++- src/decompression.rs | 24 ++++++++++++----- src/main.rs | 5 +++- src/quantization.rs | 62 +++++++++++++++++++++++++------------------- src/transform.rs | 7 ++--- 6 files changed, 69 insertions(+), 44 deletions(-) diff --git a/src/compression.rs b/src/compression.rs index 31f3263..8f9bc72 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -1,18 +1,18 @@ use std::io::Write; -use crate::{constants::{MAGIC, TILE_SZ}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform}; +use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform}; struct PixelTile { // i32: representation that supports Walsh-Hadamard - pixels: [i32; TILE_SZ * TILE_SZ] + pixels: [i32; TILE_SZ2] } struct CoefTile { - coefs: [i32; TILE_SZ * TILE_SZ] + coefs: [i32; TILE_SZ2] } struct QuantTile { - quants: [i32; TILE_SZ * TILE_SZ] + quants: [i32; TILE_SZ2] } @@ -68,7 +68,7 @@ impl PixelTile { width: usize, height: usize ) -> PixelTile { - let mut pixels = [0; TILE_SZ * TILE_SZ]; + let mut pixels = [0; TILE_SZ2]; for x in 0..TILE_SZ { for y in 0..TILE_SZ { let src_x = x0 + x; @@ -97,7 +97,7 @@ impl CoefTile { // columns for x in 0..TILE_SZ { - transform::encode(&mut coefs, x * TILE_SZ, 1); + transform::encode(&mut coefs, x, 8); } return CoefTile { coefs } diff --git a/src/constants.rs b/src/constants.rs index f53238b..2164936 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -4,4 +4,5 @@ pub const MAGIC: Header = Header { bytes: *b"rxi", version: 0 }; -pub const TILE_SZ: usize = 8; \ No newline at end of file +pub const TILE_SZ: usize = 8; +pub const TILE_SZ2: usize = TILE_SZ * TILE_SZ; \ No newline at end of file diff --git a/src/decompression.rs b/src/decompression.rs index ac37012..82d988f 100644 --- a/src/decompression.rs +++ b/src/decompression.rs @@ -1,20 +1,20 @@ use std::io::Read; -use crate::{constants::{MAGIC, TILE_SZ}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform}; +use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform}; struct PixelTile { // i32: representation that supports Walsh-Hadamard - pixels: [i32; TILE_SZ * TILE_SZ] + pixels: [i32; TILE_SZ2] } #[derive(Clone)] struct CoefTile { - coefs: [i32; TILE_SZ * TILE_SZ] + coefs: [i32; TILE_SZ2] } #[derive(Clone)] struct QuantTile { - quants: [i32; TILE_SZ * TILE_SZ] + quants: [i32; TILE_SZ2] } pub fn decompress( @@ -68,7 +68,7 @@ impl CoefTile { // columns for x in 0..TILE_SZ { - transform::decode(&mut pixels, x * TILE_SZ, 1); + transform::decode(&mut pixels, x, 8); } // rows @@ -87,16 +87,26 @@ impl PixelTile { let dst_x = x0 + x; let dst_y = y0 + y; if dst_x < width && dst_y < height { - layer[dst_y * width + dst_x] = self.pixels[y * TILE_SZ + x] as u8; + layer[dst_y * width + dst_x] = self.clip(self.pixels[y * TILE_SZ + x]); } } } } + + fn clip(&self, x: i32) -> u8 { + if x > u8::MAX as i32 { + return u8::MAX + } + if x < u8::MIN as i32 { + return u8::MIN; + } + return x as u8; + } } impl QuantTile { fn new() -> QuantTile { - QuantTile { quants: [0; TILE_SZ * TILE_SZ] } + QuantTile { quants: [0; TILE_SZ2] } } fn to_coef_tile(&self) -> CoefTile { diff --git a/src/main.rs b/src/main.rs index 6272f8d..6892323 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,7 +15,10 @@ mod transform; fn main() { for chunk in vec![ [0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 2, 3, 4, 5, 6, 7] + [0, 1, 2, 3, 4, 5, 6, 7], + [255, 255, 255, 255, 255, 255, 255, 80], + [255, 255, 255, 255, 255, 255, 255, 0], + [255, 253, 254, 252, 251, 252, 255, 254] ] { let orig = chunk; let mut enc = chunk.clone(); diff --git a/src/quantization.rs b/src/quantization.rs index 4f516cc..6e0a8d0 100644 --- a/src/quantization.rs +++ b/src/quantization.rs @@ -1,6 +1,6 @@ -use crate::constants::TILE_SZ; +use crate::constants::TILE_SZ2; -const ZIGZAG: [u8; TILE_SZ * TILE_SZ] = [ +const ZIGZAG: [u8; TILE_SZ2] = [ 0 , 1 , 5 , 6 , 14, 15, 27, 28, 2 , 4 , 7 , 13, 16, 26, 29, 42, 3 , 8 , 12, 17, 25, 30, 41, 43, @@ -10,38 +10,48 @@ const ZIGZAG: [u8; TILE_SZ * TILE_SZ] = [ 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]; + coefs: [i32; TILE_SZ2] +) -> [i32; TILE_SZ2] { + let mut quant: [i32; TILE_SZ2] = [0; TILE_SZ2]; - 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); + for cf_ix in 0..TILE_SZ2 { + let div = divisor(cf_ix); + let qval = (coefs[cf_ix] + div / 2) / div; + quant[ZIGZAG[cf_ix] as usize] = qval } + + let mut indices = [0; TILE_SZ2]; + for i in 0..indices.len() { + indices[i] = i; + } + indices.sort_by_key(|i| -quant[*i].abs()); + for i in 4..indices.len() { + quant[indices[i]] = 0; + } + 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; + quant: [i32; TILE_SZ2] +) -> [i32; TILE_SZ2] { + let mut coefs: [i32; TILE_SZ2] = [0; TILE_SZ2]; + for cf_ix in 0..TILE_SZ2 { + let div = divisor(cf_ix); + coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize] * div; } coefs +} + +pub fn divisor(cf_ix: usize) -> i32 { + if cf_ix == 0 { return 1; } + let x = cf_ix % 8; + let y = cf_ix / 8; + let div = 32 + (x as i32 + y as i32) * 12; + if div==32 && cf_ix != 0 { + dbg!(cf_ix, x, y, div); + } + return div; } \ No newline at end of file diff --git a/src/transform.rs b/src/transform.rs index 3c71400..9536300 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -4,11 +4,11 @@ use crate::constants::TILE_SZ; // Ported from: // - https://colab.research.google.com/drive/1WjtKwUcqxWafAumFO9ET74RsckZSsw6n#scrollTo=e3r9-RBpcgwH pub fn encode(data: &mut [i32], zero: usize, stride: usize) { - transform(data, zero, stride, -128, 0, 1) + transform(data, zero, stride, 0, 0, 1) } pub fn decode(data: &mut [i32], zero: usize, stride: usize) { - transform(data, zero, stride, 0, 128, 8) + transform(data, zero, stride, 0, 0, 8) } fn transform(data: &mut [i32], zero: usize, stride: usize, offset_in: i32, offset_out: i32, divisor: i32) { @@ -34,8 +34,9 @@ fn transform(data: &mut [i32], zero: usize, stride: usize, offset_in: i32, offse // reorder according to gray code // basically, this moves important stuff towards the left const GRAY: [usize; TILE_SZ] = [0, 4, 6, 2, 3, 7, 5, 1]; + for i in 0..TILE_SZ { - data[ix(GRAY[i])] = row[i] / divisor + offset_out; + data[ix(i)] = row[GRAY[i]] / divisor + offset_out; } }