Improved decode speeds again, fixed padding problems

This commit is contained in:
G2-Games 2024-07-28 16:25:33 -05:00
parent cf6d32e1aa
commit 4f2f101c03
5 changed files with 222 additions and 123 deletions

View file

@ -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<f32> {
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<f32> {
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<f32> {
/// 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<Vec<i16>> {
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,8 +158,10 @@ 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 {
for x in 0..((new_height / 8) * (new_width / 8)) {
let h = x / (new_width / 8);
let w = x % (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];
@ -174,59 +174,61 @@ pub fn dct_compress(input: &[u8], parameters: DctParameters) -> DctImage {
dct_channel.extend_from_slice(&quantized_dct);
}
}
dct_channel
}).collect();
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<i16>], parameters: DctParameters) -> Vec<u8> {
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);
});

View file

@ -43,10 +43,26 @@ impl CompressionInfo {
Ok(())
}
pub fn read_from<T: Read + ReadBytesExt>(input: &mut T) -> Self {
let mut compression_info = CompressionInfo {
chunk_count: input.read_u32::<LE>().unwrap() as usize,
chunks: Vec::new(),
};
for _ in 0..compression_info.chunk_count {
compression_info.chunks.push(ChunkInfo {
size_compressed: input.read_u32::<LE>().unwrap() as usize,
size_raw: input.read_u32::<LE>().unwrap() as usize,
});
}
compression_info
}
}
#[derive(Debug, Error)]
enum CompressionError {
pub enum CompressionError {
#[error("bad compressed element \"{1}\" at byte {2}")]
BadElement(Vec<u8>, u64, usize),
@ -54,7 +70,7 @@ enum CompressionError {
NoChunks,
}
pub fn compress(data: &[u8]) -> (Vec<u8>, CompressionInfo) {
pub fn compress(data: &[u8]) -> Result<(Vec<u8>, CompressionInfo), CompressionError> {
let mut part_data;
let mut offset = 0;
@ -84,18 +100,15 @@ pub fn compress(data: &[u8]) -> (Vec<u8>, 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<u8>) -> (usize, Vec<u8>, Vec<u8>) {
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<Vec<u8>, 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();

View file

@ -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::<LE>(self.width).unwrap();
buf.write_u32::<LE>(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<T: Read + ReadBytesExt>(input: &mut T) -> Result<Self, Error> {
let mut magic = [0u8; 8];
input.read_exact(&mut magic).unwrap();
if magic != *b"dangoimg" {
return Err(Error::InvalidIdentifier(magic));
}
#[derive(Debug, Clone, Copy)]
Ok(Header {
magic,
width: input.read_u32::<LE>()?,
height: input.read_u32::<LE>()?,
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, 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<u8> for ColorFormat {
type Error = String;
fn try_from(value: u8) -> Result<Self, Self::Error> {
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<u8> for CompressionType {
type Error = String;
fn try_from(value: u8) -> Result<Self, Self::Error> {
Ok(match value {
0 => Self::None,
1 => Self::Lossless,
2 => Self::LossyDct,
v => return Err(format!("invalid compression type {v}"))
})
}
}

View file

@ -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::<Vec<u8>>()).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::<Vec<u8>>());
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();

View file

@ -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<u8>,
}
#[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<O: Write + WriteBytesExt>(&self, mut output: O) {
pub fn from_raw(
width: u32,
height: u32,
color_format: ColorFormat,
compression_type: CompressionType,
bitmap: Vec<u8>,
) -> 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<O: Write + WriteBytesExt>(&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::<LE>().unwrap(),
height: input.read_u32::<LE>().unwrap(),
};
let header = Header::read_from(&mut input)?;
let mut compression_info = CompressionInfo {
chunk_count: input.read_u32::<LE>().unwrap() as usize,
chunks: Vec::new(),
};
for _ in 0..compression_info.chunk_count {
compression_info.chunks.push(ChunkInfo {
size_compressed: input.read_u32::<LE>().unwrap() as usize,
size_raw: input.read_u32::<LE>().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]),
}