pi-frame-server/src/imageproc.rs

use image::RgbImage;

use palette::{cast::FromComponents, IntoColor, Lab, Oklch, Srgb};
use palette::color_difference::{Ciede2000, EuclideanDistance};

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

// fn octcolor_rgb(color: &OctColor) -> &Srgb<u8> {
//     match color {
//         OctColor::Black => &DISPLAY_PALETTE[0],
//         OctColor::White => &DISPLAY_PALETTE[1],
//         OctColor::Green => &DISPLAY_PALETTE[2],
//         OctColor::Blue => &DISPLAY_PALETTE[3],
//         OctColor::Red => &DISPLAY_PALETTE[4],
//         OctColor::Yellow => &DISPLAY_PALETTE[5],
//         OctColor::Orange => &DISPLAY_PALETTE[6],
//         OctColor::HiZ => &DISPLAY_PALETTE[1],
//     }
// }

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DisplayColor {
    Black,
    White,
    Green,
    Blue,
    Red,
    Yellow,
    Orange,
}

impl From<DisplayColor> for Srgb {
    fn from(value: DisplayColor) -> Self {
        DISPLAY_PALETTE[value as usize]
    }
}

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,
            _ => panic!("unexpected DisplayColor {value}"),
        }
    }

    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 {
        value as Self
    }
}
/// Buffer to be sent to the ``EInk`` display.
#[derive(Debug)]
pub struct EInkBuffer(Vec<DisplayColor>);

impl EInkBuffer {
    #[must_use]
    pub fn into_display_buffer(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(self.0.len() / 2);

        for colors in self.0.chunks_exact(2) {
            buf.push(DisplayColor::into_byte(colors[0], colors[1]));
        }

        buf
    }
    #[must_use]
    pub fn new(width: usize, height: usize) -> Self {
        let v = vec![DisplayColor::Black; width * height];
        Self(v)
    }
    // pub fn make_image(&self) -> RgbImage {
    //     RgbImage::from_fn(800, 480, |x, y| {
    //         let srgb = Srgb::from(self.0[y * 800 + x]);
    //     })    
    // }
}

// impl EInkBuffer {
//     /// Converts the EInkBuffer into data that can be sent over the SPI API
//     /// Bin-packs the two 4-bit colors into bytes.
//     pub fn into_buffer(&self) -> Vec<u8> {
//         vec![]
//     }
//
//     pub fn new(width: usize, height: usize) -> EInkBuffer {
//         EInkBuffer {
//             data: vec![DisplayColor::Black; width * height],
//             width,
//             height,
//         }
//     }
//     pub fn set(&mut self, x: usize, y: usize, value: DisplayColor) {
//         self.data[x + y * self.width] = value;
//     }
//     pub fn get(&self, x:usize, y:usize) -> DisplayColor {
//         self.data[x
//     }
// }

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

// fn color_distance(c1: LinSrgb, c2: LinSrgb) -> f32 {
//     let r2 = (c2.red - c1.red).powf(2.0);
//     let g2 = (c2.green - c1.green).powf(2.0);
//     let b2 = (c2.blue - c1.blue).powf(2.0);
//
//     (r2 + g2 + b2).sqrt()
// }

/// Find the closest approximate palette color to the given sRGB value.
/// This uses euclidian distance in linear space.
#[must_use]
pub fn nearest_neighbor(input_color: Lab) -> (DisplayColor, Lab) {
    let (nearest, _, color_diff) = DISPLAY_PALETTE
        .iter()
        .enumerate()
        .map(|(idx, p_color)| {
            let c: Lab = (*p_color).into_color();
            (idx, input_color.difference(c), input_color - c)
        })
        .min_by(|(_, a, _), (_, b, _)| a.total_cmp(b))
        .expect("could not find a color");
    (DisplayColor::from_u8(nearest as u8), color_diff)
}

pub struct NNDither();

impl Ditherer for NNDither {
    fn dither(&mut self, img: &RgbImage, output: &mut EInkBuffer) {
        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, pixel) in srgb.iter().enumerate() {
            let (n, _) = nearest_neighbor(pixel.into_format().into_color());
            output.0[idx] = n;
        }
    }
}

pub struct FloydSteinbergDither();

/// 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, plus a scalar factor.
fn get_error_adjusted(orig: &Lab, err: &Lab, scalar: f32) -> Lab {
    let (p_l, p_a, p_b) = orig.into_components();
    let (err_l, err_a, err_b) = err.into_components();
    Lab::from_components((
        p_l + err_l * scalar,
        p_a + err_a * scalar,
        p_b + err_b * scalar,
    ))
}

struct DiffusionPoint {
    xshift: i32,
    yshift: i32,
    scale: f32,
}

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

const FLOYD_STEINBERG: [DiffusionPoint; 4] = [
    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),
];
const ATKINSON_DITHER: [DiffusionPoint; 6] = [
    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),
];

impl Ditherer for FloydSteinbergDither {
    fn dither(&mut self, img: &RgbImage, output: &mut EInkBuffer) {
        // create a copy of the image in Lab space, mutable.
        let srgb = <&[Srgb<u8>]>::from_components(&**img);
        let (xsize, ysize) = img.dimensions();
        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());
        }
        // 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.
                output.0[index] = nearest;
                // take the error, and propagate it.
                for point in ATKINSON_DITHER {
                    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] = get_error_adjusted(pix, &err, point.scale);
                    }
                }
            }
        }
    }
}