pi-frame-server/src/imageproc.rs

348 lines
10 KiB
Rust
Raw Normal View History

2024-07-28 14:28:06 +00:00
use image::{GenericImageView, GrayImage, ImageBuffer, Luma, RgbImage};
use palette::FromColor;
2024-07-28 04:01:57 +00:00
use tracing::instrument;
2024-06-29 16:55:37 +00:00
2024-07-17 18:27:31 +00:00
use image::Rgb as imgRgb;
2024-07-18 20:51:07 +00:00
use palette::color_difference::Ciede2000;
use palette::{cast::FromComponents, IntoColor, Lab, Srgb};
2024-06-29 16:55:37 +00:00
/// Palette used on the display; pixels can be one of these colors.
///
/// The RGB values are slightly adjusted to improve accuracy.
const DISPLAY_PALETTE: [Srgb; 7] = [
Srgb::new(0.047, 0.047, 0.055), // Black
Srgb::new(0.824, 0.824, 0.816), // White
Srgb::new(0.118, 0.376, 0.122), // Green
Srgb::new(0.114, 0.118, 0.667), // Blue
Srgb::new(0.549, 0.106, 0.114), // Red
Srgb::new(0.827, 0.788, 0.239), // Yellow
Srgb::new(0.757, 0.443, 0.165), // Orange
];
2024-07-28 14:28:06 +00:00
pub enum Error {
DitherError,
PaletteIndexError(usize),
}
2024-06-29 16:55:37 +00:00
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DisplayColor {
Black,
White,
Green,
Blue,
Red,
Yellow,
Orange,
}
2024-07-02 15:57:29 +00:00
impl From<DisplayColor> for Srgb {
fn from(value: DisplayColor) -> Self {
DISPLAY_PALETTE[value as usize]
2024-06-29 16:55:37 +00:00
}
}
impl DisplayColor {
fn from_u8(value: u8) -> Self {
match value {
0 => Self::Black,
1 => Self::White,
2 => Self::Green,
3 => Self::Blue,
4 => Self::Red,
5 => Self::Yellow,
6 => Self::Orange,
2024-07-16 23:55:59 +00:00
_ => panic!("unexpected DisplayColor {value}"),
2024-06-29 16:55:37 +00:00
}
}
2024-07-02 15:57:29 +00:00
fn into_byte(color1: Self, color2: Self) -> u8 {
let upper: u8 = color1.into();
let lower: u8 = color2.into();
upper << 4 | lower
}
}
impl From<DisplayColor> for u8 {
fn from(value: DisplayColor) -> Self {
2024-07-16 23:55:59 +00:00
value as Self
2024-07-02 15:57:29 +00:00
}
}
2024-07-16 23:55:59 +00:00
/// Buffer to be sent to the ``EInk`` display.
2024-06-29 16:55:37 +00:00
#[derive(Debug)]
2024-07-17 18:27:31 +00:00
pub struct EInkImage {
data: Vec<DisplayColor>,
width: u32,
height: u32,
}
2024-06-29 16:55:37 +00:00
2024-07-28 14:28:06 +00:00
pub struct TestEInkImage {
buf: ImageBuffer<Luma<u8>, Vec<u8>>,
palette: Vec<Srgb>,
}
impl TestEInkImage {
pub fn into_display_buffer(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(self.buf.len() / 2);
for pix in self.buf.chunks_exact(2) {
buf.push(pix[0] << 4 | pix[1]);
}
buf
}
pub fn into_rgbimage(&self) -> RgbImage {
RgbImage::from_fn(self.buf.width(), self.buf.height(), |x, y| {
let idx = self.buf.get_pixel(x, y).0[0];
let disp_color = self.palette.get(idx as usize).unwrap();
let arr: [u8; 3] = disp_color.into_format().into();
imgRgb(arr)
})
}
/// Constructs a new EInk Image based on the given color palette for
/// color indexing.
#[must_use]
pub fn new(palette: Vec<Srgb>) -> Self {
Self {
buf: GrayImage::new(800, 480),
palette,
}
}
}
// TODO: Evaluate using Imagebuffer<Luma<u8>, Vec<u8>> instead.
// This is what the imageops index_map function does.
// advantages are we get all the 2d array helping functions for free.
2024-07-17 18:27:31 +00:00
impl EInkImage {
2024-07-16 23:55:59 +00:00
#[must_use]
2024-07-28 14:28:06 +00:00
2024-07-02 15:57:29 +00:00
pub fn into_display_buffer(&self) -> Vec<u8> {
2024-07-17 18:27:31 +00:00
let mut buf = Vec::with_capacity(self.data.len() / 2);
2024-06-29 16:55:37 +00:00
2024-07-17 18:27:31 +00:00
for colors in self.data.chunks_exact(2) {
2024-07-02 15:57:29 +00:00
buf.push(DisplayColor::into_byte(colors[0], colors[1]));
2024-06-29 16:55:37 +00:00
}
2024-07-02 15:57:29 +00:00
buf
}
2024-07-16 23:55:59 +00:00
#[must_use]
2024-07-17 18:27:31 +00:00
pub fn new(width: u32, height: u32) -> Self {
let v = vec![DisplayColor::Black; (width * height) as usize];
Self {
data: v,
width,
height,
}
2024-06-29 16:55:37 +00:00
}
2024-07-17 18:27:31 +00:00
/// Produces a regular RGB image from this image buffer using the given
/// color palette.
pub fn make_image(&self) -> RgbImage {
2024-07-17 18:27:31 +00:00
RgbImage::from_fn(self.width, self.height, |x, y| {
let srgb = Srgb::from(self.data[(y * self.width + x) as usize]);
let arr: [u8; 3] = srgb.into_format().into();
imgRgb(arr)
})
}
2024-07-17 18:27:31 +00:00
/// Returns the dimensions (width, height) of the image buffer.
pub const fn dimensions(&self) -> (u32, u32) {
(self.width, self.height)
}
2024-06-29 16:55:37 +00:00
}
pub trait Ditherer {
2024-07-28 14:28:06 +00:00
fn dither(&mut self, img: &RgbImage, output: &mut EInkImage);
2024-06-29 16:55:37 +00:00
}
2024-07-28 14:28:06 +00:00
pub type DitherFunc = dyn Fn(&RgbImage, &mut TestEInkImage) -> Result<(), Error>;
2024-06-29 16:55:37 +00:00
/// Find the closest approximate palette color to the given sRGB value.
/// This uses euclidian distance in linear space.
2024-07-17 05:24:11 +00:00
#[must_use]
2024-07-16 23:55:59 +00:00
pub fn nearest_neighbor(input_color: Lab) -> (DisplayColor, Lab) {
let (nearest, _, color_diff) = DISPLAY_PALETTE
2024-06-29 16:55:37 +00:00
.iter()
.enumerate()
.map(|(idx, p_color)| {
2024-07-02 15:57:29 +00:00
let c: Lab = (*p_color).into_color();
2024-07-17 15:20:06 +00:00
(idx, input_color.difference(c), input_color - c)
2024-06-29 16:55:37 +00:00
})
2024-07-16 23:55:59 +00:00
.min_by(|(_, a, _), (_, b, _)| a.total_cmp(b))
2024-07-28 14:28:06 +00:00
.expect("Should always find a color");
(DisplayColor::from_u8(nearest as u8), color_diff)
}
fn nearest_neighbor2(input_color: Lab, palette:&[Srgb]) -> (DisplayColor, Lab) {
let (nearest, _, color_diff) = palette
.iter()
.enumerate()
.map(|(idx, p_color)| {
let c: Lab = Lab::from_color(*p_color);
(idx, input_color.difference(c), input_color - c)
})
.min_by(|(_, a, _), (_, b, _)| a.total_cmp(b))
.expect("Should always find a color");
2024-07-16 23:55:59 +00:00
(DisplayColor::from_u8(nearest as u8), color_diff)
2024-06-29 16:55:37 +00:00
}
pub struct NNDither();
impl Ditherer for NNDither {
2024-07-28 14:28:06 +00:00
fn dither(&mut self, img: &RgbImage, output: &mut EInkImage) {
2024-06-29 16:55:37 +00:00
assert!(img.width() == 800);
assert!(img.height() == 480);
2024-07-16 23:55:59 +00:00
2024-07-02 15:57:29 +00:00
// sRGB view into the given image. zero copy!
2024-06-29 16:55:37 +00:00
let srgb = <&[Srgb<u8>]>::from_components(&**img);
2024-07-02 15:57:29 +00:00
for (idx, pixel) in srgb.iter().enumerate() {
2024-07-16 23:55:59 +00:00
let (n, _) = nearest_neighbor(pixel.into_format().into_color());
2024-07-17 18:27:31 +00:00
output.data[idx] = n;
2024-07-02 15:57:29 +00:00
}
2024-06-29 16:55:37 +00:00
}
}
2024-07-02 15:57:29 +00:00
2024-07-17 05:24:11 +00:00
/// Compute the vector index for a given image by using the size of rows. Assumes that images
/// are indexed in row-major order.
const fn coord_to_idx(x: u32, y: u32, xsize: u32) -> usize {
(y * xsize + x) as usize
}
/// Compute the error-adjusted new lab value based on the error value of the currently scanned
/// pixel multiplied by a scalar factor.
2024-07-28 14:28:06 +00:00
fn compute_error_adjusted_color(orig: &Lab, err: &Lab, weight: f32) -> Lab {
let (orig_l, orig_a, orig_b) = orig.into_components();
2024-07-17 05:24:11 +00:00
let (err_l, err_a, err_b) = err.into_components();
Lab::from_components((
2024-07-28 14:28:06 +00:00
err_l.mul_add(weight, orig_l), // scalar * err_l + p_l
err_a.mul_add(weight, orig_a),
err_b.mul_add(weight, orig_b),
2024-07-17 05:24:11 +00:00
))
}
/// ``DiffusionPoint`` is part of the diffusion matrix, represented by a shift in x and y and an error
/// scaling factor.
2024-07-17 05:24:11 +00:00
struct DiffusionPoint {
xshift: i32,
yshift: i32,
scale: f32,
}
impl DiffusionPoint {
/// Creates a new ``DiffusionPoint``
2024-07-17 05:24:11 +00:00
const fn new(xshift: i32, yshift: i32, scale: f32) -> Self {
Self {
xshift,
yshift,
scale,
}
}
}
static FLOYD_STEINBERG_POINTS: &[DiffusionPoint] = &[
2024-07-17 05:24:11 +00:00
DiffusionPoint::new(1, 0, 7.0 / 16.0),
DiffusionPoint::new(-1, 1, 3.0 / 16.0),
DiffusionPoint::new(0, 1, 5.0 / 16.0),
DiffusionPoint::new(1, 1, 1.0 / 16.0),
];
static ATKINSON_DITHER_POINTS: &[DiffusionPoint] = &[
DiffusionPoint::new(1, 0, 1.0 / 8.0),
DiffusionPoint::new(2, 0, 1.0 / 8.0),
DiffusionPoint::new(-1, 1, 1.0 / 8.0),
DiffusionPoint::new(0, 1, 1.0 / 8.0),
DiffusionPoint::new(1, 1, 1.0 / 8.0),
DiffusionPoint::new(0, 2, 1.0 / 8.0),
];
static SIERRA_DITHER_POINTS: &[DiffusionPoint] = &[
DiffusionPoint::new(1, 0, 5.0 / 32.0),
DiffusionPoint::new(2, 0, 3.0 / 32.0),
DiffusionPoint::new(-2, 1, 2.0 / 32.0),
DiffusionPoint::new(-1, 1, 4.0 / 32.0),
DiffusionPoint::new(0, 1, 5.0 / 32.0),
DiffusionPoint::new(1, 1, 4.0 / 32.0),
DiffusionPoint::new(2, 1, 2.0 / 32.0),
DiffusionPoint::new(-1, 2, 2.0 / 32.0),
DiffusionPoint::new(0, 2, 3.0 / 32.0),
DiffusionPoint::new(1, 2, 2.0 / 32.0),
];
static STUKI_DITHER_POINTS: &[DiffusionPoint] = &[
DiffusionPoint::new(1, 0, 8.0 / 42.0),
DiffusionPoint::new(2, 0, 4.0 / 42.0),
DiffusionPoint::new(-2, 1, 2.0 / 42.0),
DiffusionPoint::new(-1, 1, 4.0 / 42.0),
DiffusionPoint::new(0, 1, 8.0 / 42.0),
DiffusionPoint::new(1, 1, 4.0 / 42.0),
DiffusionPoint::new(2, 1, 2.0 / 42.0),
DiffusionPoint::new(-2, 2, 1.0 / 42.0),
DiffusionPoint::new(-1, 2, 2.0 / 42.0),
DiffusionPoint::new(0, 2, 4.0 / 42.0),
DiffusionPoint::new(1, 2, 2.0 / 42.0),
DiffusionPoint::new(1, 2, 1.0 / 42.0),
2024-07-17 15:20:06 +00:00
];
2024-07-17 05:24:11 +00:00
#[derive(Debug)]
pub enum DiffusionMatrix {
FloydSteinberg,
Atkinson,
Sierra,
Stuki,
}
impl DiffusionMatrix {
fn value(&self) -> &'static [DiffusionPoint] {
match *self {
Self::FloydSteinberg => FLOYD_STEINBERG_POINTS,
Self::Atkinson => ATKINSON_DITHER_POINTS,
Self::Sierra => SIERRA_DITHER_POINTS,
Self::Stuki => STUKI_DITHER_POINTS,
}
}
}
#[derive(Debug)]
pub struct ErrorDiffusionDither(DiffusionMatrix);
impl ErrorDiffusionDither {
#[must_use]
pub const fn new(dm: DiffusionMatrix) -> Self {
Self(dm)
}
}
impl Ditherer for ErrorDiffusionDither {
#[instrument]
2024-07-28 14:28:06 +00:00
fn dither(&mut self, img: &RgbImage, output: &mut EInkImage) {
2024-07-16 23:55:59 +00:00
// create a copy of the image in Lab space, mutable.
2024-07-28 14:28:06 +00:00
// first, a view into the rgb components
2024-07-16 23:55:59 +00:00
let srgb = <&[Srgb<u8>]>::from_components(&**img);
2024-07-17 05:24:11 +00:00
let (xsize, ysize) = img.dimensions();
2024-07-28 14:28:06 +00:00
// our destination buffer.
2024-07-17 05:24:11 +00:00
let mut temp_img: Vec<Lab> = Vec::with_capacity((xsize * ysize) as usize);
2024-07-16 23:55:59 +00:00
for pix in srgb {
temp_img.push(pix.into_format().into_color());
}
2024-07-17 05:24:11 +00:00
// now we take our units.
for y in 0..ysize {
for x in 0..xsize {
let index = coord_to_idx(x, y, xsize);
let curr_pix = temp_img[index];
let (nearest, err) = nearest_neighbor(curr_pix);
// set the color in the output buffer.
2024-07-17 18:27:31 +00:00
output.data[index] = nearest;
2024-07-17 05:24:11 +00:00
// take the error, and propagate it.
for point in self.0.value() {
2024-07-17 05:24:11 +00:00
let Some(target_x) = x.checked_add_signed(point.xshift) else {
continue;
};
let Some(target_y) = y.checked_add_signed(point.yshift) else {
continue;
};
let target = coord_to_idx(target_x, target_y, xsize);
2024-07-17 05:24:11 +00:00
if let Some(pix) = temp_img.get(target) {
2024-07-28 14:28:06 +00:00
temp_img[target] = compute_error_adjusted_color(pix, &err, point.scale);
2024-07-17 05:24:11 +00:00
}
}
}
}
2024-07-16 23:55:59 +00:00
}
}