Compare commits

...

4 Commits

Author SHA1 Message Date
f047642f37 TODOs complete. More TODOs? 2024-04-07 23:31:22 -07:00
de10e4b4b8 Better PNG interface 2024-04-07 22:52:20 -07:00
9505b5754d Select quality 2024-04-07 22:15:24 -07:00
e87b406240 More TODOs 2024-04-07 21:01:24 -07:00
10 changed files with 404 additions and 170 deletions

12
TODO.md
View File

@@ -1,9 +1,19 @@
TODO:
I still think we can interleave layers to get better progressive loading. But we need FC, FD, FE, and FF commands for "rerun 4/3/2/1 tiles ago."
Make a super simple one corresponding to quality 3 or so, controlled by constants?
Generally improve code quality.
Done:
Actually support terminating early for "progressive" loads
Support an alpha channel. Support an alpha channel.
Improve the PNG interface. Improve the PNG interface.
Interleave layers instead of having them consecutively. (Bad idea, breaks RLE. Skipping)
Global parameters to control encoder settings: Global parameters to control encoder settings:
- Quantization level: quant 0 - Quantization level: quant 0
- Quantization level: quants 1-63 - Quantization level: quants 1-63
- Constant term - Constant term
- Linear term - Linear term
- Number of quants to keep, per tile. - Number of quants to keep, per tile.
- Whether to use wide or packed representation for the thumbnail block - Whether to use wide or packed representation for the thumbnail block
Maybe have a second instance of the global parameters block for the alpha channel.

BIN
inputs/zonked.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 518 KiB

View File

