rximg/src/decompression.rs
2024-04-07 20:30:35 -07:00

130 lines
3.5 KiB
Rust

use std::io::Read;
use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform};
struct PixelTile {
// i32: representation that supports Walsh-Hadamard
pixels: [i16; TILE_SZ2]
}
#[derive(Clone)]
struct CoefTile {
coefs: [i16; TILE_SZ2]
}
#[derive(Clone)]
struct QuantTile {
quants: [i16; TILE_SZ2]
}
pub fn decompress<R: Read>(
reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> {
// read header
let header = reader.read_header()?;
if header.bytes != MAGIC.bytes {
return Err(ProtocolReaderError::WrongHeader);
}
if header.version != MAGIC.version {
return Err(ProtocolReaderError::WrongVersion);
}
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![QuantTile::new(); n_tiles];
for i in 0..n_tiles {
tiles[i].load_zero(reader)?;
}
for i in 0..n_tiles {
tiles[i].load_rest(reader)?;
}
let mut tile_i = 0;
let mut layers = vec![vec![0; width as usize * height as usize]; n_layers as usize];
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_coef_tile().to_pixel_tile();
tile_i += 1;
pixel_tile.to_layer(
x0 as usize, y0 as usize, &mut layers[layer as usize],
width as usize, height as usize
);
}
}
}
return Ok((width, height, layers))
}
impl CoefTile {
fn to_pixel_tile(&self) -> PixelTile {
let mut pixels = self.coefs.clone();
// columns
for x in 0..TILE_SZ {
transform::decode(&mut pixels, x, 8);
}
// rows
for y in 0..TILE_SZ {
transform::decode(&mut pixels, y * TILE_SZ, 1);
}
return PixelTile { pixels };
}
}
impl PixelTile {
fn to_layer(&self, x0: usize, y0: usize, layer: &mut [u8], width: usize, height: usize) {
for x in 0..TILE_SZ {
for y in 0..TILE_SZ {
let dst_x = x0 + x;
let dst_y = y0 + y;
if dst_x < width && dst_y < height {
layer[dst_y * width + dst_x] = self.clip(self.pixels[y * TILE_SZ + x]);
}
}
}
}
fn clip(&self, x: i16) -> u8 {
if x > u8::MAX as i16 {
return u8::MAX
}
if x < u8::MIN as i16 {
return u8::MIN;
}
return x as u8;
}
}
impl QuantTile {
fn new() -> QuantTile {
QuantTile { quants: [0; TILE_SZ2] }
}
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_i16_wide()?;
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_i16_packed()?;
}
Ok(())
}
}