Skip to main content

ai_image/codecs/
qoi.rs

1//! Decoding and encoding of QOI images
2
3use crate::error::{DecodingError, EncodingError, UnsupportedError, UnsupportedErrorKind};
4use crate::{
5    ColorType, ExtendedColorType, ImageDecoder, ImageEncoder, ImageError, ImageFormat, ImageResult,
6};
7use alloc::boxed::Box;
8use no_std_io::io::{Read, Write};
9
10/// QOI decoder
11pub struct QoiDecoder<R> {
12    decoder: qoi::Decoder<R>,
13}
14
15impl<R> QoiDecoder<R>
16where
17    R: Read,
18{
19    /// Creates a new decoder that decodes from the stream ```reader```
20    pub fn new(reader: R) -> ImageResult<Self> {
21        let decoder = qoi::Decoder::from_stream(reader).map_err(decoding_error)?;
22        Ok(Self { decoder })
23    }
24}
25
26impl<R: Read> ImageDecoder for QoiDecoder<R> {
27    fn dimensions(&self) -> (u32, u32) {
28        (self.decoder.header().width, self.decoder.header().height)
29    }
30
31    fn color_type(&self) -> ColorType {
32        match self.decoder.header().channels {
33            qoi::Channels::Rgb => ColorType::Rgb8,
34            qoi::Channels::Rgba => ColorType::Rgba8,
35        }
36    }
37
38    fn read_image(mut self, buf: &mut [u8]) -> ImageResult<()> {
39        self.decoder.decode_to_buf(buf).map_err(decoding_error)?;
40        Ok(())
41    }
42
43    fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
44        (*self).read_image(buf)
45    }
46}
47
48fn decoding_error(error: qoi::Error) -> ImageError {
49    ImageError::Decoding(DecodingError::new(ImageFormat::Qoi.into(), error))
50}
51
52fn encoding_error(error: qoi::Error) -> ImageError {
53    ImageError::Encoding(EncodingError::new(ImageFormat::Qoi.into(), error))
54}
55
56/// QOI encoder
57pub struct QoiEncoder<W> {
58    writer: W,
59}
60
61impl<W: Write> QoiEncoder<W> {
62    /// Creates a new encoder that writes its output to ```writer```
63    pub fn new(writer: W) -> Self {
64        Self { writer }
65    }
66}
67
68impl<W: Write> ImageEncoder for QoiEncoder<W> {
69    #[track_caller]
70    fn write_image(
71        mut self,
72        buf: &[u8],
73        width: u32,
74        height: u32,
75        color_type: ExtendedColorType,
76    ) -> ImageResult<()> {
77        if !matches!(
78            color_type,
79            ExtendedColorType::Rgba8 | ExtendedColorType::Rgb8
80        ) {
81            return Err(ImageError::Unsupported(
82                UnsupportedError::from_format_and_kind(
83                    ImageFormat::Qoi.into(),
84                    UnsupportedErrorKind::Color(color_type),
85                ),
86            ));
87        }
88
89        let expected_buffer_len = color_type.buffer_size(width, height);
90        assert_eq!(
91            expected_buffer_len,
92            buf.len() as u64,
93            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
94            buf.len(),
95        );
96
97        // Encode data in QOI
98        let data = qoi::encode_to_vec(buf, width, height).map_err(encoding_error)?;
99
100        // Write data to buffer
101        self.writer.write_all(&data[..])?;
102        self.writer.flush()?;
103
104        Ok(())
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use std::fs::File;
112
113    #[test]
114    fn decode_test_image() {
115        let decoder = QoiDecoder::new(File::open("tests/images/qoi/basic-test.qoi").unwrap())
116            .expect("Unable to read QOI file");
117
118        assert_eq!((5, 5), decoder.dimensions());
119        assert_eq!(ColorType::Rgba8, decoder.color_type());
120    }
121}