Skip to main content

zune_jpeg/
decoder.rs

1/*
2 * Copyright (c) 2023.
3 *
4 * This software is free software;
5 *
6 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
7 */
8
9//! Main image logic.
10#![allow(clippy::doc_markdown)]
11
12use alloc::string::ToString;
13use alloc::vec::Vec;
14use alloc::{format, vec};
15
16use zune_core::bytestream::{ZByteReaderTrait, ZReader};
17use zune_core::colorspace::ColorSpace;
18use zune_core::log::{error, trace, warn};
19use zune_core::options::DecoderOptions;
20
21use crate::color_convert::choose_ycbcr_to_rgb_convert_func;
22use crate::components::{Components, SampleRatios};
23use crate::errors::{DecodeErrors, UnsupportedSchemes};
24use crate::headers::{
25    parse_app1, parse_app13, parse_app14, parse_app2, parse_dqt, parse_huffman, parse_sos,
26    parse_start_of_frame
27};
28use crate::huffman::HuffmanTable;
29use crate::idct::{choose_idct_func, choose_idct_1x1_func, choose_idct_4x4_func};
30use crate::marker::Marker;
31use crate::misc::SOFMarkers;
32use crate::upsampler::{
33    choose_horizontal_samp_function, choose_hv_samp_function, choose_v_samp_function,
34    generic_sampler, upsample_no_op
35};
36
37/// Maximum components
38pub(crate) const MAX_COMPONENTS: usize = 4;
39
40/// Maximum image dimensions supported.
41pub(crate) const MAX_DIMENSIONS: usize = 1 << 27;
42
43/// Color conversion function that can convert YCbCr colorspace to RGB(A/X) for
44/// 16 values
45///
46/// The following are guarantees to the following functions
47///
48/// 1. The `&[i16]` slices passed contain 16 items
49///
50/// 2. The slices passed are in the following order
51///     `y,cb,cr`
52///
53/// 3. `&mut [u8]` is zero initialized
54///
55/// 4. `&mut usize` points to the position in the array where new values should
56/// be used
57///
58/// The pointer should
59/// 1. Carry out color conversion
60/// 2. Update `&mut usize` with the new position
61
62pub type ColorConvert16Ptr = fn(&[i16; 16], &[i16; 16], &[i16; 16], &mut [u8], &mut usize);
63
64/// IDCT  function prototype
65///
66/// This encapsulates a dequantize and IDCT function which will carry out the
67/// following functions
68///
69/// Multiply each 64 element block of `&mut [i16]` with `&Aligned32<[i32;64]>`
70/// Carry out IDCT (type 3 dct) on ach block of 64 i16's
71pub type IDCTPtr = fn(&mut [i32; 64], &mut [i16], usize);
72
73/// An encapsulation of an ICC chunk
74pub(crate) struct ICCChunk {
75    pub(crate) seq_no:      u8,
76    pub(crate) num_markers: u8,
77    pub(crate) data:        Vec<u8>
78}
79
80/// A JPEG Decoder Instance.
81#[allow(clippy::upper_case_acronyms, clippy::struct_excessive_bools)]
82pub struct JpegDecoder<T> {
83    /// Struct to hold image information from SOI
84    pub(crate) info:              ImageInfo,
85    ///  Quantization tables, will be set to none and the tables will
86    /// be moved to `components` field
87    pub(crate) qt_tables:         [Option<[i32; 64]>; MAX_COMPONENTS],
88    /// DC Huffman Tables with a maximum of 4 tables for each  component
89    pub(crate) dc_huffman_tables: [Option<HuffmanTable>; MAX_COMPONENTS],
90    /// AC Huffman Tables with a maximum of 4 tables for each component
91    pub(crate) ac_huffman_tables: [Option<HuffmanTable>; MAX_COMPONENTS],
92    /// Image components, holds information like DC prediction and quantization
93    /// tables of a component
94    pub(crate) components:        Vec<Components>,
95    /// maximum horizontal component of all channels in the image
96    pub(crate) h_max:             usize,
97    // maximum vertical component of all channels in the image
98    pub(crate) v_max:             usize,
99    /// mcu's  width (interleaved scans)
100    pub(crate) mcu_width:         usize,
101    /// MCU height(interleaved scans
102    pub(crate) mcu_height:        usize,
103    /// Number of MCU's in the x plane
104    pub(crate) mcu_x:             usize,
105    /// Number of MCU's in the y plane
106    pub(crate) mcu_y:             usize,
107    /// Is the image interleaved?
108    pub(crate) is_interleaved:    bool,
109    /// Image input colorspace, should be YCbCr for a sane image, might be
110    /// grayscale too
111    pub(crate) input_colorspace:  ColorSpace,
112    // Progressive image details
113    /// Is the image progressive?
114    pub(crate) is_progressive:    bool,
115
116    /// Start of spectral scan
117    pub(crate) spec_start:       u8,
118    /// End of spectral scan
119    pub(crate) spec_end:         u8,
120    /// Successive approximation bit position high
121    pub(crate) succ_high:        u8,
122    /// Successive approximation bit position low
123    pub(crate) succ_low:         u8,
124    /// Number of components.
125    pub(crate) num_scans:        u8,
126    /// For a scan, check if any component has vertical/horizontal sampling.
127    pub(crate) scan_subsampled:  bool,
128    // Function pointers, for pointy stuff.
129    /// Dequantize and idct function
130    // This is determined at runtime which function to run, statically it's
131    // initialized to a platform independent one and during initialization
132    // of this struct, we check if we can switch to a faster one which
133    // depend on certain CPU extensions.
134    pub(crate) idct_func: IDCTPtr,
135    /// Specialized IDCT when we can guarantee only few coefficients are non-zero.
136    ///
137    /// **The callee must uphold a contract**. See [`choose_idct_4x4_func`].
138    pub(crate) idct_4x4_func: IDCTPtr,
139    pub(crate) idct_1x1_func: IDCTPtr,
140    // Color convert function which acts on 16 YCbCr values
141    pub(crate) color_convert_16: ColorConvert16Ptr,
142    pub(crate) z_order:          [usize; MAX_COMPONENTS],
143    /// restart markers
144    pub(crate) restart_interval: usize,
145    pub(crate) todo:             usize,
146    // decoder options
147    pub(crate) options:          DecoderOptions,
148    // byte-stream
149    pub(crate) stream:           ZReader<T>,
150    // Indicate whether headers have been decoded
151    pub(crate) headers_decoded:  bool,
152    pub(crate) seen_sof:         bool,
153
154    // exif data, lifted from app2
155    pub(crate) icc_data: Vec<ICCChunk>,
156    pub(crate) is_mjpeg: bool,
157    pub(crate) coeff:    usize, // Solves some weird bug :)
158    /// Extended XMP segments
159    pub(crate) extended_xmp_segments: Vec<ExtendedXmpSegment>,
160}
161
162impl<T> JpegDecoder<T>
163where
164    T: ZByteReaderTrait
165{
166    #[allow(clippy::redundant_field_names)]
167    fn default(options: DecoderOptions, buffer: T) -> Self {
168        let color_convert = choose_ycbcr_to_rgb_convert_func(ColorSpace::RGB, &options).unwrap();
169        JpegDecoder {
170            info:              ImageInfo::default(),
171            qt_tables:         [None, None, None, None],
172            dc_huffman_tables: [None, None, None, None],
173            ac_huffman_tables: [None, None, None, None],
174            components:        vec![],
175            // Interleaved information
176            h_max:             1,
177            v_max:             1,
178            mcu_height:        0,
179            mcu_width:         0,
180            mcu_x:             0,
181            mcu_y:             0,
182            is_interleaved:    false,
183            is_progressive:    false,
184            spec_start:        0,
185            spec_end:          0,
186            succ_high:         0,
187            succ_low:          0,
188            num_scans:         0,
189            scan_subsampled:   false, 
190            idct_func:         choose_idct_func(&options),
191            idct_4x4_func:     choose_idct_4x4_func(&options),
192            idct_1x1_func:     choose_idct_1x1_func(&options),
193            color_convert_16:  color_convert,
194            input_colorspace:  ColorSpace::YCbCr,
195            z_order:           [0; MAX_COMPONENTS],
196            restart_interval:  0,
197            todo:              0x7fff_ffff,
198            options:           options,
199            stream:            ZReader::new(buffer),
200            headers_decoded:   false,
201            seen_sof:          false,
202            icc_data:          vec![],
203            is_mjpeg:          false,
204            coeff:             1,
205            extended_xmp_segments: vec![],
206        }
207    }
208    /// Decode a buffer already in memory
209    ///
210    /// The buffer should be a valid jpeg file, perhaps created by the command
211    /// `std:::fs::read()` or a JPEG file downloaded from the internet.
212    ///
213    /// # Errors
214    /// See DecodeErrors for an explanation
215    pub fn decode(&mut self) -> Result<Vec<u8>, DecodeErrors> {
216        self.decode_headers()?;
217        let size = self.output_buffer_size().unwrap();
218        let mut out = vec![0; size];
219        self.decode_into(&mut out)?;
220        Ok(out)
221    }
222
223    /// Create a new Decoder instance
224    ///
225    /// # Arguments
226    ///  - `stream`: The raw bytes of a jpeg file.
227    #[must_use]
228    #[allow(clippy::new_without_default)]
229    pub fn new(stream: T) -> JpegDecoder<T> {
230        JpegDecoder::default(DecoderOptions::default(), stream)
231    }
232
233    /// Returns the image information
234    ///
235    /// This **must** be called after a subsequent call to [`decode`] or [`decode_headers`]
236    /// it will return `None`
237    ///
238    /// # Returns
239    /// - `Some(info)`: Image information,width, height, number of components
240    /// - None: Indicates image headers haven't been decoded
241    ///
242    /// [`decode`]: JpegDecoder::decode
243    /// [`decode_headers`]: JpegDecoder::decode_headers
244    #[must_use]
245    pub fn info(&self) -> Option<ImageInfo> {
246        // we check for fails to that call by comparing what we have to the default, if
247        // it's default we assume that the caller failed to uphold the
248        // guarantees. We can be sure that an image cannot be the default since
249        // its a hard panic in-case width or height are set to zero.
250        if !self.headers_decoded {
251            return None;
252        }
253
254        return Some(self.info.clone());
255    }
256
257    /// Return the number of bytes required to hold a decoded image frame
258    /// decoded using the given input transformations
259    ///
260    /// # Returns
261    ///  - `Some(usize)`: Minimum size for a buffer needed to decode the image
262    ///  - `None`: Indicates the image was not decoded, or image dimensions would overflow a usize
263    ///
264    #[must_use]
265    pub fn output_buffer_size(&self) -> Option<usize> {
266        return if self.headers_decoded {
267            Some(
268                usize::from(self.width())
269                    .checked_mul(usize::from(self.height()))?
270                    .checked_mul(self.options.jpeg_get_out_colorspace().num_components())?
271            )
272        } else {
273            None
274        };
275    }
276
277    /// Get an immutable reference to the decoder options
278    /// for the decoder instance
279    ///
280    /// This can be used to modify options before actual decoding
281    /// but after initial creation
282    ///
283    /// # Example
284    /// ```no_run
285    /// use zune_core::bytestream::ZCursor;
286    /// use zune_jpeg::JpegDecoder;
287    ///
288    /// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
289    /// // get current options
290    /// let mut options = decoder.options();
291    /// // modify it
292    ///  let new_options = options.set_max_width(10);
293    /// // set it back
294    /// decoder.set_options(new_options);
295    ///
296    /// ```
297    #[must_use]
298    pub const fn options(&self) -> &DecoderOptions {
299        &self.options
300    }
301    /// Return the input colorspace of the image
302    ///
303    /// This indicates the colorspace that is present in
304    /// the image, but this may be different to the colorspace that
305    /// the output will be transformed to
306    ///
307    /// # Returns
308    /// -`Some(Colorspace)`: Input colorspace
309    /// - None : Indicates the headers weren't decoded
310    #[must_use]
311    pub fn input_colorspace(&self) -> Option<ColorSpace> {
312        return if self.headers_decoded { Some(self.input_colorspace) } else { None };
313    }
314    /// Set decoder options
315    ///
316    /// This can be used to set new options even after initialization
317    /// but before decoding.
318    ///
319    /// This does not bear any significance after decoding an image
320    ///
321    /// # Arguments
322    /// - `options`: New decoder options
323    ///
324    /// # Example
325    /// Set maximum jpeg progressive passes to be 4
326    ///
327    /// ```no_run
328    /// use zune_core::bytestream::ZCursor;
329    /// use zune_jpeg::JpegDecoder;
330    /// let mut decoder =JpegDecoder::new(ZCursor::new(&[]));
331    /// // this works also because DecoderOptions implements `Copy`
332    /// let options = decoder.options().jpeg_set_max_scans(4);
333    /// // set the new options
334    /// decoder.set_options(options);
335    /// // now decode
336    /// decoder.decode().unwrap();
337    /// ```
338    pub fn set_options(&mut self, options: DecoderOptions) {
339        self.options = options;
340    }
341    fn reassemble_extended_xmp(&mut self) {
342        if self.extended_xmp_segments.is_empty() {
343            return;
344        }
345
346        // Sort by offset
347        self.extended_xmp_segments.sort_by(|a, b| a.offset.cmp(&b.offset));
348
349        let guid = &self.extended_xmp_segments[0].guid;
350        let total_size = self.extended_xmp_segments[0].total_size;
351
352        // Check for consistency
353        for segment in &self.extended_xmp_segments {
354            if &segment.guid != guid || segment.total_size != total_size {
355                error!("Inconsistent Extended XMP segments");
356                self.extended_xmp_segments.clear();
357                return;
358            }
359        }
360
361        let mut rolling_offset = 0;
362        let mut complete = true;
363
364        for segment in &self.extended_xmp_segments {
365            if segment.offset != rolling_offset {
366                // Gap or overlap
367                complete = false;
368                break;
369            }
370            rolling_offset += segment.data.len() as u32;
371        }
372
373        if complete && rolling_offset == total_size {
374            let mut result = Vec::with_capacity(total_size as usize);
375            for segment in &self.extended_xmp_segments {
376                result.extend_from_slice(&segment.data);
377            }
378            self.info.extended_xmp = Some(result);
379            self.info.extended_xmp_guid = Some(guid.clone());
380            self.extended_xmp_segments.clear();
381        } else if rolling_offset > total_size {
382            error!("Extended XMP overflow");
383            self.extended_xmp_segments.clear();
384        }
385        // Else: Incomplete, wait for more.
386    }
387    /// Decode Decoder headers
388    ///
389    /// This routine takes care of parsing supported headers from a Decoder
390    /// image
391    ///
392    /// # Supported Headers
393    ///  - APP(0)
394    ///  - SOF(O)
395    ///  - DQT -> Quantization tables
396    ///  - DHT -> Huffman tables
397    ///  - SOS -> Start of Scan
398    /// # Unsupported Headers
399    ///  - SOF(n) -> Decoder images which are not baseline/progressive
400    ///  - DAC -> Images using Arithmetic tables
401    ///  - JPG(n)
402    fn decode_headers_internal(&mut self) -> Result<(), DecodeErrors> {
403        if self.headers_decoded {
404            trace!("Headers decoded!");
405            return Ok(());
406        }
407
408        // match output colorspace here
409        // we know this will only be called once per image
410        // so makes sense
411        // We only care for ycbcr to rgb/rgba here
412        // in case one is using another colorspace.
413        // May god help you
414        let out_colorspace = self.options.jpeg_get_out_colorspace();
415
416        if matches!(
417            out_colorspace,
418            ColorSpace::BGR | ColorSpace::BGRA | ColorSpace::RGB | ColorSpace::RGBA
419        ) {
420            self.color_convert_16 = choose_ycbcr_to_rgb_convert_func(
421                self.options.jpeg_get_out_colorspace(),
422                &self.options
423            )
424            .unwrap();
425        }
426        // First two bytes should be jpeg soi marker
427        let magic_bytes = self.stream.get_u16_be_err()?;
428
429        let mut last_byte = 0;
430        let mut bytes_before_marker = 0;
431
432        if magic_bytes != 0xffd8 {
433            return Err(DecodeErrors::IllegalMagicBytes(magic_bytes));
434        }
435
436        loop {
437            // read a byte
438            let mut m = self.stream.read_u8_err()?;
439
440            // AND OF COURSE some images will have fill bytes in their marker
441            // bitstreams because why not.
442            //
443            // I am disappointed as a man.
444            if (m == 0xFF || m == 0) && last_byte == 0xFF {
445                // This handles the edge case where
446                // images have markers with fill bytes(0xFF)
447                // or byte stuffing (0)
448                // I.e 0xFF 0xFF 0xDA
449                // and
450                // 0xFF 0 0xDA
451                // It should ignore those fill bytes and take 0xDA
452                // I don't know why such images exist
453                // but they do.
454                // so this is for you (with love)
455                while m == 0xFF || m == 0x0 {
456                    last_byte = m;
457                    m = self.stream.read_u8_err()?;
458                }
459            }
460            // Last byte should be 0xFF to confirm existence of a marker since markers look
461            // like OxFF(some marker data)
462            if last_byte == 0xFF {
463                let marker = Marker::from_u8(m);
464                if let Some(n) = marker {
465                    if bytes_before_marker > 3 {
466                        if self.options.strict_mode()
467                        /*No reason to use this*/
468                        {
469                            return Err(DecodeErrors::FormatStatic(
470                                "[strict-mode]: Extra bytes between headers"
471                            ));
472                        }
473
474                        error!(
475                            "Extra bytes {} before marker 0xFF{:X}",
476                            bytes_before_marker - 3,
477                            m
478                        );
479                    }
480
481                    bytes_before_marker = 0;
482
483                    self.parse_marker_inner(n)?;
484
485                    if !self.extended_xmp_segments.is_empty() {
486                        self.reassemble_extended_xmp();
487                    }
488
489                    // break after reading the start of scan.
490                    // what follows is the image data
491                    if n == Marker::SOS {
492                        self.headers_decoded = true;
493                        trace!("Input colorspace {:?}", self.input_colorspace);
494
495                        // Check if image is RGB
496                        // The check is weird, we need to check if ID
497                        // represents R, G and B in ascii,
498                        //
499                        // I am not sure if this is even specified in any standard,
500                        // but jpegli https://github.com/google/jpegli does encode
501                        // its images that way, so this will check for that. and handle it appropriately
502                        // It is spefified here so that on a successful header decode,we can at least
503                        // try to attribute image colorspace  correctly.
504                        //
505                        // It was first the issue in https://github.com/etemesi254/zune-image/issues/291
506                        // that brought it to light
507                        //
508                        let mut is_rgb = self.components.len() == 3;
509                        let chars = ['R', 'G', 'B'];
510                        for (comp, single_char) in self.components.iter().zip(chars.iter()) {
511                            is_rgb &= comp.id == (*single_char) as u8
512                        }
513                        // Image is RGB, change colorspace
514                        if is_rgb {
515                            self.input_colorspace = ColorSpace::RGB;
516                        }
517
518                        return Ok(());
519                    }
520                } else {
521                    bytes_before_marker = 0;
522
523                    warn!("Marker 0xFF{:X} not known", m);
524
525                    let length = self.stream.get_u16_be_err()?;
526
527                    if length < 2 {
528                        return Err(DecodeErrors::Format(format!(
529                            "Found a marker with invalid length : {length}"
530                        )));
531                    }
532
533                    warn!("Skipping {} bytes", length - 2);
534                    self.stream.skip((length - 2) as usize)?;
535                }
536            }
537            last_byte = m;
538            bytes_before_marker += 1;
539        }
540        // Check if image is RGB
541    }
542    #[allow(clippy::too_many_lines)]
543    pub(crate) fn parse_marker_inner(&mut self, m: Marker) -> Result<(), DecodeErrors> {
544        match m {
545            Marker::SOF(0..=2) => {
546                let marker = {
547                    // choose marker
548                    if m == Marker::SOF(0) || m == Marker::SOF(1) {
549                        SOFMarkers::BaselineDct
550                    } else {
551                        self.is_progressive = true;
552                        SOFMarkers::ProgressiveDctHuffman
553                    }
554                };
555
556                trace!("Image encoding scheme =`{:?}`", marker);
557                // get components
558                parse_start_of_frame(marker, self)?;
559            }
560            // Start of Frame Segments not supported
561            Marker::SOF(v) => {
562                let feature = UnsupportedSchemes::from_int(v);
563
564                if let Some(feature) = feature {
565                    return Err(DecodeErrors::Unsupported(feature));
566                }
567
568                return Err(DecodeErrors::Format("Unsupported image format".to_string()));
569            }
570            //APP(0) segment
571            Marker::APP(0) => {
572                let mut length = self.stream.get_u16_be_err()?;
573
574                if length < 2 {
575                    return Err(DecodeErrors::Format(format!(
576                        "Found a marker with invalid length:{length}\n"
577                    )));
578                }
579                // skip for now
580                if length > 5 {
581                    let mut buffer = [0u8; 5];
582                    self.stream.read_exact_bytes(&mut buffer)?;
583                    if &buffer == b"AVI1\0" {
584                        self.is_mjpeg = true;
585                    }
586                    length -= 5;
587                }
588
589                self.stream.skip(length.saturating_sub(2) as usize)?;
590
591                //parse_app(buf, m, &mut self.info)?;
592            }
593            Marker::APP(1) => {
594                parse_app1(self)?;
595            }
596
597            Marker::APP(2) => {
598                parse_app2(self)?;
599            }
600            // Quantization tables
601            Marker::DQT => {
602                parse_dqt(self)?;
603            }
604            // Huffman tables
605            Marker::DHT => {
606                parse_huffman(self)?;
607            }
608            // Start of Scan Data
609            Marker::SOS => {
610                parse_sos(self)?;
611            }
612            Marker::EOI => return Err(DecodeErrors::FormatStatic("Premature End of image")),
613
614            Marker::DAC | Marker::DNL => {
615                return Err(DecodeErrors::Format(format!(
616                    "Parsing of the following header `{m:?}` is not supported,\
617                                cannot continue"
618                )));
619            }
620            Marker::DRI => {
621                if self.stream.get_u16_be_err()? != 4 {
622                    return Err(DecodeErrors::Format(
623                        "Bad DRI length, Corrupt JPEG".to_string()
624                    ));
625                }
626
627                self.restart_interval = usize::from(self.stream.get_u16_be_err()?);
628                trace!("DRI marker present ({})", self.restart_interval);
629
630                self.todo = self.restart_interval;
631            }
632            Marker::APP(14) => {
633                parse_app14(self)?;
634            }
635            Marker::APP(13) => {
636                parse_app13(self)?;
637            }
638            _ => {
639                warn!(
640                    "Capabilities for processing marker \"{:?}\" not implemented",
641                    m
642                );
643
644                let length = self.stream.get_u16_be_err()?;
645
646                if length < 2 {
647                    return Err(DecodeErrors::Format(format!(
648                        "Found a marker with invalid length:{length}\n"
649                    )));
650                }
651                warn!("Skipping {} bytes", length - 2);
652                self.stream.skip((length - 2) as usize)?;
653            }
654        }
655        Ok(())
656    }
657    /// Get the embedded ICC profile if it exists
658    /// and is correct
659    ///
660    /// One needs not to decode the whole image to extract this,
661    /// calling [`decode_headers`] for an image with an ICC profile
662    /// allows you to decode this
663    ///
664    /// # Returns
665    /// - `Some(Vec<u8>)`: The raw ICC profile of the image
666    /// - `None`: May indicate an error  in the ICC profile , non-existence of
667    /// an ICC profile, or that the headers weren't decoded.
668    ///
669    /// [`decode_headers`]:Self::decode_headers
670    #[must_use]
671    pub fn icc_profile(&self) -> Option<Vec<u8>> {
672        let mut marker_present: [Option<&ICCChunk>; 256] = [None; 256];
673
674        if !self.headers_decoded {
675            return None;
676        }
677        let num_markers = self.icc_data.len();
678
679        if num_markers == 0 || num_markers >= 255 {
680            return None;
681        }
682        // check validity
683        for chunk in &self.icc_data {
684            if usize::from(chunk.num_markers) != num_markers {
685                // all the lengths must match
686                return None;
687            }
688            if chunk.seq_no == 0 {
689                warn!("Zero sequence number in ICC, corrupt ICC chunk");
690                return None;
691            }
692            if marker_present[usize::from(chunk.seq_no)].is_some() {
693                // duplicate seq_no
694                warn!("Duplicate sequence number in ICC, corrupt chunk");
695                return None;
696            }
697
698            marker_present[usize::from(chunk.seq_no)] = Some(chunk);
699        }
700        let mut data = Vec::with_capacity(1000);
701        // assemble the data now
702        for chunk in marker_present.get(1..=num_markers).unwrap() {
703            if let Some(ch) = chunk {
704                data.extend_from_slice(&ch.data);
705            } else {
706                warn!("Missing icc sequence number, corrupt ICC chunk ");
707                return None;
708            }
709        }
710
711        Some(data)
712    }
713    /// Return the exif data for the file
714    ///
715    /// This returns the raw exif data starting at the
716    /// TIFF header
717    ///
718    /// # Returns
719    /// -`Some(data)`: The raw exif data, if present in the image
720    /// - None: May indicate the following
721    ///
722    ///    1. The image doesn't have exif data
723    ///    2. The image headers haven't been decoded
724    #[must_use]
725    pub fn exif(&self) -> Option<&Vec<u8>> {
726        return self.info.exif_data.as_ref();
727    }
728    /// Return the XMP data for the file
729    ///
730    /// This returns raw XMP data starting at the XML header
731    /// One needs an XML/XMP decoder to extract valuable metadata
732    ///
733    ///
734    /// # Returns
735    ///  - `Some(data)`: Raw xmp data
736    ///  - `None`: May indicate the following
737    ///     1. The image does not have xmp data
738    ///     2. The image headers have not been decoded
739    ///
740    /// # Example
741    ///
742    /// ```no_run
743    /// use zune_core::bytestream::ZCursor;
744    /// use zune_jpeg::JpegDecoder;
745    /// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
746    /// // decode headers to extract xmp metadata if present
747    /// decoder.decode_headers().unwrap();
748    /// if let Some(data) = decoder.xmp(){
749    ///     let stringified = String::from_utf8_lossy(data);
750    ///     println!("XMP")
751    /// } else{
752    ///     println!("No XMP Found")
753    /// }
754    ///
755    /// ```
756    pub fn xmp(&self) -> Option<&Vec<u8>> {
757        return self.info.xmp_data.as_ref();
758    }
759    /// Return the IPTC data for the file
760    ///
761    /// This returns the raw IPTC data.
762    ///
763    /// # Returns
764    /// -`Some(data)`: The raw IPTC data, if present in the image
765    /// - None: May indicate the following
766    ///
767    ///    1. The image doesn't have IPTC data
768    ///    2. The image headers haven't been decoded
769    #[must_use]
770    pub fn iptc(&self) -> Option<&Vec<u8>> {
771        return self.info.iptc_data.as_ref();
772    }
773    /// Get the output colorspace the image pixels will be decoded into
774    ///
775    ///
776    /// # Note.
777    /// This field can only be regarded after decoding headers,
778    /// as markers such as Adobe APP14 may dictate different colorspaces
779    /// than requested.
780    ///
781    /// Calling `decode_headers` is sufficient to know what colorspace the
782    /// output is, if this is called after `decode` it indicates the colorspace
783    /// the output is currently in
784    ///
785    /// Additionally not all input->output colorspace mappings are supported
786    /// but all input colorspaces can map to RGB colorspace, so that's a safe bet
787    /// if one is handling image formats
788    ///
789    ///# Returns
790    /// - `Some(Colorspace)`: If headers have been decoded, the colorspace the
791    ///output array will be in
792    ///- `None
793    #[must_use]
794    pub fn output_colorspace(&self) -> Option<ColorSpace> {
795        return if self.headers_decoded {
796            Some(self.options.jpeg_get_out_colorspace())
797        } else {
798            None
799        };
800    }
801
802    /// Decode into a pre-allocated buffer
803    ///
804    /// It is an error if the buffer size is smaller than
805    /// [`output_buffer_size()`](Self::output_buffer_size)
806    ///
807    /// If the buffer is bigger than expected, we ignore the end padding bytes
808    ///
809    /// # Example
810    ///
811    /// - Read  headers and then alloc a buffer big enough to hold the image
812    ///
813    /// ```no_run
814    /// use zune_core::bytestream::ZCursor;
815    /// use zune_jpeg::JpegDecoder;
816    /// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
817    /// // before we get output, we must decode the headers to get width
818    /// // height, and input colorspace
819    /// decoder.decode_headers().unwrap();
820    ///
821    /// let mut out = vec![0;decoder.output_buffer_size().unwrap()];
822    /// // write into out
823    /// decoder.decode_into(&mut out).unwrap();
824    /// ```
825    ///
826    ///
827    pub fn decode_into(&mut self, out: &mut [u8]) -> Result<(), DecodeErrors> {
828        self.decode_headers_internal()?;
829
830        let expected_size = self.output_buffer_size().unwrap();
831
832        if out.len() < expected_size {
833            // too small of a size
834            return Err(DecodeErrors::TooSmallOutput(expected_size, out.len()));
835        }
836
837        // ensure we don't touch anyone else's scratch space
838        let out_len = core::cmp::min(out.len(), expected_size);
839        let out = &mut out[0..out_len];
840
841        if self.is_progressive {
842            self.decode_mcu_ycbcr_progressive(out)
843        } else {
844            self.decode_mcu_ycbcr_baseline(out)
845        }
846    }
847
848    /// Read only headers from a jpeg image buffer
849    ///
850    /// This allows you to extract important information like
851    /// image width and height without decoding the full image
852    ///
853    /// # Examples
854    /// ```no_run
855    /// use zune_core::bytestream::ZCursor;
856    /// use zune_jpeg::{JpegDecoder};
857    ///
858    /// let img_data = std::fs::read("a_valid.jpeg").unwrap();
859    /// let mut decoder = JpegDecoder::new(ZCursor::new(&img_data));
860    /// decoder.decode_headers().unwrap();
861    ///
862    /// println!("Total decoder dimensions are : {:?} pixels",decoder.dimensions());
863    /// println!("Number of components in the image are {}", decoder.info().unwrap().components);
864    /// ```
865    /// # Errors
866    /// See DecodeErrors enum for list of possible errors during decoding
867    pub fn decode_headers(&mut self) -> Result<(), DecodeErrors> {
868        self.decode_headers_internal()?;
869        Ok(())
870    }
871    /// Create a new decoder with the specified options to be used for decoding
872    /// an image
873    ///
874    /// # Arguments
875    /// - `buf`: The input buffer from where we will pull in compressed jpeg bytes from
876    /// - `options`: Options specific to this decoder instance
877    #[must_use]
878    pub fn new_with_options(buf: T, options: DecoderOptions) -> JpegDecoder<T> {
879        JpegDecoder::default(options, buf)
880    }
881
882    /// Set up-sampling routines in case an image is down sampled
883    pub(crate) fn set_upsampling(&mut self) -> Result<(), DecodeErrors> {
884        // no sampling, return early
885        // check if horizontal max ==1
886        if self.h_max == self.v_max && self.h_max == 1 {
887            return Ok(());
888        }
889
890        for comp in &mut self.components {
891            let hs = self.h_max / comp.horizontal_sample;
892            let vs = self.v_max / comp.vertical_sample;
893
894            let samp_factor = match (hs, vs) {
895                (1, 1) => {
896                    comp.sample_ratio = SampleRatios::None;
897                    upsample_no_op
898                }
899                (2, 1) => {
900                    comp.sample_ratio = SampleRatios::H;
901                    choose_horizontal_samp_function(&self.options)
902                }
903                (1, 2) => {
904                    comp.sample_ratio = SampleRatios::V;
905                    choose_v_samp_function(&self.options)
906                }
907                (2, 2) => {
908                    comp.sample_ratio = SampleRatios::HV;
909                    choose_hv_samp_function(&self.options)
910                }
911                (hs, vs) => {
912                    comp.sample_ratio = SampleRatios::Generic(hs, vs);
913                    generic_sampler()
914                }
915            };
916            comp.setup_upsample_scanline();
917            comp.up_sampler = samp_factor;
918        }
919
920        return Ok(());
921    }
922    #[must_use]
923    /// Get the width of the image as a u16
924    ///
925    /// The width lies between 1 and 65535
926    pub(crate) fn width(&self) -> u16 {
927        self.info.width
928    }
929
930    /// Get the height of the image as a u16
931    ///
932    /// The height lies between 1 and 65535
933    #[must_use]
934    pub(crate) fn height(&self) -> u16 {
935        self.info.height
936    }
937
938    /// Get image dimensions as a tuple of width and height
939    /// or `None` if the image hasn't been decoded.
940    ///
941    /// # Returns
942    /// - `Some(width,height)`: Image dimensions
943    /// -  None : The image headers haven't been decoded
944    #[must_use]
945    pub const fn dimensions(&self) -> Option<(usize, usize)> {
946        return if self.headers_decoded {
947            Some((self.info.width as usize, self.info.height as usize))
948        } else {
949            None
950        };
951    }
952}
953
954#[derive(Default, Clone, Eq, PartialEq, Debug)]
955pub struct GainMapInfo {
956    pub data: Vec<u8>
957}
958
959#[derive(Default, Clone, Eq, PartialEq, Debug)]
960pub(crate) struct ExtendedXmpSegment {
961    pub(crate) offset: u32,
962    pub(crate) total_size: u32,
963    pub(crate) guid: Vec<u8>,
964    pub(crate) data: Vec<u8>,
965}
966
967/// A struct representing Image Information
968#[derive(Default, Clone, Eq, PartialEq)]
969#[allow(clippy::module_name_repetitions)]
970pub struct ImageInfo {
971    /// Width of the image
972    pub width: u16,
973    /// Height of image
974    pub height: u16,
975    /// PixelDensity
976    pub pixel_density: u8,
977    /// Start of frame markers
978    pub sof: SOFMarkers,
979    /// Horizontal sample
980    pub x_density: u16,
981    /// Vertical sample
982    pub y_density: u16,
983    /// Number of components
984    pub components: u8,
985    /// Gain Map information, useful for
986    /// UHDR images
987    pub gain_map_info: Vec<GainMapInfo>,
988    /// Multi picture information, useful for
989    /// UHDR images
990    pub multi_picture_information: Option<Vec<u8>>,
991    /// Exif Data
992    pub exif_data: Option<Vec<u8>>,
993    /// XMP Data
994    pub xmp_data: Option<Vec<u8>>,
995    /// IPTC Data
996    pub iptc_data: Option<Vec<u8>>,
997    /// Extended XMP Data
998    pub extended_xmp: Option<Vec<u8>>,
999    /// Extended XMP Guid
1000    pub extended_xmp_guid: Option<Vec<u8>>,
1001    /// Image sub-sampling ratio
1002    pub sample_ratio: SampleRatios,
1003    /// The offset at which Multi picture information was found
1004    pub multi_picture_information_offset: Option<u64>,
1005}
1006
1007impl ImageInfo {
1008    /// Set width of the image
1009    ///
1010    /// Found in the start of frame
1011
1012    pub(crate) fn set_width(&mut self, width: u16) {
1013        self.width = width;
1014    }
1015
1016    /// Set height of the image
1017    ///
1018    /// Found in the start of frame
1019
1020    pub(crate) fn set_height(&mut self, height: u16) {
1021        self.height = height;
1022    }
1023
1024    /// Set the image density
1025    ///
1026    /// Found in the start of frame
1027
1028    pub(crate) fn set_density(&mut self, density: u8) {
1029        self.pixel_density = density;
1030    }
1031
1032    /// Set image Start of frame marker
1033    ///
1034    /// found in the Start of frame header
1035
1036    pub(crate) fn set_sof_marker(&mut self, marker: SOFMarkers) {
1037        self.sof = marker;
1038    }
1039
1040    /// Set image x-density(dots per pixel)
1041    ///
1042    /// Found in the APP(0) marker
1043    #[allow(dead_code)]
1044    pub(crate) fn set_x(&mut self, sample: u16) {
1045        self.x_density = sample;
1046    }
1047
1048    /// Set image y-density
1049    ///
1050    /// Found in the APP(0) marker
1051    #[allow(dead_code)]
1052    pub(crate) fn set_y(&mut self, sample: u16) {
1053        self.y_density = sample;
1054    }
1055}