diff --git a/TODO.md b/TODO.md index 9126cf5..5622bb1 100644 --- a/TODO.md +++ b/TODO.md @@ -1,8 +1,8 @@ -Support an alpha channel. -Improve the PNG interface. Actually support terminating early for "progressive" loads Done: +Support an alpha channel. +Improve the PNG interface. Interleave layers instead of having them consecutively. (Bad idea, breaks RLE. Skipping) Global parameters to control encoder settings: - Quantization level: quant 0 diff --git a/inputs/zonked.png b/inputs/zonked.png new file mode 100644 index 0000000..ce4f73c Binary files /dev/null and b/inputs/zonked.png differ diff --git a/src/compression.rs b/src/compression.rs index bddcd5b..f68b0be 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -1,6 +1,6 @@ use std::io::Write; -use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quality_settings::{self, QualitySettings}, quantization, transform}; +use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, image::Image, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quality_settings::{self, QualitySettings}, quantization, transform}; struct PixelTile { @@ -20,37 +20,34 @@ struct QuantTile { pub fn compress( - width: u32, - height: u32, - // n_components: u16, - layers: &[(QualitySettings, &[u8])], + image: Image, writer: &mut protocol::ProtocolWriter ) -> ProtocolWriterResult<()> { // validation - for l in 0..layers.len() { - assert!(layers[l].1.len() == width as usize * height as usize); + for l in 0..image.layers.len() { + assert!(image.layers[l].1.len() == image.width as usize * image.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)?; + writer.write_u32_wide(image.width as u32)?; + writer.write_u32_wide(image.height as u32)?; + writer.write_u32_wide(image.layers.len() as u32)?; // write quality settings for each layout - for (quality, _) in layers.iter() { + for (quality, _) in image.layers.iter() { writer.write_quality_settings(*quality)?; } // build list of tiles let mut tiles = vec![]; - for (quality, layer) in layers.iter() { - for x0 in (0..width).step_by(TILE_SZ) { - for y0 in (0..height).step_by(TILE_SZ) { + for (quality, layer) in image.layers.iter() { + for x0 in (0..image.width).step_by(TILE_SZ) { + for y0 in (0..image.height).step_by(TILE_SZ) { let pixel_tile = PixelTile::from_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 quant_tile = QuantTile::from_coef_tile(&coef_tile); diff --git a/src/decompression.rs b/src/decompression.rs index 554f8d6..6ac57d9 100644 --- a/src/decompression.rs +++ b/src/decompression.rs @@ -1,6 +1,6 @@ use std::io::Read; -use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quality_settings::{self, QualitySettings}, quantization, transform}; +use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, image::Image, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quality_settings::{self, QualitySettings}, quantization, transform}; struct PixelTile { pixels: [i16; TILE_SZ2] @@ -19,7 +19,7 @@ struct QuantTile { pub fn decompress( reader: &mut ProtocolReader -) -> ProtocolReaderResult<(u32, u32, Vec>)> { +) -> ProtocolReaderResult { // read header let header = reader.read_header()?; if header.bytes != MAGIC.bytes { @@ -78,7 +78,14 @@ pub fn decompress( } } } - 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(image) } impl CoefTile { diff --git a/src/image.rs b/src/image.rs new file mode 100644 index 0000000..d2a9947 --- /dev/null +++ b/src/image.rs @@ -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)> +} + +impl Image { + pub fn new(width: usize, height: usize) -> Self { + Image { + width, height, + layers: vec![] + } + } + + pub fn add(&mut self, quality: QualitySettings, pixels: Vec) -> Result<(), ImageError> { + if pixels.len() != self.width * self.height { + return Err(ImageError::WrongSize(self.width, self.height)); + } + self.layers.push((quality, pixels)); + + Ok(()) + } +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index b92c536..0db6758 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,104 +1,44 @@ use std::{fs::File, io::{Cursor, Write}}; -use png::{BitDepth, ColorType}; - -use crate::protocol::{ProtocolReader, ProtocolWriter}; +use crate::{protocol::{ProtocolReader, ProtocolWriter}, quality_settings::QualitySettings}; mod compression; mod constants; mod decompression; +mod png_utils; mod protocol; mod quality_settings; mod quantization; mod transform; +mod image; fn main() { - for chunk in vec![ - [0, 0, 0, 0, 0, 0, 0, 0], - [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(); - transform::encode(&mut enc, 0, 1); - let mut dec = enc.clone(); - transform::decode(&mut dec, 0, 1); - dbg!(orig, enc, dec); - } - - hard_main(); + run_for("zonked".to_string()); + run_for("avatar2".to_string()); } -fn hard_main() { - let (width, height, r, g, b) = load_image(); +fn run_for(name: String) { + let image = png_utils::load_image( + format!("inputs/{}.png", name), + QualitySettings::new(2), + QualitySettings::new(2), + ).unwrap(); let mut writer = ProtocolWriter::new(vec![]); - let quality_settings = quality_settings::QualitySettings::new(3); - compression::compress( - width as u32, height as u32, &[(quality_settings, &r), (quality_settings, &g), (quality_settings, &b)], + image, &mut writer ).unwrap(); let compressed = writer.destroy(); - let mut output_file = File::create("outputs/avatar2.rxi").unwrap(); + let mut output_file = File::create(format!("outputs/{}.rxi", name)).unwrap(); output_file.write_all(&compressed).unwrap(); let mut reader = ProtocolReader::new(Cursor::new(compressed)); - let (width2, height2, decompressed) = - decompression::decompress(&mut reader).unwrap(); + let decompressed = decompression::decompress(&mut reader).unwrap(); - assert_eq!(3, decompressed.len()); - save_image(width2 as usize, height2 as usize, decompressed); + png_utils::save_image( + format!("outputs/{}_out.png", name), + decompressed + ).unwrap(); } - - -fn load_image() -> (usize, usize, Vec, Vec, Vec) { - 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 = bytes[0..].iter().cloned().step_by(3).collect(); - let g: Vec = bytes[1..].iter().cloned().step_by(3).collect(); - let b: Vec = 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>) { - 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 = 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; - } - } - - let mut writer = encoder.write_header().unwrap(); - writer.write_image_data(&idata).unwrap(); - -} \ No newline at end of file diff --git a/src/png_utils.rs b/src/png_utils.rs new file mode 100644 index 0000000..04c94a3 --- /dev/null +++ b/src/png_utils.rs @@ -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 { + 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 = bytes[0..].iter().cloned().step_by(3).collect(); + let g: Vec = bytes[1..].iter().cloned().step_by(3).collect(); + let b: Vec = 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 = bytes[0..].iter().cloned().step_by(4).collect(); + let g: Vec = bytes[1..].iter().cloned().step_by(4).collect(); + let b: Vec = bytes[2..].iter().cloned().step_by(4).collect(); + let a: Vec = 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 = 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(()) +} \ No newline at end of file