use anyhow::{bail, Context, Result};
use crate::timed;
use image::{DynamicImage, ImageDecoder, ImageFormat};
use std::io::Cursor;
use std::path::Path;
#[derive(Clone, Debug)]
pub struct Gray {
pub w: usize,
pub h: usize,
pub px: Vec<f32>,
}
impl Gray {
pub fn new(w: usize, h: usize) -> Self {
Gray { w, h, px: vec![0.0; w * h] }
}
}
pub struct Decoded {
pub work: Gray,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
Image(ImageFormat),
Jxl,
Heif,
Unknown,
}
pub const EXTENSIONS: [&str; 25] = [
"jpg", "jpeg", "jpe", "jfif", "png", "gif", "webp", "bmp", "tif", "tiff", "avif", "heic",
"heif", "hif", "jxl", "ico", "pnm", "pbm", "pgm", "ppm", "tga", "dds", "qoi", "exr", "ff",
];
pub const NOT_AN_IMAGE: &str = "not an image";
pub fn sniff(b: &[u8]) -> Kind {
if b.len() < 12 {
return Kind::Unknown;
}
if b.starts_with(&[0xFF, 0xD8, 0xFF]) {
return Kind::Image(ImageFormat::Jpeg);
}
if b.starts_with(b"\x89PNG\r\n\x1a\n") {
return Kind::Image(ImageFormat::Png);
}
if b.starts_with(b"GIF8") {
return Kind::Image(ImageFormat::Gif);
}
if b.starts_with(b"RIFF") && &b[8..12] == b"WEBP" {
return Kind::Image(ImageFormat::WebP);
}
if b.starts_with(b"BM") {
return Kind::Image(ImageFormat::Bmp);
}
if b.starts_with(b"II*\0") || b.starts_with(b"MM\0*") {
return Kind::Image(ImageFormat::Tiff);
}
if b.starts_with(&[0xFF, 0x0A]) || b.starts_with(b"\0\0\0\x0cJXL \r\n\x87\n") {
return Kind::Jxl;
}
if &b[4..8] == b"ftyp" {
let brand = &b[8..12];
if matches!(
brand,
b"heic" | b"heix" | b"hevc" | b"hevx" | b"heim" | b"heis" | b"mif1" | b"msf1"
| b"avif" | b"avis"
) {
return Kind::Heif;
}
}
if b.starts_with(b"qoif") {
return Kind::Image(ImageFormat::Qoi);
}
if b.starts_with(&[0x76, 0x2F, 0x31, 0x01]) {
return Kind::Image(ImageFormat::OpenExr);
}
if b.starts_with(b"farbfeld") {
return Kind::Image(ImageFormat::Farbfeld);
}
if b[0] == b'P' && (b'1'..=b'7').contains(&b[1]) {
return Kind::Image(ImageFormat::Pnm);
}
if b.starts_with(&[0, 0, 1, 0]) {
return Kind::Image(ImageFormat::Ico);
}
Kind::Unknown
}
fn decode_budget() -> usize {
let available = std::fs::read_to_string("/proc/meminfo").ok().and_then(|s| {
let line = s.lines().find(|l| l.starts_with("MemAvailable:"))?;
line.split_whitespace().nth(1)?.parse::<usize>().ok()
});
match available {
Some(kb) => (kb / 8 * 1024).max(64 << 20),
None => 256 << 20,
}
}
fn working_bytes(w: usize, h: usize, work: usize) -> u64 {
let k = box_factor(w, h, work);
((w / k).max(1) as u64) * ((h / k).max(1) as u64) * 4
}
struct Budget {
state: std::sync::Mutex<Queue>,
room: std::sync::Condvar,
limit: usize,
}
struct Queue {
held: usize,
issued: u64,
serving: u64,
}
static BUDGET: std::sync::LazyLock<Budget> = std::sync::LazyLock::new(|| Budget {
state: std::sync::Mutex::new(Queue { held: 0, issued: 0, serving: 0 }),
room: std::sync::Condvar::new(),
limit: decode_budget(),
});
thread_local! {
static HELD_HERE: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
}
struct Permit(usize);
impl Drop for Permit {
fn drop(&mut self) {
let mut q = BUDGET.state.lock().unwrap_or_else(|e| e.into_inner());
q.held -= self.0;
drop(q);
HELD_HERE.with(|h| h.set(h.get() - 1));
BUDGET.room.notify_all();
}
}
fn reserve(bytes: u64) -> Permit {
timed!(29, reserve_inner(bytes))
}
fn reserve_inner(bytes: u64) -> Permit {
let want = bytes.min(isize::MAX as u64) as usize;
if HELD_HERE.with(|h| h.get()) > 0 {
let mut q = BUDGET.state.lock().unwrap_or_else(|e| e.into_inner());
q.held += want;
drop(q);
HELD_HERE.with(|h| h.set(h.get() + 1));
return Permit(want);
}
let mut q = BUDGET.state.lock().unwrap_or_else(|e| e.into_inner());
let ticket = q.issued;
q.issued += 1;
loop {
if q.serving == ticket && (q.held == 0 || q.held + want <= BUDGET.limit) {
q.held += want;
q.serving += 1;
break;
}
q = BUDGET.room.wait(q).unwrap_or_else(|e| e.into_inner());
}
drop(q);
HELD_HERE.with(|h| h.set(h.get() + 1));
BUDGET.room.notify_all();
Permit(want)
}
pub fn decode(path: &Path, work_size: usize) -> Result<Decoded> {
let bytes = timed!(0, std::fs::read(path).with_context(|| format!("read {}", path.display()))?);
let kind = sniff(&bytes);
let (w, h, gray) = match kind {
Kind::Image(ImageFormat::Jpeg) => timed!(22, decode_image_crate(&bytes, ImageFormat::Jpeg, work_size)?),
Kind::Image(ImageFormat::Png) => timed!(23, decode_image_crate(&bytes, ImageFormat::Png, work_size)?),
Kind::Image(ImageFormat::WebP) => timed!(24, decode_image_crate(&bytes, ImageFormat::WebP, work_size)?),
Kind::Image(ImageFormat::Tiff) => timed!(25, decode_image_crate(&bytes, ImageFormat::Tiff, work_size)?),
Kind::Image(fmt) => decode_image_crate(&bytes, fmt, work_size)?,
Kind::Jxl => timed!(26, decode_jxl(&bytes, work_size)?),
Kind::Heif => timed!(27, decode_heif(&bytes, work_size)?),
Kind::Unknown => {
match image::guess_format(&bytes) {
Ok(fmt) => decode_image_crate(&bytes, fmt, work_size)?,
Err(_) => bail!(NOT_AN_IMAGE),
}
}
};
let _ = (w, h);
Ok(Decoded { work: gray })
}
pub struct Probe {
pub kind: Kind,
pub w: u32,
pub h: u32,
}
pub fn probe(path: &Path) -> Option<Probe> {
use std::io::{BufRead, BufReader};
let mut r = BufReader::with_capacity(16 << 10, std::fs::File::open(path).ok()?);
let head = r.fill_buf().ok()?;
let kind = sniff(head);
let (w, h) = match kind {
Kind::Image(ImageFormat::Jpeg) => jpeg_size(&mut r)?,
Kind::Image(ImageFormat::Png) => {
let be = |i: usize| head.get(i..i + 4).map(|b| u32::from_be_bytes(b.try_into().unwrap()));
(be(16)?, be(20)?)
}
Kind::Image(fmt) => image::ImageReader::with_format(r, fmt).into_dimensions().ok()?,
Kind::Jxl => jxl_size(&mut r)?,
Kind::Heif => {
let ctx = libheif_rs::HeifContext::read_from_file(path.to_str()?).ok()?;
let handle = ctx.primary_image_handle().ok()?;
(handle.width(), handle.height())
}
Kind::Unknown => image::ImageReader::new(r).with_guessed_format().ok()?.into_dimensions().ok()?,
};
Some(Probe { kind, w, h })
}
fn jpeg_size(r: &mut std::io::BufReader<std::fs::File>) -> Option<(u32, u32)> {
fn byte(r: &mut impl std::io::Read) -> Option<u8> {
let mut b = [0u8];
r.read_exact(&mut b).ok().map(|_| b[0])
}
let be16 = |r: &mut std::io::BufReader<std::fs::File>| Some(u16::from_be_bytes([byte(r)?, byte(r)?]));
be16(r)?;
loop {
if byte(r)? != 0xFF {
return None;
}
let mut m = byte(r)?;
while m == 0xFF {
m = byte(r)?;
}
if matches!(m, 0x01 | 0xD0..=0xD8) {
continue;
}
let len = be16(r)? as i64;
if (0xC0..=0xCF).contains(&m) && !matches!(m, 0xC4 | 0xC8 | 0xCC) {
let _precision = byte(r)?;
let h = be16(r)?;
let w = be16(r)?;
return Some((w as u32, h as u32));
}
if len < 2 {
return None;
}
r.seek_relative(len - 2).ok()?;
}
}
fn jxl_size(r: &mut impl std::io::Read) -> Option<(u32, u32)> {
let mut uninit = jxl_oxide::JxlImage::builder().pool(jxl_oxide::JxlThreadPool::none()).build_uninit();
let mut buf = vec![0u8; 4096];
let mut valid = 0usize;
for _ in 0..256 {
let n = r.read(&mut buf[valid..]).ok()?;
if n == 0 {
return None;
}
valid += n;
let used = uninit.feed_bytes(&buf[..valid]).ok()?;
buf.copy_within(used..valid, 0);
valid -= used;
match uninit.try_init().ok()? {
jxl_oxide::InitializeResult::NeedMoreData(u) => uninit = u,
jxl_oxide::InitializeResult::Initialized(img) => return Some((img.width(), img.height())),
}
}
None
}
fn decode_image_crate(bytes: &[u8], fmt: ImageFormat, work: usize) -> Result<(u32, u32, Gray)> {
if fmt == ImageFormat::Png {
if let Some(done) = decode_png_rows(bytes, work) {
return Ok(done);
}
}
decode_whole(bytes, fmt, work)
}
fn decode_whole(bytes: &[u8], fmt: ImageFormat, work: usize) -> Result<(u32, u32, Gray)> {
let reader = image::ImageReader::with_format(Cursor::new(bytes), fmt);
let mut decoder = reader.into_decoder()?;
let orientation = decoder.orientation().unwrap_or(image::metadata::Orientation::NoTransforms);
let rotates = !matches!(
orientation,
image::metadata::Orientation::NoTransforms
| image::metadata::Orientation::FlipHorizontal
| image::metadata::Orientation::FlipVertical
| image::metadata::Orientation::Rotate180
);
let converts = !matches!(
decoder.color_type(),
image::ColorType::Rgb8 | image::ColorType::Rgba8 | image::ColorType::L8 | image::ColorType::La8
);
let (dw, dh) = decoder.dimensions();
let _permit = reserve(
bytes.len() as u64
+ decoder.total_bytes().saturating_mul(1 + rotates as u64)
+ if converts { dw as u64 * dh as u64 * 4 } else { 0 }
+ working_bytes(dw as usize, dh as usize, work),
);
let mut img = timed!(1, DynamicImage::from_decoder(decoder)?);
if orientation != image::metadata::Orientation::NoTransforms {
img.apply_orientation(orientation);
}
let (w, h) = (img.width(), img.height());
let gray = timed!(2, dynamic_to_gray(&img, work));
Ok((w, h, gray))
}
fn decode_png_rows(bytes: &[u8], work: usize) -> Option<(u32, u32, Gray)> {
const IMAGE_MAX_ALLOC: usize = 512 << 20;
let mut dec = png::Decoder::new_with_limits(Cursor::new(bytes), png::Limits { bytes: IMAGE_MAX_ALLOC });
dec.set_ignore_text_chunk(false);
dec.set_transformations(png::Transformations::EXPAND);
let mut reader = dec.read_info().ok()?;
let info = reader.info();
if info.interlaced || info.animation_control.is_some() || info.exif_metadata.is_some() {
return None;
}
let (w, h) = (info.width, info.height);
let (color, depth) = reader.output_color_type();
if depth != png::BitDepth::Eight || reader.output_buffer_size()? > IMAGE_MAX_ALLOC {
return None;
}
let (wu, hu) = (w as usize, h as usize);
let _permit = reserve(bytes.len() as u64 + working_bytes(wu, hu, work) + 2 * reader.output_line_size(w)? as u64);
let reduced = timed!(1, match color {
png::ColorType::Rgb => png_rows::<_, 3, false>(&mut reader, wu, hu, work)?,
png::ColorType::Rgba => png_rows::<_, 4, true>(&mut reader, wu, hu, work)?,
png::ColorType::Grayscale => png_rows::<_, 1, false>(&mut reader, wu, hu, work)?,
png::ColorType::GrayscaleAlpha => png_rows::<_, 2, true>(&mut reader, wu, hu, work)?,
png::ColorType::Indexed => return None,
});
Some((w, h, timed!(3, fit_to(reduced, work))))
}
fn png_rows<R: std::io::BufRead + std::io::Seek, const CH: usize, const ALPHA: bool>(
reader: &mut png::Reader<R>,
w: usize,
h: usize,
work: usize,
) -> Option<Gray> {
let mut r = Reducer::<CH, ALPHA>::new(w, h, work);
for _ in 0..h {
let row = reader.next_row().ok()??;
if row.data().len() != w * CH {
return None;
}
r.push(row.data());
}
if reader.next_row().ok()?.is_some() {
return None;
}
Some(r.finish())
}
fn dynamic_to_gray(img: &DynamicImage, work: usize) -> Gray {
let (w, h) = (img.width() as usize, img.height() as usize);
match img {
DynamicImage::ImageRgb8(b) => reduce_to_gray(w, h, b.as_raw(), 3, false, work),
DynamicImage::ImageRgba8(b) => reduce_to_gray(w, h, b.as_raw(), 4, true, work),
DynamicImage::ImageLuma8(b) => reduce_to_gray(w, h, b.as_raw(), 1, false, work),
DynamicImage::ImageLumaA8(b) => reduce_to_gray(w, h, b.as_raw(), 2, true, work),
_ => {
let b = img.to_rgba8();
reduce_to_gray(w, h, b.as_raw(), 4, true, work)
}
}
}
fn box_factor(w: usize, h: usize, work: usize) -> usize {
let long = w.max(h);
if work == 0 {
return 1;
}
(long / (2 * work)).max(1)
}
pub fn reduce_to_gray(w: usize, h: usize, data: &[u8], ch: usize, alpha: bool, work: usize) -> Gray {
let g = match (ch, alpha) {
(3, false) => reduce::<3, false>(w, h, data, work),
(4, true) => reduce::<4, true>(w, h, data, work),
(1, false) => reduce::<1, false>(w, h, data, work),
(2, true) => reduce::<2, true>(w, h, data, work),
(4, false) => reduce::<4, false>(w, h, data, work),
_ => return reduce_dyn(w, h, data, ch, alpha, work),
};
timed!(3, fit_to(g, work))
}
fn reduce_rows(w: usize, h: usize, ch: usize, alpha: bool, work: usize, row: impl FnMut(&mut [u8])) -> Gray {
fn rows<const CH: usize, const ALPHA: bool>(w: usize, h: usize, work: usize, mut row: impl FnMut(&mut [u8])) -> Gray {
let mut r = Reducer::<CH, ALPHA>::new(w, h, work);
let mut line = vec![0u8; w * CH];
for _ in 0..r.rows_read() {
row(&mut line);
r.push(&line);
}
r.finish()
}
let g = match (ch, alpha) {
(3, false) => rows::<3, false>(w, h, work, row),
(4, true) => rows::<4, true>(w, h, work, row),
(1, false) => rows::<1, false>(w, h, work, row),
(2, true) => rows::<2, true>(w, h, work, row),
(4, false) => rows::<4, false>(w, h, work, row),
_ => {
let mut row = row;
let mut data = vec![0u8; w * h * ch];
for line in data.chunks_exact_mut((w * ch).max(1)) {
row(line);
}
return reduce_to_gray(w, h, &data, ch, alpha, work);
}
};
timed!(3, fit_to(g, work))
}
#[inline(always)]
fn grey_of<const CH: usize, const ALPHA: bool>(p: &[u8]) -> f32 {
let color_ch = if ALPHA { CH - 1 } else { CH };
let mut v = 0u32;
for c in 0..color_ch {
v += p[c] as u32;
}
let mut g = if color_ch == 1 { v as f32 } else { v as f32 / color_ch as f32 };
if ALPHA {
let a = p[CH - 1] as f32 / 255.0;
g = g * a + 128.0 * (1.0 - a);
}
g
}
fn reduce<const CH: usize, const ALPHA: bool>(w: usize, h: usize, data: &[u8], work: usize) -> Gray {
let mut r = Reducer::<CH, ALPHA>::new(w, h, work);
for y in 0..r.rows_read() {
r.push(&data[y * w * CH..(y + 1) * w * CH]);
}
r.finish()
}
struct Reducer<const CH: usize, const ALPHA: bool> {
w: usize,
h: usize,
k: usize,
ow: usize,
oh: usize,
inv: f32,
px: Vec<f32>,
grey: Vec<f32>,
row: Vec<f32>,
sy: usize,
}
impl<const CH: usize, const ALPHA: bool> Reducer<CH, ALPHA> {
fn new(w: usize, h: usize, work: usize) -> Self {
let k = box_factor(w, h, work);
let ow = (w / k).max(1);
let oh = (h / k).max(1);
Reducer {
w,
h,
k,
ow,
oh,
inv: 1.0 / (255.0 * (k * k) as f32),
px: Vec::with_capacity(ow * oh),
grey: vec![0.0; w],
row: if k == 1 { Vec::new() } else { vec![0.0f32; ow] },
sy: 0,
}
}
fn rows_read(&self) -> usize {
(self.oh * self.k).min(self.h)
}
fn push(&mut self, line: &[u8]) {
let (w, k, sy) = (self.w, self.k, self.sy);
self.sy += 1;
if sy >= self.rows_read() {
return;
}
if k == 1 {
grey_row::<CH, ALPHA>(&line[..w * CH], &mut self.grey[..w]);
let inv = self.inv;
self.px.extend(self.grey[..self.ow].iter().map(|g| g * inv));
return;
}
if sy % k == 0 {
self.row.fill(0.0);
}
grey_row::<CH, ALPHA>(&line[..w * CH], &mut self.grey[..w]);
box_row(k, &self.grey[..w], &mut self.row);
if sy % k == k - 1 || sy + 1 == self.h {
let inv = self.inv;
self.px.extend(self.row.iter().map(|v| v * inv));
}
}
fn finish(self) -> Gray {
debug_assert_eq!(self.px.len(), self.ow * self.oh);
Gray { w: self.ow, h: self.oh, px: self.px }
}
}
#[inline]
fn box_row(k: usize, grey: &[f32], row: &mut [f32]) {
if row.len() * k <= grey.len() {
match k {
2 => return box_k::<2>(&grey[..row.len() * 2], row),
3 => return box_k::<3>(&grey[..row.len() * 3], row),
4 => return box_k::<4>(&grey[..row.len() * 4], row),
_ => {}
}
}
let w = grey.len();
for (ox, r) in row.iter_mut().enumerate() {
let mut acc = 0.0f32;
for &g in grey[ox * k..(ox * k + k).min(w)].iter() {
acc += g;
}
*r += acc;
}
}
#[inline(always)]
fn box_k<const K: usize>(grey: &[f32], row: &mut [f32]) {
debug_assert_eq!(grey.len(), row.len() * K);
for (r, g) in row.iter_mut().zip(grey.chunks_exact(K)) {
let mut acc = 0.0f32;
for t in 0..K {
acc += g[t];
}
*r += acc;
}
}
#[inline]
fn grey_row<const CH: usize, const ALPHA: bool>(line: &[u8], out: &mut [f32]) {
let n = out.len();
debug_assert!(line.len() >= n * CH);
#[allow(unused_mut)]
let mut x = 0usize;
#[cfg(target_feature = "avx2")]
if CH == 3 || CH == 4 {
const STEP: usize = 8;
while x + STEP <= n && x * CH + 32 <= line.len() {
unsafe { grey_eight::<CH, ALPHA>(&line[x * CH..], &mut out[x..x + STEP]) };
x += STEP;
}
}
for x in x..n {
out[x] = grey_of::<CH, ALPHA>(&line[x * CH..x * CH + CH]);
}
}
#[cfg(target_feature = "avx2")]
#[inline]
unsafe fn grey_eight<const CH: usize, const ALPHA: bool>(src: &[u8], out: &mut [f32]) {
use std::arch::x86_64::*;
if CH != 3 && CH != 4 {
for (i, o) in out.iter_mut().enumerate() {
*o = grey_of::<CH, ALPHA>(&src[i * CH..i * CH + CH]);
}
return;
}
unsafe {
let v = _mm256_loadu_si256(src.as_ptr() as *const __m256i);
let p = if CH == 3 {
let moved = _mm256_permutevar8x32_epi32(v, _mm256_setr_epi32(0, 1, 2, 3, 3, 4, 5, 6));
#[rustfmt::skip]
let shuf = _mm256_setr_epi8(
0, 1, 2, -1, 3, 4, 5, -1, 6, 7, 8, -1, 9, 10, 11, -1,
0, 1, 2, -1, 3, 4, 5, -1, 6, 7, 8, -1, 9, 10, 11, -1,
);
_mm256_shuffle_epi8(moved, shuf)
} else {
v
};
let mask = _mm256_set1_epi32(0xff);
let c0 = _mm256_and_si256(p, mask);
let c1 = _mm256_and_si256(_mm256_srli_epi32(p, 8), mask);
let c2 = _mm256_and_si256(_mm256_srli_epi32(p, 16), mask);
let mut sum = _mm256_add_epi32(_mm256_add_epi32(c0, c1), c2);
let colour_ch = if ALPHA { CH - 1 } else { CH };
if colour_ch == 4 {
sum = _mm256_add_epi32(sum, _mm256_srli_epi32(p, 24));
}
let mut g = _mm256_cvtepi32_ps(sum);
if colour_ch != 1 {
g = _mm256_div_ps(g, _mm256_set1_ps(colour_ch as f32));
}
if ALPHA {
let a = _mm256_div_ps(_mm256_cvtepi32_ps(_mm256_srli_epi32(p, 24)), _mm256_set1_ps(255.0));
let t = _mm256_sub_ps(_mm256_set1_ps(1.0), a);
g = _mm256_add_ps(_mm256_mul_ps(g, a), _mm256_mul_ps(_mm256_set1_ps(128.0), t));
}
_mm256_storeu_ps(out.as_mut_ptr(), g);
}
}
fn reduce_dyn(w: usize, h: usize, data: &[u8], ch: usize, alpha: bool, work: usize) -> Gray {
let k = box_factor(w, h, work);
let ow = (w / k).max(1);
let oh = (h / k).max(1);
let mut out = Gray::new(ow, oh);
let color_ch = if alpha { ch - 1 } else { ch };
let inv = 1.0 / (255.0 * (k * k) as f32);
for oy in 0..oh {
let row = &mut out.px[oy * ow..(oy + 1) * ow];
for sy in oy * k..(oy * k + k).min(h) {
let line = &data[sy * w * ch..(sy + 1) * w * ch];
for ox in 0..ow {
let mut acc = 0.0f32;
for sx in ox * k..(ox * k + k).min(w) {
let p = &line[sx * ch..sx * ch + ch];
let mut v = 0u32;
for c in 0..color_ch {
v += p[c] as u32;
}
let mut g = v as f32 / color_ch as f32;
if alpha {
let a = p[ch - 1] as f32 / 255.0;
g = g * a + 128.0 * (1.0 - a);
}
acc += g;
}
row[ox] += acc;
}
}
for v in row.iter_mut() {
*v *= inv;
}
}
fit_to(out, work)
}
pub fn fit_to(g: Gray, work: usize) -> Gray {
let long = g.w.max(g.h);
if work == 0 || long <= work {
return g;
}
let s = work as f32 / long as f32;
let tw = ((g.w as f32 * s).round() as usize).max(1);
let th = ((g.h as f32 * s).round() as usize).max(1);
resize_area(&g, tw, th)
}
pub fn resize_area(g: &Gray, tw: usize, th: usize) -> Gray {
if tw == g.w && th == g.h {
return g.clone();
}
let xw = weights(g.w, tw);
let yw = weights(g.h, th);
let mut tmp: Vec<f32> = Vec::with_capacity(tw * g.h);
for y in 0..g.h {
let src = &g.px[y * g.w..(y + 1) * g.w];
tmp.extend((0..tw).map(|ox| {
let start = xw.start[ox] as usize;
let (a, b) = (xw.at[ox] as usize, xw.at[ox + 1] as usize);
let ws = &xw.w[a..b];
let ss = &src[start..start + ws.len()];
let mut acc = 0.0;
for (sv, wgt) in ss.iter().zip(ws) {
acc += sv * wgt;
}
acc
}));
}
let mut px: Vec<f32> = Vec::with_capacity(tw * th);
for oy in 0..th {
let start = yw.start[oy] as usize;
let (a, b) = (yw.at[oy] as usize, yw.at[oy + 1] as usize);
let base = px.len();
let w0 = yw.w[a];
let s0 = &tmp[start * tw..(start + 1) * tw];
px.extend(s0.iter().map(|v| v * w0));
let dst = &mut px[base..base + tw];
for (i, wgt) in yw.w[a + 1..b].iter().enumerate() {
let src = &tmp[(start + i + 1) * tw..(start + i + 2) * tw];
for x in 0..tw {
dst[x] += src[x] * wgt;
}
}
}
Gray { w: tw, h: th, px }
}
struct Taps {
start: Vec<u32>,
at: Vec<u32>,
w: Vec<f32>,
}
fn weights(src: usize, dst: usize) -> Taps {
let s = src as f64 / dst as f64;
let mut t = Taps { start: Vec::with_capacity(dst), at: Vec::with_capacity(dst + 1), w: Vec::new() };
for o in 0..dst {
let a = o as f64 * s;
let b = ((o + 1) as f64 * s).min(src as f64);
let i0 = a.floor() as usize;
let i1 = (b.ceil() as usize).min(src).max(i0 + 1);
t.start.push(i0 as u32);
t.at.push(t.w.len() as u32);
let mut total = 0.0;
for i in i0..i1 {
let lo = (i as f64).max(a);
let hi = ((i + 1) as f64).min(b);
let wgt = (hi - lo).max(0.0);
t.w.push(wgt as f32);
total += wgt;
}
let inv = if total > 0.0 { (1.0 / total) as f32 } else { 0.0 };
let from = t.at[o] as usize;
for w in t.w[from..].iter_mut() {
*w *= inv;
}
}
t.at.push(t.w.len() as u32);
t
}
fn decode_jxl(bytes: &[u8], work: usize) -> Result<(u32, u32, Gray)> {
let image = jxl_oxide::JxlImage::builder()
.pool(jxl_oxide::JxlThreadPool::none())
.read(Cursor::new(bytes))
.map_err(|e| anyhow::anyhow!("jxl: {e}"))?;
let header = image.image_header();
let nch = if header.metadata.grayscale() { 1 } else { 3 } + header.metadata.alpha().is_some() as u64;
let _permit = reserve(
bytes.len() as u64
+ (image.width() as u64) * (image.height() as u64) * nch * 8
+ working_bytes(image.width() as usize, image.height() as usize, work),
);
let render = image
.render_frame(0)
.map_err(|e| anyhow::anyhow!("jxl render: {e}"))?;
let mut stream = render.stream();
let (w, h, ch) = (stream.width() as usize, stream.height() as usize, stream.channels() as usize);
let color = if image.image_header().metadata.grayscale() { 1 } else { 3 };
let has_alpha = image.image_header().metadata.alpha().is_some();
let out_ch = color + has_alpha as usize;
let mut rowf = vec![0f32; w * ch];
let g = reduce_rows(w, h, out_ch, has_alpha, work, |line| {
let got = stream.write_to_buffer(&mut rowf);
rowf[got..].fill(0.0);
for (px, out) in rowf.chunks_exact(ch).zip(line.chunks_exact_mut(out_ch)) {
for c in 0..color {
out[c] = (px[c].clamp(0.0, 1.0) * 255.0 + 0.5) as u8;
}
if has_alpha {
out[color] = (px[color].clamp(0.0, 1.0) * 255.0 + 0.5) as u8;
}
}
});
Ok((w as u32, h as u32, g))
}
fn decode_heif(bytes: &[u8], work: usize) -> Result<(u32, u32, Gray)> {
use libheif_rs::{ColorSpace, HeifContext, LibHeif, RgbChroma};
let lib = LibHeif::new();
let ctx = HeifContext::read_from_bytes(bytes).map_err(|e| anyhow::anyhow!("heif: {e}"))?;
let handle = ctx.primary_image_handle().map_err(|e| anyhow::anyhow!("heif: {e}"))?;
let has_alpha = handle.has_alpha_channel();
let chroma = if has_alpha { RgbChroma::Rgba } else { RgbChroma::Rgb };
let _permit = reserve(
bytes.len() as u64
+ (handle.width() as u64) * (handle.height() as u64) * if has_alpha { 8 } else { 6 }
+ working_bytes(handle.width() as usize, handle.height() as usize, work),
);
let img = lib
.decode(&handle, ColorSpace::Rgb(chroma), None)
.map_err(|e| anyhow::anyhow!("heif decode: {e}"))?;
let planes = img.planes();
let plane = planes.interleaved.context("heif: no interleaved plane")?;
let (w, h) = (plane.width as usize, plane.height as usize);
let ch = if has_alpha { 4 } else { 3 };
let mut y = 0;
let g = reduce_rows(w, h, ch, has_alpha, work, |line| {
line.copy_from_slice(&plane.data[y * plane.stride..y * plane.stride + w * ch]);
y += 1;
});
Ok((w as u32, h as u32, g))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn area_resize_preserves_mean() {
let mut g = Gray::new(10, 7);
for (i, v) in g.px.iter_mut().enumerate() {
*v = (i % 13) as f32 / 12.0;
}
let mean: f32 = g.px.iter().sum::<f32>() / g.px.len() as f32;
let r = resize_area(&g, 3, 2);
let m2: f32 = r.px.iter().sum::<f32>() / r.px.len() as f32;
assert!((mean - m2).abs() < 0.05, "{mean} vs {m2}");
}
#[test]
fn specialised_reduction_matches_the_general_one() {
for &(ch, alpha) in &[(1, false), (2, true), (3, false), (4, true), (4, false)] {
for &(w, h, work) in
&[(37usize, 23usize, 64usize), (200, 150, 32), (64, 64, 0), (100, 80, 64), (121, 97, 64), (900, 5, 64), (5, 900, 64), (301, 7, 32)]
{
let data: Vec<u8> = (0..w * h * ch)
.map(|i| ((i * 37 + i / 17 * 11) % 251) as u8)
.collect();
let fast = reduce_to_gray(w, h, &data, ch, alpha, work);
let slow = reduce_dyn(w, h, &data, ch, alpha, work);
assert_eq!((fast.w, fast.h), (slow.w, slow.h), "{ch} {alpha} {w}x{h}@{work}");
assert_eq!(fast.px, slow.px, "{ch} {alpha} {w}x{h}@{work}");
}
}
}
#[test]
fn reduce_rows_matches_the_whole_buffer() {
for &(ch, alpha) in &[(1, false), (2, true), (3, false), (4, true), (4, false), (5, false)] {
for &(w, h, work) in &[(37usize, 23usize, 64usize), (200, 150, 32), (64, 64, 0), (121, 97, 64), (900, 5, 64), (5, 900, 64)] {
let data: Vec<u8> = (0..w * h * ch).map(|i| ((i * 37 + i / 17 * 11) % 251) as u8).collect();
let want = reduce_to_gray(w, h, &data, ch, alpha, work);
let mut y = 0;
let got = reduce_rows(w, h, ch, alpha, work, |line| {
line.copy_from_slice(&data[y * w * ch..(y + 1) * w * ch]);
y += 1;
});
assert_eq!((got.w, got.h), (want.w, want.h), "{ch} {alpha} {w}x{h}@{work}");
assert!(got.px.iter().zip(want.px.iter()).all(|(a, b)| a.to_bits() == b.to_bits()), "{ch} {alpha} {w}x{h}@{work}");
}
}
}
fn png_of(w: u32, h: u32, color: png::ColorType, depth: png::BitDepth) -> Vec<u8> {
let mut out = Vec::new();
{
let mut e = png::Encoder::new(&mut out, w, h);
e.set_color(color);
e.set_depth(depth);
if color == png::ColorType::Indexed {
e.set_palette((0..256u32).flat_map(|i| [i as u8, (i * 7) as u8, (255 - i) as u8]).collect::<Vec<u8>>());
e.set_trns((0..256u32).map(|i| (i * 3) as u8).collect::<Vec<u8>>());
}
let mut wr = e.write_header().unwrap();
let bits = color.samples() * depth as usize;
let row = (w as usize * bits).div_ceil(8);
let data: Vec<u8> = (0..row * h as usize).map(|i| ((i * 37 + i / 13 * 11 + i / 997) % 251) as u8).collect();
wr.write_image_data(&data).unwrap();
}
out
}
#[test]
fn png_rows_decode_exactly_as_the_whole_picture_does() {
use png::{BitDepth as D, ColorType as C};
let layouts = [
(C::Rgb, D::Eight, true),
(C::Rgba, D::Eight, true),
(C::Grayscale, D::Eight, true),
(C::GrayscaleAlpha, D::Eight, true),
(C::Grayscale, D::One, true),
(C::Grayscale, D::Four, true),
(C::Indexed, D::Eight, true),
(C::Indexed, D::Two, true),
(C::Rgb, D::Sixteen, false),
(C::Rgba, D::Sixteen, false),
];
for &(color, depth, streams) in layouts.iter() {
for &(w, h, work) in &[(37u32, 23u32, 64usize), (300, 200, 64), (257, 511, 64), (900, 5, 64), (5, 700, 64), (64, 64, 0)] {
let bytes = png_of(w, h, color, depth);
let tag = format!("{color:?} {depth:?} {w}x{h}@{work}");
assert_eq!(decode_png_rows(&bytes, work).is_some(), streams, "{tag}");
let got = decode_image_crate(&bytes, ImageFormat::Png, work).unwrap().2;
let want = decode_whole(&bytes, ImageFormat::Png, work).unwrap().2;
assert_eq!((got.w, got.h), (want.w, want.h), "{tag}");
assert!(got.px.iter().zip(want.px.iter()).all(|(a, b)| a.to_bits() == b.to_bits()), "{tag}");
}
}
let bytes = png_of(300, 200, C::Rgb, D::Eight);
for cut in [40, 200, bytes.len() / 2, bytes.len() - 20, bytes.len() - 5] {
let got = decode_image_crate(&bytes[..cut], ImageFormat::Png, 64).map(|r| r.2.px);
let want = decode_whole(&bytes[..cut], ImageFormat::Png, 64).map(|r| r.2.px);
match (got, want) {
(Ok(a), Ok(b)) => assert_eq!(a, b, "cut {cut}"),
(Err(a), Err(b)) => assert_eq!(a.to_string(), b.to_string(), "cut {cut}"),
(a, b) => panic!("cut {cut}: {:?} against {:?}", a.is_ok(), b.is_ok()),
}
}
}
#[test]
fn sniff_basics() {
assert_eq!(sniff(b"\xFF\xD8\xFF\xE0\0\x10JFIF\0\x01\x01"), Kind::Image(ImageFormat::Jpeg));
assert_eq!(sniff(b"\0\0\0\x18ftypavif\0\0\0\0"), Kind::Heif);
assert_eq!(sniff(b"RIFF\0\0\0\0WEBPVP8 "), Kind::Image(ImageFormat::WebP));
}
}
#[cfg(test)]
mod bench {
use super::*;
fn ms(f: impl Fn()) -> f64 {
let mut best = f64::MAX;
for _ in 0..9 {
let t = std::time::Instant::now();
f();
best = best.min(t.elapsed().as_secs_f64() * 1000.0);
}
best
}
#[test]
#[ignore]
fn reduce_timings() {
for &(w, h) in &[(1200usize, 900usize), (4000, 3000)] {
for &(ch, alpha, name) in &[(3usize, false, "rgb8"), (4, true, "rgba8"), (1, false, "l8")] {
let data: Vec<u8> = (0..w * h * ch).map(|i| ((i * 37 + i / 101 * 7) % 251) as u8).collect();
let t = ms(|| {
std::hint::black_box(reduce_to_gray(w, h, &data, ch, alpha, 640));
});
let k = box_factor(w, h, 640);
println!("reduce {name} {w}x{h} (k={k}): {t:8.3} ms -> {:5.1} Mpx/s", (w * h) as f64 / t / 1000.0);
}
}
let mut g = Gray::new(1280, 960);
for (i, v) in g.px.iter_mut().enumerate() {
*v = ((i * 37) % 251) as f32 / 251.0;
}
let t = ms(|| {
std::hint::black_box(resize_area(&g, 640, 480));
});
println!("resize_area 1280x960 -> 640x480: {t:8.3} ms");
}
}