use crate::color_space::{srgb_channel_to_linear, srgb_fraction_to_linear};
use crate::color_space_lab::{PaletteLab, match_pixel_oklab, rgb_to_oklab, WAB};
use crate::palettes::Palette;
use rayon::prelude::*;
pub struct KernelOffset {
pub dx: i32,
pub dy: i32,
pub weight: f64,
}
pub struct Kernel {
pub offsets: &'static [KernelOffset],
}
pub static FLOYD_STEINBERG: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 7.0 / 16.0 },
KernelOffset { dx: -1, dy: 1, weight: 3.0 / 16.0 },
KernelOffset { dx: 0, dy: 1, weight: 5.0 / 16.0 },
KernelOffset { dx: 1, dy: 1, weight: 1.0 / 16.0 },
],
};
pub static ATKINSON: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 1.0 / 8.0 },
KernelOffset { dx: 2, dy: 0, weight: 1.0 / 8.0 },
KernelOffset { dx: -1, dy: 1, weight: 1.0 / 8.0 },
KernelOffset { dx: 0, dy: 1, weight: 1.0 / 8.0 },
KernelOffset { dx: 1, dy: 1, weight: 1.0 / 8.0 },
KernelOffset { dx: 0, dy: 2, weight: 1.0 / 8.0 },
],
};
pub static BURKES: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 8.0 / 32.0 },
KernelOffset { dx: 2, dy: 0, weight: 4.0 / 32.0 },
KernelOffset { dx: -2, dy: 1, weight: 2.0 / 32.0 },
KernelOffset { dx: -1, dy: 1, weight: 4.0 / 32.0 },
KernelOffset { dx: 0, dy: 1, weight: 8.0 / 32.0 },
KernelOffset { dx: 1, dy: 1, weight: 4.0 / 32.0 },
KernelOffset { dx: 2, dy: 1, weight: 2.0 / 32.0 },
],
};
pub static STUCKI: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 8.0 / 42.0 },
KernelOffset { dx: 2, dy: 0, weight: 4.0 / 42.0 },
KernelOffset { dx: -2, dy: 1, weight: 2.0 / 42.0 },
KernelOffset { dx: -1, dy: 1, weight: 4.0 / 42.0 },
KernelOffset { dx: 0, dy: 1, weight: 8.0 / 42.0 },
KernelOffset { dx: 1, dy: 1, weight: 4.0 / 42.0 },
KernelOffset { dx: 2, dy: 1, weight: 2.0 / 42.0 },
KernelOffset { dx: -2, dy: 2, weight: 1.0 / 42.0 },
KernelOffset { dx: -1, dy: 2, weight: 2.0 / 42.0 },
KernelOffset { dx: 0, dy: 2, weight: 4.0 / 42.0 },
KernelOffset { dx: 1, dy: 2, weight: 2.0 / 42.0 },
KernelOffset { dx: 2, dy: 2, weight: 1.0 / 42.0 },
],
};
pub static SIERRA: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 5.0 / 32.0 },
KernelOffset { dx: 2, dy: 0, weight: 3.0 / 32.0 },
KernelOffset { dx: -2, dy: 1, weight: 2.0 / 32.0 },
KernelOffset { dx: -1, dy: 1, weight: 4.0 / 32.0 },
KernelOffset { dx: 0, dy: 1, weight: 5.0 / 32.0 },
KernelOffset { dx: 1, dy: 1, weight: 4.0 / 32.0 },
KernelOffset { dx: 2, dy: 1, weight: 2.0 / 32.0 },
KernelOffset { dx: -1, dy: 2, weight: 2.0 / 32.0 },
KernelOffset { dx: 0, dy: 2, weight: 3.0 / 32.0 },
KernelOffset { dx: 1, dy: 2, weight: 2.0 / 32.0 },
],
};
pub static SIERRA_LITE: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 2.0 / 4.0 },
KernelOffset { dx: -1, dy: 1, weight: 1.0 / 4.0 },
KernelOffset { dx: 0, dy: 1, weight: 1.0 / 4.0 },
],
};
pub static JARVIS_JUDICE_NINKE: Kernel = Kernel {
offsets: &[
KernelOffset { dx: 1, dy: 0, weight: 7.0 / 48.0 },
KernelOffset { dx: 2, dy: 0, weight: 5.0 / 48.0 },
KernelOffset { dx: -2, dy: 1, weight: 3.0 / 48.0 },
KernelOffset { dx: -1, dy: 1, weight: 5.0 / 48.0 },
KernelOffset { dx: 0, dy: 1, weight: 7.0 / 48.0 },
KernelOffset { dx: 1, dy: 1, weight: 5.0 / 48.0 },
KernelOffset { dx: 2, dy: 1, weight: 3.0 / 48.0 },
KernelOffset { dx: -2, dy: 2, weight: 1.0 / 48.0 },
KernelOffset { dx: -1, dy: 2, weight: 3.0 / 48.0 },
KernelOffset { dx: 0, dy: 2, weight: 5.0 / 48.0 },
KernelOffset { dx: 1, dy: 2, weight: 3.0 / 48.0 },
KernelOffset { dx: 2, dy: 2, weight: 1.0 / 48.0 },
],
};
fn build_palette_lab(palette: &Palette) -> (Vec<[f64; 3]>, PaletteLab) {
let palette_linear: Vec<[f64; 3]> = palette
.colors
.iter()
.map(|&[r, g, b]| {
[
srgb_channel_to_linear(r),
srgb_channel_to_linear(g),
srgb_channel_to_linear(b),
]
})
.collect();
let palette_lab = PaletteLab::from_linear_rgb(&palette_linear);
(palette_linear, palette_lab)
}
pub fn error_diffusion_dither(
pixels: &[u8],
width: usize,
height: usize,
palette: &Palette,
kernel: &Kernel,
serpentine: bool,
) -> Vec<u8> {
error_diffusion_dither_impl(pixels, width, height, palette, None, kernel, serpentine)
}
pub fn error_diffusion_dither_with_canonical(
pixels: &[u8],
width: usize,
height: usize,
palette: &Palette,
canonical_palette: &Palette,
kernel: &Kernel,
serpentine: bool,
) -> Vec<u8> {
error_diffusion_dither_impl(
pixels,
width,
height,
palette,
Some(canonical_palette),
kernel,
serpentine,
)
}
fn error_diffusion_dither_impl(
pixels: &[u8],
width: usize,
height: usize,
palette: &Palette,
canonical_palette: Option<&Palette>,
kernel: &Kernel,
serpentine: bool,
) -> Vec<u8> {
let (_palette_linear, palette_lab) = build_palette_lab(palette);
let palette_srgb_f: Vec<[f64; 3]> = palette
.colors
.iter()
.map(|&[r, g, b]| [r as f64, g as f64, b as f64])
.collect();
let mut buf: Vec<f64> = pixels.iter().map(|&v| v as f64).collect();
let lut: Vec<f64> = (0u8..=255).map(srgb_channel_to_linear).collect();
let mut output = vec![0u8; width * height];
for y in 0..height {
let reverse = serpentine && y % 2 == 1;
for xi in 0..width {
let x = if reverse { width - 1 - xi } else { xi };
let idx = (y * width + x) * 3;
if let Some(canonical_palette) = canonical_palette
&& let Some(exact_idx) =
exact_palette_index(&pixels[idx..idx + 3], canonical_palette)
{
output[y * width + x] = exact_idx;
continue;
}
let rs = buf[idx].clamp(0.0, 255.0);
let gs = buf[idx + 1].clamp(0.0, 255.0);
let bs = buf[idx + 2].clamp(0.0, 255.0);
let r_lin = lut[rs.round() as usize];
let g_lin = lut[gs.round() as usize];
let b_lin = lut[bs.round() as usize];
let pixel_lab = rgb_to_oklab(r_lin, g_lin, b_lin);
let best_idx = match_pixel_oklab(pixel_lab, &palette_lab, WAB);
output[y * width + x] = best_idx as u8;
let err_r = rs - palette_srgb_f[best_idx][0];
let err_g = gs - palette_srgb_f[best_idx][1];
let err_b = bs - palette_srgb_f[best_idx][2];
for offset in kernel.offsets {
let effective_dx = if reverse { -offset.dx } else { offset.dx };
let nx = x as i64 + effective_dx as i64;
let ny = y as i64 + offset.dy as i64;
if nx >= 0 && nx < width as i64 && ny >= 0 && ny < height as i64 {
let ni = (ny as usize * width + nx as usize) * 3;
buf[ni] += err_r * offset.weight;
buf[ni + 1] += err_g * offset.weight;
buf[ni + 2] += err_b * offset.weight;
}
}
}
}
output
}
pub fn floyd_steinberg(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &FLOYD_STEINBERG, serpentine)
}
pub fn atkinson(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &ATKINSON, serpentine)
}
pub fn burkes(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &BURKES, serpentine)
}
pub fn stucki(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &STUCKI, serpentine)
}
pub fn sierra(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &SIERRA, serpentine)
}
pub fn sierra_lite(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &SIERRA_LITE, serpentine)
}
pub fn jarvis_judice_ninke(pixels: &[u8], w: usize, h: usize, palette: &Palette, serpentine: bool) -> Vec<u8> {
error_diffusion_dither(pixels, w, h, palette, &JARVIS_JUDICE_NINKE, serpentine)
}
fn exact_palette_index(rgb: &[u8], palette: &Palette) -> Option<u8> {
palette
.colors
.iter()
.position(|&color| rgb == color)
.and_then(|idx| u8::try_from(idx).ok())
}
pub fn try_exact_palette_map(pixels: &[u8], canonical_palette: &Palette) -> Option<Vec<u8>> {
pixels
.par_chunks(3)
.map(|rgb| exact_palette_index(rgb, canonical_palette))
.collect()
}
pub fn direct_map(pixels: &[u8], palette: &Palette, canonical_palette: &Palette) -> Vec<u8> {
let (_, palette_lab) = build_palette_lab(palette);
pixels
.par_chunks(3)
.map(|rgb| {
if let Some(idx) = exact_palette_index(rgb, canonical_palette) {
return idx;
}
let r = srgb_channel_to_linear(rgb[0]);
let g = srgb_channel_to_linear(rgb[1]);
let b = srgb_channel_to_linear(rgb[2]);
let lab = rgb_to_oklab(r, g, b);
match_pixel_oklab(lab, &palette_lab, WAB) as u8
})
.collect()
}
const BAYER_4X4: [[f64; 4]; 4] = [
[-0.46875, 0.03125, -0.34375, 0.15625],
[ 0.28125, -0.21875, 0.40625, -0.09375],
[-0.28125, 0.21875, -0.40625, 0.09375],
[ 0.46875, -0.03125, 0.34375, -0.15625],
];
fn ordered_spread(palette: &Palette) -> f64 {
let is_grayscale =
palette.colors.len() >= 3 && palette.colors.iter().all(|&[r, g, b]| r == g && g == b);
if is_grayscale {
1.0 / (palette.colors.len() - 1) as f64
} else {
1.0
}
}
pub fn ordered_dither(pixels: &[u8], width: usize, palette: &Palette) -> Vec<u8> {
ordered_dither_impl(pixels, width, palette, None)
}
pub fn ordered_dither_with_canonical(
pixels: &[u8],
width: usize,
palette: &Palette,
canonical_palette: &Palette,
) -> Vec<u8> {
ordered_dither_impl(pixels, width, palette, Some(canonical_palette))
}
fn ordered_dither_impl(
pixels: &[u8],
width: usize,
palette: &Palette,
canonical_palette: Option<&Palette>,
) -> Vec<u8> {
let (_palette_linear, palette_lab) = build_palette_lab(palette);
let spread = ordered_spread(palette);
let lut: Vec<[f64; 256]> = (0..16)
.map(|t| {
let threshold = BAYER_4X4[t / 4][t % 4] * spread;
let mut row = [0.0_f64; 256];
for (v, slot) in row.iter_mut().enumerate() {
*slot = srgb_fraction_to_linear((v as f64 / 255.0 + threshold).clamp(0.0, 1.0));
}
row
})
.collect();
pixels
.par_chunks(3)
.enumerate()
.map(|(i, rgb)| {
if let Some(canonical_palette) = canonical_palette
&& let Some(idx) = exact_palette_index(rgb, canonical_palette)
{
return idx;
}
let x = i % width;
let y = i / width;
let t = (y % 4) * 4 + (x % 4);
let lab = rgb_to_oklab(
lut[t][rgb[0] as usize],
lut[t][rgb[1] as usize],
lut[t][rgb[2] as usize],
);
match_pixel_oklab(lab, &palette_lab, WAB) as u8
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::palettes::ColorScheme;
fn solid_image(r: u8, g: u8, b: u8, w: usize, h: usize) -> Vec<u8> {
vec![r, g, b].into_iter().cycle().take(w * h * 3).collect()
}
#[test]
fn solid_white_maps_to_white() {
let pixels = solid_image(255, 255, 255, 4, 4);
let out = floyd_steinberg(&pixels, 4, 4, ColorScheme::Mono.palette(), true);
assert!(out.iter().all(|&v| v == 1));
}
#[test]
fn solid_black_maps_to_black() {
let pixels = solid_image(0, 0, 0, 4, 4);
let out = floyd_steinberg(&pixels, 4, 4, ColorScheme::Mono.palette(), true);
assert!(out.iter().all(|&v| v == 0));
}
#[test]
fn output_length_matches_pixels() {
let pixels = solid_image(128, 64, 200, 10, 7);
let out = floyd_steinberg(&pixels, 10, 7, ColorScheme::Bwr.palette(), true);
assert_eq!(out.len(), 10 * 7);
}
#[test]
fn all_algorithms_produce_correct_length() {
let pixels = solid_image(128, 64, 200, 10, 7);
let pal = ColorScheme::Bwr.palette();
assert_eq!(burkes(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(atkinson(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(stucki(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(sierra(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(sierra_lite(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(jarvis_judice_ninke(&pixels, 10, 7, pal, false).len(), 70);
assert_eq!(ordered_dither(&pixels, 10, pal).len(), 70);
assert_eq!(direct_map(&pixels, pal, pal).len(), 70);
}
#[test]
fn direct_map_pure_red_maps_to_red_in_bwr() {
let pixels = vec![255u8, 0, 0];
let out = direct_map(&pixels, ColorScheme::Bwr.palette(), ColorScheme::Bwr.palette());
assert_eq!(out[0], 2, "pure sRGB red should map to red ink (index 2) in BWR");
}
#[test]
fn direct_map_pure_blue_maps_to_blue_in_bwgbry() {
let pixels = vec![0u8, 0, 255];
let out = direct_map(&pixels, ColorScheme::Bwgbry.palette(), ColorScheme::Bwgbry.palette());
assert_eq!(out[0], 4, "pure sRGB blue should map to blue ink (index 4) in BWGBRY");
}
#[test]
fn serpentine_differs_from_raster_on_gradient() {
let pixels: Vec<u8> = (0u8..=255)
.flat_map(|v| [v, v / 2, 255 - v])
.cycle()
.take(16 * 16 * 3)
.collect();
let raster = floyd_steinberg(&pixels, 16, 16, ColorScheme::Mono.palette(), false);
let serpentine = floyd_steinberg(&pixels, 16, 16, ColorScheme::Mono.palette(), true);
assert_ne!(raster, serpentine, "serpentine and raster should differ on a gradient");
}
#[test]
fn ordered_dither_activity_is_perceptually_uniform() {
const W: usize = 256;
const H: usize = 16; const BANDS: usize = 8;
const BAND_W: usize = W / BANDS;
let mut pixels = Vec::with_capacity(W * H * 3);
for _ in 0..H {
for x in 0..W {
let v = x as u8;
pixels.extend_from_slice(&[v, v, v]);
}
}
let out = ordered_dither(&pixels, W, ColorScheme::Mono.palette());
let mut transitions = [0usize; BANDS];
for y in 0..H {
for x in 0..W - 1 {
if out[y * W + x] != out[y * W + x + 1] {
transitions[(x / BAND_W).min(BANDS - 1)] += 1;
}
}
}
let mid = &transitions[1..BANDS - 2];
let max = *mid.iter().max().unwrap();
let min = *mid.iter().min().unwrap();
assert!(min > 0, "every mid-tone band should have at least one transition: {transitions:?}");
let ratio = max as f64 / min as f64;
assert!(
ratio < 5.0,
"ordered-dither activity is perceptually skewed (max/min = {ratio:.2}); \
expected sRGB-space ordered dither to spread roughly uniformly across \
mid-tones. Per-band transitions: {transitions:?}"
);
}
#[test]
fn bayer_matrix_is_zero_mean() {
let sum: f64 = BAYER_4X4.iter().flatten().sum();
assert!(sum.abs() < 1e-12, "Bayer thresholds must sum to zero, got {sum}");
}
#[test]
fn ordered_dither_is_brightness_neutral_at_midpoint() {
let pal = ColorScheme::Mono.palette();
let (_, pal_lab) = build_palette_lab(pal);
let midpoint = (0u16..=255)
.find(|&v| {
let lin = srgb_channel_to_linear(v as u8);
let lab = rgb_to_oklab(lin, lin, lin);
match_pixel_oklab(lab, &pal_lab, WAB) == 1
})
.expect("mono palette must have a black→white transition") as u8;
const N: usize = 32;
let pixels = solid_image(midpoint, midpoint, midpoint, N, N);
let out = ordered_dither(&pixels, N, pal);
let white = out.iter().filter(|&&v| v == 1).count() as f64 / out.len() as f64;
assert!(
(white - 0.5).abs() <= 1.0 / 16.0,
"ordered dither at the mono midpoint should be ~50% white, got {white}"
);
}
#[test]
fn ordered_dither_grayscale16_tracks_ramp_and_uses_many_levels() {
const W: usize = 256;
const H: usize = 4;
let pal = ColorScheme::Grayscale16.palette();
let mut pixels = Vec::with_capacity(W * H * 3);
for _ in 0..H {
for x in 0..W {
let v = x as u8;
pixels.extend_from_slice(&[v, v, v]);
}
}
let out = ordered_dither(&pixels, W, pal);
let levels: Vec<f64> = pal.colors.iter().map(|&[r, _, _]| r as f64).collect();
let mut max_err = 0.0_f64;
for bx in 0..W / 4 {
let mut sum = 0.0;
for y in 0..H {
for dx in 0..4 {
let x = bx * 4 + dx;
sum += levels[out[y * W + x] as usize];
}
}
let avg = sum / (H * 4) as f64;
let input = (bx * 4 + 1) as f64; max_err = max_err.max((avg - input).abs());
}
assert!(max_err < 24.0, "block-averaged output should track the ramp, max_err={max_err}");
let mut seen = [false; 16];
for &idx in &out {
seen[idx as usize] = true;
}
let used = seen.iter().filter(|&&s| s).count();
assert!(used >= 12, "GS16 ordered dither should use ≥12 levels, used {used}");
}
#[test]
fn serpentine_first_row_matches_raster() {
let pixels: Vec<u8> = (0u8..=255)
.flat_map(|v| [v, 128u8.wrapping_add(v), 64u8])
.cycle()
.take(8 * 4 * 3)
.collect();
let raster = floyd_steinberg(&pixels, 8, 4, ColorScheme::Mono.palette(), false);
let serpentine = floyd_steinberg(&pixels, 8, 4, ColorScheme::Mono.palette(), true);
assert_eq!(
&raster[0..8],
&serpentine[0..8],
"first row must be identical (both scan left-to-right)"
);
}
}