From 521bb01c87bc046961035c23b869593e1fdc38bc Mon Sep 17 00:00:00 2001 From: G2-Games Date: Thu, 25 Jul 2024 17:55:58 -0500 Subject: [PATCH] Initial work on Discrete Cosine Transform --- src/compression/dct.rs | 176 +++++++++++++++++++++++++++++++++++++++++ src/main.rs | 64 ++++++++++++++- 2 files changed, 237 insertions(+), 3 deletions(-) diff --git a/src/compression/dct.rs b/src/compression/dct.rs index 8b13789..2cec060 100644 --- a/src/compression/dct.rs +++ b/src/compression/dct.rs @@ -1 +1,177 @@ +use std::f32::consts::{PI, SQRT_2}; +/// Perform a Discrete Cosine Transform on the input matrix. +pub fn dct(input: &[u8], width: usize, height: usize) -> Vec { + if input.len() != width * height { + panic!("Input matrix size must be width×height") + } + + let sqrt_width_zero = 1.0 / (width as f32).sqrt(); + let sqrt_width = SQRT_2 / (width as f32).sqrt(); + + let sqrt_height_zero = 1.0 / (height as f32).sqrt(); + let sqrt_height = SQRT_2 / (height as f32).sqrt(); + + let mut output = Vec::new(); + for u in 0..width { + for v in 0..height { + + // according to the formula of DCT + let cu = if u == 0 { + sqrt_width_zero + } else { + sqrt_width + }; + + let cv = if v == 0 { + sqrt_height_zero + } else { + sqrt_height + }; + + // calculate DCT + let mut tmp_sum = 0.0; + for x in 0..width { + for y in 0..height { + let dct = (input[x * width + y] as f32 - 128.0) * + f32::cos((2.0 * x as f32 + 1.0) * u as f32 * PI / (2.0 * width as f32)) * + f32::cos((2.0 * y as f32 + 1.0) * v as f32 * PI / (2.0 * height as f32)); + + tmp_sum += dct; + } + } + + output.push(cu * cv * tmp_sum) + } + } + + output +} + +/// Perform an inverse Discrete Cosine Transform on the input matrix. +pub fn idct(input: &[f32], width: usize, height: usize) -> Vec { + if input.len() != width * height { + panic!("Input matrix size must be width×height") + } + + let sqrt_width_zero = 1.0 / (width as f32).sqrt(); + let sqrt_width = SQRT_2 / (width as f32).sqrt(); + + let sqrt_height_zero = 1.0 / (height as f32).sqrt(); + let sqrt_height = SQRT_2 / (height as f32).sqrt(); + + let mut output = Vec::new(); + for x in 0..width { + for y in 0..height { + + let mut tmp_sum = 0.0; + for u in 0..width { + for v in 0..height { + let cu = if u == 0 { + sqrt_width_zero + } else { + sqrt_width + }; + + let cv = if v == 0 { + sqrt_height_zero + } else { + sqrt_height + }; + + let idct = input[u * width + v] as f32 * + f32::cos((2.0 * x as f32 + 1.0) * u as f32 * PI / (2.0 * width as f32)) * + f32::cos((2.0 * y as f32 + 1.0) * v as f32 * PI / (2.0 * height as f32)); + + tmp_sum += cu * cv * idct + } + } + + output.push((tmp_sum + 128.0) as u8) + } + } + + output +} + +/// JPEG 8x8 Base Quantization Matrix for a quality level of 50. +/// +/// Instead of using this, utilize the [`quantization_matrix`] function to +/// get a quantization matrix corresponding to the image quality value. +const BASE_QUANTIZATION_MATRIX: [u16; 64] = [ + 16, 11, 10, 16, 24, 40, 51, 61, + 12, 12, 14, 19, 26, 58, 60, 55, + 14, 13, 16, 24, 40, 57, 69, 56, + 14, 17, 22, 29, 51, 87, 80, 62, + 18, 22, 37, 56, 68, 109, 103, 77, + 24, 35, 55, 64, 81, 104, 113, 92, + 49, 64, 78, 87, 103, 121, 120, 101, + 72, 92, 95, 98, 112, 100, 103, 99, +]; + +/// Generate the 8x8 quantization matrix for the given quality level. +pub fn quantization_matrix(quality: u32) -> [u16; 64] { + let factor = if quality < 50 { + 5000.0 / quality as f32 + } else { + 200.0 - 2.0 * quality as f32 + }; + + let new_matrix = BASE_QUANTIZATION_MATRIX.map(|i| + f32::floor((factor * i as f32 + 50.0) / 100.0) as u16 + ); + new_matrix.map(|i| if i == 0 { 1 } else { i }) +} + +/// Quantize an input matrix, returning the result. +pub fn quantize(input: &[f32], quant_matrix: [u16; 64]) -> Vec { + 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 { + input.iter().zip(quant_matrix).map(|(v, q)| (*v as i16 * q as i16) as f32).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quantization_matrix_q80() { + let result = gen_quantization_matrix(80); + + assert_eq!( + result, + [ + 6, 4, 4, 6, 10, 16, 20, 24, + 5, 5, 6, 8, 10, 23, 24, 22, + 6, 5, 6, 10, 16, 23, 28, 22, + 6, 7, 9, 12, 20, 35, 32, 25, + 7, 9, 15, 22, 27, 44, 41, 31, + 10, 14, 22, 26, 32, 42, 45, 37, + 20, 26, 31, 35, 41, 48, 48, 40, + 29, 37, 38, 39, 45, 40, 41, 40 + ] + ); + } + + #[test] + fn quantization_matrix_q100() { + let result = gen_quantization_matrix(100); + + assert_eq!( + result, + [ + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0 + ] + ); + } +} diff --git a/src/main.rs b/src/main.rs index 59475a4..3523cf6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,7 +5,7 @@ mod compression { mod binio; mod header; mod operations; -mod picture; +pub mod picture; use picture::DangoPicture; use std::{ @@ -13,18 +13,75 @@ use std::{ io::{BufReader, BufWriter}, time::Instant, }; +use compression::dct::{dct, dequantize, idct, quantization_matrix, quantize}; -use image::RgbaImage; +use image::{ColorType, DynamicImage, GenericImage, GrayImage, Rgba}; fn main() { - let image_data = image::open("transparent2.png").unwrap().to_rgba8(); + let input = image::open("test_input.png").unwrap().to_luma8(); + input.save("original.png").unwrap(); + + let mut dct_image = Vec::new(); + for h in 0..input.height() as usize / 8 { + for w in 0..input.width() as usize / 8 { + let mut chunk = Vec::new(); + for i in 0..8 { + let start = (w * 8) + (h * 8) + (i * input.width() as usize); + let row = &input.as_raw()[start..start + 8]; + chunk.extend_from_slice(&row); + } + + if h + w == 0 { + println!("{:?}", chunk); + } + + // Perform the DCT on the image section + let dct: Vec = dct(&chunk, 8, 8); + let quantzied_dct = quantize(&dct, quantization_matrix(50)); + + dct_image.push(quantzied_dct); + } + } + + let mut decoded_image = DynamicImage::new(input.width(), input.height(), ColorType::L8); + for (i, chunk) in dct_image.iter().enumerate() { + let dequantized_dct = dequantize(chunk, quantization_matrix(50)); + let original = idct(&dequantized_dct, 8, 8); + + // Write rows of blocks + let start_x = (i * 8) % (input.width() as usize - 2); + let start_y = (i * 8) / input.width() as usize * 8; + dbg!(start_x); + dbg!(start_y); + let mut sub = decoded_image.sub_image(start_x as u32, start_y as u32, 8, 8); + for y in 0..8 { + for x in 0..8 { + let value = original[(y as usize * 8) + x as usize]; + sub.put_pixel(x, y, Rgba([value, value, value, 255])) + } + } + } + decoded_image.save(format!("test.png")).unwrap(); + + /* + // Reverse the DCT + let idct: Vec = idct(&dct, 8, 8).iter().map(|c| *c as u8).collect(); + + let img = GrayImage::from_raw(input.width(), input.height(), idct).unwrap(); + img.save("test.png").unwrap(); + */ + + /* + 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()); encoded_dpf.encode(&mut outfile); println!("Encoding took {}ms", timer.elapsed().as_millis()); + println!("DECODING ---"); let timer = Instant::now(); let mut infile = BufReader::new(File::open("test.dpf").unwrap()); let decoded_dpf = DangoPicture::decode(&mut infile).unwrap(); @@ -37,4 +94,5 @@ fn main() { ) .unwrap(); out_image.save("test.png").unwrap(); + */ }