125 lines
3.1 KiB
Rust
125 lines
3.1 KiB
Rust
use std::io::Write;
|
|
|
|
use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform};
|
|
|
|
|
|
struct PixelTile {
|
|
// i32: representation that supports Walsh-Hadamard
|
|
pixels: [i16; TILE_SZ2]
|
|
}
|
|
struct CoefTile {
|
|
coefs: [i16; TILE_SZ2]
|
|
}
|
|
|
|
struct QuantTile {
|
|
quants: [i16; TILE_SZ2]
|
|
}
|
|
|
|
|
|
pub fn compress<W: Write>(
|
|
width: u32,
|
|
height: u32,
|
|
// n_components: u16,
|
|
layers: &[&[u8]],
|
|
writer: &mut protocol::ProtocolWriter<W>
|
|
) -> ProtocolWriterResult<()> {
|
|
|
|
// validation
|
|
for l in 0..layers.len() {
|
|
assert!(layers[l].len() == width as usize * height as usize);
|
|
}
|
|
|
|
// write header
|
|
writer.write_header(MAGIC)?;
|
|
writer.write_u32_wide(width)?;
|
|
writer.write_u32_wide(height)?;
|
|
writer.write_u32_wide(layers.len() as u32)?;
|
|
|
|
let mut tiles = vec![];
|
|
for layer in layers.iter() {
|
|
for x0 in (0..width).step_by(TILE_SZ) {
|
|
for y0 in (0..height).step_by(TILE_SZ) {
|
|
let pixel_tile = PixelTile::from_layer(
|
|
x0 as usize, y0 as usize, layer,
|
|
width as usize, height as usize
|
|
);
|
|
let coef_tile = CoefTile::from_pixel_tile(&pixel_tile);
|
|
let quant_tile = QuantTile::from_coef_tile(&coef_tile);
|
|
tiles.push(quant_tile);
|
|
}
|
|
}
|
|
}
|
|
|
|
for t in tiles.iter() {
|
|
t.write_zero(writer)?;
|
|
}
|
|
for t in tiles.iter() {
|
|
t.write_rest(writer)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
impl PixelTile {
|
|
fn from_layer(
|
|
x0: usize,
|
|
y0: usize,
|
|
layer: &[u8],
|
|
width: usize,
|
|
height: usize
|
|
) -> PixelTile {
|
|
let mut pixels = [0; TILE_SZ2];
|
|
for x in 0..TILE_SZ {
|
|
for y in 0..TILE_SZ {
|
|
let src_x = x0 + x;
|
|
let src_y = y0 + y;
|
|
pixels[y * TILE_SZ + x] =
|
|
if src_x < width && src_y < height {
|
|
layer[src_y * width + src_x] as i16
|
|
} else {
|
|
0
|
|
};
|
|
}
|
|
}
|
|
|
|
return PixelTile { pixels };
|
|
}
|
|
}
|
|
|
|
impl CoefTile {
|
|
fn from_pixel_tile(pt: &PixelTile) -> CoefTile {
|
|
let mut coefs = pt.pixels.clone();
|
|
|
|
// rows
|
|
for y in 0..TILE_SZ {
|
|
transform::encode(&mut coefs, y * TILE_SZ, 1);
|
|
}
|
|
|
|
// columns
|
|
for x in 0..TILE_SZ {
|
|
transform::encode(&mut coefs, x, 8);
|
|
}
|
|
|
|
return CoefTile { coefs }
|
|
}
|
|
}
|
|
|
|
impl QuantTile {
|
|
fn from_coef_tile(pt: &CoefTile) -> QuantTile {
|
|
QuantTile {
|
|
quants: quantization::to_quantized(pt.coefs)
|
|
}
|
|
}
|
|
|
|
fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
|
|
writer.write_i16_wide(self.quants[0])?;
|
|
Ok(())
|
|
}
|
|
|
|
fn write_rest<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
|
|
for i in 1..self.quants.len() {
|
|
writer.write_i16_packed(self.quants[i])?;
|
|
}
|
|
Ok(())
|
|
}
|
|
} |