@@ -1,48 +1,53 @@
use std::io::Write; use std::io::Write;
use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quantization, transform}; use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, image::Image, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quality_settings::{self, QualitySettings}, quantization, transform};
struct PixelTile { struct PixelTile {
// i32: representation that supports Walsh-Hadamard quality: QualitySettings,
pixels: [i16; TILE_SZ2] pixels: [i16; TILE_SZ2]
} }
struct CoefTile { struct CoefTile {
quality: QualitySettings,
coefs: [i16; TILE_SZ2] coefs: [i16; TILE_SZ2]
} }
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
struct QuantTile { struct QuantTile {
quality: QualitySettings,
quants: [i16; TILE_SZ2] quants: [i16; TILE_SZ2]
} }
pub fn compress<W: Write>( pub fn compress<W: Write>(
width: u32, image: Image,
height: u32,
// n_components: u16,
layers: &[&[u8]],
writer: &mut protocol::ProtocolWriter<W> writer: &mut protocol::ProtocolWriter<W>
) -> ProtocolWriterResult<()> { ) -> ProtocolWriterResult<()> {
// validation // validation
for l in 0..layers.len() { for l in 0..image.layers.len() {
assert!(layers[l].len() == width as usize * height as usize); assert!(image.layers[l].1.len() == image.width as usize * image.height as usize);
} }
// write header // write header
writer.write_header(MAGIC)?; writer.write_header(MAGIC)?;
writer.write_u32_wide(width)?; writer.write_u32_wide(image.width as u32)?;
writer.write_u32_wide(height)?; writer.write_u32_wide(image.height as u32)?;
writer.write_u32_wide(layers.len() as u32)?; writer.write_u32_wide(image.layers.len() as u32)?;
// write quality settings for each layout
for (quality, _) in image.layers.iter() {
writer.write_quality_settings(*quality)?;
}
// build list of tiles
let mut tiles = vec![]; let mut tiles = vec![];
for layer in layers.iter() { for (quality, layer) in image.layers.iter() {
for x0 in (0..width).step_by(TILE_SZ) { for y0 in (0..image.height).step_by(TILE_SZ) {
for y0 in (0..height).step_by(TILE_SZ) { for x0 in (0..image.width).step_by(TILE_SZ) {
let pixel_tile = PixelTile::from_layer( let pixel_tile = PixelTile::from_layer(
x0 as usize, y0 as usize, layer, x0 as usize, y0 as usize, *quality, layer,
width as usize, height as usize image.width as usize, image.height as usize
); );
let coef_tile = CoefTile::from_pixel_tile(&pixel_tile); let coef_tile = CoefTile::from_pixel_tile(&pixel_tile);
let quant_tile = QuantTile::from_coef_tile(&coef_tile); let quant_tile = QuantTile::from_coef_tile(&coef_tile);
@@ -51,13 +56,16 @@ pub fn compress<W: Write>(
} }
} }
// write the thumbnail block
for t in tiles.iter() { for t in tiles.iter() {
t.write_zero(writer)?; t.write_zero(writer)?;
} }
for ti in 0..tiles.len() {
let prev = if ti > 0 { Some(tiles[ti - 1]) } else { None }; // write the tiles
let t = tiles[ti]; let mut prev: Option<QuantTile> = None;
for t in tiles {
t.write_rest(prev, writer)?; t.write_rest(prev, writer)?;
prev = Some(t);
} }
Ok(()) Ok(())
@@ -67,13 +75,14 @@ impl PixelTile {
fn from_layer( fn from_layer(
x0: usize, x0: usize,
y0: usize, y0: usize,
quality: QualitySettings,
layer: &[u8], layer: &[u8],
width: usize, width: usize,
height: usize height: usize
) -> PixelTile { ) -> PixelTile {
let mut pixels = [0; TILE_SZ2]; let mut pixels = [0; TILE_SZ2];
for x in 0..TILE_SZ { for y in 0..TILE_SZ {
for y in 0..TILE_SZ { for x in 0..TILE_SZ {
let src_x = x0 + x; let src_x = x0 + x;
let src_y = y0 + y; let src_y = y0 + y;
pixels[y * TILE_SZ + x] = pixels[y * TILE_SZ + x] =
@@ -85,7 +94,7 @@ impl PixelTile {
} }
} }
return PixelTile { pixels }; return PixelTile { quality, pixels };
} }
} }
@@ -103,19 +112,24 @@ impl CoefTile {
transform::encode(&mut coefs, x, 8); transform::encode(&mut coefs, x, 8);
} }
return CoefTile { coefs } return CoefTile { quality: pt.quality, coefs }
} }
} }
impl QuantTile { impl QuantTile {
fn from_coef_tile(pt: &CoefTile) -> QuantTile { fn from_coef_tile(pt: &CoefTile) -> QuantTile {
QuantTile { QuantTile {
quants: quantization::to_quantized(pt.coefs) quality: pt.quality,
quants: quantization::to_quantized(pt.quality, pt.coefs)
} }
} }
fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> { fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
writer.write_i16_wide(self.quants[0])?; if self.quality.values[quality_settings::THUMBNAIL_IS_WIDE] != 0 {
writer.write_i16_wide(self.quants[0])?;
} else {
writer.write_i16_packed(self.quants[0])?;
}
Ok(()) Ok(())
} }

View File

@@ -1,9 +1,8 @@
use std::io::Read; use std::io::{ErrorKind, Read};
use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform}; use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, image::Image, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quality_settings::{self, QualitySettings}, quantization, transform};
struct PixelTile { struct PixelTile {
// i32: representation that supports Walsh-Hadamard
pixels: [i16; TILE_SZ2] pixels: [i16; TILE_SZ2]
} }
@@ -14,46 +13,102 @@ struct CoefTile {
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
struct QuantTile { struct QuantTile {
quality: QualitySettings,
quants: [i16; TILE_SZ2] quants: [i16; TILE_SZ2]
} }
pub enum ProgressiveImage {
Present { image: Image, complete: bool },
NotHereYet
}
pub fn decompress<R: Read>( pub fn decompress<R: Read>(
reader: &mut ProtocolReader<R> reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> { ) -> ProtocolReaderResult<ProgressiveImage> {
// read header // read header
let header = reader.read_header()?; let read_header_and_stuff: Result<_, ProtocolReaderError> = (|| {
if header.bytes != MAGIC.bytes { let header = reader.read_header()?;
return Err(ProtocolReaderError::WrongHeader); if header.bytes != MAGIC.bytes {
} return Err(ProtocolReaderError::WrongHeader);
if header.version != MAGIC.version { }
return Err(ProtocolReaderError::WrongVersion); if header.version != MAGIC.version {
} return Err(ProtocolReaderError::WrongVersion);
}
let width = reader.read_u32_wide()?; let width = reader.read_u32_wide()?;
let height = reader.read_u32_wide()?; let height = reader.read_u32_wide()?;
let n_layers = reader.read_u32_wide()?; let n_layers = reader.read_u32_wide()?;
// read quality settings for each layout
let mut quality_settings = vec![];
for _ in 0..n_layers {
quality_settings.push(reader.read_quality_settings()?);
}
Ok((width, height, n_layers, quality_settings))
})();
let Ok((width, height, n_layers, quality_settings)) = read_header_and_stuff else {
match read_header_and_stuff {
Ok(_) => unreachable!(),
Err(ProtocolReaderError::Io(io)) => {
if io.kind() == ErrorKind::UnexpectedEof {
return ProtocolReaderResult::Ok(ProgressiveImage::NotHereYet)
}
return Err(ProtocolReaderError::Io(io))
}
Err(e) => return Err(e),
}
};
let width_in_tiles = (width as usize + TILE_SZ - 1) / TILE_SZ; 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 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 n_tiles_per_layer = width_in_tiles * height_in_tiles;
let mut tiles = vec![QuantTile::new(); n_tiles];
for i in 0..n_tiles { let mut tiles = vec![QuantTile::new(); n_tiles_per_layer * n_layers as usize];
tiles[i].load_zero(reader)?; let read_tiles_and_stuff: Result<(), ProtocolReaderError> = (|tiles: &mut [QuantTile]| {
} // read thumbnail block
for i in 0..n_tiles { for l in 0..n_layers {
let prev = if i > 0 { Some(tiles[i - 1]) } else { None }; for i in 0..n_tiles_per_layer {
tiles[i].load_rest( tiles[l as usize * n_tiles_per_layer + i].load_zero(
prev, quality_settings[l as usize], reader
reader )?;
)?; }
} }
// read remaining tiles
let mut prev: Option<QuantTile> = None;
for i in 0..n_tiles_per_layer * n_layers as usize {
tiles[i].load_rest(
prev,
reader
)?;
prev = Some(tiles[i]);
}
Ok(())
})(&mut tiles);
let complete;
match read_tiles_and_stuff {
Err(ProtocolReaderError::Io(i)) => {
if i.kind() == ErrorKind::UnexpectedEof {
complete = false
} else {
return Err(ProtocolReaderError::Io(i))
}
}
Err(e) => return Err(e),
Ok(()) => {
complete = true;
}
};
let mut tile_i = 0; let mut tile_i = 0;
let mut layers = vec![vec![0; width as usize * height as usize]; n_layers as usize]; let mut layers = vec![vec![0; width as usize * height as usize]; n_layers as usize];
for layer in 0..n_layers { for layer in 0..n_layers {
for x0 in (0..width).step_by(TILE_SZ) { for y0 in (0..height).step_by(TILE_SZ) {
for y0 in (0..height).step_by(TILE_SZ) { for x0 in (0..width).step_by(TILE_SZ) {
let pixel_tile = tiles[tile_i].to_coef_tile().to_pixel_tile(); let pixel_tile = tiles[tile_i].to_coef_tile().to_pixel_tile();
tile_i += 1; tile_i += 1;
pixel_tile.to_layer( pixel_tile.to_layer(
@@ -63,7 +118,13 @@ pub fn decompress<R: Read>(
} }
} }
} }
return Ok((width, height, layers))
let mut image = Image::new(width as usize, height as usize);
for (l, layer) in layers.into_iter().enumerate() {
image.add(quality_settings[l], layer).expect("by construction, this layer should be OK");
}
return Ok(ProgressiveImage::Present { image: image, complete });
} }
impl CoefTile { impl CoefTile {
@@ -86,8 +147,8 @@ impl CoefTile {
impl PixelTile { impl PixelTile {
fn to_layer(&self, x0: usize, y0: usize, layer: &mut [u8], width: usize, height: usize) { 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 {
for y in 0..TILE_SZ { for x in 0..TILE_SZ {
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 {
@@ -110,17 +171,29 @@ impl PixelTile {
impl QuantTile { impl QuantTile {
fn new() -> QuantTile { fn new() -> QuantTile {
QuantTile { quants: [0; TILE_SZ2] } QuantTile {
quality: QualitySettings { values: [0, 0, 0, 0, 0, 0, 0, 0] },
quants: [0; TILE_SZ2]
}
} }
fn to_coef_tile(&self) -> CoefTile { fn to_coef_tile(&self) -> CoefTile {
CoefTile { CoefTile {
coefs: quantization::from_quantized(self.quants) coefs: quantization::from_quantized(self.quality, self.quants)
} }
} }
fn load_zero<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> { fn load_zero<R: Read>(
self.quants[0] = reader.read_i16_wide()?; &mut self,
quality: QualitySettings,
reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<()> {
self.quality = quality;
if quality.values[quality_settings::THUMBNAIL_IS_WIDE] != 0 {
self.quants[0] = reader.read_i16_wide()?;
} else {
self.quants[0] = reader.read_i16_packed()?;
}
Ok(()) Ok(())
} }

47
src/image.rs Normal file
View File

@@ -0,0 +1,47 @@
use std::io;
use png::ColorType;
use thiserror::Error;
use crate::quality_settings::QualitySettings;
#[derive(Error, Debug)]
pub enum ImageError {
#[error("can't make the layer fit {0}x{1}")]
WrongSize(usize, usize),
#[error("unsupported color type {0:?}")]
InvalidInputColorType(ColorType),
#[error("can't save an image with {0} layers as png")]
WrongNumberOfLayersToSave(usize),
#[error("general IO error")]
IoError(#[from] io::Error),
#[error("error in encoding PNG")]
PngEncodingError(#[from] png::EncodingError),
#[error("error in decoding PNG")]
PngDecodingError(#[from] png::DecodingError),
}
pub struct Image {
pub width: usize,
pub height: usize,
pub layers: Vec<(QualitySettings, Vec<u8>)>
}
impl Image {
pub fn new(width: usize, height: usize) -> Self {
Image {
width, height,
layers: vec![]
}
}
pub fn add(&mut self, quality: QualitySettings, pixels: Vec<u8>) -> Result<(), ImageError> {
if pixels.len() != self.width * self.height {
return Err(ImageError::WrongSize(self.width, self.height));
}
self.layers.push((quality, pixels));
Ok(())
}
}

View File

@@ -1,100 +1,60 @@
use std::{fs::File, io::{Cursor, Write}}; use std::{fs::File, io::{Cursor, Write}};
use png::{BitDepth, ColorType}; use crate::{protocol::{ProtocolReader, ProtocolWriter}, quality_settings::QualitySettings};
use crate::protocol::{ProtocolReader, ProtocolWriter};
mod compression; mod compression;
mod constants; mod constants;
mod decompression; mod decompression;
mod png_utils;
mod protocol; mod protocol;
mod quality_settings;
mod quantization; mod quantization;
mod transform; mod transform;
mod image;
fn main() { fn main() {
for chunk in vec![ // run_for("zonked".to_string(), 9, 2, Some(330000));
[0, 0, 0, 0, 0, 0, 0, 0], run_for("zonked".to_string(), 9, 2, Some(37000));
[0, 1, 2, 3, 4, 5, 6, 7], run_for("avatar2".to_string(), 9, 2, Some(100000));
[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();
transform::encode(&mut enc, 0, 1);
let mut dec = enc.clone();
transform::decode(&mut dec, 0, 1);
dbg!(orig, enc, dec);
}
hard_main();
} }
fn hard_main() { fn run_for(name: String, quality_rgb: u8, quality_alpha: u8, bytes_to_chop: Option<usize>) {
let (width, height, r, g, b) = load_image(); let image = png_utils::load_image(
format!("inputs/{}.png", name),
QualitySettings::new(quality_rgb),
QualitySettings::new(quality_alpha),
).unwrap();
let mut writer = ProtocolWriter::new(vec![]); let mut writer = ProtocolWriter::new(vec![]);
compression::compress( compression::compress(
width as u32, height as u32, &[&r, &g, &b], image,
&mut writer &mut writer
).unwrap(); ).unwrap();
let compressed = writer.destroy(); let mut compressed = writer.destroy();
let mut output_file = File::create("outputs/avatar2.rxi").unwrap(); if let Some(b2c) = bytes_to_chop {
compressed.drain(b2c..);
}
let mut output_file = File::create(format!("outputs/{}.rxi", name)).unwrap();
output_file.write_all(&compressed).unwrap(); output_file.write_all(&compressed).unwrap();
let mut reader = ProtocolReader::new(Cursor::new(compressed)); let mut reader = ProtocolReader::new(Cursor::new(compressed));
let (width2, height2, decompressed) = let decompressed = decompression::decompress(&mut reader).unwrap();
decompression::decompress(&mut reader).unwrap();
assert_eq!(3, decompressed.len()); let decompressed_image = match decompressed {
save_image(width2 as usize, height2 as usize, decompressed); decompression::ProgressiveImage::Present { image, complete: _ } => {
} image
fn load_image() -> (usize, usize, Vec<u8>, Vec<u8>, Vec<u8>) {
let decoder = png::Decoder::new(File::open("inputs/avatar2.png").unwrap());
let mut reader = decoder.read_info().unwrap();
let mut buf = vec![0; reader.output_buffer_size()];
let info = reader.next_frame(&mut buf).unwrap();
let bytes = &buf[..info.buffer_size()];
let width = reader.info().width;
let height = reader.info().height;
assert_eq!(BitDepth::Eight, reader.info().bit_depth);
assert_eq!(3, reader.info().bytes_per_pixel());
assert_eq!(ColorType::Rgb, reader.info().color_type);
let r: Vec<u8> = bytes[0..].iter().cloned().step_by(3).collect();
let g: Vec<u8> = bytes[1..].iter().cloned().step_by(3).collect();
let b: Vec<u8> = bytes[2..].iter().cloned().step_by(3).collect();
assert_eq!(r.len(), (width * height) as usize);
assert_eq!(g.len(), (width * height) as usize);
assert_eq!(b.len(), (width * height) as usize);
(width as usize, height as usize, r, g, b)
}
fn save_image(width: usize, height: usize, image: Vec<Vec<u8>>) {
assert_eq!(image.len(), 3);
let mut encoder = png::Encoder::new(
File::create("outputs/avatar2_out.png").unwrap(),
width as u32, height as u32,
);
encoder.set_color(ColorType::Rgb);
let mut idata: Vec<u8> = vec![0; width * height * 3];
for i in [0, 1, 2] {
for (dst, src) in
idata[i..].iter_mut().step_by(3).zip(image[i].iter())
{
*dst = *src;
} }
} decompression::ProgressiveImage::NotHereYet => {
// image not here yet
println!("not here yet: {}", name);
return
}
};
let mut writer = encoder.write_header().unwrap(); png_utils::save_image(
writer.write_image_data(&idata).unwrap(); format!("outputs/{}_out.png", name),
decompressed_image
} ).unwrap();
}

83
src/png_utils.rs Normal file
View File

@@ -0,0 +1,83 @@
use std::fs::File;
use png::{BitDepth, ColorType};
use crate::{image::{Image, ImageError}, quality_settings::QualitySettings};
pub fn load_image(filename: String, quality_rgb: QualitySettings, quality_alpha: QualitySettings) -> Result<Image, ImageError> {
let decoder = png::Decoder::new(File::open(filename).unwrap());
let mut reader = decoder.read_info().unwrap();
let mut buf = vec![0; reader.output_buffer_size()];
let info = reader.next_frame(&mut buf).unwrap();
let bytes = &buf[..info.buffer_size()];
assert_eq!(BitDepth::Eight, reader.info().bit_depth);
assert_eq!(ColorType::Rgb, reader.info().color_type);
let width = reader.info().width;
let height = reader.info().height;
match reader.info().color_type {
ColorType::Rgb => {
assert_eq!(3, reader.info().bytes_per_pixel());
let r: Vec<u8> = bytes[0..].iter().cloned().step_by(3).collect();
let g: Vec<u8> = bytes[1..].iter().cloned().step_by(3).collect();
let b: Vec<u8> = bytes[2..].iter().cloned().step_by(3).collect();
let mut image = Image::new(width as usize, height as usize);
image.add(quality_rgb, r)?;
image.add(quality_rgb, g)?;
image.add(quality_rgb, b)?;
return Ok(image);
}
ColorType::Rgba => {
assert_eq!(4, reader.info().bytes_per_pixel());
let r: Vec<u8> = bytes[0..].iter().cloned().step_by(4).collect();
let g: Vec<u8> = bytes[1..].iter().cloned().step_by(4).collect();
let b: Vec<u8> = bytes[2..].iter().cloned().step_by(4).collect();
let a: Vec<u8> = bytes[3..].iter().cloned().step_by(4).collect();
let mut image = Image::new(width as usize, height as usize);
image.add(quality_rgb, r)?;
image.add(quality_rgb, g)?;
image.add(quality_rgb, b)?;
image.add(quality_alpha, a)?;
return Ok(image);
}
ct => {
return Err(ImageError::InvalidInputColorType(ct))
}
}
}
pub fn save_image(filename: String, image: Image) -> Result<(), ImageError> {
let n_layers = image.layers.len();
let color_type = match n_layers {
3 => ColorType::Rgb,
4 => ColorType::Rgba,
_ => return Err(ImageError::WrongNumberOfLayersToSave(n_layers)),
};
let mut encoder = png::Encoder::new(
File::create(filename)?,
image.width as u32, image.height as u32,
);
encoder.set_color(color_type);
let mut idata: Vec<u8> = vec![0; image.width * image.height * n_layers];
for i in 0..n_layers {
for (dst, src) in
idata[i..].iter_mut().step_by(n_layers).zip(image.layers[i].1.iter())
{
*dst = *src;
}
}
let mut writer = encoder.write_header()?;
writer.write_image_data(&idata)?;
Ok(())
}

View File

@@ -1,7 +1,9 @@
use std::io::{self, ErrorKind, Read, Write}; use std::io::{self, Read, Write};
use thiserror::Error; use thiserror::Error;
use crate::quality_settings::QualitySettings;
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
pub struct Header { pub struct Header {
pub bytes: [u8; 3], pub bytes: [u8; 3],
@@ -23,10 +25,6 @@ pub struct ProtocolWriter<W: Write> {
#[derive(Error, Debug)] #[derive(Error, Debug)]
pub enum ProtocolReaderError { pub enum ProtocolReaderError {
// this is explicitly supported: rxi images are progressive
#[error("EOF before end of image")]
EarlyEof,
#[error("wrong header")] #[error("wrong header")]
WrongHeader, WrongHeader,
@@ -57,6 +55,12 @@ impl<W: Write> ProtocolWriter<W> {
self.writer.write_all(&[value.version])?; self.writer.write_all(&[value.version])?;
Ok(()) Ok(())
} }
pub fn write_quality_settings(&mut self, quality: QualitySettings) -> ProtocolWriterResult<()> {
self.writer.write_all(&quality.values)?;
Ok(())
}
pub fn write_u32_wide(&mut self, value: u32) -> ProtocolWriterResult<()> { pub fn write_u32_wide(&mut self, value: u32) -> ProtocolWriterResult<()> {
self.writer.write_all(&value.to_le_bytes())?; self.writer.write_all(&value.to_le_bytes())?;
Ok(()) Ok(())
@@ -89,10 +93,10 @@ impl<R: Read> ProtocolReader<R> {
pub fn read_header(&mut self) -> ProtocolReaderResult<Header> { pub fn read_header(&mut self) -> ProtocolReaderResult<Header> {
let mut bytes_buf = [0; 3]; let mut bytes_buf = [0; 3];
self.read_exact(&mut bytes_buf)?; self.reader.read_exact(&mut bytes_buf)?;
let mut version_buf = [0; 1]; let mut version_buf = [0; 1];
self.read_exact(&mut version_buf)?; self.reader.read_exact(&mut version_buf)?;
Ok(Header { Ok(Header {
bytes: bytes_buf, bytes: bytes_buf,
@@ -100,50 +104,43 @@ impl<R: Read> ProtocolReader<R> {
}) })
} }
pub fn read_quality_settings(&mut self) -> ProtocolReaderResult<QualitySettings> {
let mut qsettings_buf = [0; 8];
self.reader.read_exact(&mut qsettings_buf)?;
Ok(QualitySettings { values: qsettings_buf })
}
pub fn read_u32_wide(&mut self) -> ProtocolReaderResult<u32> { pub fn read_u32_wide(&mut self) -> ProtocolReaderResult<u32> {
let mut u32_buf = [0; 4]; let mut u32_buf = [0; 4];
self.read_exact(&mut u32_buf)?; self.reader.read_exact(&mut u32_buf)?;
Ok(u32::from_le_bytes(u32_buf)) Ok(u32::from_le_bytes(u32_buf))
} }
pub fn read_i16_wide(&mut self) -> ProtocolReaderResult<i16> { pub fn read_i16_wide(&mut self) -> ProtocolReaderResult<i16> {
let mut i16_buf = [0; 2]; let mut i16_buf = [0; 2];
self.read_exact(&mut i16_buf)?; self.reader.read_exact(&mut i16_buf)?;
Ok(i16::from_le_bytes(i16_buf)) Ok(i16::from_le_bytes(i16_buf))
} }
pub fn read_u8_wide(&mut self) -> ProtocolReaderResult<u8> { pub fn read_u8_wide(&mut self) -> ProtocolReaderResult<u8> {
let mut u8_buf = [0; 1]; let mut u8_buf = [0; 1];
self.read_exact(&mut u8_buf)?; self.reader.read_exact(&mut u8_buf)?;
Ok(u8::from_le_bytes(u8_buf)) Ok(u8::from_le_bytes(u8_buf))
} }
pub fn read_i16_packed(&mut self) -> ProtocolReaderResult<i16> { pub fn read_i16_packed(&mut self) -> ProtocolReaderResult<i16> {
let mut i8_buf = [0; 1]; let mut i8_buf = [0; 1];
self.read_exact(&mut i8_buf)?; self.reader.read_exact(&mut i8_buf)?;
let i8_value = i8::from_le_bytes(i8_buf); let i8_value = i8::from_le_bytes(i8_buf);
if i8_value != i8::MAX { return Ok(i8_value as i16) } if i8_value != i8::MAX { return Ok(i8_value as i16) }
let mut i16_buf = [0; 2]; let mut i16_buf = [0; 2];
self.read_exact(&mut i16_buf)?; self.reader.read_exact(&mut i16_buf)?;
let i16_value = i16::from_le_bytes(i16_buf); let i16_value = i16::from_le_bytes(i16_buf);
Ok(i16_value as i16) Ok(i16_value as i16)
} }
// wrap UnexpectedEof, since that's an error for other Io clients,
// but it's specifically not an error for us
fn read_exact(&mut self, buf: &mut [u8]) -> ProtocolReaderResult<()> {
match self.reader.read_exact(buf) {
Ok(_) => Ok(()),
Err(e) => {
if e.kind() == ErrorKind::UnexpectedEof {
return Err(ProtocolReaderError::EarlyEof)
}
return Err(ProtocolReaderError::Io(e));
}
}
}
} }

40
src/quality_settings.rs Normal file
View File

@@ -0,0 +1,40 @@
#[derive(Clone, Copy)]
pub struct QualitySettings {
pub values: [u8; 8],
}
pub const QUANTIZATION_QUANT0_DIVISOR: usize = 0;
pub const QUANTIZATION_CONSTANT: usize = 1;
pub const QUANTIZATION_LINEAR: usize = 2;
pub const QUANTIZATION_N_QUANTS_TO_KEEP: usize = 3;
pub const THUMBNAIL_IS_WIDE: usize = 4;
// 5, 6, 7: reserved
impl QualitySettings {
pub fn new(level: u8) -> Self {
// Concept:
// Levels 0, 1, and 2 are incredibly deep fried
// Levels 3, 4, 5, 6 are usable with artifacting
// Levels 7, 8, 9 are above what you would seriously use
if level == 0 {
return QualitySettings { values: [255, 255, 32, 1, 0, 0, 0, 0] };
} else if level == 1 {
return QualitySettings { values: [192, 192, 24, 2, 0, 0, 0, 0] };
} else if level == 2 {
return QualitySettings { values: [128, 128, 16, 3, 0, 0, 0, 0] };
} else if level == 3 {
return QualitySettings { values: [96, 96, 12, 4, 0, 0, 0, 0] };
// TODO: Levels 4, 5, 6
} else if level == 7 {
return QualitySettings { values: [1, 24, 8, 24, 1, 0, 0, 0] };
} else if level == 8 {
return QualitySettings { values: [1, 12, 4, 48, 1, 0, 0, 0] };
} else { // level >= 9
return QualitySettings { values: [1, 1, 1, 96, 1, 0, 0, 0] };
}
}
}

View File

@@ -1,4 +1,4 @@
use crate::constants::TILE_SZ2; use crate::{constants::TILE_SZ2, quality_settings::{self, QualitySettings}};
const ZIGZAG: [u8; TILE_SZ2] = [ const ZIGZAG: [u8; TILE_SZ2] = [
0 , 1 , 5 , 6 , 14, 15, 27, 28, 0 , 1 , 5 , 6 , 14, 15, 27, 28,
@@ -12,12 +12,13 @@ const ZIGZAG: [u8; TILE_SZ2] = [
]; ];
pub fn to_quantized( pub fn to_quantized(
quality_settings: QualitySettings,
coefs: [i16; TILE_SZ2] coefs: [i16; TILE_SZ2]
) -> [i16; TILE_SZ2] { ) -> [i16; TILE_SZ2] {
let mut quant: [i16; TILE_SZ2] = [0; TILE_SZ2]; let mut quant: [i16; TILE_SZ2] = [0; TILE_SZ2];
for cf_ix in 0..TILE_SZ2 { for cf_ix in 0..TILE_SZ2 {
let div = divisor(cf_ix); let div = divisor(quality_settings, cf_ix);
let qval = (coefs[cf_ix] + div / 2) / div; let qval = (coefs[cf_ix] + div / 2) / div;
quant[ZIGZAG[cf_ix] as usize] = qval quant[ZIGZAG[cf_ix] as usize] = qval
} }
@@ -27,7 +28,7 @@ pub fn to_quantized(
indices[i] = i; indices[i] = i;
} }
indices.sort_by_key(|i| -quant[*i].abs()); indices.sort_by_key(|i| -quant[*i].abs());
for i in 4..indices.len() { for i in (quality_settings.values[quality_settings::QUANTIZATION_N_QUANTS_TO_KEEP] as usize)..indices.len() {
quant[indices[i]] = 0; quant[indices[i]] = 0;
} }
@@ -35,21 +36,30 @@ pub fn to_quantized(
} }
pub fn from_quantized( pub fn from_quantized(
quality_settings: QualitySettings,
quant: [i16; TILE_SZ2] quant: [i16; TILE_SZ2]
) -> [i16; TILE_SZ2] { ) -> [i16; TILE_SZ2] {
let mut coefs: [i16; TILE_SZ2] = [0; TILE_SZ2]; let mut coefs: [i16; TILE_SZ2] = [0; TILE_SZ2];
for cf_ix in 0..TILE_SZ2 { for cf_ix in 0..TILE_SZ2 {
let div = divisor(cf_ix); let div = divisor(quality_settings, cf_ix);
coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize].wrapping_mul(div); coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize].wrapping_mul(div);
} }
coefs coefs
} }
pub fn divisor(cf_ix: usize) -> i16 { pub fn divisor(
if cf_ix == 0 { return 1; } quality_settings: QualitySettings,
cf_ix: usize
) -> i16 {
if cf_ix == 0 { return quality_settings.values[quality_settings::QUANTIZATION_QUANT0_DIVISOR] as i16; }
let x = cf_ix % 8; let x = cf_ix % 8;
let y = cf_ix / 8; let y = cf_ix / 8;
let div = 32 + (x as i16 + y as i16) * 12; let div =
(quality_settings.values[quality_settings::QUANTIZATION_CONSTANT] as i16).wrapping_add(
(x as i16 + y as i16).wrapping_mul(
quality_settings.values[quality_settings::QUANTIZATION_LINEAR] as i16
)
);
if div==32 && cf_ix != 0 { if div==32 && cf_ix != 0 {
dbg!(cf_ix, x, y, div); dbg!(cf_ix, x, y, div);
} }