pi-frame-server/src/imageproc.rs

use image::{GrayImage, ImageBuffer, Luma, RgbImage};
use palette::FromColor;
use serde::{Deserialize, Serialize};
use tracing::instrument;

use image::Rgb as imgRgb;
use palette::color_difference::Ciede2000;
use palette::{cast::FromComponents, IntoColor, Lab, Srgb};

/// 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
];

const SIMPLE_PALETTE: [Srgb; 7] = [
    Srgb::new(0.0, 0.0, 0.0),       // Black
    Srgb::new(1.0, 1.0, 1.0),       // White
    Srgb::new(0.0, 1.0, 0.0),       // Green
    Srgb::new(0.0, 0.0, 1.0),       // Blue
    Srgb::new(1.0, 0.0, 0.0),       // Red
    Srgb::new(1.0, 1.0, 0.0),       // Yellow
    Srgb::new(0.757, 0.443, 0.165), // Orange
];

#[derive(strum::EnumString, Serialize, Deserialize, PartialEq, Eq, Debug)]
pub enum DitherPalette {
    Default,
    Simple,
}

impl DitherPalette {
    #[must_use]
    pub const fn value(&self) -> &[Srgb] {
        match self {
            Self::Default => &DISPLAY_PALETTE,
            Self::Simple => &SIMPLE_PALETTE, // FIXME: use simple pallete based on binary.
        }
    }
}

#[derive(strum::EnumString, Serialize, Deserialize, PartialEq, Eq, Debug)]
pub enum DitherMethod {
    NearestNeighbor,
    FloydSteinberg,
    Atkinson,
    Stuki,
    Sierra,
}

impl DitherMethod {
    #[must_use]
    pub fn get_ditherer(&self) -> Box<dyn Ditherer> {
        match self {
            Self::NearestNeighbor => Box::new(NNDither {}),
            Self::Atkinson => Box::new(ErrorDiffusionDither::new(ATKINSON_DITHER_POINTS)),
            Self::FloydSteinberg => Box::new(ErrorDiffusionDither::new(FLOYD_STEINBERG_POINTS)),
            Self::Stuki => Box::new(ErrorDiffusionDither::new(STUKI_DITHER_POINTS)),
            Self::Sierra => Box::new(ErrorDiffusionDither::new(SIERRA_DITHER_POINTS)),
        }
    }
}

pub enum ProcessingError {
    DitherError,
    PaletteIndexError(usize),
}

/// Buffer to be sent to the ``EInk`` display.
#[derive(Debug)]
pub struct DitheredImage {
    buf: ImageBuffer<Luma<u8>, Vec<u8>>,
    palette: Vec<Srgb>,
}

impl DitheredImage {
    #[must_use]
    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
    }

    /// Convert the index-based image back into an RGB image using the color palette.
    #[must_use]
    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)
                .expect("Palette will never be out of bounds");
            let arr: [u8; 3] = disp_color.into_format().into();
            imgRgb(arr)
        })
    }

    /// Constructs a new Dithered Image based on the given dimensions and color palette for
    /// color indexing.
    #[must_use]
    pub fn new(width: u32, height: u32, palette: Vec<Srgb>) -> Self {
        Self {
            buf: GrayImage::new(width, height),
            palette,
        }
    }
}

impl Default for DitheredImage {
    fn default() -> Self {
        Self::new(800, 480, DISPLAY_PALETTE.to_vec())
    }
}

pub trait Ditherer {
    fn dither(&mut self, img: &RgbImage, output: &mut DitheredImage);
}

/// Find the closest approximate palette color to the given sRGB value.
/// This uses euclidian distance in linear space.
fn nearest_neighbor(input_color: Lab, palette: &[Srgb]) -> (u8, 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");
    (nearest as u8, color_diff)
}

pub struct NNDither();

impl Ditherer for NNDither {
    fn dither(&mut self, img: &RgbImage, output: &mut DitheredImage) {
        assert!(img.width() == 800);
        assert!(img.height() == 480);

        // sRGB view into the given image. zero copy!
        let srgb = <&[Srgb<u8>]>::from_components(&**img);

        for (idx, pix) in output.buf.iter_mut().enumerate() {
            let (n, _) = nearest_neighbor(srgb[idx].into_format().into_color(), &output.palette);
            *pix = n;
        }
    }
}

/// 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.
fn compute_error_adjusted_color(orig: &Lab, err: &Lab, weight: f32) -> Lab {
    let (orig_l, orig_a, orig_b) = orig.into_components();
    let (err_l, err_a, err_b) = err.into_components();
    Lab::from_components((
        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),
    ))
}

/// ``DiffusionPoint`` is part of the diffusion matrix, represented by a shift in x and y and an error
/// scaling factor.
#[derive(Debug)]
pub struct DiffusionPoint {
    xshift: i32,
    yshift: i32,
    scale: f32,
}

impl DiffusionPoint {
    /// Creates a new ``DiffusionPoint``
    const fn new(xshift: i32, yshift: i32, scale: f32) -> Self {
        Self {
            xshift,
            yshift,
            scale,
        }
    }
}

static FLOYD_STEINBERG_POINTS: &[DiffusionPoint] = &[
    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),
];

pub type DiffusionMatrix<'a> = &'a [DiffusionPoint];

#[derive(Debug)]
pub struct ErrorDiffusionDither<'a>(&'a [DiffusionPoint]);

impl<'a> ErrorDiffusionDither<'a> {
    #[must_use]
    pub const fn new(dm: DiffusionMatrix<'a>) -> Self {
        Self(dm)
    }
}

impl<'a> Ditherer for ErrorDiffusionDither<'a> {
    #[instrument]
    fn dither(&mut self, img: &RgbImage, output: &mut DitheredImage) {
        // create a copy of the image in Lab space, mutable.
        // first, a view into the rgb components
        let srgb = <&[Srgb<u8>]>::from_components(&**img);
        let (xsize, ysize) = img.dimensions();
        // our destination buffer.
        let mut temp_img: Vec<Lab> = Vec::with_capacity((xsize * ysize) as usize);
        for pix in srgb {
            temp_img.push(pix.into_format().into_color());
        }

        // TODO: rework this to make more sense.
        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, &output.palette);
                // set the color in the output buffer.
                *output.buf.get_mut(index).unwrap() = nearest;
                // take the error, and propagate it.
                for point in self.0 {
                    // bounds checking.
                    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);
                    if let Some(pix) = temp_img.get(target) {
                        temp_img[target] = compute_error_adjusted_color(pix, &err, point.scale);
                    }
                }
            }
        }
    }
}