361 lines
13 KiB
Rust
361 lines
13 KiB
Rust
use crate::colors::ColorType;
|
|
use crate::deflate;
|
|
use crate::error::PngError;
|
|
use crate::filters::*;
|
|
use crate::headers::*;
|
|
use crate::interlace::{deinterlace_image, interlace_image};
|
|
use byteorder::{BigEndian, WriteBytesExt};
|
|
use crc::crc32;
|
|
use rgb::ComponentSlice;
|
|
use rgb::RGBA8;
|
|
use std::collections::HashMap;
|
|
use std::fs::File;
|
|
use std::io::{Read, Seek, SeekFrom};
|
|
use std::iter::Iterator;
|
|
use std::path::Path;
|
|
use std::sync::Arc;
|
|
|
|
pub(crate) const STD_COMPRESSION: u8 = 6;
|
|
/// Must use normal compression, as faster ones (Huffman/RLE-only) are not representative
|
|
pub(crate) const STD_STRATEGY: u8 = 0;
|
|
/// OK to use a bit smalller window for evaluation
|
|
pub(crate) const STD_WINDOW: u8 = 13;
|
|
pub(crate) const STD_FILTERS: [u8; 2] = [0, 5];
|
|
|
|
pub(crate) mod scan_lines;
|
|
|
|
use self::scan_lines::{ScanLines, ScanLinesMut};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct PngImage {
|
|
/// The headers stored in the IHDR chunk
|
|
pub ihdr: IhdrData,
|
|
/// The uncompressed, optionally filtered data from the IDAT chunk
|
|
pub data: Vec<u8>,
|
|
/// The palette containing colors used in an Indexed image
|
|
/// Contains 3 bytes per color (R+G+B), up to 768
|
|
pub palette: Option<Vec<RGBA8>>,
|
|
/// The pixel value that should be rendered as transparent
|
|
pub transparency_pixel: Option<Vec<u8>>,
|
|
/// All non-critical headers from the PNG are stored here
|
|
pub aux_headers: HashMap<[u8; 4], Vec<u8>>,
|
|
}
|
|
|
|
/// Contains all data relevant to a PNG image
|
|
#[derive(Debug, Clone)]
|
|
pub struct PngData {
|
|
/// Uncompressed image data
|
|
pub raw: Arc<PngImage>,
|
|
/// The filtered and compressed data of the IDAT chunk
|
|
pub idat_data: Vec<u8>,
|
|
}
|
|
|
|
type PaletteWithTrns = (Option<Vec<RGBA8>>, Option<Vec<u8>>);
|
|
|
|
impl PngData {
|
|
/// Create a new `PngData` struct by opening a file
|
|
#[inline]
|
|
pub fn new(filepath: &Path, fix_errors: bool) -> Result<Self, PngError> {
|
|
let byte_data = Self::read_file(filepath)?;
|
|
|
|
Self::from_slice(&byte_data, fix_errors)
|
|
}
|
|
|
|
pub fn read_file(filepath: &Path) -> Result<Vec<u8>, PngError> {
|
|
let mut file = match File::open(filepath) {
|
|
Ok(f) => f,
|
|
Err(_) => return Err(PngError::new("Failed to open file for reading")),
|
|
};
|
|
// Check file for PNG header
|
|
let mut header = [0; 8];
|
|
if file.read_exact(&mut header).is_err() {
|
|
return Err(PngError::new("Not a PNG file: too small"));
|
|
}
|
|
if !file_header_is_valid(&header) {
|
|
return Err(PngError::new("Invalid PNG header detected"));
|
|
}
|
|
if file.seek(SeekFrom::Start(0)).is_err() {
|
|
return Err(PngError::new("Failed to read from file"));
|
|
}
|
|
// Read raw png data into memory
|
|
let mut byte_data: Vec<u8> =
|
|
Vec::with_capacity(file.metadata().map(|m| m.len() as usize).unwrap_or(0));
|
|
match file.read_to_end(&mut byte_data) {
|
|
Ok(_) => (),
|
|
Err(_) => return Err(PngError::new("Failed to read from file")),
|
|
}
|
|
Ok(byte_data)
|
|
}
|
|
|
|
/// Create a new `PngData` struct by reading a slice
|
|
pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<Self, PngError> {
|
|
let mut byte_offset: usize = 0;
|
|
// Test that png header is valid
|
|
let header = byte_data.get(0..8).ok_or(PngError::TruncatedData)?;
|
|
if !file_header_is_valid(header) {
|
|
return Err(PngError::NotPNG);
|
|
}
|
|
byte_offset += 8;
|
|
// Read the data headers
|
|
let mut aux_headers: HashMap<[u8; 4], Vec<u8>> = HashMap::new();
|
|
let mut idat_headers: Vec<u8> = Vec::new();
|
|
while let Some(header) = parse_next_header(byte_data, &mut byte_offset, fix_errors)? {
|
|
match &header.name {
|
|
b"IDAT" => idat_headers.extend(header.data),
|
|
b"acTL" => return Err(PngError::APNGNotSupported),
|
|
_ => {
|
|
aux_headers.insert(header.name, header.data.to_owned());
|
|
}
|
|
}
|
|
}
|
|
// Parse the headers into our PngData
|
|
if idat_headers.is_empty() {
|
|
return Err(PngError::ChunkMissing("IDAT"));
|
|
}
|
|
let ihdr = match aux_headers.remove(b"IHDR") {
|
|
Some(ihdr) => ihdr,
|
|
None => return Err(PngError::ChunkMissing("IHDR")),
|
|
};
|
|
let ihdr_header = parse_ihdr_header(&ihdr)?;
|
|
let raw_data = deflate::inflate(idat_headers.as_ref())?;
|
|
|
|
let (palette, transparency_pixel) = Self::palette_to_rgba(
|
|
ihdr_header.color_type,
|
|
aux_headers.remove(b"PLTE"),
|
|
aux_headers.remove(b"tRNS"),
|
|
)?;
|
|
|
|
let mut raw = PngImage {
|
|
ihdr: ihdr_header,
|
|
data: raw_data,
|
|
palette,
|
|
transparency_pixel,
|
|
aux_headers,
|
|
};
|
|
raw.data = raw.unfilter_image();
|
|
// Return the PngData
|
|
Ok(Self {
|
|
idat_data: idat_headers,
|
|
raw: Arc::new(raw),
|
|
})
|
|
}
|
|
|
|
/// Handle transparency header
|
|
fn palette_to_rgba(
|
|
color_type: ColorType,
|
|
palette_data: Option<Vec<u8>>,
|
|
trns_data: Option<Vec<u8>>,
|
|
) -> Result<PaletteWithTrns, PngError> {
|
|
if color_type == ColorType::Indexed {
|
|
let palette_data =
|
|
palette_data.ok_or_else(|| PngError::new("no palette in indexed image"))?;
|
|
let mut palette: Vec<_> = palette_data
|
|
.chunks(3)
|
|
.map(|color| RGBA8::new(color[0], color[1], color[2], 255))
|
|
.collect();
|
|
|
|
if let Some(trns_data) = trns_data {
|
|
for (color, trns) in palette.iter_mut().zip(trns_data) {
|
|
color.a = trns;
|
|
}
|
|
}
|
|
Ok((Some(palette), None))
|
|
} else {
|
|
Ok((None, trns_data))
|
|
}
|
|
}
|
|
|
|
/// Format the `PngData` struct into a valid PNG bytestream
|
|
pub fn output(&self) -> Vec<u8> {
|
|
// PNG header
|
|
let mut output = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
|
|
// IHDR
|
|
let mut ihdr_data = Vec::with_capacity(13);
|
|
let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.width);
|
|
let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.height);
|
|
let _ = ihdr_data.write_u8(self.raw.ihdr.bit_depth.as_u8());
|
|
let _ = ihdr_data.write_u8(self.raw.ihdr.color_type.png_header_code());
|
|
let _ = ihdr_data.write_u8(0); // Compression -- deflate
|
|
let _ = ihdr_data.write_u8(0); // Filter method -- 5-way adaptive filtering
|
|
let _ = ihdr_data.write_u8(self.raw.ihdr.interlaced);
|
|
write_png_block(b"IHDR", &ihdr_data, &mut output);
|
|
// Ancillary headers
|
|
for (key, header) in self
|
|
.raw
|
|
.aux_headers
|
|
.iter()
|
|
.filter(|&(key, _)| !(key == b"bKGD" || key == b"hIST" || key == b"tRNS"))
|
|
{
|
|
write_png_block(key, header, &mut output);
|
|
}
|
|
// Palette
|
|
if let Some(ref palette) = self.raw.palette {
|
|
let mut palette_data = Vec::with_capacity(palette.len() * 3);
|
|
let max_palette_size = 1 << (self.raw.ihdr.bit_depth.as_u8() as usize);
|
|
for px in palette.iter().take(max_palette_size) {
|
|
palette_data.extend_from_slice(px.rgb().as_slice());
|
|
}
|
|
write_png_block(b"PLTE", &palette_data, &mut output);
|
|
let num_transparent =
|
|
palette
|
|
.iter()
|
|
.take(max_palette_size)
|
|
.enumerate()
|
|
.fold(
|
|
0,
|
|
|prev, (index, px)| {
|
|
if px.a != 255 {
|
|
index + 1
|
|
} else {
|
|
prev
|
|
}
|
|
},
|
|
);
|
|
if num_transparent > 0 {
|
|
let trns_data: Vec<_> = palette[0..num_transparent].iter().map(|px| px.a).collect();
|
|
write_png_block(b"tRNS", &trns_data, &mut output);
|
|
}
|
|
} else if let Some(ref transparency_pixel) = self.raw.transparency_pixel {
|
|
// Transparency pixel
|
|
write_png_block(b"tRNS", transparency_pixel, &mut output);
|
|
}
|
|
// Special ancillary headers that need to come after PLTE but before IDAT
|
|
for (key, header) in self
|
|
.raw
|
|
.aux_headers
|
|
.iter()
|
|
.filter(|&(key, _)| key == b"bKGD" || key == b"hIST" || key == b"tRNS")
|
|
{
|
|
write_png_block(key, header, &mut output);
|
|
}
|
|
// IDAT data
|
|
write_png_block(b"IDAT", &self.idat_data, &mut output);
|
|
// Stream end
|
|
write_png_block(b"IEND", &[], &mut output);
|
|
|
|
output
|
|
}
|
|
}
|
|
|
|
impl PngImage {
|
|
/// Convert the image to the specified interlacing type
|
|
/// Returns true if the interlacing was changed, false otherwise
|
|
/// The `interlace` parameter specifies the *new* interlacing mode
|
|
/// Assumes that the data has already been de-filtered
|
|
#[inline]
|
|
#[must_use]
|
|
pub fn change_interlacing(&self, interlace: u8) -> Option<PngImage> {
|
|
if interlace == self.ihdr.interlaced {
|
|
return None;
|
|
}
|
|
|
|
Some(if interlace == 1 {
|
|
// Convert progressive to interlaced data
|
|
interlace_image(self)
|
|
} else {
|
|
// Convert interlaced to progressive data
|
|
deinterlace_image(self)
|
|
})
|
|
}
|
|
|
|
/// Return the number of channels in the image, based on color type
|
|
#[inline]
|
|
pub fn channels_per_pixel(&self) -> u8 {
|
|
self.ihdr.color_type.channels_per_pixel()
|
|
}
|
|
|
|
/// Return an iterator over the scanlines of the image
|
|
#[inline]
|
|
pub fn scan_lines(&self) -> ScanLines<'_> {
|
|
ScanLines::new(self)
|
|
}
|
|
|
|
/// Return an iterator over the scanlines of the image
|
|
#[inline]
|
|
pub fn scan_lines_mut(&mut self) -> ScanLinesMut<'_> {
|
|
ScanLinesMut::new(self)
|
|
}
|
|
|
|
/// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
|
|
fn unfilter_image(&self) -> Vec<u8> {
|
|
let mut unfiltered = Vec::with_capacity(self.data.len());
|
|
let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
|
|
let mut last_line: Vec<u8> = Vec::new();
|
|
let mut last_pass = 1;
|
|
for line in self.scan_lines() {
|
|
if let Some(pass) = line.pass {
|
|
if pass != last_pass {
|
|
last_line = Vec::new();
|
|
last_pass = pass;
|
|
}
|
|
}
|
|
let unfiltered_line = unfilter_line(line.filter, bpp, &line.data, &last_line);
|
|
unfiltered.push(0);
|
|
unfiltered.extend_from_slice(&unfiltered_line);
|
|
last_line = unfiltered_line;
|
|
}
|
|
unfiltered
|
|
}
|
|
|
|
/// Apply the specified filter type to all rows in the image
|
|
/// 0: None
|
|
/// 1: Sub
|
|
/// 2: Up
|
|
/// 3: Average
|
|
/// 4: Paeth
|
|
/// 5: All (heuristically pick the best filter for each line)
|
|
pub fn filter_image(&self, filter: u8) -> Vec<u8> {
|
|
let mut filtered = Vec::with_capacity(self.data.len());
|
|
let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
|
|
let mut last_line: &[u8] = &[];
|
|
let mut last_pass: Option<u8> = None;
|
|
for line in self.scan_lines() {
|
|
match filter {
|
|
0 | 1 | 2 | 3 | 4 => {
|
|
let filter = if last_pass == line.pass || filter <= 1 {
|
|
filter
|
|
} else {
|
|
0
|
|
};
|
|
filtered.push(filter);
|
|
filtered.extend_from_slice(&filter_line(filter, bpp, &line.data, last_line));
|
|
}
|
|
5 => {
|
|
// Heuristically guess best filter per line
|
|
// Uses MSAD algorithm mentioned in libpng reference docs
|
|
// http://www.libpng.org/pub/png/book/chapter09.html
|
|
let mut trials: Vec<(u8, Vec<u8>)> = Vec::with_capacity(5);
|
|
// Avoid vertical filtering on first line of each interlacing pass
|
|
for filter in if last_pass == line.pass { 0..5 } else { 0..2 } {
|
|
trials.push((filter, filter_line(filter, bpp, &line.data, last_line)));
|
|
}
|
|
let (best_filter, best_line) = trials
|
|
.iter()
|
|
.min_by_key(|(_, line)| {
|
|
line.iter().fold(0u64, |acc, &x| {
|
|
let signed = x as i8;
|
|
acc + i16::from(signed).abs() as u64
|
|
})
|
|
})
|
|
.unwrap();
|
|
filtered.push(*best_filter);
|
|
filtered.extend_from_slice(best_line);
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
last_line = line.data;
|
|
last_pass = line.pass;
|
|
}
|
|
filtered
|
|
}
|
|
}
|
|
fn write_png_block(key: &[u8], header: &[u8], output: &mut Vec<u8>) {
|
|
let mut header_data = Vec::with_capacity(header.len() + 4);
|
|
header_data.extend_from_slice(key);
|
|
header_data.extend_from_slice(header);
|
|
output.reserve(header_data.len() + 8);
|
|
let _ = output.write_u32::<BigEndian>(header_data.len() as u32 - 4);
|
|
let crc = crc32::checksum_ieee(&header_data);
|
|
output.append(&mut header_data);
|
|
let _ = output.write_u32::<BigEndian>(crc);
|
|
}
|