From 3b31b3f027ab2c99f7fe0b8d07d225c549bc38a8 Mon Sep 17 00:00:00 2001
From: saji <saji@saji.dev>
Date: Wed, 17 Jul 2024 13:27:31 -0500
Subject: [PATCH] add dimensions data to EInkImage

---
 src/imageproc.rs | 46 +++++++++++++++++++++++++++++++---------------
 src/main.rs      |  4 ++--
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/imageproc.rs b/src/imageproc.rs
index d5ea99c..fd23107 100644
--- a/src/imageproc.rs
+++ b/src/imageproc.rs
@@ -1,8 +1,8 @@
 use image::RgbImage;
 
+use image::Rgb as imgRgb;
 use palette::color_difference::{Ciede2000, EuclideanDistance};
 use palette::{cast::FromComponents, IntoColor, Lab, Oklch, Srgb};
-use image::Rgb as imgRgb;
 
 /// Palette used on the display; pixels can be one of these colors.
 ///
@@ -63,35 +63,51 @@ impl From<DisplayColor> for u8 {
 }
 /// Buffer to be sent to the ``EInk`` display.
 #[derive(Debug)]
-pub struct EInkBuffer(Vec<DisplayColor>);
+pub struct EInkImage {
+    data: Vec<DisplayColor>,
+    width: u32,
+    height: u32,
+}
 
-impl EInkBuffer {
+impl EInkImage {
     #[must_use]
     pub fn into_display_buffer(&self) -> Vec<u8> {
-        let mut buf = Vec::with_capacity(self.0.len() / 2);
+        let mut buf = Vec::with_capacity(self.data.len() / 2);
 
-        for colors in self.0.chunks_exact(2) {
+        for colors in self.data.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 new(width: u32, height: u32) -> Self {
+        let v = vec![DisplayColor::Black; (width * height) as usize];
+        Self {
+            data: v,
+            width,
+            height,
+        }
     }
+
+    /// Produces a regular RGB image from this image buffer using the given
+    /// color palette.
     pub fn make_image(&self) -> RgbImage {
-        RgbImage::from_fn(800, 480, |x, y| {
-            let srgb = Srgb::from(self.0[(y * 800 + x) as usize]);
+        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)
         })
     }
+
+    /// Returns the dimensions (width, height) of the image buffer.
+    pub const fn dimensions(&self) -> (u32, u32) {
+        (self.width, self.height)
+    }
 }
 
 pub trait Ditherer {
-    fn dither(&self, img: &RgbImage, output: &mut EInkBuffer);
+    fn dither(&self, img: &RgbImage, output: &mut EInkImage);
 }
 
 /// Find the closest approximate palette color to the given sRGB value.
@@ -113,7 +129,7 @@ pub fn nearest_neighbor(input_color: Lab) -> (DisplayColor, Lab) {
 pub struct NNDither();
 
 impl Ditherer for NNDither {
-    fn dither(&self, img: &RgbImage, output: &mut EInkBuffer) {
+    fn dither(&self, img: &RgbImage, output: &mut EInkImage) {
         assert!(img.width() == 800);
         assert!(img.height() == 480);
 
@@ -122,7 +138,7 @@ impl Ditherer for NNDither {
 
         for (idx, pixel) in srgb.iter().enumerate() {
             let (n, _) = nearest_neighbor(pixel.into_format().into_color());
-            output.0[idx] = n;
+            output.data[idx] = n;
         }
     }
 }
@@ -235,7 +251,7 @@ impl ErrorDiffusionDither {
 }
 
 impl Ditherer for ErrorDiffusionDither {
-    fn dither(&self, img: &RgbImage, output: &mut EInkBuffer) {
+    fn dither(&self, img: &RgbImage, output: &mut EInkImage) {
         // create a copy of the image in Lab space, mutable.
         let srgb = <&[Srgb<u8>]>::from_components(&**img);
         let (xsize, ysize) = img.dimensions();
@@ -251,7 +267,7 @@ impl Ditherer for ErrorDiffusionDither {
                 let curr_pix = temp_img[index];
                 let (nearest, err) = nearest_neighbor(curr_pix);
                 // set the color in the output buffer.
-                output.0[index] = nearest;
+                output.data[index] = nearest;
                 // take the error, and propagate it.
                 for point in self.0.value() {
                     let Some(target_x) = x.checked_add_signed(point.xshift) else {
diff --git a/src/main.rs b/src/main.rs
index 8b7dd8c..0800518 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,7 @@ pub mod display;
 pub mod imageproc;
 
 use crate::display::Wrapper;
-use crate::imageproc::{DiffusionMatrix, Ditherer, EInkBuffer, ErrorDiffusionDither};
+use crate::imageproc::{DiffusionMatrix, Ditherer, EInkImage, ErrorDiffusionDither};
 use clap::{Parser, Subcommand};
 use image::RgbImage;
 
@@ -33,7 +33,7 @@ fn main() -> anyhow::Result<()> {
         let img: RgbImage = image::io::Reader::open("myimage.png")?.decode()?.into();
         let mut display = Wrapper::new()?;
 
-        let mut eink_buf = EInkBuffer::new(800, 480);
+        let mut eink_buf = EInkImage::new(800, 480);
         let dither = ErrorDiffusionDither::new(DiffusionMatrix::Atkinson);
 
         dither.dither(&img, &mut eink_buf);