png.rs
11.7 kB · rust · 377 lines
1use crate::errors::{value_error, MrlyError, Result};2use crate::resample::block;3use png::{4 AdaptiveFilterType, BitDepth, ColorType, Compression, Decoder, Encoder, FilterType,5 Transformations,6};78impl From<png::EncodingError> for MrlyError {9 fn from(error: png::EncodingError) -> MrlyError {10 MrlyError::Value(error.to_string())11 }12}1314impl From<png::DecodingError> for MrlyError {15 fn from(error: png::DecodingError) -> MrlyError {16 MrlyError::Value(error.to_string())17 }18}1920/// Encodes rgba colors as a png, drawing each source pixel as a scale by scale block.21///22/// Images of 256 colors or fewer are written as a palette png at the smallest bit depth that23/// fits; everything else stays 8-bit rgba. Both forms read back through [`unpng`].24pub fn png(colors: &[[u8; 4]], width: usize, height: usize, scale: usize) -> Result<Vec<u8>> {25 if scale < 1 {26 return value_error("scale must be at least 1.");27 }28 if colors.len() != width * height {29 return value_error("colors length must equal width * height.");30 }31 let pixels = block(colors, width, height, scale);32 let (width, height) = (width * scale, height * scale);33 let mut bytes = Vec::with_capacity(pixels.len() + 128);34 {35 let mut encoder = Encoder::new(&mut bytes, width as u32, height as u32);36 encoder.set_compression(Compression::Best);37 encoder.set_filter(FilterType::Paeth);38 encoder.set_adaptive_filter(AdaptiveFilterType::Adaptive);39 let data = match table(&pixels) {40 Some(table) => {41 let depth = depth(table.len());42 encoder.set_color(ColorType::Indexed);43 encoder.set_depth(depth);44 encoder.set_palette(plte(&table));45 let veils = trns(&table);46 if !veils.is_empty() {47 encoder.set_trns(veils);48 }49 indices(&pixels, &table, width, height, depth)50 }51 None => {52 encoder.set_color(ColorType::Rgba);53 encoder.set_depth(BitDepth::Eight);54 pixels.concat()55 }56 };57 let mut writer = encoder.write_header()?;58 writer.write_image_data(&data)?;59 writer.finish()?;60 }61 Ok(bytes)62}6364fn table(pixels: &[[u8; 4]]) -> Option<Vec<[u8; 4]>> {65 let mut keys: Vec<u32> = Vec::with_capacity(256);66 for pixel in pixels {67 let key = u32::from_be_bytes(*pixel);68 if let Err(slot) = keys.binary_search(&key) {69 if keys.len() == 256 {70 return None;71 }72 keys.insert(slot, key);73 }74 }75 let mut table: Vec<[u8; 4]> = keys.into_iter().map(u32::to_be_bytes).collect();76 table.sort_by_key(|color| color[3]);77 Some(table)78}7980fn depth(count: usize) -> BitDepth {81 match count {82 0..=2 => BitDepth::One,83 3..=4 => BitDepth::Two,84 5..=16 => BitDepth::Four,85 _ => BitDepth::Eight,86 }87}8889fn plte(table: &[[u8; 4]]) -> Vec<u8> {90 table91 .iter()92 .flat_map(|color| [color[0], color[1], color[2]])93 .collect()94}9596fn trns(table: &[[u8; 4]]) -> Vec<u8> {97 let opaque = table98 .iter()99 .position(|color| color[3] == 255)100 .unwrap_or(table.len());101 table[..opaque].iter().map(|color| color[3]).collect()102}103104fn indices(105 pixels: &[[u8; 4]],106 table: &[[u8; 4]],107 width: usize,108 height: usize,109 depth: BitDepth,110) -> Vec<u8> {111 let mut lookup: Vec<(u32, u8)> = table112 .iter()113 .enumerate()114 .map(|(slot, color)| (u32::from_be_bytes(*color), slot as u8))115 .collect();116 lookup.sort_unstable();117 let bits = depth as usize;118 let per_byte = 8 / bits;119 let stride = width.div_ceil(per_byte);120 let mut data = vec![0u8; stride * height];121 if width == 0 {122 return data;123 }124 for (y, row) in pixels.chunks_exact(width).enumerate() {125 for (x, pixel) in row.iter().enumerate() {126 let key = u32::from_be_bytes(*pixel);127 let slot = lookup.binary_search_by_key(&key, |entry| entry.0).unwrap();128 let shift = 8 - bits * (x % per_byte + 1);129 data[y * stride + x / per_byte] |= lookup[slot].1 << shift;130 }131 }132 data133}134135/// Decodes a png to its width, height, and rgba colors, or an error for a broken file.136///137/// Grayscale, rgb, palette and 16-bit files all come back as 8-bit rgba.138pub fn unpng(bytes: &[u8]) -> Result<(usize, usize, Vec<[u8; 4]>)> {139 let mut decoder = Decoder::new(bytes);140 decoder.set_transformations(Transformations::normalize_to_color8());141 let mut reader = decoder.read_info()?;142 let mut data = vec![0u8; reader.output_buffer_size()];143 let info = reader.next_frame(&mut data)?;144 let data = &data[..info.buffer_size()];145 let colors = match info.color_type {146 ColorType::Grayscale => data.iter().map(|&g| [g, g, g, 255]).collect(),147 ColorType::GrayscaleAlpha => data148 .chunks_exact(2)149 .map(|p| [p[0], p[0], p[0], p[1]])150 .collect(),151 ColorType::Rgb => data152 .chunks_exact(3)153 .map(|p| [p[0], p[1], p[2], 255])154 .collect(),155 ColorType::Rgba => data156 .chunks_exact(4)157 .map(|p| [p[0], p[1], p[2], p[3]])158 .collect(),159 ColorType::Indexed => return value_error("png palette did not expand."),160 };161 Ok((info.width as usize, info.height as usize, colors))162}163164#[cfg(test)]165mod tests {166 use super::*;167 fn raw(168 color: ColorType,169 depth: BitDepth,170 width: u32,171 height: u32,172 data: &[u8],173 palette: &[u8],174 trns: &[u8],175 ) -> Vec<u8> {176 let mut bytes = Vec::new();177 let mut encoder = Encoder::new(&mut bytes, width, height);178 encoder.set_color(color);179 encoder.set_depth(depth);180 if !palette.is_empty() {181 encoder.set_palette(palette.to_vec());182 }183 if !trns.is_empty() {184 encoder.set_trns(trns.to_vec());185 }186 let mut writer = encoder.write_header().unwrap();187 writer.write_image_data(data).unwrap();188 writer.finish().unwrap();189 bytes190 }191 #[test]192 fn png_signature_and_scaled_size() {193 let colors = vec![194 [255, 0, 0, 255],195 [0, 255, 0, 255],196 [0, 0, 255, 255],197 [255, 255, 0, 255],198 ];199 let bytes = png(&colors, 2, 2, 4).unwrap();200 assert_eq!(&bytes[0..8], &[137, 80, 78, 71, 13, 10, 26, 10]);201 assert_eq!(&bytes[16..20], &8u32.to_be_bytes());202 assert_eq!(&bytes[20..24], &8u32.to_be_bytes());203 }204 #[test]205 fn png_rejects_bad_inputs() {206 let colors = vec![[0, 0, 0, 255]];207 assert!(png(&colors, 1, 1, 0).is_err());208 assert!(png(&colors, 2, 2, 1).is_err());209 assert!(png(&[], 0, 0, 1).is_err());210 }211 #[test]212 fn unpng_round_trips_the_encoder() {213 let colors = vec![214 [255, 0, 0, 255],215 [0, 255, 0, 128],216 [0, 0, 255, 0],217 [7, 8, 9, 10],218 [250, 251, 252, 253],219 [1, 1, 1, 255],220 ];221 let bytes = png(&colors, 3, 2, 1).unwrap();222 let (w, h, out) = unpng(&bytes).unwrap();223 assert_eq!((w, h), (3, 2));224 assert_eq!(out, colors);225 }226 #[test]227 fn unpng_round_trips_scaled_output() {228 let colors = vec![[9, 9, 9, 255], [0, 0, 0, 0]];229 let bytes = png(&colors, 2, 1, 3).unwrap();230 let (w, h, out) = unpng(&bytes).unwrap();231 assert_eq!((w, h), (6, 3));232 assert_eq!(out[0], [9, 9, 9, 255]);233 assert_eq!(out[5], [0, 0, 0, 0]);234 assert_eq!(out.len(), 18);235 }236 #[test]237 fn unpng_expands_gray_and_rgb() {238 let gray = raw(239 ColorType::Grayscale,240 BitDepth::Eight,241 3,242 1,243 &[0, 128, 255],244 &[],245 &[],246 );247 let (w, h, out) = unpng(&gray).unwrap();248 assert_eq!((w, h), (3, 1));249 assert_eq!(250 out,251 [[0, 0, 0, 255], [128, 128, 128, 255], [255, 255, 255, 255]]252 );253 let bits = raw(254 ColorType::Grayscale,255 BitDepth::One,256 3,257 1,258 &[0b1010_0000],259 &[],260 &[],261 );262 let (_, _, out) = unpng(&bits).unwrap();263 assert_eq!(264 out,265 [[255, 255, 255, 255], [0, 0, 0, 255], [255, 255, 255, 255]]266 );267 let veiled = raw(268 ColorType::GrayscaleAlpha,269 BitDepth::Eight,270 1,271 1,272 &[7, 9],273 &[],274 &[],275 );276 assert_eq!(unpng(&veiled).unwrap().2, [[7, 7, 7, 9]]);277 let rgb = raw(278 ColorType::Rgb,279 BitDepth::Eight,280 1,281 2,282 &[1, 2, 3, 4, 5, 6],283 &[],284 &[],285 );286 let (w, h, out) = unpng(&rgb).unwrap();287 assert_eq!((w, h), (1, 2));288 assert_eq!(out, [[1, 2, 3, 255], [4, 5, 6, 255]]);289 }290 #[test]291 fn unpng_expands_palette_and_transparency() {292 let palette = [10, 20, 30, 40, 50, 60, 70, 80, 90];293 let paletted = raw(294 ColorType::Indexed,295 BitDepth::Two,296 4,297 1,298 &[0b0001_1001],299 &palette,300 &[0],301 );302 let (w, h, out) = unpng(&paletted).unwrap();303 assert_eq!((w, h), (4, 1));304 assert_eq!(305 out,306 [307 [10, 20, 30, 0],308 [40, 50, 60, 255],309 [70, 80, 90, 255],310 [40, 50, 60, 255]311 ]312 );313 let opaque = raw(314 ColorType::Indexed,315 BitDepth::Eight,316 2,317 1,318 &[2, 0],319 &palette,320 &[],321 );322 assert_eq!(323 unpng(&opaque).unwrap().2,324 [[70, 80, 90, 255], [10, 20, 30, 255]]325 );326 }327 #[test]328 fn unpng_strips_sixteen_bit_samples() {329 let deep = [0x12, 0x34, 0xff, 0xff, 0x80, 0x00, 0x00, 0xff];330 let bytes = raw(ColorType::Rgba, BitDepth::Sixteen, 1, 1, &deep, &[], &[]);331 assert_eq!(unpng(&bytes).unwrap().2, [[0x12, 0xff, 0x80, 0x00]]);332 }333 #[test]334 fn unpng_rejects_garbage() {335 assert!(unpng(&[]).is_err());336 assert!(unpng(b"not a png at all").is_err());337 let mut bytes = png(&[[1, 2, 3, 4]], 1, 1, 1).unwrap();338 bytes[25] = 16;339 assert!(unpng(&bytes).is_err());340 bytes.truncate(40);341 assert!(unpng(&bytes).is_err());342 }343344 fn form(bytes: &[u8]) -> (ColorType, BitDepth) {345 let reader = Decoder::new(bytes).read_info().unwrap();346 let info = reader.info();347 (info.color_type, info.bit_depth)348 }349 #[test]350 fn png_packs_few_colors_into_a_palette() {351 let mut colors = Vec::with_capacity(64 * 64);352 for y in 0..64usize {353 for x in 0..64usize {354 colors.push(if (x / 8 + y / 8) % 2 == 0 {355 [17, 34, 51, 255]356 } else {357 [238, 221, 204, 64]358 });359 }360 }361 let bytes = png(&colors, 64, 64, 1).unwrap();362 assert_eq!(form(&bytes), (ColorType::Indexed, BitDepth::One));363 assert!(bytes.len() < 1024, "{} bytes", bytes.len());364 let (w, h, out) = unpng(&bytes).unwrap();365 assert_eq!((w, h), (64, 64));366 assert_eq!(out, colors);367 }368 #[test]369 fn png_keeps_rgba_past_the_palette() {370 let colors: Vec<[u8; 4]> = (0..4096u32)371 .map(|i| [(i >> 4) as u8, (i & 15) as u8, (i % 251) as u8, 255])372 .collect();373 let bytes = png(&colors, 64, 64, 1).unwrap();374 assert_eq!(form(&bytes), (ColorType::Rgba, BitDepth::Eight));375 assert_eq!(unpng(&bytes).unwrap().2, colors);376 }377}