hyperion/
image.rs

1use std::convert::TryFrom;
2
3use base64::Engine;
4use image::ImageEncoder;
5use thiserror::Error;
6
7use crate::models::Color;
8
9mod reducer;
10pub use reducer::*;
11
12pub trait Image: Sized {
13    /// Get the width of the image, in pixels
14    fn width(&self) -> u16;
15
16    /// Get the height of the image, in pixels
17    fn height(&self) -> u16;
18
19    /// Get the color at the given coordinates
20    fn color_at(&self, x: u16, y: u16) -> Option<Color>;
21
22    /// Get the color at the given coordinates skipping bound checks
23    ///
24    /// # Safety
25    ///
26    /// The caller is responsible to ensure that 0 <= x < width and 0 <= y < height.
27    unsafe fn color_at_unchecked(&self, x: u16, y: u16) -> Color;
28
29    /// Convert this image trait object to a raw image
30    fn to_raw_image(&self) -> RawImage;
31}
32
33#[derive(Debug, Error)]
34pub enum RawImageError {
35    #[error("invalid data ({data} bytes) for the given dimensions ({width} x {height} x {channels} = {expected})")]
36    InvalidData {
37        data: usize,
38        width: u32,
39        height: u32,
40        channels: u32,
41        expected: usize,
42    },
43    #[error("invalid width")]
44    InvalidWidth,
45    #[error("invalid height")]
46    InvalidHeight,
47    #[error("raw image data missing")]
48    RawImageMissing,
49    #[error("image width is zero")]
50    ZeroWidth,
51    #[error("image height is zero")]
52    ZeroHeight,
53    #[error("i/o error")]
54    Io(#[from] std::io::Error),
55    #[error("encoding error")]
56    Encoding(#[from] image::ImageError),
57}
58
59#[derive(Clone)]
60pub struct RawImage {
61    data: Vec<u8>,
62    width: u16,
63    height: u16,
64}
65
66impl RawImage {
67    pub const CHANNELS: u16 = 3;
68
69    pub fn write_to_kitty(&self, out: &mut dyn std::io::Write) -> Result<(), RawImageError> {
70        // Buffer for raw PNG data
71        let mut buf = Vec::new();
72        // PNG encoder
73        let encoder = image::codecs::png::PngEncoder::new(&mut buf);
74        // Write PNG to buffer
75        encoder.write_image(
76            &self.data[..],
77            self.width as _,
78            self.height as _,
79            image::ColorType::Rgb8.into(),
80        )?;
81        // Encode to base64
82        let encoded = base64::engine::general_purpose::STANDARD_NO_PAD.encode(&buf);
83        // Split into chunks
84        let chunks = encoded.as_bytes().chunks(4096).collect::<Vec<_>>();
85        // Transmit chunks
86        for (i, chunk) in chunks.iter().enumerate() {
87            let last = if i == chunks.len() - 1 { b"0" } else { b"1" };
88
89            match i {
90                0 => {
91                    // First chunk
92                    out.write_all(b"\x1B_Gf=100,a=T,m=")?;
93                }
94                _ => {
95                    // Other chunks
96                    out.write_all(b"\x1B_Gm=")?;
97                }
98            }
99
100            out.write_all(last)?;
101            out.write_all(b";")?;
102            out.write_all(chunk)?;
103            out.write_all(b"\x1B\\")?;
104        }
105
106        // Finish with new-line
107        out.write_all(b"\n")?;
108
109        Ok(())
110    }
111}
112
113impl Image for RawImage {
114    fn width(&self) -> u16 {
115        self.width
116    }
117
118    fn height(&self) -> u16 {
119        self.height
120    }
121
122    fn color_at(&self, x: u16, y: u16) -> Option<Color> {
123        if x < self.width && y < self.height {
124            unsafe { Some(self.color_at_unchecked(x, y)) }
125        } else {
126            None
127        }
128    }
129
130    unsafe fn color_at_unchecked(&self, x: u16, y: u16) -> Color {
131        let idx = (y as usize * self.width as usize + x as usize) * Self::CHANNELS as usize;
132        Color::new(
133            *self.data.get_unchecked(idx),
134            *self.data.get_unchecked(idx + 1),
135            *self.data.get_unchecked(idx + 2),
136        )
137    }
138
139    fn to_raw_image(&self) -> RawImage {
140        self.clone()
141    }
142}
143
144impl std::fmt::Debug for RawImage {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        let mut f = f.debug_struct("RawImage");
147        f.field("width", &self.width);
148        f.field("height", &self.height);
149        f.field("channels", &Self::CHANNELS);
150
151        if self.data.len() > 32 {
152            f.field("data", &format!("[{} bytes]", self.data.len()));
153        } else {
154            f.field("data", &self.data);
155        }
156
157        f.finish()
158    }
159}
160
161impl TryFrom<(Vec<u8>, u32, u32)> for RawImage {
162    type Error = RawImageError;
163
164    fn try_from((data, width, height): (Vec<u8>, u32, u32)) -> Result<Self, Self::Error> {
165        let expected = width as usize * height as usize * Self::CHANNELS as usize;
166
167        if data.len() != expected {
168            return Err(RawImageError::InvalidData {
169                data: data.len(),
170                width,
171                height,
172                channels: Self::CHANNELS as _,
173                expected,
174            });
175        } else if width == 0 {
176            return Err(RawImageError::ZeroWidth);
177        } else if height == 0 {
178            return Err(RawImageError::ZeroHeight);
179        } else if width >= u16::MAX as u32 {
180            return Err(RawImageError::InvalidWidth);
181        } else if height >= u16::MAX as u32 {
182            return Err(RawImageError::InvalidHeight);
183        }
184
185        Ok(Self {
186            data,
187            width: width as _,
188            height: height as _,
189        })
190    }
191}
192
193pub struct ImageView<'i, T: Image> {
194    inner: &'i T,
195    xmin: u16,
196    xmax: u16,
197    ymin: u16,
198    ymax: u16,
199}
200
201impl<'i, T: Image> Image for ImageView<'i, T> {
202    fn width(&self) -> u16 {
203        self.xmax - self.xmin
204    }
205
206    fn height(&self) -> u16 {
207        self.ymax - self.ymin
208    }
209
210    fn color_at(&self, x: u16, y: u16) -> Option<Color> {
211        self.inner.color_at(x + self.xmin, y + self.ymin)
212    }
213
214    unsafe fn color_at_unchecked(&self, x: u16, y: u16) -> Color {
215        self.inner.color_at_unchecked(x + self.xmin, y + self.ymin)
216    }
217
218    fn to_raw_image(&self) -> RawImage {
219        let w = self.width();
220        let h = self.height();
221        let mut data = Vec::with_capacity(w as usize * h as usize * RawImage::CHANNELS as usize);
222
223        unsafe {
224            for y in 0..h {
225                for x in 0..w {
226                    let (r, g, b) = self.color_at_unchecked(x, y).into_components();
227                    data.push(r);
228                    data.push(g);
229                    data.push(b);
230                }
231            }
232        }
233
234        RawImage {
235            data,
236            width: w,
237            height: h,
238        }
239    }
240}
241
242pub trait ImageViewExt: Image {
243    fn wrap(&self, x: std::ops::Range<u16>, y: std::ops::Range<u16>) -> ImageView<'_, Self>;
244}
245
246impl<T: Image> ImageViewExt for T {
247    fn wrap(&self, x: std::ops::Range<u16>, y: std::ops::Range<u16>) -> ImageView<'_, Self> {
248        ImageView {
249            inner: self,
250            xmin: x.start,
251            xmax: x.end,
252            ymin: y.start,
253            ymax: y.end,
254        }
255    }
256}
257
258pub mod prelude {
259    pub use super::{Image, ImageViewExt};
260}