diff --git a/src/compression/dct.rs b/src/compression/dct.rs index f382175..06e21b7 100644 --- a/src/compression/dct.rs +++ b/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 { if input.len() != width * height { @@ -133,7 +135,10 @@ pub fn dequantize(input: &[i16], quant_matrix: [u16; 64]) -> Vec { 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>, 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> = input.windows(width as usize).step_by(width as usize).map(|r| r.to_vec()).collect(); @@ -141,23 +146,65 @@ 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 h in 0..new_height / 8 { - for w in 0..new_width / 8 { - let mut chunk = Vec::new(); - for i in 0..8 { - let row = &img_2d[(h * 8) + i][w * 8..(w * 8) + 8]; - chunk.extend_from_slice(&row); + 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(); + for i in 0..8 { + let row = &img_2d[(h * 8) + i][w * 8..(w * 8) + 8]; + chunk.extend_from_slice(&row); + } + + // Perform the DCT on the image section + let dct: Vec = dct(&chunk, 8, 8); + let quantzied_dct = quantize(&dct, quantization_matrix(parameters.quality)); + + dct_channel.extend_from_slice(&quantzied_dct); } - - // Perform the DCT on the image section - let dct: Vec = dct(&chunk, 8, 8); - let quantzied_dct = quantize(&dct, quantization_matrix(quality)); - - dct_image.push(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>, + + /// New width after padding. + pub width: u32, + + /// New height after padding. + pub height: u32, } #[cfg(test)] diff --git a/src/header.rs b/src/header.rs index 29f025a..245d8c3 100644 --- a/src/header.rs +++ b/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, + } + } +} diff --git a/src/main.rs b/src/main.rs index 45a1217..23f844b 100644 --- a/src/main.rs +++ b/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::>()); - 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();