use imageproc::image::{ RgbImage, RgbaImage }; use itertools::Itertools; use crate::common::*; impl Header { pub fn read (header: &[u8]) -> Option { let magic = header[..4].iter().eq(MAGIC.iter()); if !magic { return None; } let width = u32::from_be_bytes( header[4..8].try_into().ok()?); let height = u32::from_be_bytes(header[8..12].try_into().ok()?); let channels = Self::channels(u8::from_be_bytes(header[12..13].try_into().ok()?))?; let colorspace = Self::colorspace(u8::from_be_bytes(header[13..14].try_into().ok()?))?; Some(Self { width, height, channels, colorspace }) } fn colorspace (space: u8) -> Option { match space { 0 => Some(Colorspace::Srgb), 1 => Some(Colorspace::Linear), _ => None } } fn channels (ch: u8) -> Option { match ch { 3 => Some(Channels::RGB), 4 => Some(Channels::RGBA), _ => None } } } pub trait Decoder { fn decode (header: Header, bytes: &[u8]) -> Option where Self: Sized; } impl Decoder for Vec<[u8; 4]> { /// decode a complete slice into array of pixel color values fn decode (header: Header, bytes: &[u8]) -> Option { if bytes.len() < 8 { return None; } let end = [0, 0, 0, 0, 0, 0, 0, 1]; if bytes[(bytes.len() - 8)..] != end { return None; } let bytes = &bytes[..bytes.len() - 8]; let mut known = [[0; 4]; 64]; let mut prev = [0, 0, 0, 255]; let mut out = vec![]; let mut iter = bytes.iter(); loop { let Some(&byte) = iter.next() else { break; }; let (color, len) = match byte { 254 => { // QOI_OP_RGB let (r, g, b) = iter.next_tuple().unwrap(); ([*r, *g, *b, prev[3]], 1) }, 255 => { // QOI_OP_RGBA let (r, g, b, a) = iter.next_tuple().unwrap(); ([*r, *g, *b, *a], 1) }, q if q >> 6 == 0 => { // QOI_OP_INDEX (known[q as usize & 63], 1) }, q if q >> 6 == 1 => { // QOI_OP_DIFF let dr = ((q & 48) >> 4) as i8 - 2; let dg = ((q & 12) >> 2) as i8 - 2; let db = (q & 3) as i8 - 2; ([ prev[0].wrapping_add_signed(dr), prev[1].wrapping_add_signed(dg), prev[2].wrapping_add_signed(db), prev[3]], 1) }, q if q >> 6 == 2 => { // QOI_OP_LUMA let extra = *iter.next()?; let dg = (q & 63) as i8 - 32; let dr = ((extra & 240) >> 4) as i8 - 8; let db = (extra & 15) as i8 - 8; ([ prev[0].wrapping_add_signed(dr + dg), prev[1].wrapping_add_signed(dg), prev[2].wrapping_add_signed(db + dg), prev[3]], 1) }, q if q >> 6 == 3 => { // QOI_OP_RUN let len = ((q & 63) + 1) as usize; (prev.clone(), len) }, _ => return None }; known[hash(color)] = color; prev = color.clone(); (0..len).for_each(|_| out.push(color)); } if header.width as u64 * header.height as u64 != out.len() as u64 { return None; } Some(out) } } impl Decoder for Vec<[u8; 3]> { fn decode (header: Header, bytes: &[u8]) -> Option where Self: Sized { let pixels: Vec<[u8; 4]> = Decoder::decode(header, bytes)?; let pixels: Vec<[u8; 3]> = pixels.iter().filter_map(|e| (&e[0..3]).try_into().ok()).collect(); if pixels.len() as u64 != header.width as u64 * header.height as u64 { return None; } Some(pixels) } } impl Decoder for Vec { fn decode (header: Header, bytes: &[u8]) -> Option where Self: Sized { let pixels: Vec<[u8; 4]> = Decoder::decode(header, bytes)?; if header.channels == Channels::RGBA { bytemuck::try_cast_vec(pixels).ok() } else { let pixels: Vec<[u8; 3]> = pixels.iter().filter_map(|e| (&e[0..3]).try_into().ok()).collect(); if pixels.len() as u64 != header.width as u64 * header.height as u64 { return None; } bytemuck::try_cast_vec(pixels).ok() } } } impl Decoder for RgbaImage { fn decode (mut header: Header, bytes: &[u8]) -> Option where Self: Sized { header.channels = Channels::RGBA; let pixels = Decoder::decode(header, bytes)?; RgbaImage::from_raw(header.width, header.height, pixels) } } impl Decoder for RgbImage { fn decode (mut header: Header, bytes: &[u8]) -> Option where Self: Sized { header.channels = Channels::RGB; let pixels = Decoder::decode(header, bytes)?; RgbImage::from_raw(header.width, header.height, pixels) } } /// a generic decoder over a type which implements Decoder trait pub fn decode (bytes: &[u8]) -> Option where T: Decoder { let header = Header::read(&bytes[..14])?; let rest = &bytes[14..]; Decoder::decode(header, rest) } #[cfg(test)] mod test { use std::time::Instant; use imageproc::image::RgbaImage; use super::*; #[test] fn basic () { let file = std::fs::read("in/dice.qoi").unwrap(); let now = Instant::now(); let header = Header::read(&file[..14]).unwrap(); assert_eq!(Header { width: 800, height: 600, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header); let rest = &file[14..]; let pixels: Vec<[u8; 4]> = Decoder::decode(header, rest).unwrap(); println!("dice decoding done in {}ms", now.elapsed().as_millis()); let now = Instant::now(); let pixels = bytemuck::try_cast_vec(pixels).unwrap(); println!("dice cast done in {}ns", now.elapsed().as_nanos()); let rgba = RgbaImage::from_raw(header.width, header.height, pixels).unwrap(); rgba.save("out/dice.png").unwrap(); } #[test] fn edge () { let file = std::fs::read("in/edgecase.qoi").unwrap(); let now = Instant::now(); let header = Header::read(&file[..14]).unwrap(); assert_eq!(Header { width: 256, height: 64, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header); let rest = &file[14..]; let pixels: Vec<[u8; 4]> = Decoder::decode(header, rest).unwrap(); println!("edgecase decoding done in {}ms", now.elapsed().as_millis()); let pixels = bytemuck::try_cast_vec(pixels).unwrap(); let rgba = RgbaImage::from_raw(header.width, header.height, pixels).unwrap(); rgba.save("out/edgecase.png").unwrap(); } #[test] fn edge_as_jpeg () { let file = std::fs::read("in/edgecase.qoi").unwrap(); let header = Header::read(&file[..14]).unwrap(); assert_eq!(Header { width: 256, height: 64, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header); let rest = &file[14..]; let rgb = RgbImage::decode(header, rest).unwrap(); rgb.save_with_format("out/edgecase.jpg", imageproc::image::ImageFormat::Jpeg).unwrap(); } }