Improved DCT decoding and ergonomics

src/compression/dct.rs

@@ -1,5 +1,7 @@
 use std::f32::consts::{PI, SQRT_2};
 
+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 {
@@ -133,7 +135,10 @@ 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()
 }
 
-pub fn dct_compress(input: &[u8], width: u32, height: u32, quality: u32) -> (Vec<Vec<i16>>, usize, usize) {
+/// Take in an image encoded in some [`ColorFormat`] and perform DCT on it,
+/// returning the modified data. This function also pads the image dimensions
+/// to a multiple of 8, which must be reversed when decoding.
+pub fn dct_compress(input: &[u8], width: u32, height: u32, parameters: DctParameters) -> DctImage {
     let new_width = width as usize + (8 - width % 8) as usize;
     let new_height = height as usize + (8 - height % 8) as usize;
     let mut img_2d: Vec<Vec<u8>> = input.windows(width as usize).step_by(width as usize).map(|r| r.to_vec()).collect();
@@ -141,6 +146,8 @@ pub fn dct_compress(input: &[u8], width: u32, height: u32, quality: u32) -> (Vec
     img_2d.resize(new_height, vec![0u8; new_width]);
 
     let mut dct_image = Vec::new();
+    for _ in 0..1 {
+        let mut dct_channel = Vec::new();
         for h in 0..new_height / 8 {
             for w in 0..new_width / 8 {
                 let mut chunk = Vec::new();
@@ -151,13 +158,53 @@ pub fn dct_compress(input: &[u8], width: u32, height: u32, quality: u32) -> (Vec
 
                 // Perform the DCT on the image section
                 let dct: Vec<f32> = dct(&chunk, 8, 8);
-            let quantzied_dct = quantize(&dct, quantization_matrix(quality));
+                let quantzied_dct = quantize(&dct, quantization_matrix(parameters.quality));
 
-            dct_image.push(quantzied_dct);
+                dct_channel.extend_from_slice(&quantzied_dct);
             }
         }
+        dct_image.push(dct_channel);
+    }
 
-    (dct_image, new_width, new_height)
+    DctImage {
+        channels: dct_image,
+        width: new_width as u32,
+        height: new_height as u32
+    }
+}
+
+/// Parameters to pass to the [`dct_compress`] function.
+pub struct DctParameters {
+    /// A quality level from 1-100. Higher values provide better results.
+    /// Default value is 80.
+    pub quality: u32,
+
+    /// The color format of the input bytes.
+    ///
+    /// Since DCT can only process one channel at a time, knowing the format
+    /// is important.
+    pub format: ColorFormat,
+}
+
+impl Default for DctParameters {
+    fn default() -> Self {
+        Self {
+            quality: 80,
+            format: ColorFormat::Rgba32
+        }
+    }
+}
+
+/// The results of DCT compression
+pub struct DctImage {
+    /// The DCT encoded version of each channel.
+    pub channels: Vec<Vec<i16>>,
+
+    /// New width after padding.
+    pub width: u32,
+
+    /// New height after padding.
+    pub height: u32,
 }
 
 #[cfg(test)]

src/header.rs

@@ -39,3 +39,35 @@ pub enum ColorFormat {
     /// RGB, 8 bits per channel
     Rgb24,
 }
+
+impl ColorFormat {
+    /// Bits per color channel.
+    ///
+    /// Ex. Rgba32 has `8bpc`
+    pub fn bpc(&self) -> u8 {
+        match self {
+            ColorFormat::Rgba32 => 8,
+            ColorFormat::Rgb24 => 8,
+        }
+    }
+
+    /// Bits per pixel.
+    ///
+    /// Ex. Rgba32 has `32bpp`
+    pub fn bpp(&self) -> u16 {
+        match self {
+            ColorFormat::Rgba32 => 32,
+            ColorFormat::Rgb24 => 24,
+        }
+    }
+
+    /// Number of color channels.
+    ///
+    /// Ex. Rgba32 has `4` channels
+    pub fn channels(self) -> u16 {
+        match self {
+            ColorFormat::Rgba32 => 4,
+            ColorFormat::Rgb24 => 3,
+        }
+    }
+}

src/main.rs

@@ -7,33 +7,39 @@ mod header;
 mod operations;
 pub mod picture;
 
+use header::ColorFormat;
 use picture::DangoPicture;
 use std::{
     fs::File,
     io::{BufReader, BufWriter, Write},
     time::Instant,
 };
-use compression::{dct::{dct, dct_compress, dequantize, idct, quantization_matrix, quantize}, lossless};
+use compression::{dct::{dct, dct_compress, dequantize, idct, quantization_matrix, quantize, DctParameters}, lossless};
 
-use image::{ColorType, DynamicImage, GenericImage, GrayImage, Luma, Rgba};
+use image::{GenericImage, GrayImage, Luma};
 
 fn main() {
-    let input = image::open("transparent.png").unwrap().to_luma8();
+    let input = image::open("test_input.png").unwrap().to_luma8();
     input.save("original.png").unwrap();
 
-    let (dct_image, new_width, new_height) = dct_compress(input.as_raw(), input.width(), input.height(), 100);
-    let compressed_dct = lossless::compress(&dct_image.iter().flatten().flat_map(|b| b.to_le_bytes()).collect::<Vec<u8>>());
-    let mut dct_save = File::create("dct_raw.dct").unwrap();
-    dct_save.write_all(&compressed_dct.0).unwrap();
+    let dct_result = dct_compress(
+        input.as_raw(),
+        input.width(),
+        input.height(),
+        DctParameters {
+            quality: 100,
+            format: ColorFormat::Rgba32,
+        }
+    );
 
-    let mut decoded_image = GrayImage::new(new_width as u32, new_height as u32);
-    for (i, chunk) in dct_image.iter().enumerate() {
+    let mut decoded_image = GrayImage::new(dct_result.width, dct_result.height);
+    for (i, chunk) in dct_result.channels[0].windows(64).step_by(64).enumerate() {
         let dequantized_dct = dequantize(chunk, quantization_matrix(100));
         let original = idct(&dequantized_dct, 8, 8);
 
         // Write rows of blocks
-        let start_x = (i * 8) % (new_width as usize);
-        let start_y = ((i * 8) / new_width as usize) * 8;
+        let start_x = (i * 8) % dct_result.width as usize;
+        let start_y = ((i * 8) / dct_result.width as usize) * 8;
 
         let mut sub = decoded_image.sub_image(start_x as u32, start_y as u32, 8, 8);
         for y in 0..8 {
@@ -42,6 +48,8 @@ fn main() {
                 sub.put_pixel(x, y, Luma([value]))
             }
         }
+
+        decoded_image.save(format!("test.png")).unwrap();
     }
     decoded_image.save(format!("test.png")).unwrap();