blurhash-update/src/lib.rs
asonix b1cefd7863
All checks were successful
/ tests (push) Successful in 2m32s
/ check (x86_64-unknown-linux-musl) (push) Successful in 6s
/ clippy (push) Successful in 8s
/ check (aarch64-unknown-linux-musl) (push) Successful in 7s
/ check (armv7-unknown-linux-musleabihf) (push) Successful in 6s
Add length assertion
2024-02-18 13:25:39 -06:00

356 lines
10 KiB
Rust

mod base83;
mod srgb_lookup;
use std::f32::consts::PI;
use srgb_lookup::srgb_to_linear;
const BYTES_PER_PIXEL: usize = 4;
/// How many components should be used in blurhash creation
///
/// More components will increase the definition of the blurhash, but also increase processing
/// time.
pub struct Components {
pub x: u32,
pub y: u32,
}
/// Width and height of the input image
#[derive(Clone, Copy, Debug)]
pub struct ImageBounds {
pub width: u32,
pub height: u32,
}
struct ComponentState {
x: u32,
y: u32,
basis: f32,
}
/// Error raised when too many components are requested
#[derive(Debug)]
pub struct ComponentError;
/// Encoder type used to produce blurhashes
pub struct Encoder {
index: usize,
components: Components,
factors: Box<[(ComponentState, [f32; BYTES_PER_PIXEL])]>,
bounds: ImageBounds,
}
/// A simple "encode this image please" function
///
/// The input image must be in the sRGB colorspace and be formatted as 8bit RGBA
pub fn encode(
components: Components,
bounds: ImageBounds,
rgba8_image: &[u8],
) -> Result<String, ComponentError> {
let mut encoder = Encoder::new(components, bounds)?;
encoder.update(rgba8_image);
Ok(encoder.finalize())
}
impl Encoder {
/// Create a new Encoder to produce a blurhash
///
/// Errors if too many components are requested
pub fn new(
Components { x, y }: Components,
bounds: ImageBounds,
) -> Result<Self, ComponentError> {
if !(1..=9).contains(&x) || !(1..=9).contains(&y) {
return Err(ComponentError);
}
Ok(Self {
index: 0,
components: Components { x, y },
factors: Box::from(
(0..y)
.flat_map(|y| {
(0..x).map(move |x| (ComponentState { x, y, basis: 0. }, [0., 0., 0., 0.]))
})
.collect::<Vec<_>>(),
),
bounds,
})
}
/// Update the encoder with bytes from an image
///
/// The input image must be in the sRGB colorspace and be formatted as 8bit RGBA
/// The input doesn't need to contain whole pixels, the encoder is capable of handling partial
/// pixels
pub fn update(&mut self, rgba8_image: &[u8]) {
// get offset in terms of already-processed bytes
let offset = self.index % BYTES_PER_PIXEL;
// get offset in terms of remaining bytes on head of rgba8_image
let offset = (BYTES_PER_PIXEL - offset) % BYTES_PER_PIXEL;
let basis_scale_x = PI / self.bounds.width as f32;
let basis_scale_y = PI / self.bounds.height as f32;
for (ComponentState { basis, .. }, [_, g, b, _]) in self.factors.iter_mut() {
for (val, slot) in rgba8_image[..offset]
.iter()
.map(|byte| *basis * srgb_to_linear(*byte))
.zip(
[b, g][..offset.saturating_sub(BYTES_PER_PIXEL - 2)]
.iter_mut()
.rev(),
)
{
**slot += val;
}
}
let pixels = ((self.index + offset) / BYTES_PER_PIXEL) as u32;
let mut chunks = rgba8_image[offset..].chunks_exact(BYTES_PER_PIXEL);
for (i, chunk) in (&mut chunks).enumerate() {
let px = pixels + i as u32;
let px_x = px % self.bounds.width;
let px_y = px / self.bounds.width;
let scale_x = px_x as f32 * basis_scale_x;
let scale_y = px_y as f32 * basis_scale_y;
for (ComponentState { x, y, .. }, rgb) in self.factors.iter_mut() {
let basis = f32::cos(*x as f32 * scale_x) * f32::cos(*y as f32 * scale_y);
assert_eq!(chunk.len(), rgb.len());
for (val, slot) in chunk
.iter()
.map(|byte| basis * srgb_to_linear(*byte))
.zip(rgb)
{
*slot += val;
}
}
}
if !chunks.remainder().is_empty() {
let px = pixels + (rgba8_image[offset..].len() / BYTES_PER_PIXEL) as u32;
let px_x = px % self.bounds.width;
let px_y = px / self.bounds.width;
let scale_x = px_x as f32 * basis_scale_x;
let scale_y = px_y as f32 * basis_scale_y;
for (ComponentState { x, y, basis }, rgb) in self.factors.iter_mut() {
*basis = f32::cos(*x as f32 * scale_x) * f32::cos(*y as f32 * scale_y);
for (val, slot) in chunks
.remainder()
.iter()
.map(|byte| *basis * srgb_to_linear(*byte))
.zip(rgb)
{
*slot += val;
}
}
}
self.index += rgba8_image.len();
}
/// Produce a blurhash from the provided encoder
pub fn finalize(mut self) -> String {
for (ComponentState { x, y, .. }, rgb) in self.factors.iter_mut() {
let normalisation = if *x == 0 && *y == 0 { 1. } else { 2. };
let scale = normalisation / (self.bounds.width * self.bounds.height) as f32;
for slot in rgb {
*slot *= scale;
}
}
let mut blurhash = String::with_capacity(30);
let (_, dc) = self.factors[0];
let ac = &self.factors[1..];
let size_flag = self.components.x - 1 + (self.components.y - 1) * 9;
base83::encode(size_flag, 1, &mut blurhash);
let maximum = ac.iter().fold(0.0_f32, |maximum, (_, [r, g, b, _])| {
maximum.max(r.abs()).max(g.abs()).max(b.abs())
});
let quantized_maximum = (maximum * 166. - 0.5).floor().max(0.) as u32;
base83::encode(quantized_maximum, 1, &mut blurhash);
let maximum_value = (quantized_maximum + 1) as f32 / 166.;
base83::encode(encode_dc(dc), 4, &mut blurhash);
for (_, rgb) in ac {
base83::encode(encode_ac(*rgb, maximum_value), 2, &mut blurhash);
}
blurhash
}
}
fn encode_dc(rgb: [f32; BYTES_PER_PIXEL]) -> u32 {
let [r, g, b, _] = rgb.map(linear_to_srgb);
(r << 16) + (g << 8) + b
}
fn encode_ac(rgb: [f32; BYTES_PER_PIXEL], maximum_value: f32) -> u32 {
let [r, g, b, _] = rgb.map(|c| encode_ac_digit(c, maximum_value));
r * 19 * 19 + g * 19 + b
}
fn encode_ac_digit(d: f32, maximum_value: f32) -> u32 {
((sign_pow(d / maximum_value, 0.5) * 9. + 9.5) as i32).clamp(0, 18) as u32
}
fn linear_to_srgb(value: f32) -> u32 {
let v = f32::max(0., f32::min(1., value));
if v <= 0.003_130_8 {
(v * 12.92 * 255. + 0.5).round() as u32
} else {
((1.055 * f32::powf(v, 1. / 2.4) - 0.055) * 255. + 0.5).round() as u32
}
}
fn sign(n: f32) -> f32 {
if n < 0. {
-1.
} else {
1.
}
}
fn sign_pow(val: f32, exp: f32) -> f32 {
sign(val) * val.abs().powf(exp)
}
impl std::fmt::Display for ComponentError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Components out of bounds")
}
}
impl std::error::Error for ComponentError {}
#[cfg(test)]
mod tests {
use image::{EncodableLayout, GenericImageView};
#[test]
fn contrived() {
let input = [
0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120,
0, 0, 60, 120, 0, 0, 60, 120, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60,
0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0,
];
let width = 4;
let height = 4;
let hash = super::encode(
crate::Components { x: 4, y: 3 },
crate::ImageBounds { width, height },
&input,
)
.unwrap();
assert_eq!(hash, "LQ9~?d$,fQ$,G1S%fQS%A{SPfQSP");
}
#[test]
fn one_component() {
let inputs = [
("data/19dd1c444d1c7939.png", "00AQtR"),
("data/f73d2ee39133d871.jpg", "00E{R{"),
("data/shenzi.png", "0039[D"),
];
for (input, output) in inputs {
let img = image::open(input).unwrap();
let (width, height) = img.dimensions();
let hash = super::encode(
crate::Components { x: 1, y: 1 },
crate::ImageBounds { width, height },
img.to_rgba8().as_bytes(),
)
.unwrap();
assert_eq!(hash, output, "expected output wrong");
}
}
#[test]
fn matches_blurhash() {
let inputs = [
("data/19dd1c444d1c7939.png", "L3AQtR2FSz6NrsOCW:ODR*,EE};h"),
("data/f73d2ee39133d871.jpg", "LJE{R{Z}V?N#0JR*Rit7^htTfkaI"),
("data/shenzi.png", "L239[DQ.91t,rJX9Qns+8zt5.PR6"),
];
for (input, output) in inputs {
let img = image::open(input).unwrap();
let (width, height) = img.dimensions();
let hash = super::encode(
crate::Components { x: 4, y: 3 },
crate::ImageBounds { width, height },
img.to_rgba8().as_bytes(),
)
.unwrap();
assert_eq!(hash, output, "expected output wrong");
}
}
#[test]
fn matches_self_when_split() {
let inputs = [
"data/19dd1c444d1c7939.png",
"data/f73d2ee39133d871.jpg",
"data/shenzi.png",
];
for input in inputs {
let img = image::open(input).unwrap();
let (width, height) = img.dimensions();
let rgba8_img = img.to_rgba8();
let bytes = rgba8_img.as_bytes();
let b1 = super::encode(
crate::Components { x: 4, y: 3 },
crate::ImageBounds { width, height },
bytes,
)
.unwrap();
for chunk_count in 2..20 {
let mut encoder = super::Encoder::new(
crate::Components { x: 4, y: 3 },
crate::ImageBounds { width, height },
)
.unwrap();
let chunk_size = bytes.len() / chunk_count;
for chunk in bytes.chunks(chunk_size) {
encoder.update(chunk);
}
let b2 = encoder.finalize();
assert_eq!(b1, b2);
}
}
}
}