Skip to main content

ai_image/codecs/jpeg/
decoder.rs

1use alloc::{boxed::Box, format, string::String, vec::Vec};
2use core::fmt;
3use core::marker::PhantomData;
4use no_std_io::io::{BufRead, Seek};
5
6use zune_core::bytestream::ZCursor;
7
8use crate::color::ColorType;
9use crate::error::{
10    DecodingError, ImageError, ImageResult, LimitError, UnsupportedError, UnsupportedErrorKind,
11};
12use crate::metadata::Orientation;
13use crate::{ImageDecoder, ImageFormat, Limits};
14
15type ZuneColorSpace = zune_core::colorspace::ColorSpace;
16
17/// JPEG decoder
18pub struct JpegDecoder<R> {
19    input: Vec<u8>,
20    orig_color_space: ZuneColorSpace,
21    width: u16,
22    height: u16,
23    limits: Limits,
24    orientation: Option<Orientation>,
25    // For API compatibility with the previous jpeg_decoder wrapper.
26    // Can be removed later, which would be an API break.
27    phantom: PhantomData<R>,
28}
29
30impl<R: BufRead + Seek> JpegDecoder<R> {
31    /// Create a new decoder that decodes from the stream ```r```
32    pub fn new(r: R) -> ImageResult<JpegDecoder<R>> {
33        let mut input = Vec::new();
34        let mut r = r;
35        r.read_to_end(&mut input)?;
36        let options = zune_core::options::DecoderOptions::default()
37            .set_strict_mode(false)
38            .set_max_width(usize::MAX)
39            .set_max_height(usize::MAX);
40        let mut decoder =
41            zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(input.as_slice()), options);
42        decoder.decode_headers().map_err(ImageError::from_jpeg)?;
43        // now that we've decoded the headers we can `.unwrap()`
44        // all these functions that only fail if called before decoding the headers
45        let (width, height) = decoder.dimensions().unwrap();
46        // JPEG can only express dimensions up to 65535x65535, so this conversion cannot fail
47        let width: u16 = width.try_into().unwrap();
48        let height: u16 = height.try_into().unwrap();
49        let orig_color_space = decoder.input_colorspace().expect("headers were decoded");
50
51        // Now configure the decoder color output.
52        decoder.set_options({
53            let requested_color = match orig_color_space {
54                ZuneColorSpace::RGB
55                | ZuneColorSpace::RGBA
56                | ZuneColorSpace::Luma
57                | ZuneColorSpace::LumaA => orig_color_space,
58                // Late failure
59                _ => ZuneColorSpace::RGB,
60            };
61
62            decoder.options().jpeg_set_out_colorspace(requested_color)
63        });
64
65        // Limits are disabled by default in the constructor for all decoders
66        let limits = Limits::no_limits();
67        Ok(JpegDecoder {
68            input,
69            orig_color_space,
70            width,
71            height,
72            limits,
73            orientation: None,
74            phantom: PhantomData,
75        })
76    }
77}
78
79impl<R: BufRead + Seek> ImageDecoder for JpegDecoder<R> {
80    fn dimensions(&self) -> (u32, u32) {
81        (u32::from(self.width), u32::from(self.height))
82    }
83
84    fn color_type(&self) -> ColorType {
85        ColorType::from_jpeg(self.orig_color_space)
86    }
87
88    fn icc_profile(&mut self) -> ImageResult<Option<Vec<u8>>> {
89        let options = zune_core::options::DecoderOptions::default()
90            .set_strict_mode(false)
91            .set_max_width(usize::MAX)
92            .set_max_height(usize::MAX);
93        let mut decoder =
94            zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(&self.input), options);
95        decoder.decode_headers().map_err(ImageError::from_jpeg)?;
96        Ok(decoder.icc_profile())
97    }
98
99    fn exif_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
100        let options = zune_core::options::DecoderOptions::default()
101            .set_strict_mode(false)
102            .set_max_width(usize::MAX)
103            .set_max_height(usize::MAX);
104        let mut decoder =
105            zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(&self.input), options);
106        decoder.decode_headers().map_err(ImageError::from_jpeg)?;
107        let exif = decoder.exif().cloned();
108
109        self.orientation = Some(
110            exif.as_ref()
111                .and_then(|exif| Orientation::from_exif_chunk(exif))
112                .unwrap_or(Orientation::NoTransforms),
113        );
114
115        Ok(exif)
116    }
117
118    fn xmp_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
119        let options = zune_core::options::DecoderOptions::default()
120            .set_strict_mode(false)
121            .set_max_width(usize::MAX)
122            .set_max_height(usize::MAX);
123        let mut decoder =
124            zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(&self.input), options);
125        decoder.decode_headers().map_err(ImageError::from_jpeg)?;
126
127        Ok(decoder.xmp().cloned())
128    }
129
130    fn iptc_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
131        let options = zune_core::options::DecoderOptions::default()
132            .set_strict_mode(false)
133            .set_max_width(usize::MAX)
134            .set_max_height(usize::MAX);
135        let mut decoder =
136            zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(&self.input), options);
137        decoder.decode_headers().map_err(ImageError::from_jpeg)?;
138
139        Ok(decoder.iptc().cloned())
140    }
141
142    fn orientation(&mut self) -> ImageResult<Orientation> {
143        // `exif_metadata` caches the orientation, so call it if `orientation` hasn't been set yet.
144        if self.orientation.is_none() {
145            let _ = self.exif_metadata()?;
146        }
147        Ok(self.orientation.unwrap())
148    }
149
150    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
151        let advertised_len = self.total_bytes();
152        let actual_len = buf.len() as u64;
153
154        if actual_len != advertised_len {
155            return Err(ImageError::Decoding(DecodingError::new(
156                ImageFormat::Jpeg.into(),
157                format!(
158                    "Length of the decoded data {actual_len} \
159                    doesn't match the advertised dimensions of the image \
160                    that imply length {advertised_len}"
161                ),
162            )));
163        }
164
165        let mut decoder = new_zune_decoder(&self.input, self.orig_color_space, self.limits);
166        decoder.decode_into(buf).map_err(ImageError::from_jpeg)?;
167        Ok(())
168    }
169
170    fn set_limits(&mut self, limits: Limits) -> ImageResult<()> {
171        limits.check_support(&crate::LimitSupport::default())?;
172        let (width, height) = self.dimensions();
173        limits.check_dimensions(width, height)?;
174        self.limits = limits;
175        Ok(())
176    }
177
178    fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
179        (*self).read_image(buf)
180    }
181}
182
183impl ColorType {
184    fn from_jpeg(colorspace: ZuneColorSpace) -> ColorType {
185        let colorspace = to_supported_color_space(colorspace);
186        use zune_core::colorspace::ColorSpace::*;
187        match colorspace {
188            // As of zune-jpeg 0.3.13 the output is always 8-bit,
189            // but support for 16-bit JPEG might be added in the future.
190            RGB => ColorType::Rgb8,
191            RGBA => ColorType::Rgba8,
192            Luma => ColorType::L8,
193            LumaA => ColorType::La8,
194            // to_supported_color_space() doesn't return any of the other variants
195            _ => unreachable!(),
196        }
197    }
198}
199
200fn to_supported_color_space(orig: ZuneColorSpace) -> ZuneColorSpace {
201    use zune_core::colorspace::ColorSpace::*;
202    match orig {
203        RGB | RGBA | Luma | LumaA => orig,
204        // the rest is not supported by `image` so it will be converted to RGB during decoding
205        _ => RGB,
206    }
207}
208
209fn new_zune_decoder(
210    input: &[u8],
211    orig_color_space: ZuneColorSpace,
212    limits: Limits,
213) -> zune_jpeg::JpegDecoder<ZCursor<&[u8]>> {
214    let target_color_space = to_supported_color_space(orig_color_space);
215    let mut options = zune_core::options::DecoderOptions::default()
216        .jpeg_set_out_colorspace(target_color_space)
217        .set_strict_mode(false);
218    options = options.set_max_width(match limits.max_image_width {
219        Some(max_width) => max_width as usize, // u32 to usize never truncates
220        None => usize::MAX,
221    });
222    options = options.set_max_height(match limits.max_image_height {
223        Some(max_height) => max_height as usize, // u32 to usize never truncates
224        None => usize::MAX,
225    });
226    zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(input), options)
227}
228
229/// Wrapper to provide `core::error::Error` for zune-jpeg's `DecodeErrors` in no_std.
230struct JpegDecodeError(String);
231
232impl fmt::Display for JpegDecodeError {
233    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
234        f.write_str(&self.0)
235    }
236}
237
238impl fmt::Debug for JpegDecodeError {
239    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
240        f.write_str(&self.0)
241    }
242}
243
244impl core::error::Error for JpegDecodeError {}
245
246impl ImageError {
247    fn from_jpeg(err: zune_jpeg::errors::DecodeErrors) -> ImageError {
248        use zune_jpeg::errors::DecodeErrors::*;
249        match err {
250            Unsupported(desc) => ImageError::Unsupported(UnsupportedError::from_format_and_kind(
251                ImageFormat::Jpeg.into(),
252                UnsupportedErrorKind::GenericFeature(format!("{desc:?}")),
253            )),
254            LargeDimensions(_) => ImageError::Limits(LimitError::from_kind(
255                crate::error::LimitErrorKind::DimensionError,
256            )),
257            err => ImageError::Decoding(DecodingError::new(
258                ImageFormat::Jpeg.into(),
259                JpegDecodeError(format!("{err:?}")),
260            )),
261        }
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use std::{fs, io::Cursor};
269
270    #[test]
271    fn test_exif_orientation() {
272        let data = fs::read("tests/images/jpg/portrait_2.jpg").unwrap();
273        let mut decoder = JpegDecoder::new(Cursor::new(data)).unwrap();
274        assert_eq!(decoder.orientation().unwrap(), Orientation::FlipHorizontal);
275    }
276}