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

View File

@ -5,3 +5,4 @@ pub const MAGIC: Header = Header {
version: 0
};
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 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<R: Read>(
@ -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 {

View File

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

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,
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;
}

View File

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