Better PNG interface

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

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<W: Write>(
-    width: u32,
-    height: u32,
-    // n_components: u16,
-    layers: &[(QualitySettings, &[u8])],
+    image: Image,
     writer: &mut protocol::ProtocolWriter<W>
 ) -> 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);

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<R: Read>(
     reader: &mut ProtocolReader<R>
-) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> {
+) -> ProtocolReaderResult<Image> {
     // read header
     let header = reader.read_header()?;
     if header.bytes != MAGIC.bytes {
@@ -78,7 +78,14 @@ 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(image)
 }
 
 impl CoefTile {

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<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(())
+    }
+}

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);
+    run_for("zonked".to_string());
+    run_for("avatar2".to_string());
 }
-
-    hard_main();
-}
-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();
-
-    assert_eq!(3, decompressed.len());
-    save_image(width2 as usize, height2 as usize, decompressed);
-}
-
-
-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;
-        }
-    }
-
-    let mut writer = encoder.write_header().unwrap();
-    writer.write_image_data(&idata).unwrap();
+    let decompressed = decompression::decompress(&mut reader).unwrap();
 
+    png_utils::save_image(
+        format!("outputs/{}_out.png", name), 
+        decompressed
+    ).unwrap();
 }

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<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(())
+}