ai_image/codecs/webp/
decoder.rs

1use alloc::{boxed::Box, vec::Vec};
2use no_std_io::io::{BufRead, Read, Seek};
3
4use image_webp::LoopCount;
5
6use crate::buffer::ConvertBuffer;
7use crate::error::{DecodingError, ImageError, ImageResult};
8use crate::metadata::Orientation;
9use crate::{
10    AnimationDecoder, ColorType, Delay, Frame, Frames, ImageDecoder, ImageFormat, RgbImage, Rgba,
11    RgbaImage,
12};
13
14/// WebP Image format decoder.
15///
16/// Supports both lossless and lossy WebP images.
17pub struct WebPDecoder<R> {
18    inner: image_webp::WebPDecoder<R>,
19    orientation: Option<Orientation>,
20}
21
22impl<R: BufRead + Seek> WebPDecoder<R> {
23    /// Create a new `WebPDecoder` from the Reader `r`.
24    pub fn new(r: R) -> ImageResult<Self> {
25        Ok(Self {
26            inner: image_webp::WebPDecoder::new(r).map_err(ImageError::from_webp_decode)?,
27            orientation: None,
28        })
29    }
30
31    /// Returns true if the image as described by the bitstream is animated.
32    pub fn has_animation(&self) -> bool {
33        self.inner.is_animated()
34    }
35
36    /// Sets the background color if the image is an extended and animated webp.
37    pub fn set_background_color(&mut self, color: Rgba<u8>) -> ImageResult<()> {
38        self.inner
39            .set_background_color(color.0)
40            .map_err(ImageError::from_webp_decode)
41    }
42}
43
44impl<R: BufRead + Seek> ImageDecoder for WebPDecoder<R> {
45    fn dimensions(&self) -> (u32, u32) {
46        self.inner.dimensions()
47    }
48
49    fn color_type(&self) -> ColorType {
50        if self.inner.has_alpha() {
51            ColorType::Rgba8
52        } else {
53            ColorType::Rgb8
54        }
55    }
56
57    fn read_image(mut self, buf: &mut [u8]) -> ImageResult<()> {
58        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
59
60        self.inner
61            .read_image(buf)
62            .map_err(ImageError::from_webp_decode)
63    }
64
65    fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
66        (*self).read_image(buf)
67    }
68
69    fn icc_profile(&mut self) -> ImageResult<Option<Vec<u8>>> {
70        self.inner
71            .icc_profile()
72            .map_err(ImageError::from_webp_decode)
73    }
74
75    fn exif_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
76        let exif = self
77            .inner
78            .exif_metadata()
79            .map_err(ImageError::from_webp_decode)?;
80
81        self.orientation = Some(
82            exif.as_ref()
83                .and_then(|exif| Orientation::from_exif_chunk(exif))
84                .unwrap_or(Orientation::NoTransforms),
85        );
86
87        Ok(exif)
88    }
89
90    fn xmp_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
91        self.inner
92            .xmp_metadata()
93            .map_err(ImageError::from_webp_decode)
94    }
95
96    fn orientation(&mut self) -> ImageResult<Orientation> {
97        // `exif_metadata` caches the orientation, so call it if `orientation` hasn't been set yet.
98        if self.orientation.is_none() {
99            let _ = self.exif_metadata()?;
100        }
101        Ok(self.orientation.unwrap())
102    }
103}
104
105impl<'a, R: 'a + BufRead + Seek> AnimationDecoder<'a> for WebPDecoder<R> {
106    fn loop_count(&self) -> crate::metadata::LoopCount {
107        match self.inner.loop_count() {
108            LoopCount::Forever => crate::metadata::LoopCount::Infinite,
109            LoopCount::Times(n) => crate::metadata::LoopCount::Finite(n.into()),
110        }
111    }
112
113    fn into_frames(self) -> Frames<'a> {
114        struct FramesInner<R: Read + Seek> {
115            decoder: WebPDecoder<R>,
116            current: u32,
117        }
118        impl<R: BufRead + Seek> Iterator for FramesInner<R> {
119            type Item = ImageResult<Frame>;
120
121            fn next(&mut self) -> Option<Self::Item> {
122                if self.current == self.decoder.inner.num_frames() {
123                    return None;
124                }
125                self.current += 1;
126                let (width, height) = self.decoder.inner.dimensions();
127
128                let (img, delay) = if self.decoder.inner.has_alpha() {
129                    let mut img = RgbaImage::new(width, height);
130                    match self.decoder.inner.read_frame(&mut img) {
131                        Ok(delay) => (img, delay),
132                        Err(image_webp::DecodingError::NoMoreFrames) => return None,
133                        Err(e) => return Some(Err(ImageError::from_webp_decode(e))),
134                    }
135                } else {
136                    let mut img = RgbImage::new(width, height);
137                    match self.decoder.inner.read_frame(&mut img) {
138                        Ok(delay) => (img.convert(), delay),
139                        Err(image_webp::DecodingError::NoMoreFrames) => return None,
140                        Err(e) => return Some(Err(ImageError::from_webp_decode(e))),
141                    }
142                };
143
144                Some(Ok(Frame::from_parts(
145                    img,
146                    0,
147                    0,
148                    Delay::from_numer_denom_ms(delay, 1),
149                )))
150            }
151        }
152
153        Frames::new(Box::new(FramesInner {
154            decoder: self,
155            current: 0,
156        }))
157    }
158}
159
160impl ImageError {
161    fn from_webp_decode(e: image_webp::DecodingError) -> Self {
162        match e {
163            image_webp::DecodingError::IoError(e) => ImageError::IoError(e),
164            _ => ImageError::Decoding(DecodingError::new(ImageFormat::WebP.into(), e)),
165        }
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn add_with_overflow_size() {
175        let bytes = vec![
176            0x52, 0x49, 0x46, 0x46, 0xaf, 0x37, 0x80, 0x47, 0x57, 0x45, 0x42, 0x50, 0x6c, 0x64,
177            0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xfb, 0x7e, 0x73, 0x00, 0x06, 0x00, 0x00, 0x00,
178            0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65,
179            0x40, 0xfb, 0xff, 0xff, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65,
180            0x00, 0x00, 0x00, 0x00, 0x62, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x49,
181            0x49, 0x54, 0x55, 0x50, 0x4c, 0x54, 0x59, 0x50, 0x45, 0x33, 0x37, 0x44, 0x4d, 0x46,
182        ];
183
184        let data = no_std_io::io::Cursor::new(bytes);
185
186        let _ = WebPDecoder::new(data);
187    }
188}
ai_image/codecs/webp/decoder.rs

ai_image/codecs/webp/
decoder.rs
