From e1a52c7077cad6935981193829c5d93aa6870105 Mon Sep 17 00:00:00 2001 From: G2-Games Date: Sun, 28 Jul 2024 18:29:02 -0500 Subject: [PATCH] Added proper encoding and decoding, fixed minor issues --- Cargo.toml | 1 + src/compression/dct.rs | 30 ++++++------ src/header.rs | 21 +++++++- src/main.rs | 106 ++++++++++++++++++----------------------- src/picture.rs | 90 +++++++++++++++++++++++++++------- 5 files changed, 155 insertions(+), 93 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 380e456..ac83169 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/compression/dct.rs b/src/compression/dct.rs index 008c8ea..92af693 100644 --- a/src/compression/dct.rs +++ b/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 { 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 { /// 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") + 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 { - 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 { - 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 { /// to a multiple of 8, which must be reversed when decoding. pub fn dct_compress(input: &[u8], parameters: DctParameters) -> Vec> { 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> { .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> = 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> { /// 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], parameters: DctParameters) -> Vec { +pub fn dct_decompress(input: &[i16], parameters: DctParameters) -> Vec { 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_iter().enumerate().for_each(|(chan_num, channel)| { - println!("Decoding channel {chan_num}"); - + input.par_chunks(new_width * new_height).enumerate().for_each(|(chan_num, channel)| { 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); diff --git a/src/header.rs b/src/header.rs index 042fabe..a2eaad3 100644 --- a/src/header.rs +++ b/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::(self.width).unwrap(); buf.write_u32::(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 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 for CompressionType { }) } } + +impl Into for CompressionType { + fn into(self) -> u8 { + match self { + CompressionType::None => 0, + CompressionType::Lossless => 1, + CompressionType::LossyDct => 2, + } + } +} diff --git a/src/main.rs b/src/main.rs index 0d78686..0130dbc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,80 +7,66 @@ mod header; mod operations; pub mod picture; -use std::{fs::File, io::Write, time::Instant}; - -use header::ColorFormat; -use compression::{dct::{dct_compress, dct_decompress, DctParameters}, lossless}; +use std::{fs::File, io::{BufReader, BufWriter}, time::Instant}; +use header::{ColorFormat, CompressionType}; +use image::RgbaImage; +use picture::DangoPicture; fn main() { let input = image::open("transparent2.png").unwrap().to_rgba8(); input.save("original.png").unwrap(); - let timer = Instant::now(); - 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::>()).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( + let dpf_lossy = DangoPicture::from_raw( input.width(), input.height(), - decoded_dct - ).unwrap().save("dct-final.png").unwrap(); + ColorFormat::Rgba32, + CompressionType::LossyDct, + Some(10), + input.as_raw().clone() + ); - /* - // Reverse the DCT - let idct: Vec = idct(&dct, 8, 8).iter().map(|c| *c as u8).collect(); + let dpf_lossless = DangoPicture::from_raw( + input.width(), + input.height(), + ColorFormat::Rgba32, + CompressionType::Lossless, + None, + input.as_raw().clone() + ); - 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 ---"); + println!("\n--- LOSSY ---"); + println!("Encoding"); let timer = Instant::now(); - let mut outfile = BufWriter::new(File::create("test.dpf").unwrap()); - encoded_dpf.encode(&mut outfile); + let mut outfile = BufWriter::new(std::fs::File::create("test-lossy.dpf").unwrap()); + 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( - decoded_dpf.header.width, - decoded_dpf.header.height, - decoded_dpf.bitmap.into(), - ) - .unwrap(); - out_image.save("test.png").unwrap(); - */ + println!("\n--- LOSSLESS ---"); + println!("Encoding"); + let timer = Instant::now(); + let mut outfile = BufWriter::new(std::fs::File::create("test-lossless.dpf").unwrap()); + dpf_lossless.encode(&mut outfile).unwrap(); + println!("Encoding took {}ms", timer.elapsed().as_millis()); + + 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()); } diff --git a/src/picture.rs b/src/picture.rs index 6b2b001..3e00289 100644 --- a/src/picture.rs +++ b/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, bitmap: Vec, ) -> 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>(&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(mut input: I) -> Result { - 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 }) }