| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| use crate::resize::{Rgb, PAD}; |
| use crate::{Result, VisionError}; |
| use std::io::Cursor; |
|
|
| #[derive(Clone, Copy, Debug, PartialEq, Eq)] |
| pub enum Format { |
| Png, |
| Jpeg, |
| Gif, |
| Bmp, |
| WebP, |
| Tiff, |
| Unknown, |
| } |
|
|
| pub fn sniff(b: &[u8]) -> Format { |
| if b.len() >= 8 && b[..8] == [0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a] { |
| Format::Png |
| } else if b.len() >= 3 && b[0] == 0xff && b[1] == 0xd8 && b[2] == 0xff { |
| Format::Jpeg |
| } else if b.len() >= 6 && (&b[..6] == b"GIF87a" || &b[..6] == b"GIF89a") { |
| Format::Gif |
| } else if b.len() >= 2 && &b[..2] == b"BM" { |
| Format::Bmp |
| } else if b.len() >= 12 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" { |
| Format::WebP |
| } else if b.len() >= 4 && (b[..4] == [b'I', b'I', 42, 0] || b[..4] == [b'M', b'M', 0, 42]) { |
| Format::Tiff |
| } else { |
| Format::Unknown |
| } |
| } |
|
|
| fn err<T: std::fmt::Display>(e: T) -> VisionError { |
| VisionError::Decode(e.to_string()) |
| } |
|
|
| |
|
|
| |
| pub fn probe(bytes: &[u8]) -> Result<(usize, usize)> { |
| match sniff(bytes) { |
| Format::Png => { |
| if bytes.len() < 24 || &bytes[12..16] != b"IHDR" { |
| return Err(err("bad PNG header")); |
| } |
| let w = u32::from_be_bytes(bytes[16..20].try_into().unwrap()) as usize; |
| let h = u32::from_be_bytes(bytes[20..24].try_into().unwrap()) as usize; |
| Ok((w, h)) |
| } |
| Format::Jpeg => probe_jpeg(bytes), |
| Format::Gif => { |
| if bytes.len() < 10 { |
| return Err(err("bad GIF header")); |
| } |
| Ok(( |
| u16::from_le_bytes([bytes[6], bytes[7]]) as usize, |
| u16::from_le_bytes([bytes[8], bytes[9]]) as usize, |
| )) |
| } |
| Format::Unknown => Err(VisionError::Unsupported), |
| _ => { |
| let r = image::ImageReader::new(Cursor::new(bytes)).with_guessed_format().map_err(err)?; |
| let (w, h) = r.into_dimensions().map_err(err)?; |
| Ok((w as usize, h as usize)) |
| } |
| } |
| } |
|
|
| fn probe_jpeg(b: &[u8]) -> Result<(usize, usize)> { |
| let mut i = 2; |
| while i + 4 <= b.len() { |
| if b[i] != 0xff { |
| i += 1; |
| continue; |
| } |
| let m = b[i + 1]; |
| if m == 0xff { |
| i += 1; |
| continue; |
| } |
| if m == 0xd8 || (0xd0..=0xd7).contains(&m) || m == 0x01 || m == 0x00 { |
| i += 2; |
| continue; |
| } |
| let len = u16::from_be_bytes([b[i + 2], b[i + 3]]) as usize; |
| |
| if (0xc0..=0xcf).contains(&m) && m != 0xc4 && m != 0xc8 && m != 0xcc { |
| if i + 9 > b.len() { |
| break; |
| } |
| let h = u16::from_be_bytes([b[i + 5], b[i + 6]]) as usize; |
| let w = u16::from_be_bytes([b[i + 7], b[i + 8]]) as usize; |
| return Ok((w, h)); |
| } |
| if m == 0xda { |
| break; |
| } |
| i += 2 + len; |
| } |
| Err(err("JPEG: no SOF marker")) |
| } |
|
|
| |
|
|
| |
| |
| pub fn decode_rgb(bytes: &[u8], max_pixels: usize) -> Result<Rgb> { |
| let fmt = sniff(bytes); |
| if fmt == Format::Unknown { |
| return Err(VisionError::Unsupported); |
| } |
| let (w, h) = probe(bytes)?; |
| if w == 0 || h == 0 { |
| return Err(err("zero-sized image")); |
| } |
| if w.checked_mul(h).map_or(true, |p| p > max_pixels) { |
| return Err(VisionError::TooLarge { w, h, max_pixels }); |
| } |
| match fmt { |
| Format::Png => decode_png(bytes), |
| Format::Jpeg => decode_jpeg(bytes), |
| _ => decode_other(bytes), |
| } |
| } |
|
|
| |
| pub fn decode_to_rgb(bytes: &[u8]) -> Result<(usize, usize, Vec<u8>)> { |
| let mut r = decode_rgb(bytes, crate::DEFAULT_MAX_PIXELS)?; |
| let mut v = std::mem::take(&mut r.data); |
| v.truncate(r.w * r.h * 3); |
| Ok((r.w, r.h, v)) |
| } |
|
|
| |
|
|
| fn decode_png(bytes: &[u8]) -> Result<Rgb> { |
| if let Some(p) = crate::png_fast::parse(bytes).filter(crate::png_fast::worthwhile) { |
| let mut out = Rgb::scratch(p.w, p.h); |
| let prog = std::sync::atomic::AtomicUsize::new(0); |
| crate::png_fast::decode_into(&p, out.data.as_mut_ptr(), out.data.len(), &prog)?; |
| return Ok(out); |
| } |
| decode_png_slow(bytes) |
| } |
|
|
| fn decode_png_slow(bytes: &[u8]) -> Result<Rgb> { |
| use png::{BitDepth, ColorType, Transformations}; |
| let limits = png::Limits { bytes: 1usize << 40 }; |
| |
| let probe_dec = png::Decoder::new_with_limits(Cursor::new(bytes), limits.clone()); |
| let probe_reader = probe_dec.read_info().map_err(err)?; |
| let (w, h) = { |
| let i = probe_reader.info(); |
| (i.width as usize, i.height as usize) |
| }; |
| let gray16 = probe_reader.info().color_type == ColorType::Grayscale |
| && probe_reader.info().bit_depth == BitDepth::Sixteen; |
| drop(probe_reader); |
|
|
| let mut dec = png::Decoder::new_with_limits(Cursor::new(bytes), limits); |
| dec.set_transformations(if gray16 { Transformations::IDENTITY } else { Transformations::EXPAND | Transformations::STRIP_16 }); |
| let mut reader = dec.read_info().map_err(err)?; |
| let (ct, bd) = reader.output_color_type(); |
| let size = reader.output_buffer_size().ok_or_else(|| err("PNG buffer size overflow"))?; |
| |
| let direct = ct == ColorType::Rgb && bd == BitDepth::Eight; |
| let mut buf = vec![0u8; if direct { size + PAD } else { size }]; |
| reader.next_frame(&mut buf[..size]).map_err(err)?; |
| if direct { |
| return Ok(Rgb { w, h, data: buf }); |
| } |
| let mut out = vec![0u8; w * h * 3 + PAD]; |
| let n = w * h; |
| match (ct, bd) { |
| (ColorType::Rgba, BitDepth::Eight) => rgba_to_rgb(&buf, &mut out[..n * 3]), |
| (ColorType::Grayscale, BitDepth::Eight) => { |
| for (i, &g) in buf[..n].iter().enumerate() { |
| out[i * 3] = g; |
| out[i * 3 + 1] = g; |
| out[i * 3 + 2] = g; |
| } |
| } |
| (ColorType::GrayscaleAlpha, BitDepth::Eight) => { |
| for i in 0..n { |
| let g = buf[i * 2]; |
| out[i * 3] = g; |
| out[i * 3 + 1] = g; |
| out[i * 3 + 2] = g; |
| } |
| } |
| (ColorType::Grayscale, BitDepth::Sixteen) => { |
| for i in 0..n { |
| let (hi, lo) = (buf[i * 2], buf[i * 2 + 1]); |
| let g = if hi != 0 { 255 } else { lo }; |
| out[i * 3] = g; |
| out[i * 3 + 1] = g; |
| out[i * 3 + 2] = g; |
| } |
| } |
| _ => return Err(err(format!("unsupported PNG output {:?}/{:?}", ct, bd))), |
| } |
| Ok(Rgb { w, h, data: out }) |
| } |
|
|
| pub fn rgba_to_rgb(src: &[u8], dst: &mut [u8]) { |
| use rayon::prelude::*; |
| dst.par_chunks_mut(3 * 4096).zip(src.par_chunks(4 * 4096)).for_each(|(d, s)| { |
| for (dp, sp) in d.chunks_exact_mut(3).zip(s.chunks_exact(4)) { |
| dp.copy_from_slice(&sp[..3]); |
| } |
| }); |
| } |
|
|
| |
|
|
| fn decode_jpeg(bytes: &[u8]) -> Result<Rgb> { |
| use turbojpeg_sys as tj; |
| unsafe { |
| let h = tj::tj3Init(tj::TJINIT_TJINIT_DECOMPRESS as i32); |
| if h.is_null() { |
| return Err(err("tj3Init failed")); |
| } |
| struct Guard(tj::tjhandle); |
| impl Drop for Guard { |
| fn drop(&mut self) { |
| unsafe { tj::tj3Destroy(self.0) } |
| } |
| } |
| let _g = Guard(h); |
| let msg = |h: tj::tjhandle| -> VisionError { |
| let p = tj::tj3GetErrorStr(h); |
| let s = if p.is_null() { "jpeg error".to_string() } else { std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned() }; |
| VisionError::Decode(s) |
| }; |
| |
| tj::tj3Set(h, tj::TJPARAM_TJPARAM_FASTDCT as i32, 0); |
| tj::tj3Set(h, tj::TJPARAM_TJPARAM_FASTUPSAMPLE as i32, 0); |
| if tj::tj3DecompressHeader(h, bytes.as_ptr(), bytes.len() as _) != 0 { |
| return Err(msg(h)); |
| } |
| let w = tj::tj3Get(h, tj::TJPARAM_TJPARAM_JPEGWIDTH as i32) as usize; |
| let ht = tj::tj3Get(h, tj::TJPARAM_TJPARAM_JPEGHEIGHT as i32) as usize; |
| let cs = tj::tj3Get(h, tj::TJPARAM_TJPARAM_COLORSPACE as i32); |
| let prec = tj::tj3Get(h, tj::TJPARAM_TJPARAM_PRECISION as i32); |
| if prec != 8 { |
| return Err(err(format!("unsupported JPEG precision {prec}"))); |
| } |
| if cs == tj::TJCS_TJCS_CMYK as i32 || cs == tj::TJCS_TJCS_YCCK as i32 { |
| let mut cmyk = vec![0u8; w * ht * 4]; |
| if tj::tj3Decompress8(h, bytes.as_ptr(), bytes.len() as _, cmyk.as_mut_ptr(), (w * 4) as i32, tj::TJPF_TJPF_CMYK as i32) != 0 { |
| return Err(msg(h)); |
| } |
| let mut out = vec![0u8; w * ht * 3 + PAD]; |
| |
| for (o, p) in out.chunks_exact_mut(3).zip(cmyk.chunks_exact(4)) { |
| let (c, m, y, k) = (255 - p[0] as u32, 255 - p[1] as u32, 255 - p[2] as u32, 255 - p[3] as u32); |
| let nk = 255 - k; |
| let md = |a: u32| -> u8 { |
| let t = a * nk + 128; |
| let v = ((t >> 8) + t) >> 8; |
| nk.saturating_sub(v).min(255) as u8 |
| }; |
| o[0] = md(c); |
| o[1] = md(m); |
| o[2] = md(y); |
| } |
| return Ok(Rgb { w, h: ht, data: out }); |
| } |
| let mut out = vec![0u8; w * ht * 3 + PAD]; |
| if tj::tj3Decompress8(h, bytes.as_ptr(), bytes.len() as _, out.as_mut_ptr(), (w * 3) as i32, tj::TJPF_TJPF_RGB as i32) != 0 { |
| return Err(msg(h)); |
| } |
| Ok(Rgb { w, h: ht, data: out }) |
| } |
| } |
|
|
| |
|
|
| fn decode_other(bytes: &[u8]) -> Result<Rgb> { |
| let img = image::load_from_memory(bytes).map_err(err)?; |
| let rgb = img.to_rgb8(); |
| let (w, h) = (rgb.width() as usize, rgb.height() as usize); |
| Ok(Rgb::from_vec(w, h, rgb.into_raw())) |
| } |
|
|
