Fixes to bugs involving transform code

This commit is contained in:
Pyrex 2024-04-07 20:17:36 -07:00
parent f61ec3dddc
commit f76be1b112
6 changed files with 69 additions and 44 deletions

View File

@ -1,18 +1,18 @@
use std::io::Write; 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 { struct PixelTile {
// i32: representation that supports Walsh-Hadamard // i32: representation that supports Walsh-Hadamard
pixels: [i32; TILE_SZ * TILE_SZ] pixels: [i32; TILE_SZ2]
} }
struct CoefTile { struct CoefTile {
coefs: [i32; TILE_SZ * TILE_SZ] coefs: [i32; TILE_SZ2]
} }
struct QuantTile { struct QuantTile {
quants: [i32; TILE_SZ * TILE_SZ] quants: [i32; TILE_SZ2]
} }
@ -68,7 +68,7 @@ impl PixelTile {
width: usize, width: usize,
height: usize height: usize
) -> PixelTile { ) -> PixelTile {
let mut pixels = [0; TILE_SZ * TILE_SZ]; let mut pixels = [0; TILE_SZ2];
for x in 0..TILE_SZ { for x in 0..TILE_SZ {
for y in 0..TILE_SZ { for y in 0..TILE_SZ {
let src_x = x0 + x; let src_x = x0 + x;
@ -97,7 +97,7 @@ impl CoefTile {
// columns // columns
for x in 0..TILE_SZ { for x in 0..TILE_SZ {
transform::encode(&mut coefs, x * TILE_SZ, 1); transform::encode(&mut coefs, x, 8);
} }
return CoefTile { coefs } return CoefTile { coefs }

View File

@ -4,4 +4,5 @@ pub const MAGIC: Header = Header {
bytes: *b"rxi", bytes: *b"rxi",
version: 0 version: 0
}; };
pub const TILE_SZ: usize = 8; pub const TILE_SZ: usize = 8;
pub const TILE_SZ2: usize = TILE_SZ * TILE_SZ;

View File

@ -1,20 +1,20 @@
use std::io::Read; 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 { struct PixelTile {
// i32: representation that supports Walsh-Hadamard // i32: representation that supports Walsh-Hadamard
pixels: [i32; TILE_SZ * TILE_SZ] pixels: [i32; TILE_SZ2]
} }
#[derive(Clone)] #[derive(Clone)]
struct CoefTile { struct CoefTile {
coefs: [i32; TILE_SZ * TILE_SZ] coefs: [i32; TILE_SZ2]
} }
#[derive(Clone)] #[derive(Clone)]
struct QuantTile { struct QuantTile {
quants: [i32; TILE_SZ * TILE_SZ] quants: [i32; TILE_SZ2]
} }
pub fn decompress<R: Read>( pub fn decompress<R: Read>(
@ -68,7 +68,7 @@ impl CoefTile {
// columns // columns
for x in 0..TILE_SZ { for x in 0..TILE_SZ {
transform::decode(&mut pixels, x * TILE_SZ, 1); transform::decode(&mut pixels, x, 8);
} }
// rows // rows
@ -87,16 +87,26 @@ impl PixelTile {
let dst_x = x0 + x; let dst_x = x0 + x;
let dst_y = y0 + y; let dst_y = y0 + y;
if dst_x < width && dst_y < height { 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 { impl QuantTile {
fn new() -> QuantTile { fn new() -> QuantTile {
QuantTile { quants: [0; TILE_SZ * TILE_SZ] } QuantTile { quants: [0; TILE_SZ2] }
} }
fn to_coef_tile(&self) -> CoefTile { fn to_coef_tile(&self) -> CoefTile {

View File

@ -15,7 +15,10 @@ mod transform;
fn main() { fn main() {
for chunk in vec![ for chunk in vec![
[0, 0, 0, 0, 0, 0, 0, 0], [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 orig = chunk;
let mut enc = chunk.clone(); let mut enc = chunk.clone();

View File

@ -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, 0 , 1 , 5 , 6 , 14, 15, 27, 28,
2 , 4 , 7 , 13, 16, 26, 29, 42, 2 , 4 , 7 , 13, 16, 26, 29, 42,
3 , 8 , 12, 17, 25, 30, 41, 43, 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, 21, 34, 37, 47, 50, 56, 59, 61,
35, 36, 48, 49, 57, 58, 62, 63, 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( pub fn to_quantized(
coefs: [i32; TILE_SZ * TILE_SZ] coefs: [i32; TILE_SZ2]
) -> [i32; TILE_SZ * TILE_SZ] { ) -> [i32; TILE_SZ2] {
let mut quant: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ]; let mut quant: [i32; TILE_SZ2] = [0; TILE_SZ2];
for cf_ix in 0..TILE_SZ * TILE_SZ { for cf_ix in 0..TILE_SZ2 {
quant[ZIGZAG[cf_ix] as usize] = let div = divisor(cf_ix);
coefs[cf_ix] let qval = (coefs[cf_ix] + div / 2) / div;
// (coefs[cf_ix] + DIVISORS[cf_ix] as i32 / 2) / quant[ZIGZAG[cf_ix] as usize] = qval
// (DIVISORS[cf_ix] as i32);
} }
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 quant
} }
pub fn from_quantized( pub fn from_quantized(
quant: [i32; TILE_SZ* TILE_SZ] quant: [i32; TILE_SZ2]
) -> [i32; TILE_SZ * TILE_SZ] { ) -> [i32; TILE_SZ2] {
let mut coefs: [i32; TILE_SZ * TILE_SZ] = [0; TILE_SZ * TILE_SZ]; let mut coefs: [i32; TILE_SZ2] = [0; TILE_SZ2];
for cf_ix in 0..TILE_SZ * TILE_SZ { for cf_ix in 0..TILE_SZ2 {
coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize]; // * DIVISORS[cf_ix] as i32; let div = divisor(cf_ix);
coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize] * div;
} }
coefs 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;
} }

View File

@ -4,11 +4,11 @@ use crate::constants::TILE_SZ;
// Ported from: // Ported from:
// - https://colab.research.google.com/drive/1WjtKwUcqxWafAumFO9ET74RsckZSsw6n#scrollTo=e3r9-RBpcgwH // - https://colab.research.google.com/drive/1WjtKwUcqxWafAumFO9ET74RsckZSsw6n#scrollTo=e3r9-RBpcgwH
pub fn encode(data: &mut [i32], zero: usize, stride: usize) { 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) { 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) { 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 // reorder according to gray code
// basically, this moves important stuff towards the left // basically, this moves important stuff towards the left
const GRAY: [usize; TILE_SZ] = [0, 4, 6, 2, 3, 7, 5, 1]; const GRAY: [usize; TILE_SZ] = [0, 4, 6, 2, 3, 7, 5, 1];
for i in 0..TILE_SZ { for i in 0..TILE_SZ {
data[ix(GRAY[i])] = row[i] / divisor + offset_out; data[ix(i)] = row[GRAY[i]] / divisor + offset_out;
} }
} }