Fixes to bugs involving transform code

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 }

src/constants.rs

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

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

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

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]
+    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
 }
 
 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
 }
+
+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;
+}

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