From 4f2f101c03580b53570e3d4e4e50de8a35ef527d Mon Sep 17 00:00:00 2001 From: G2-Games Date: Sun, 28 Jul 2024 16:25:33 -0500 Subject: [PATCH] Improved decode speeds again, fixed padding problems --- src/compression/dct.rs | 90 +++++++++++++++--------------- src/compression/lossless.rs | 29 +++++++--- src/header.rs | 93 ++++++++++++++++++++++++++++--- src/main.rs | 27 ++++----- src/picture.rs | 106 ++++++++++++++++++++---------------- 5 files changed, 222 insertions(+), 123 deletions(-) diff --git a/src/compression/dct.rs b/src/compression/dct.rs index e16b8b2..008c8ea 100644 --- a/src/compression/dct.rs +++ b/src/compression/dct.rs @@ -1,6 +1,6 @@ use std::{f32::consts::{PI, SQRT_2}, sync::{Arc, Mutex}}; -use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, ParallelIterator}; +use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator}; use crate::header::ColorFormat; @@ -20,7 +20,6 @@ pub fn dct(input: &[u8], width: usize, height: usize) -> Vec { 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 { @@ -33,7 +32,6 @@ pub fn dct(input: &[u8], width: usize, height: usize) -> Vec { sqrt_height }; - // calculate DCT let mut tmp_sum = 0.0; for x in 0..width { for y in 0..height { @@ -140,7 +138,7 @@ pub fn dequantize(input: &[i16], quant_matrix: [u16; 64]) -> Vec { /// 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], parameters: DctParameters) -> DctImage { +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 quantization_matrix = quantization_matrix(parameters.quality); @@ -160,20 +158,21 @@ pub fn dct_compress(input: &[u8], parameters: DctParameters) -> DctImage { img_2d.resize(new_height, vec![0u8; new_width]); 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); - } + for x in 0..((new_height / 8) * (new_width / 8)) { + let h = x / (new_width / 8); + let w = x % (new_width / 8); - // Perform the DCT on the image section - let dct: Vec = dct(&chunk, 8, 8); - let quantized_dct = quantize(&dct, quantization_matrix); - - dct_channel.extend_from_slice(&quantized_dct); + 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 quantized_dct = quantize(&dct, quantization_matrix); + + dct_channel.extend_from_slice(&quantized_dct); } dct_channel @@ -181,52 +180,55 @@ pub fn dct_compress(input: &[u8], parameters: DctParameters) -> DctImage { channels.into_iter().for_each(|c| dct_image.push(c)); - DctImage { - channels: dct_image, - width: new_width as u32, - height: new_height as u32 - } + dct_image } -/// Take in +/// 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 { + let new_width = parameters.width + (8 - parameters.width % 8); + let new_height = parameters.height + (8 - parameters.width % 8); + // Precalculate the quantization matrix let quantization_matrix = quantization_matrix(parameters.quality); - // Number of bytes per row of chunks - let chunk_row = parameters.width * 8; + let final_img = Arc::new(Mutex::new(vec![0u8; (new_width * new_height) * parameters.format.channels() as usize])); - let final_img = Arc::new(Mutex::new(vec![0u8; (parameters.width * parameters.height) * 4])); input.par_iter().enumerate().for_each(|(chan_num, channel)| { println!("Decoding channel {chan_num}"); - let mut decoded_image = vec![]; - for (i, chunk) in channel.windows(64).step_by(64).enumerate() { - - // Allocate a new row of the image for every new row of chunks - if i % (parameters.width / 8) == 0 { - decoded_image.extend_from_slice(&vec![0u8; chunk_row]); - } - - let dequantized_dct = dequantize(chunk, quantization_matrix); + 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)| { + let dequantized_dct = dequantize(&chunk, quantization_matrix); let original = idct(&dequantized_dct, 8, 8); // Write rows of blocks - let start_x = (i * 8) % parameters.width; - let start_y = ((i * 8) / parameters.width) * 8; + let start_x = (i * 8) % new_width; + let start_y = ((i * 8) / new_width) * 8; let start = start_x + (start_y * parameters.width); for row_num in 0..8 { - let row_offset = row_num * parameters.width; - let row_data = &original[row_num * 8..(row_num * 8) + 8]; - decoded_image[start + row_offset..start + row_offset + 8].copy_from_slice(row_data); - } - } + if start_y + row_num >= parameters.height { + break; + } - final_img.lock().unwrap().iter_mut() + let row_offset = row_num * parameters.width; + + let offset = if start_x + 8 >= parameters.width { + parameters.width % 8 + } else { + 8 + }; + + let row_data = &original[row_num * 8..(row_num * 8) + offset]; + decoded_image.lock().unwrap()[start + row_offset..start + row_offset + offset].copy_from_slice(row_data); + } + }); + + final_img.lock().unwrap().par_iter_mut() .skip(chan_num) - .step_by(4) - .zip(decoded_image.iter()) + .step_by(parameters.format.channels() as usize) + .zip(decoded_image.lock().unwrap().par_iter()) .for_each(|(c, n)| *c = *n); }); diff --git a/src/compression/lossless.rs b/src/compression/lossless.rs index 08ed7e7..0233201 100644 --- a/src/compression/lossless.rs +++ b/src/compression/lossless.rs @@ -43,10 +43,26 @@ impl CompressionInfo { Ok(()) } + + pub fn read_from(input: &mut T) -> Self { + let mut compression_info = CompressionInfo { + chunk_count: input.read_u32::().unwrap() as usize, + chunks: Vec::new(), + }; + + for _ in 0..compression_info.chunk_count { + compression_info.chunks.push(ChunkInfo { + size_compressed: input.read_u32::().unwrap() as usize, + size_raw: input.read_u32::().unwrap() as usize, + }); + } + + compression_info + } } #[derive(Debug, Error)] -enum CompressionError { +pub enum CompressionError { #[error("bad compressed element \"{1}\" at byte {2}")] BadElement(Vec, u64, usize), @@ -54,7 +70,7 @@ enum CompressionError { NoChunks, } -pub fn compress(data: &[u8]) -> (Vec, CompressionInfo) { +pub fn compress(data: &[u8]) -> Result<(Vec, CompressionInfo), CompressionError> { let mut part_data; let mut offset = 0; @@ -84,18 +100,15 @@ pub fn compress(data: &[u8]) -> (Vec, CompressionInfo) { } if output_info.chunk_count == 0 { - panic!("No chunks compressed!") + return Err(CompressionError::NoChunks) } - (output_buf, output_info) + Ok((output_buf, output_info)) } fn compress_lzw(data: &[u8], last: Vec) -> (usize, Vec, Vec) { let mut count = 0; - let mut dictionary = HashMap::new(); - for i in 0..=255 { - dictionary.insert(vec![i], i as u64); - } + let mut dictionary: HashMap, u64> = HashMap::from_iter((0..=255).into_iter().map(|i| (vec![i], i as u64))); let mut dictionary_count = (dictionary.len() + 1) as u64; let mut element = Vec::new(); diff --git a/src/header.rs b/src/header.rs index f18cc3f..042fabe 100644 --- a/src/header.rs +++ b/src/header.rs @@ -1,13 +1,29 @@ -use byteorder::{WriteBytesExt, LE}; -use std::io::{Cursor, Write}; +use byteorder::{ReadBytesExt, WriteBytesExt, LE}; +use std::io::{Cursor, Read, Write}; +use crate::picture::Error; + +/// A DPF file header. This must be included at the beginning +/// of a valid DPF file. pub struct Header { + /// Identifier. Must be set to "dangoimg". pub magic: [u8; 8], - /// Width of the image in pixels + /// Width of the image in pixels. pub width: u32, - /// Height of the image in pixels + + /// Height of the image in pixels. pub height: u32, + + /// Type of compression used on the data. + pub compression_type: CompressionType, + + /// Level of compression. Only applies in Lossy mode, otherwise this value + /// should be set to -1. + pub compression_level: i8, + + /// Format of color data in the image. + pub color_format: ColorFormat, } impl Default for Header { @@ -16,6 +32,9 @@ impl Default for Header { magic: *b"dangoimg", width: 0, height: 0, + compression_type: CompressionType::Lossless, + compression_level: -1, + color_format: ColorFormat::Rgba32, } } } @@ -28,17 +47,39 @@ impl Header { buf.write_u32::(self.width).unwrap(); buf.write_u32::(self.height).unwrap(); + buf.write_u8(self.compression_type as u8).unwrap(); + buf.write_i8(self.compression_level).unwrap(); + buf.into_inner().try_into().unwrap() } + + pub fn read_from(input: &mut T) -> Result { + let mut magic = [0u8; 8]; + input.read_exact(&mut magic).unwrap(); + + if magic != *b"dangoimg" { + return Err(Error::InvalidIdentifier(magic)); + } + + Ok(Header { + magic, + width: input.read_u32::()?, + height: input.read_u32::()?, + + compression_type: input.read_u8()?.try_into().unwrap(), + compression_level: input.read_i8()?, + color_format: input.read_u8()?.try_into().unwrap(), + }) + } } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ColorFormat { /// RGBA, 8 bits per channel - Rgba32, + Rgba32 = 0, /// RGB, 8 bits per channel - Rgb24, + Rgb24 = 1, } impl ColorFormat { @@ -72,3 +113,41 @@ impl ColorFormat { } } } + +impl TryFrom for ColorFormat { + type Error = String; + + fn try_from(value: u8) -> Result { + Ok(match value { + 0 => Self::Rgba32, + 1 => Self::Rgb24, + v => return Err(format!("invalid color format {v}")), + }) + } +} + +/// The type of compression used in the image +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompressionType { + /// No compression at all, raw bitmap + None = 0, + + /// Lossless compression + Lossless = 1, + + /// Lossy Discrete Cosine Transform compression + LossyDct = 2, +} + +impl TryFrom for CompressionType { + type Error = String; + + fn try_from(value: u8) -> Result { + Ok(match value { + 0 => Self::None, + 1 => Self::Lossless, + 2 => Self::LossyDct, + v => return Err(format!("invalid compression type {v}")) + }) + } +} diff --git a/src/main.rs b/src/main.rs index 6f7b08f..0d78686 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,47 +12,42 @@ use std::{fs::File, io::Write, time::Instant}; use header::ColorFormat; use compression::{dct::{dct_compress, dct_decompress, DctParameters}, lossless}; -use image::RgbaImage; -use picture::DangoPicture; - fn main() { - let input = image::open("kirara_motorbike.jpg").unwrap().to_rgba8(); + let input = image::open("transparent2.png").unwrap().to_rgba8(); input.save("original.png").unwrap(); - let dct_output = File::create("test.dpf").unwrap(); - DangoPicture::from_raw(input.width(), input.height(), &input.as_raw().clone()).encode(&dct_output); - let timer = Instant::now(); let dct_result = dct_compress( input.as_raw(), DctParameters { - quality: 30, + 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(); - let compressed_dct = lossless::compress(&dct_result.channels.concat().iter().flat_map(|x| x.to_le_bytes()).collect::>()); dct_output.write_all(&compressed_dct.0).unwrap(); let timer = Instant::now(); let decoded_dct = dct_decompress( - &dct_result.channels, + &dct_result, DctParameters { - quality: 30, + quality: 68, format: ColorFormat::Rgba32, - width: dct_result.width as usize, - height: dct_result.height as usize + width: input.width() as usize, + height: input.height() as usize } ); println!("Decoding took {}ms", timer.elapsed().as_millis()); - RgbaImage::from_raw( - dct_result.width, - dct_result.height, + image::RgbaImage::from_raw( + input.width(), + input.height(), decoded_dct ).unwrap().save("dct-final.png").unwrap(); diff --git a/src/picture.rs b/src/picture.rs index 2624aa5..6b2b001 100644 --- a/src/picture.rs +++ b/src/picture.rs @@ -1,11 +1,11 @@ -use std::io::{Read, Write}; +use std::io::{self, Read, Write}; -use byteorder::{ReadBytesExt, WriteBytesExt, LE}; +use byteorder::{ReadBytesExt, WriteBytesExt}; use thiserror::Error; use crate::{ - compression::lossless::{compress, decompress, ChunkInfo, CompressionInfo}, - header::Header, + compression::{dct::{dct_compress, DctParameters}, lossless::{compress, decompress, CompressionError, CompressionInfo}}, + header::{ColorFormat, CompressionType, Header}, operations::{diff_line, line_diff}, }; @@ -14,29 +14,73 @@ pub struct DangoPicture { pub bitmap: Vec, } +#[derive(Error, Debug)] +pub enum Error { + #[error("incorrect identifier, got {0:?}")] + InvalidIdentifier([u8; 8]), + + #[error("io operation failed: {0}")] + IoError(#[from] io::Error), + + #[error("compression operation failed: {0}")] + CompressionError(#[from] CompressionError), +} + impl DangoPicture { - /// Encode the image into anything that implements [Write] - pub fn encode(&self, mut output: O) { + pub fn from_raw( + width: u32, + height: u32, + color_format: ColorFormat, + compression_type: CompressionType, + bitmap: Vec, + ) -> Self { let header = Header { - width: self.header.width, - height: self.header.height, + width, + height, + + compression_type, + color_format, ..Default::default() }; - // Write out the header - output.write_all(&header.to_bytes()).unwrap(); + DangoPicture { + header, + bitmap, + } + } - let modified_data = diff_line(header.width, header.height, &self.bitmap); + /// Encode the image into anything that implements [Write] + pub fn encode(&self, mut output: O) -> Result<(), Error> { + // Write out the header + output.write_all(&self.header.to_bytes()).unwrap(); + + 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::LossyDct => { + &dct_compress( + &self.bitmap, + DctParameters { + quality: self.header.compression_level as u32, + format: self.header.color_format, + width: self.header.width as usize, + height: self.header.height as usize, + } + ).concat().iter().flat_map(|i| i.to_le_bytes()).collect() + }, + }; // Compress the image data - let (compressed_data, compression_info) = compress(&modified_data); + let (compressed_data, compression_info) = compress(&modified_data)?; // Write out compression info compression_info.write_into(&mut output).unwrap(); // Write out compressed data output.write_all(&compressed_data).unwrap(); + + Ok(()) } /// Decode the image from anything that implements [Read] @@ -48,23 +92,9 @@ impl DangoPicture { return Err(Error::InvalidIdentifier(magic)); } - let header = Header { - magic, - width: input.read_u32::().unwrap(), - height: input.read_u32::().unwrap(), - }; + let header = Header::read_from(&mut input)?; - let mut compression_info = CompressionInfo { - chunk_count: input.read_u32::().unwrap() as usize, - chunks: Vec::new(), - }; - - for _ in 0..compression_info.chunk_count { - compression_info.chunks.push(ChunkInfo { - size_compressed: input.read_u32::().unwrap() as usize, - size_raw: input.read_u32::().unwrap() as usize, - }); - } + let compression_info = CompressionInfo::read_from(&mut input); let preprocessed_bitmap = decompress(&mut input, &compression_info); @@ -72,24 +102,4 @@ impl DangoPicture { Ok(DangoPicture { header, bitmap }) } - - pub fn from_raw(width: u32, height: u32, bitmap: &[u8]) -> Self { - let header = Header { - width, - height, - - ..Default::default() - }; - - DangoPicture { - header, - bitmap: bitmap.into(), - } - } -} - -#[derive(Error, Debug)] -pub enum Error { - #[error("incorrect identifier, got {}", 0)] - InvalidIdentifier([u8; 8]), }