Added proper encoding and decoding, fixed minor issues

Cargo.toml

@@ -6,5 +6,6 @@ edition = "2021"
 [dependencies]
 byteorder = "1.5.0"
 image = "0.25.2"
+integer-encoding = "4.0.0"
 rayon = "1.10.0"
 thiserror = "1.0.63"

src/compression/dct.rs

@@ -1,13 +1,13 @@
 use std::{f32::consts::{PI, SQRT_2}, sync::{Arc, Mutex}};
 
-use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator};
+use rayon::prelude::*;
 
 use crate::header::ColorFormat;
 
 /// Perform a Discrete Cosine Transform on the input matrix.
 pub fn dct(input: &[u8], width: usize, height: usize) -> Vec<f32> {
     if input.len() != width * height {
-        panic!("Input matrix size must be width×height")
+        panic!("Input matrix size must be width * height, got {}", input.len())
     }
 
     let sqrt_width_zero = 1.0 / (width as f32).sqrt();
@@ -53,7 +53,7 @@ pub fn dct(input: &[u8], width: usize, height: usize) -> Vec<f32> {
 /// Perform an inverse Discrete Cosine Transform on the input matrix.
 pub fn idct(input: &[f32], width: usize, height: usize) -> Vec<u8> {
     if input.len() != width * height {
-        panic!("Input matrix size must be width×height")
+        panic!("Input matrix size must be width * height, got {}", input.len())
     }
 
     let sqrt_width_zero = 1.0 / (width as f32).sqrt();
@@ -127,12 +127,18 @@ pub fn quantization_matrix(quality: u32) -> [u16; 64] {
 
 /// Quantize an input matrix, returning the result.
 pub fn quantize(input: &[f32], quant_matrix: [u16; 64]) -> Vec<i16> {
-    input.iter().zip(quant_matrix).map(|(v, q)| (v / q as f32).round() as i16).collect()
+    input.iter()
+        .zip(quant_matrix)
+        .map(|(v, q)| (v / q as f32).round() as i16)
+        .collect()
 }
 
 /// Dequantize an input matrix, returning an approximation of the original.
 pub fn dequantize(input: &[i16], quant_matrix: [u16; 64]) -> Vec<f32> {
-    input.iter().zip(quant_matrix).map(|(v, q)| (*v as i16 * q as i16) as f32).collect()
+    input.iter()
+        .zip(quant_matrix)
+        .map(|(v, q)| (*v as i16 * q as i16) as f32)
+        .collect()
 }
 
 /// Take in an image encoded in some [`ColorFormat`] and perform DCT on it,
@@ -140,7 +146,7 @@ pub fn dequantize(input: &[i16], quant_matrix: [u16; 64]) -> Vec<f32> {
 /// to a multiple of 8, which must be reversed when decoding.
 pub fn dct_compress(input: &[u8], parameters: DctParameters) -> Vec<Vec<i16>> {
     let new_width = parameters.width + (8 - parameters.width % 8);
-    let new_height = parameters.height + (8 - parameters.width % 8);
+    let new_height = parameters.height + (8 - parameters.height % 8);
     let quantization_matrix = quantization_matrix(parameters.quality);
 
     let mut dct_image = Vec::with_capacity(input.len());
@@ -150,7 +156,6 @@ pub fn dct_compress(input: &[u8], parameters: DctParameters) -> Vec<Vec<i16>> {
             .step_by(parameters.format.channels() as usize)
             .copied()
             .collect();
-        println!("Encoding channel {ch}");
 
         // Create 2d array of the channel for ease of processing
         let mut img_2d: Vec<Vec<u8>> = channel.windows(parameters.width).step_by(parameters.width).map(|r| r.to_vec()).collect();
@@ -185,20 +190,17 @@ pub fn dct_compress(input: &[u8], parameters: DctParameters) -> Vec<Vec<i16>> {
 
 /// Take in an image encoded with DCT and quantized and perform IDCT on it,
 /// returning an approximation of the original data.
-pub fn dct_decompress(input: &[Vec<i16>], parameters: DctParameters) -> Vec<u8> {
+pub fn dct_decompress(input: &[i16], parameters: DctParameters) -> Vec<u8> {
     let new_width = parameters.width + (8 - parameters.width % 8);
-    let new_height = parameters.height + (8 - parameters.width % 8);
+    let new_height = parameters.height + (8 - parameters.height % 8);
 
     // Precalculate the quantization matrix
     let quantization_matrix = quantization_matrix(parameters.quality);
 
     let final_img = Arc::new(Mutex::new(vec![0u8; (new_width * new_height) * parameters.format.channels() as usize]));
-
+    input.par_chunks(new_width * new_height).enumerate().for_each(|(chan_num, channel)| {
-    input.par_iter().enumerate().for_each(|(chan_num, channel)| {
-        println!("Decoding channel {chan_num}");
-
         let decoded_image = Arc::new(Mutex::new(vec![0u8; parameters.width * parameters.height]));
-        channel.into_par_iter().copied().chunks(64).enumerate().for_each(|(i, chunk)| {
+        channel.par_chunks(64).enumerate().for_each(|(i, chunk)| {
             let dequantized_dct = dequantize(&chunk, quantization_matrix);
             let original = idct(&dequantized_dct, 8, 8);
 

src/header.rs

@@ -40,16 +40,20 @@ impl Default for Header {
 }
 
 impl Header {
-    pub fn to_bytes(&self) -> [u8; 16] {
+    pub fn to_bytes(&self) -> [u8; 19] {
         let mut buf = Cursor::new(Vec::new());
 
         buf.write_all(&self.magic).unwrap();
         buf.write_u32::<LE>(self.width).unwrap();
         buf.write_u32::<LE>(self.height).unwrap();
 
-        buf.write_u8(self.compression_type as u8).unwrap();
+        // Write compression info
+        buf.write_u8(self.compression_type.into()).unwrap();
         buf.write_i8(self.compression_level).unwrap();
 
+        // Write color format
+        buf.write_u8(self.color_format as u8).unwrap();
+
         buf.into_inner().try_into().unwrap()
     }
 
@@ -73,6 +77,8 @@ impl Header {
     }
 }
 
+/// The format of bytes in the image.
+#[repr(u8)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ColorFormat {
     /// RGBA, 8 bits per channel
@@ -127,6 +133,7 @@ impl TryFrom<u8> for ColorFormat {
 }
 
 /// The type of compression used in the image
+#[repr(u8)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum CompressionType {
     /// No compression at all, raw bitmap
@@ -151,3 +158,13 @@ impl TryFrom<u8> for CompressionType {
         })
     }
 }
+
+impl Into<u8> for CompressionType {
+    fn into(self) -> u8 {
+        match self {
+            CompressionType::None => 0,
+            CompressionType::Lossless => 1,
+            CompressionType::LossyDct => 2,
+        }
+    }
+}

src/main.rs

@@ -7,80 +7,66 @@ mod header;
 mod operations;
 pub mod picture;
 
-use std::{fs::File, io::Write, time::Instant};
+use std::{fs::File, io::{BufReader, BufWriter}, time::Instant};
-
+use header::{ColorFormat, CompressionType};
-use header::ColorFormat;
+use image::RgbaImage;
-use compression::{dct::{dct_compress, dct_decompress, DctParameters}, lossless};
+use picture::DangoPicture;
 
 fn main() {
     let input = image::open("transparent2.png").unwrap().to_rgba8();
     input.save("original.png").unwrap();
 
-    let timer = Instant::now();
+    let dpf_lossy = DangoPicture::from_raw(
-    let dct_result = dct_compress(
-        input.as_raw(),
-        DctParameters {
-            quality: 68,
-            format: ColorFormat::Rgba32,
-            width: input.width() as usize,
-            height: input.height() as usize,
-        }
-    );
-
-    let compressed_dct = lossless::compress(&dct_result.concat().iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>()).unwrap();
-    println!("Encoding took {}ms", timer.elapsed().as_millis());
-
-    let mut dct_output = File::create("test-dct.dpf").unwrap();
-    dct_output.write_all(&compressed_dct.0).unwrap();
-
-    let timer = Instant::now();
-    let decoded_dct = dct_decompress(
-        &dct_result,
-        DctParameters {
-            quality: 68,
-            format: ColorFormat::Rgba32,
-            width: input.width() as usize,
-            height: input.height() as usize
-        }
-    );
-    println!("Decoding took {}ms", timer.elapsed().as_millis());
-
-    image::RgbaImage::from_raw(
         input.width(),
         input.height(),
-        decoded_dct
+        ColorFormat::Rgba32,
-    ).unwrap().save("dct-final.png").unwrap();
+        CompressionType::LossyDct,
+        Some(10),
+        input.as_raw().clone()
+    );
 
-    /*
+    let dpf_lossless = DangoPicture::from_raw(
-    // Reverse the DCT
+        input.width(),
-    let idct: Vec<u8> = idct(&dct, 8, 8).iter().map(|c| *c as u8).collect();
+        input.height(),
+        ColorFormat::Rgba32,
+        CompressionType::Lossless,
+        None,
+        input.as_raw().clone()
+    );
 
-    let img = GrayImage::from_raw(input.width(), input.height(), idct).unwrap();
+    println!("\n--- LOSSY ---");
-    img.save("test.png").unwrap();
+    println!("Encoding");
-    */
-
-    /*
-    let image_data = image::open("bw.jpg").unwrap().to_rgba8();
-    let encoded_dpf = DangoPicture::from_raw(image_data.width(), image_data.height(), &image_data);
-
-    println!("ENCODING ---");
     let timer = Instant::now();
-    let mut outfile = BufWriter::new(File::create("test.dpf").unwrap());
+    let mut outfile = BufWriter::new(std::fs::File::create("test-lossy.dpf").unwrap());
-    encoded_dpf.encode(&mut outfile);
+    dpf_lossy.encode(&mut outfile).unwrap();
     println!("Encoding took {}ms", timer.elapsed().as_millis());
 
-    println!("DECODING ---");
+    let mut outbuf = Vec::new();
+    dpf_lossy.encode(&mut outbuf).unwrap();
+    println!("Size is {}Mb", (((outbuf.len() as f32 / 1_000_000.0) * 100.0) as u32 as f32) / 100.0);
+
+    println!("Decoding");
     let timer = Instant::now();
-    let mut infile = BufReader::new(File::open("test.dpf").unwrap());
+    let mut infile = BufReader::new(File::open("test-lossy.dpf").unwrap());
     let decoded_dpf = DangoPicture::decode(&mut infile).unwrap();
+    RgbaImage::from_raw(decoded_dpf.header.width, decoded_dpf.header.height, decoded_dpf.bitmap.into()).unwrap().save("test-lossy.png").unwrap();
     println!("Decoding took {}ms", timer.elapsed().as_millis());
 
-    let out_image = RgbaImage::from_raw(
+    println!("\n--- LOSSLESS ---");
-        decoded_dpf.header.width,
+    println!("Encoding");
-        decoded_dpf.header.height,
+    let timer = Instant::now();
-        decoded_dpf.bitmap.into(),
+    let mut outfile = BufWriter::new(std::fs::File::create("test-lossless.dpf").unwrap());
-    )
+    dpf_lossless.encode(&mut outfile).unwrap();
-    .unwrap();
+    println!("Encoding took {}ms", timer.elapsed().as_millis());
-    out_image.save("test.png").unwrap();
+
-    */
+    let mut outbuf = Vec::new();
+    dpf_lossless.encode(&mut outbuf).unwrap();
+    println!("Size is {}Mb", (((outbuf.len() as f32 / 1_000_000.0) * 100.0) as u32 as f32) / 100.0);
+
+    println!("Decoding");
+    let timer = Instant::now();
+    let mut infile = BufReader::new(File::open("test-lossless.dpf").unwrap());
+    let decoded_dpf = DangoPicture::decode(&mut infile).unwrap();
+    RgbaImage::from_raw(decoded_dpf.header.width, decoded_dpf.header.height, decoded_dpf.bitmap.into()).unwrap().save("test-lossless.png").unwrap();
+    println!("Decoding took {}ms", timer.elapsed().as_millis());
 }

src/picture.rs

@@ -1,10 +1,12 @@
-use std::io::{self, Read, Write};
+use std::{fs::File, io::{self, BufWriter, Read, Write}};
 
 use byteorder::{ReadBytesExt, WriteBytesExt};
+use integer_encoding::VarInt;
 use thiserror::Error;
 
 use crate::{
-    compression::{dct::{dct_compress, DctParameters}, lossless::{compress, decompress, CompressionError, CompressionInfo}},
+    compression::{dct::{dct_compress, dct_decompress, DctParameters},
+    lossless::{compress, decompress, CompressionError, CompressionInfo}},
     header::{ColorFormat, CompressionType, Header},
     operations::{diff_line, line_diff},
 };
@@ -16,7 +18,7 @@ pub struct DangoPicture {
 
 #[derive(Error, Debug)]
 pub enum Error {
-    #[error("incorrect identifier, got {0:?}")]
+    #[error("incorrect identifier, got {0:02X?}")]
     InvalidIdentifier([u8; 8]),
 
     #[error("io operation failed: {0}")]
@@ -27,21 +29,35 @@ pub enum Error {
 }
 
 impl DangoPicture {
+    /// Create a DPF
     pub fn from_raw(
         width: u32,
         height: u32,
         color_format: ColorFormat,
         compression_type: CompressionType,
+        compression_level: Option<u8>,
         bitmap: Vec<u8>,
     ) -> Self {
+        let compression_level = match compression_level {
+            Some(level) => {
+                if level < 1 || level > 100 {
+                    panic!("Compression level out of range 1..100")
+                }
+                level as i8
+            },
+            None => -1,
+        };
+
         let header = Header {
+            magic: *b"dangoimg",
+
             width,
             height,
 
             compression_type,
-            color_format,
+            compression_level,
 
-            ..Default::default()
+            color_format,
         };
 
         DangoPicture {
@@ -55,9 +71,12 @@ impl DangoPicture {
         // Write out the header
         output.write_all(&self.header.to_bytes()).unwrap();
 
+        // Based on the compression type, modify the data accordingly
         let modified_data = match self.header.compression_type {
             CompressionType::None => &self.bitmap,
-            CompressionType::Lossless => &diff_line(self.header.width, self.header.height, &self.bitmap),
+            CompressionType::Lossless => {
+                &diff_line(self.header.width, self.header.height, &self.bitmap)
+            },
             CompressionType::LossyDct => {
                 &dct_compress(
                     &self.bitmap,
@@ -67,11 +86,15 @@ impl DangoPicture {
                         width: self.header.width as usize,
                         height: self.header.height as usize,
                     }
-                ).concat().iter().flat_map(|i| i.to_le_bytes()).collect()
+                )
+                .concat()
+                .into_iter()
+                .flat_map(VarInt::encode_var_vec)
+                .collect()
             },
         };
 
-        // Compress the image data
+        // Compress the final image data using the basic LZW scheme
         let (compressed_data, compression_info) = compress(&modified_data)?;
 
         // Write out compression info
@@ -83,22 +106,55 @@ impl DangoPicture {
         Ok(())
     }
 
+    /// Encode and write the image out to a file.
+    pub fn save<P: ?Sized + AsRef<std::path::Path>>(&self, path: &P) -> Result<(), Error> {
+        let mut out_file = BufWriter::new(File::create(path.as_ref())?);
+
+        self.encode(&mut out_file)?;
+
+        Ok(())
+    }
+
     /// Decode the image from anything that implements [Read]
     pub fn decode<I: Read + ReadBytesExt>(mut input: I) -> Result<DangoPicture, Error> {
-        let mut magic = [0u8; 8];
-        input.read_exact(&mut magic).unwrap();
-
-        if magic != *b"dangoimg" {
-            return Err(Error::InvalidIdentifier(magic));
-        }
-
         let header = Header::read_from(&mut input)?;
 
         let compression_info = CompressionInfo::read_from(&mut input);
 
-        let preprocessed_bitmap = decompress(&mut input, &compression_info);
+        let pre_bitmap = decompress(&mut input, &compression_info);
 
-        let bitmap = line_diff(header.width, header.height, &preprocessed_bitmap);
+        let bitmap = match header.compression_type {
+            CompressionType::None => pre_bitmap,
+            CompressionType::Lossless => {
+                line_diff(header.width, header.height, &pre_bitmap)
+            },
+            CompressionType::LossyDct => {
+                let mut decoded = Vec::new();
+                let mut offset = 0;
+                loop {
+                    if offset > pre_bitmap.len() {
+                        break;
+                    }
+
+                    if let Some(num) = i16::decode_var(&pre_bitmap[offset..]) {
+                        offset += num.1;
+                        decoded.push(num.0 as i16);
+                    } else {
+                        break;
+                    }
+                }
+
+                dct_decompress(
+                    &decoded,
+                    DctParameters {
+                        quality: header.compression_level as u32,
+                        format: header.color_format,
+                        width: header.width as usize,
+                        height: header.height as usize,
+                    }
+                )
+            },
+        };
 
         Ok(DangoPicture { header, bitmap })
     }