image/codecs/jpeg/
encoder.rs

1#![allow(clippy::too_many_arguments)]
2use std::borrow::Cow;
3use std::io::{self, Write};
4
5use crate::error::{
6    ImageError, ImageResult, ParameterError, ParameterErrorKind, UnsupportedError,
7    UnsupportedErrorKind,
8};
9use crate::traits::PixelWithColorType;
10use crate::utils::clamp;
11use crate::{
12    ColorType, DynamicImage, ExtendedColorType, GenericImageView, ImageBuffer, ImageEncoder,
13    ImageFormat, Luma, Pixel, Rgb,
14};
15
16use num_traits::ToPrimitive;
17
18use super::entropy::build_huff_lut_const;
19use super::transform;
20
21// Markers
22// Baseline DCT
23static SOF0: u8 = 0xC0;
24// Huffman Tables
25static DHT: u8 = 0xC4;
26// Start of Image (standalone)
27static SOI: u8 = 0xD8;
28// End of image (standalone)
29static EOI: u8 = 0xD9;
30// Start of Scan
31static SOS: u8 = 0xDA;
32// Quantization Tables
33static DQT: u8 = 0xDB;
34// Application segments start and end
35static APP0: u8 = 0xE0;
36static APP1: u8 = 0xE1;
37static APP2: u8 = 0xE2;
38
39// section K.1
40// table K.1
41#[rustfmt::skip]
42static STD_LUMA_QTABLE: [u8; 64] = [
43    16, 11, 10, 16,  24,  40,  51,  61,
44    12, 12, 14, 19,  26,  58,  60,  55,
45    14, 13, 16, 24,  40,  57,  69,  56,
46    14, 17, 22, 29,  51,  87,  80,  62,
47    18, 22, 37, 56,  68, 109, 103,  77,
48    24, 35, 55, 64,  81, 104, 113,  92,
49    49, 64, 78, 87, 103, 121, 120, 101,
50    72, 92, 95, 98, 112, 100, 103,  99,
51];
52
53// table K.2
54#[rustfmt::skip]
55static STD_CHROMA_QTABLE: [u8; 64] = [
56    17, 18, 24, 47, 99, 99, 99, 99,
57    18, 21, 26, 66, 99, 99, 99, 99,
58    24, 26, 56, 99, 99, 99, 99, 99,
59    47, 66, 99, 99, 99, 99, 99, 99,
60    99, 99, 99, 99, 99, 99, 99, 99,
61    99, 99, 99, 99, 99, 99, 99, 99,
62    99, 99, 99, 99, 99, 99, 99, 99,
63    99, 99, 99, 99, 99, 99, 99, 99,
64];
65
66// section K.3
67// Code lengths and values for table K.3
68static STD_LUMA_DC_CODE_LENGTHS: [u8; 16] = [
69    0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
70];
71
72static STD_LUMA_DC_VALUES: [u8; 12] = [
73    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
74];
75
76static STD_LUMA_DC_HUFF_LUT: [(u8, u16); 256] =
77    build_huff_lut_const(&STD_LUMA_DC_CODE_LENGTHS, &STD_LUMA_DC_VALUES);
78
79// Code lengths and values for table K.4
80static STD_CHROMA_DC_CODE_LENGTHS: [u8; 16] = [
81    0x00, 0x03, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
82];
83
84static STD_CHROMA_DC_VALUES: [u8; 12] = [
85    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
86];
87
88static STD_CHROMA_DC_HUFF_LUT: [(u8, u16); 256] =
89    build_huff_lut_const(&STD_CHROMA_DC_CODE_LENGTHS, &STD_CHROMA_DC_VALUES);
90
91// Code lengths and values for table k.5
92static STD_LUMA_AC_CODE_LENGTHS: [u8; 16] = [
93    0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04, 0x00, 0x00, 0x01, 0x7D,
94];
95
96static STD_LUMA_AC_VALUES: [u8; 162] = [
97    0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
98    0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42, 0xB1, 0xC1, 0x15, 0x52, 0xD1, 0xF0,
99    0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28,
100    0x29, 0x2A, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
101    0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
102    0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
103    0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7,
104    0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5,
105    0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1, 0xE2,
106    0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
107    0xF9, 0xFA,
108];
109
110static STD_LUMA_AC_HUFF_LUT: [(u8, u16); 256] =
111    build_huff_lut_const(&STD_LUMA_AC_CODE_LENGTHS, &STD_LUMA_AC_VALUES);
112
113// Code lengths and values for table k.6
114static STD_CHROMA_AC_CODE_LENGTHS: [u8; 16] = [
115    0x00, 0x02, 0x01, 0x02, 0x04, 0x04, 0x03, 0x04, 0x07, 0x05, 0x04, 0x04, 0x00, 0x01, 0x02, 0x77,
116];
117static STD_CHROMA_AC_VALUES: [u8; 162] = [
118    0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
119    0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1, 0x09, 0x23, 0x33, 0x52, 0xF0,
120    0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24, 0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A, 0x26,
121    0x27, 0x28, 0x29, 0x2A, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
122    0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
123    0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
124    0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5,
125    0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3,
126    0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA,
127    0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
128    0xF9, 0xFA,
129];
130
131static STD_CHROMA_AC_HUFF_LUT: [(u8, u16); 256] =
132    build_huff_lut_const(&STD_CHROMA_AC_CODE_LENGTHS, &STD_CHROMA_AC_VALUES);
133
134static DCCLASS: u8 = 0;
135static ACCLASS: u8 = 1;
136
137static LUMADESTINATION: u8 = 0;
138static CHROMADESTINATION: u8 = 1;
139
140static LUMAID: u8 = 1;
141static CHROMABLUEID: u8 = 2;
142static CHROMAREDID: u8 = 3;
143
144/// The permutation of dct coefficients.
145#[rustfmt::skip]
146static UNZIGZAG: [u8; 64] = [
147     0,  1,  8, 16,  9,  2,  3, 10,
148    17, 24, 32, 25, 18, 11,  4,  5,
149    12, 19, 26, 33, 40, 48, 41, 34,
150    27, 20, 13,  6,  7, 14, 21, 28,
151    35, 42, 49, 56, 57, 50, 43, 36,
152    29, 22, 15, 23, 30, 37, 44, 51,
153    58, 59, 52, 45, 38, 31, 39, 46,
154    53, 60, 61, 54, 47, 55, 62, 63,
155];
156
157// E x i f \0 \0
158/// The header for an EXIF APP1 segment
159static EXIF_HEADER: [u8; 6] = [0x45, 0x78, 0x69, 0x66, 0x00, 0x00];
160
161/// A representation of a JPEG component
162#[derive(Copy, Clone)]
163struct Component {
164    /// The Component's identifier
165    id: u8,
166
167    /// Horizontal sampling factor
168    h: u8,
169
170    /// Vertical sampling factor
171    v: u8,
172
173    /// The quantization table selector
174    tq: u8,
175
176    /// Index to the Huffman DC Table
177    dc_table: u8,
178
179    /// Index to the AC Huffman Table
180    ac_table: u8,
181
182    /// The dc prediction of the component
183    _dc_pred: i32,
184}
185
186pub(crate) struct BitWriter<W> {
187    w: W,
188    accumulator: u32,
189    nbits: u8,
190}
191
192impl<W: Write> BitWriter<W> {
193    fn new(w: W) -> Self {
194        BitWriter {
195            w,
196            accumulator: 0,
197            nbits: 0,
198        }
199    }
200
201    fn write_bits(&mut self, bits: u16, size: u8) -> io::Result<()> {
202        if size == 0 {
203            return Ok(());
204        }
205
206        self.nbits += size;
207        self.accumulator |= u32::from(bits) << (32 - self.nbits) as usize;
208
209        while self.nbits >= 8 {
210            let byte = self.accumulator >> 24;
211            self.w.write_all(&[byte as u8])?;
212
213            if byte == 0xFF {
214                self.w.write_all(&[0x00])?;
215            }
216
217            self.nbits -= 8;
218            self.accumulator <<= 8;
219        }
220
221        Ok(())
222    }
223
224    fn pad_byte(&mut self) -> io::Result<()> {
225        self.write_bits(0x7F, 7)
226    }
227
228    fn huffman_encode(&mut self, val: u8, table: &[(u8, u16); 256]) -> io::Result<()> {
229        let (size, code) = table[val as usize];
230
231        assert!(size <= 16, "bad huffman value");
232
233        self.write_bits(code, size)
234    }
235
236    fn write_block(
237        &mut self,
238        block: &[i32; 64],
239        prevdc: i32,
240        dctable: &[(u8, u16); 256],
241        actable: &[(u8, u16); 256],
242    ) -> io::Result<i32> {
243        // Differential DC encoding
244        let dcval = block[0];
245        let diff = dcval - prevdc;
246        let (size, value) = encode_coefficient(diff);
247
248        self.huffman_encode(size, dctable)?;
249        self.write_bits(value, size)?;
250
251        // Figure F.2
252        let mut zero_run = 0;
253
254        for &k in &UNZIGZAG[1..] {
255            if block[k as usize] == 0 {
256                zero_run += 1;
257            } else {
258                while zero_run > 15 {
259                    self.huffman_encode(0xF0, actable)?;
260                    zero_run -= 16;
261                }
262
263                let (size, value) = encode_coefficient(block[k as usize]);
264                let symbol = (zero_run << 4) | size;
265
266                self.huffman_encode(symbol, actable)?;
267                self.write_bits(value, size)?;
268
269                zero_run = 0;
270            }
271        }
272
273        if block[UNZIGZAG[63] as usize] == 0 {
274            self.huffman_encode(0x00, actable)?;
275        }
276
277        Ok(dcval)
278    }
279
280    fn write_marker(&mut self, marker: u8) -> io::Result<()> {
281        self.w.write_all(&[0xFF, marker])
282    }
283
284    fn write_segment(&mut self, marker: u8, data: &[u8]) -> io::Result<()> {
285        self.w.write_all(&[0xFF, marker])?;
286        self.w.write_all(&(data.len() as u16 + 2).to_be_bytes())?;
287        self.w.write_all(data)
288    }
289}
290
291/// Represents a unit in which the density of an image is measured
292#[derive(Clone, Copy, Debug, Eq, PartialEq)]
293pub enum PixelDensityUnit {
294    /// Represents the absence of a unit, the values indicate only a
295    /// [pixel aspect ratio](https://en.wikipedia.org/wiki/Pixel_aspect_ratio)
296    PixelAspectRatio,
297
298    /// Pixels per inch (2.54 cm)
299    Inches,
300
301    /// Pixels per centimeter
302    Centimeters,
303}
304
305/// Represents the pixel density of an image
306///
307/// For example, a 300 DPI image is represented by:
308///
309/// ```rust
310/// use image::codecs::jpeg::*;
311/// let hdpi = PixelDensity::dpi(300);
312/// assert_eq!(hdpi, PixelDensity {density: (300,300), unit: PixelDensityUnit::Inches})
313/// ```
314#[derive(Clone, Copy, Debug, Eq, PartialEq)]
315pub struct PixelDensity {
316    /// A couple of values for (Xdensity, Ydensity)
317    pub density: (u16, u16),
318    /// The unit in which the density is measured
319    pub unit: PixelDensityUnit,
320}
321
322impl PixelDensity {
323    /// Creates the most common pixel density type:
324    /// the horizontal and the vertical density are equal,
325    /// and measured in pixels per inch.
326    #[must_use]
327    pub fn dpi(density: u16) -> Self {
328        PixelDensity {
329            density: (density, density),
330            unit: PixelDensityUnit::Inches,
331        }
332    }
333}
334
335impl Default for PixelDensity {
336    /// Returns a pixel density with a pixel aspect ratio of 1
337    fn default() -> Self {
338        PixelDensity {
339            density: (1, 1),
340            unit: PixelDensityUnit::PixelAspectRatio,
341        }
342    }
343}
344
345/// The representation of a JPEG encoder
346pub struct JpegEncoder<W> {
347    writer: BitWriter<W>,
348
349    components: Vec<Component>,
350    tables: Vec<[u8; 64]>,
351
352    luma_dctable: Cow<'static, [(u8, u16); 256]>,
353    luma_actable: Cow<'static, [(u8, u16); 256]>,
354    chroma_dctable: Cow<'static, [(u8, u16); 256]>,
355    chroma_actable: Cow<'static, [(u8, u16); 256]>,
356
357    pixel_density: PixelDensity,
358
359    icc_profile: Vec<u8>,
360    exif: Vec<u8>,
361}
362
363impl<W: Write> JpegEncoder<W> {
364    /// Create a new encoder that writes its output to ```w```
365    pub fn new(w: W) -> JpegEncoder<W> {
366        JpegEncoder::new_with_quality(w, 75)
367    }
368
369    /// Create a new encoder that writes its output to ```w```, and has
370    /// the quality parameter ```quality``` with a value in the range 1-100
371    /// where 1 is the worst and 100 is the best.
372    pub fn new_with_quality(w: W, quality: u8) -> JpegEncoder<W> {
373        let components = vec![
374            Component {
375                id: LUMAID,
376                h: 1,
377                v: 1,
378                tq: LUMADESTINATION,
379                dc_table: LUMADESTINATION,
380                ac_table: LUMADESTINATION,
381                _dc_pred: 0,
382            },
383            Component {
384                id: CHROMABLUEID,
385                h: 1,
386                v: 1,
387                tq: CHROMADESTINATION,
388                dc_table: CHROMADESTINATION,
389                ac_table: CHROMADESTINATION,
390                _dc_pred: 0,
391            },
392            Component {
393                id: CHROMAREDID,
394                h: 1,
395                v: 1,
396                tq: CHROMADESTINATION,
397                dc_table: CHROMADESTINATION,
398                ac_table: CHROMADESTINATION,
399                _dc_pred: 0,
400            },
401        ];
402
403        // Derive our quantization table scaling value using the libjpeg algorithm
404        let scale = u32::from(clamp(quality, 1, 100));
405        let scale = if scale < 50 {
406            5000 / scale
407        } else {
408            200 - scale * 2
409        };
410
411        let mut tables = vec![STD_LUMA_QTABLE, STD_CHROMA_QTABLE];
412        for t in tables.iter_mut() {
413            for v in t.iter_mut() {
414                *v = clamp((u32::from(*v) * scale + 50) / 100, 1, u32::from(u8::MAX)) as u8;
415            }
416        }
417
418        JpegEncoder {
419            writer: BitWriter::new(w),
420
421            components,
422            tables,
423
424            luma_dctable: Cow::Borrowed(&STD_LUMA_DC_HUFF_LUT),
425            luma_actable: Cow::Borrowed(&STD_LUMA_AC_HUFF_LUT),
426            chroma_dctable: Cow::Borrowed(&STD_CHROMA_DC_HUFF_LUT),
427            chroma_actable: Cow::Borrowed(&STD_CHROMA_AC_HUFF_LUT),
428
429            pixel_density: PixelDensity::default(),
430
431            icc_profile: Vec::new(),
432            exif: Vec::new(),
433        }
434    }
435
436    /// Set the pixel density of the images the encoder will encode.
437    /// If this method is not called, then a default pixel aspect ratio of 1x1 will be applied,
438    /// and no DPI information will be stored in the image.
439    pub fn set_pixel_density(&mut self, pixel_density: PixelDensity) {
440        self.pixel_density = pixel_density;
441    }
442
443    /// Encodes the image stored in the raw byte buffer ```image```
444    /// that has dimensions ```width``` and ```height```
445    /// and ```ColorType``` ```c```
446    ///
447    /// The Image in encoded with subsampling ratio 4:2:2
448    ///
449    /// # Panics
450    ///
451    /// Panics if `width * height * color_type.bytes_per_pixel() != image.len()`.
452    #[track_caller]
453    pub fn encode(
454        &mut self,
455        image: &[u8],
456        width: u32,
457        height: u32,
458        color_type: ExtendedColorType,
459    ) -> ImageResult<()> {
460        let expected_buffer_len = color_type.buffer_size(width, height);
461        assert_eq!(
462            expected_buffer_len,
463            image.len() as u64,
464            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
465            image.len(),
466        );
467
468        match color_type {
469            ExtendedColorType::L8 => {
470                let image: ImageBuffer<Luma<_>, _> =
471                    ImageBuffer::from_raw(width, height, image).unwrap();
472                self.encode_image(&image)
473            }
474            ExtendedColorType::Rgb8 => {
475                let image: ImageBuffer<Rgb<_>, _> =
476                    ImageBuffer::from_raw(width, height, image).unwrap();
477                self.encode_image(&image)
478            }
479            _ => Err(ImageError::Unsupported(
480                UnsupportedError::from_format_and_kind(
481                    ImageFormat::Jpeg.into(),
482                    UnsupportedErrorKind::Color(color_type),
483                ),
484            )),
485        }
486    }
487
488    fn write_exif(&mut self) -> ImageResult<()> {
489        if !self.exif.is_empty() {
490            let mut formatted = EXIF_HEADER.to_vec();
491            formatted.extend_from_slice(&self.exif);
492            self.writer.write_segment(APP1, &formatted)?;
493        }
494
495        Ok(())
496    }
497
498    /// Encodes the given image.
499    ///
500    /// As a special feature this does not require the whole image to be present in memory at the
501    /// same time such that it may be computed on the fly, which is why this method exists on this
502    /// encoder but not on others. Instead the encoder will iterate over 8-by-8 blocks of pixels at
503    /// a time, inspecting each pixel exactly once. You can rely on this behaviour when calling
504    /// this method.
505    ///
506    /// The Image in encoded with subsampling ratio 4:2:2
507    pub fn encode_image<I: GenericImageView>(&mut self, image: &I) -> ImageResult<()>
508    where
509        I::Pixel: PixelWithColorType,
510    {
511        let n = I::Pixel::CHANNEL_COUNT;
512        let color_type = I::Pixel::COLOR_TYPE;
513        let num_components = if n == 1 || n == 2 { 1 } else { 3 };
514
515        self.writer.write_marker(SOI)?;
516
517        let mut buf = Vec::new();
518
519        build_jfif_header(&mut buf, self.pixel_density);
520        self.writer.write_segment(APP0, &buf)?;
521        self.write_exif()?;
522
523        // Write ICC profile chunks if present
524        self.write_icc_profile_chunks()?;
525
526        build_frame_header(
527            &mut buf,
528            8,
529            // TODO: not idiomatic yet. Should be an EncodingError and mention jpg. Further it
530            // should check dimensions prior to writing.
531            u16::try_from(image.width()).map_err(|_| {
532                ImageError::Parameter(ParameterError::from_kind(
533                    ParameterErrorKind::DimensionMismatch,
534                ))
535            })?,
536            u16::try_from(image.height()).map_err(|_| {
537                ImageError::Parameter(ParameterError::from_kind(
538                    ParameterErrorKind::DimensionMismatch,
539                ))
540            })?,
541            &self.components[..num_components],
542        );
543        self.writer.write_segment(SOF0, &buf)?;
544
545        assert_eq!(self.tables.len(), 2);
546        let numtables = if num_components == 1 { 1 } else { 2 };
547
548        for (i, table) in self.tables[..numtables].iter().enumerate() {
549            build_quantization_segment(&mut buf, 8, i as u8, table);
550            self.writer.write_segment(DQT, &buf)?;
551        }
552
553        build_huffman_segment(
554            &mut buf,
555            DCCLASS,
556            LUMADESTINATION,
557            &STD_LUMA_DC_CODE_LENGTHS,
558            &STD_LUMA_DC_VALUES,
559        );
560        self.writer.write_segment(DHT, &buf)?;
561
562        build_huffman_segment(
563            &mut buf,
564            ACCLASS,
565            LUMADESTINATION,
566            &STD_LUMA_AC_CODE_LENGTHS,
567            &STD_LUMA_AC_VALUES,
568        );
569        self.writer.write_segment(DHT, &buf)?;
570
571        if num_components == 3 {
572            build_huffman_segment(
573                &mut buf,
574                DCCLASS,
575                CHROMADESTINATION,
576                &STD_CHROMA_DC_CODE_LENGTHS,
577                &STD_CHROMA_DC_VALUES,
578            );
579            self.writer.write_segment(DHT, &buf)?;
580
581            build_huffman_segment(
582                &mut buf,
583                ACCLASS,
584                CHROMADESTINATION,
585                &STD_CHROMA_AC_CODE_LENGTHS,
586                &STD_CHROMA_AC_VALUES,
587            );
588            self.writer.write_segment(DHT, &buf)?;
589        }
590
591        build_scan_header(&mut buf, &self.components[..num_components]);
592        self.writer.write_segment(SOS, &buf)?;
593
594        if ExtendedColorType::Rgb8 == color_type || ExtendedColorType::Rgba8 == color_type {
595            self.encode_rgb(image)
596        } else {
597            self.encode_gray(image)
598        }?;
599
600        self.writer.pad_byte()?;
601        self.writer.write_marker(EOI)?;
602        Ok(())
603    }
604
605    fn encode_gray<I: GenericImageView>(&mut self, image: &I) -> io::Result<()> {
606        let mut yblock = [0u8; 64];
607        let mut y_dcprev = 0;
608        let mut dct_yblock = [0i32; 64];
609
610        for y in (0..image.height()).step_by(8) {
611            for x in (0..image.width()).step_by(8) {
612                copy_blocks_gray(image, x, y, &mut yblock);
613
614                // Level shift and fdct
615                // Coeffs are scaled by 8
616                transform::fdct(&yblock, &mut dct_yblock);
617
618                // Quantization
619                for (i, dct) in dct_yblock.iter_mut().enumerate() {
620                    *dct = ((*dct / 8) as f32 / f32::from(self.tables[0][i])).round() as i32;
621                }
622
623                let la = &*self.luma_actable;
624                let ld = &*self.luma_dctable;
625
626                y_dcprev = self.writer.write_block(&dct_yblock, y_dcprev, ld, la)?;
627            }
628        }
629
630        Ok(())
631    }
632
633    fn encode_rgb<I: GenericImageView>(&mut self, image: &I) -> io::Result<()> {
634        let mut y_dcprev = 0;
635        let mut cb_dcprev = 0;
636        let mut cr_dcprev = 0;
637
638        let mut dct_yblock = [0i32; 64];
639        let mut dct_cb_block = [0i32; 64];
640        let mut dct_cr_block = [0i32; 64];
641
642        let mut yblock = [0u8; 64];
643        let mut cb_block = [0u8; 64];
644        let mut cr_block = [0u8; 64];
645
646        for y in (0..image.height()).step_by(8) {
647            for x in (0..image.width()).step_by(8) {
648                // RGB -> YCbCr
649                copy_blocks_ycbcr(image, x, y, &mut yblock, &mut cb_block, &mut cr_block);
650
651                // Level shift and fdct
652                // Coeffs are scaled by 8
653                transform::fdct(&yblock, &mut dct_yblock);
654                transform::fdct(&cb_block, &mut dct_cb_block);
655                transform::fdct(&cr_block, &mut dct_cr_block);
656
657                // Quantization
658                for i in 0usize..64 {
659                    dct_yblock[i] =
660                        ((dct_yblock[i] / 8) as f32 / f32::from(self.tables[0][i])).round() as i32;
661                    dct_cb_block[i] = ((dct_cb_block[i] / 8) as f32 / f32::from(self.tables[1][i]))
662                        .round() as i32;
663                    dct_cr_block[i] = ((dct_cr_block[i] / 8) as f32 / f32::from(self.tables[1][i]))
664                        .round() as i32;
665                }
666
667                let la = &*self.luma_actable;
668                let ld = &*self.luma_dctable;
669                let cd = &*self.chroma_dctable;
670                let ca = &*self.chroma_actable;
671
672                y_dcprev = self.writer.write_block(&dct_yblock, y_dcprev, ld, la)?;
673                cb_dcprev = self.writer.write_block(&dct_cb_block, cb_dcprev, cd, ca)?;
674                cr_dcprev = self.writer.write_block(&dct_cr_block, cr_dcprev, cd, ca)?;
675            }
676        }
677
678        Ok(())
679    }
680
681    fn write_icc_profile_chunks(&mut self) -> io::Result<()> {
682        if self.icc_profile.is_empty() {
683            return Ok(());
684        }
685
686        const MAX_CHUNK_SIZE: usize = 65533 - 14;
687        const MAX_CHUNK_COUNT: usize = 255;
688        const MAX_ICC_PROFILE_SIZE: usize = MAX_CHUNK_SIZE * MAX_CHUNK_COUNT;
689
690        if self.icc_profile.len() > MAX_ICC_PROFILE_SIZE {
691            return Err(io::Error::new(
692                io::ErrorKind::InvalidInput,
693                "ICC profile too large",
694            ));
695        }
696
697        let chunk_iter = self.icc_profile.chunks(MAX_CHUNK_SIZE);
698        let num_chunks = chunk_iter.len() as u8;
699        let mut segment = Vec::new();
700
701        for (i, chunk) in chunk_iter.enumerate() {
702            let chunk_number = (i + 1) as u8;
703            let length = 14 + chunk.len();
704
705            segment.clear();
706            segment.reserve(length);
707            segment.extend_from_slice(b"ICC_PROFILE\0");
708            segment.push(chunk_number);
709            segment.push(num_chunks);
710            segment.extend_from_slice(chunk);
711
712            self.writer.write_segment(APP2, &segment)?;
713        }
714
715        Ok(())
716    }
717}
718
719impl<W: Write> ImageEncoder for JpegEncoder<W> {
720    #[track_caller]
721    fn write_image(
722        mut self,
723        buf: &[u8],
724        width: u32,
725        height: u32,
726        color_type: ExtendedColorType,
727    ) -> ImageResult<()> {
728        self.encode(buf, width, height, color_type)
729    }
730
731    fn set_icc_profile(&mut self, icc_profile: Vec<u8>) -> Result<(), UnsupportedError> {
732        self.icc_profile = icc_profile;
733        Ok(())
734    }
735
736    fn set_exif_metadata(&mut self, exif: Vec<u8>) -> Result<(), UnsupportedError> {
737        self.exif = exif;
738        Ok(())
739    }
740
741    fn make_compatible_img(
742        &self,
743        _: crate::io::encoder::MethodSealedToImage,
744        img: &DynamicImage,
745    ) -> Option<DynamicImage> {
746        use ColorType::*;
747        match img.color() {
748            L8 | Rgb8 => None,
749            La8 | L16 | La16 => Some(img.to_luma8().into()),
750            Rgba8 | Rgb16 | Rgb32F | Rgba16 | Rgba32F => Some(img.to_rgb8().into()),
751        }
752    }
753}
754
755fn build_jfif_header(m: &mut Vec<u8>, density: PixelDensity) {
756    m.clear();
757    m.extend_from_slice(b"JFIF");
758    m.extend_from_slice(&[
759        0,
760        0x01,
761        0x02,
762        match density.unit {
763            PixelDensityUnit::PixelAspectRatio => 0x00,
764            PixelDensityUnit::Inches => 0x01,
765            PixelDensityUnit::Centimeters => 0x02,
766        },
767    ]);
768    m.extend_from_slice(&density.density.0.to_be_bytes());
769    m.extend_from_slice(&density.density.1.to_be_bytes());
770    m.extend_from_slice(&[0, 0]);
771}
772
773fn build_frame_header(
774    m: &mut Vec<u8>,
775    precision: u8,
776    width: u16,
777    height: u16,
778    components: &[Component],
779) {
780    m.clear();
781
782    m.push(precision);
783    m.extend_from_slice(&height.to_be_bytes());
784    m.extend_from_slice(&width.to_be_bytes());
785    m.push(components.len() as u8);
786
787    for &comp in components {
788        let hv = (comp.h << 4) | comp.v;
789        m.extend_from_slice(&[comp.id, hv, comp.tq]);
790    }
791}
792
793fn build_scan_header(m: &mut Vec<u8>, components: &[Component]) {
794    m.clear();
795
796    m.push(components.len() as u8);
797
798    for &comp in components {
799        let tables = (comp.dc_table << 4) | comp.ac_table;
800        m.extend_from_slice(&[comp.id, tables]);
801    }
802
803    // spectral start and end, approx. high and low
804    m.extend_from_slice(&[0, 63, 0]);
805}
806
807fn build_huffman_segment(
808    m: &mut Vec<u8>,
809    class: u8,
810    destination: u8,
811    numcodes: &[u8; 16],
812    values: &[u8],
813) {
814    m.clear();
815
816    let tcth = (class << 4) | destination;
817    m.push(tcth);
818
819    m.extend_from_slice(numcodes);
820
821    let sum: usize = numcodes.iter().map(|&x| x as usize).sum();
822
823    assert_eq!(sum, values.len());
824
825    m.extend_from_slice(values);
826}
827
828fn build_quantization_segment(m: &mut Vec<u8>, precision: u8, identifier: u8, qtable: &[u8; 64]) {
829    m.clear();
830
831    let p = if precision == 8 { 0 } else { 1 };
832
833    let pqtq = (p << 4) | identifier;
834    m.push(pqtq);
835
836    for &i in &UNZIGZAG[..] {
837        m.push(qtable[i as usize]);
838    }
839}
840
841fn encode_coefficient(coefficient: i32) -> (u8, u16) {
842    let mut magnitude = coefficient.unsigned_abs() as u16;
843    let mut num_bits = 0u8;
844
845    while magnitude > 0 {
846        magnitude >>= 1;
847        num_bits += 1;
848    }
849
850    let mask = (1 << num_bits as usize) - 1;
851
852    let val = if coefficient < 0 {
853        (coefficient - 1) as u16 & mask
854    } else {
855        coefficient as u16 & mask
856    };
857
858    (num_bits, val)
859}
860
861#[inline]
862fn rgb_to_ycbcr<P: Pixel>(pixel: P) -> (u8, u8, u8) {
863    let [r, g, b] = pixel.to_rgb().0;
864    let r: i32 = i32::from(r.to_u8().unwrap());
865    let g: i32 = i32::from(g.to_u8().unwrap());
866    let b: i32 = i32::from(b.to_u8().unwrap());
867
868    /*
869       JPEG RGB -> YCbCr is defined as following equations using Bt.601 Full Range matrix:
870       Y  =  0.29900 * R + 0.58700 * G + 0.11400 * B
871       Cb = -0.16874 * R - 0.33126 * G + 0.50000 * B  + 128
872       Cr =  0.50000 * R - 0.41869 * G - 0.08131 * B  + 128
873
874       To avoid using slow floating point conversion is done in fixed point,
875       using following coefficients with rounding to nearest integer mode:
876    */
877
878    const C_YR: i32 = 19595; // 0.29900 = 19595 * 2^-16
879    const C_YG: i32 = 38469; // 0.58700 = 38469 * 2^-16
880    const C_YB: i32 = 7471; // 0.11400 = 7471 * 2^-16
881    const Y_ROUNDING: i32 = (1 << 15) - 1; // + 0.5 to perform rounding shift right in-place
882    const C_UR: i32 = 11059; // 0.16874 = 11059 * 2^-16
883    const C_UG: i32 = 21709; // 0.33126 = 21709 * 2^-16
884    const C_UB: i32 = 32768; // 0.5 = 32768 * 2^-16
885    const UV_BIAS_ROUNDING: i32 = (128 * (1 << 16)) + ((1 << 15) - 1); // 128 + 0.5 = ((128 * (1 << 16)) + ((1 << 15) - 1)) * 2^-16 ; + 0.5 to perform rounding shift right in-place
886    const C_VR: i32 = C_UB; // 0.5 = 32768 * 2^-16
887    const C_VG: i32 = 27439; // 0.41869 = 27439 * 2^-16
888    const C_VB: i32 = 5329; // 0.08131409 = 5329 * 2^-16
889
890    let y = (C_YR * r + C_YG * g + C_YB * b + Y_ROUNDING) >> 16;
891    let cb = (-C_UR * r - C_UG * g + C_UB * b + UV_BIAS_ROUNDING) >> 16;
892    let cr = (C_VR * r - C_VG * g - C_VB * b + UV_BIAS_ROUNDING) >> 16;
893
894    (y as u8, cb as u8, cr as u8)
895}
896
897/// Returns the pixel at (x,y) if (x,y) is in the image,
898/// otherwise the closest pixel in the image
899#[inline]
900fn pixel_at_or_near<I: GenericImageView>(source: &I, x: u32, y: u32) -> I::Pixel {
901    if source.in_bounds(x, y) {
902        source.get_pixel(x, y)
903    } else {
904        source.get_pixel(x.min(source.width() - 1), y.min(source.height() - 1))
905    }
906}
907
908fn copy_blocks_ycbcr<I: GenericImageView>(
909    source: &I,
910    x0: u32,
911    y0: u32,
912    yb: &mut [u8; 64],
913    cbb: &mut [u8; 64],
914    crb: &mut [u8; 64],
915) {
916    for y in 0..8 {
917        for x in 0..8 {
918            let pixel = pixel_at_or_near(source, x + x0, y + y0);
919            let (yc, cb, cr) = rgb_to_ycbcr(pixel);
920
921            yb[(y * 8 + x) as usize] = yc;
922            cbb[(y * 8 + x) as usize] = cb;
923            crb[(y * 8 + x) as usize] = cr;
924        }
925    }
926}
927
928fn copy_blocks_gray<I: GenericImageView>(source: &I, x0: u32, y0: u32, gb: &mut [u8; 64]) {
929    use num_traits::cast::ToPrimitive;
930    for y in 0..8 {
931        for x in 0..8 {
932            let pixel = pixel_at_or_near(source, x0 + x, y0 + y);
933            let [luma] = pixel.to_luma().0;
934            gb[(y * 8 + x) as usize] = luma.to_u8().unwrap();
935        }
936    }
937}
938
939#[cfg(test)]
940mod tests {
941    use std::io::Cursor;
942
943    #[cfg(feature = "benchmarks")]
944    extern crate test;
945    #[cfg(feature = "benchmarks")]
946    use test::Bencher;
947
948    use crate::error::ParameterErrorKind::DimensionMismatch;
949    use crate::{ColorType, DynamicImage, ExtendedColorType, ImageEncoder, ImageError};
950    use crate::{ImageDecoder as _, ImageFormat};
951
952    use super::super::JpegDecoder;
953    use super::{
954        build_frame_header, build_huffman_segment, build_jfif_header, build_quantization_segment,
955        build_scan_header, Component, JpegEncoder, PixelDensity, DCCLASS, LUMADESTINATION,
956        STD_LUMA_DC_CODE_LENGTHS, STD_LUMA_DC_VALUES,
957    };
958
959    fn decode(encoded: &[u8]) -> Vec<u8> {
960        let decoder = JpegDecoder::new(Cursor::new(encoded)).expect("Could not decode image");
961
962        let mut decoded = vec![0; decoder.total_bytes() as usize];
963        decoder
964            .read_image(&mut decoded)
965            .expect("Could not decode image");
966        decoded
967    }
968
969    #[test]
970    fn roundtrip_sanity_check() {
971        // create a 1x1 8-bit image buffer containing a single red pixel
972        let img = [255u8, 0, 0];
973
974        // encode it into a memory buffer
975        let mut encoded_img = Vec::new();
976        {
977            let encoder = JpegEncoder::new_with_quality(&mut encoded_img, 100);
978            encoder
979                .write_image(&img, 1, 1, ExtendedColorType::Rgb8)
980                .expect("Could not encode image");
981        }
982
983        // decode it from the memory buffer
984        {
985            let decoded = decode(&encoded_img);
986            // note that, even with the encode quality set to 100, we do not get the same image
987            // back. Therefore, we're going to assert that it's at least red-ish:
988            assert_eq!(3, decoded.len());
989            assert!(decoded[0] > 0x80);
990            assert!(decoded[1] < 0x80);
991            assert!(decoded[2] < 0x80);
992        }
993    }
994
995    #[test]
996    fn grayscale_roundtrip_sanity_check() {
997        // create a 2x2 8-bit image buffer containing a white diagonal
998        let img = [255u8, 0, 0, 255];
999
1000        // encode it into a memory buffer
1001        let mut encoded_img = Vec::new();
1002        {
1003            let encoder = JpegEncoder::new_with_quality(&mut encoded_img, 100);
1004            encoder
1005                .write_image(&img[..], 2, 2, ExtendedColorType::L8)
1006                .expect("Could not encode image");
1007        }
1008
1009        // decode it from the memory buffer
1010        {
1011            let decoded = decode(&encoded_img);
1012            // note that, even with the encode quality set to 100, we do not get the same image
1013            // back. Therefore, we're going to assert that the diagonal is at least white-ish:
1014            assert_eq!(4, decoded.len());
1015            assert!(decoded[0] > 0x80);
1016            assert!(decoded[1] < 0x80);
1017            assert!(decoded[2] < 0x80);
1018            assert!(decoded[3] > 0x80);
1019        }
1020    }
1021
1022    #[test]
1023    fn jfif_header_density_check() {
1024        let mut buffer = Vec::new();
1025        build_jfif_header(&mut buffer, PixelDensity::dpi(300));
1026        assert_eq!(
1027            buffer,
1028            vec![
1029                b'J',
1030                b'F',
1031                b'I',
1032                b'F',
1033                0,
1034                1,
1035                2, // JFIF version 1.2
1036                1, // density is in dpi
1037                300u16.to_be_bytes()[0],
1038                300u16.to_be_bytes()[1],
1039                300u16.to_be_bytes()[0],
1040                300u16.to_be_bytes()[1],
1041                0,
1042                0, // No thumbnail
1043            ]
1044        );
1045    }
1046
1047    #[test]
1048    fn test_image_too_large() {
1049        // JPEG cannot encode images larger than 65,535×65,535
1050        // create a 65,536×1 8-bit black image buffer
1051        let img = [0; 65_536];
1052        // Try to encode an image that is too large
1053        let mut encoded = Vec::new();
1054        let encoder = JpegEncoder::new_with_quality(&mut encoded, 100);
1055        let result = encoder.write_image(&img, 65_536, 1, ExtendedColorType::L8);
1056        match result {
1057            Err(ImageError::Parameter(err)) => {
1058                assert_eq!(err.kind(), DimensionMismatch);
1059            }
1060            other => {
1061                panic!(
1062                    "Encoding an image that is too large should return a DimensionError \
1063                                it returned {other:?} instead"
1064                )
1065            }
1066        }
1067    }
1068
1069    #[test]
1070    fn test_build_jfif_header() {
1071        let mut buf = vec![];
1072        let density = PixelDensity::dpi(100);
1073        build_jfif_header(&mut buf, density);
1074        assert_eq!(
1075            buf,
1076            [0x4A, 0x46, 0x49, 0x46, 0x00, 0x01, 0x02, 0x01, 0, 100, 0, 100, 0, 0]
1077        );
1078    }
1079
1080    #[test]
1081    fn test_build_frame_header() {
1082        let mut buf = vec![];
1083        let components = vec![
1084            Component {
1085                id: 1,
1086                h: 1,
1087                v: 1,
1088                tq: 5,
1089                dc_table: 5,
1090                ac_table: 5,
1091                _dc_pred: 0,
1092            },
1093            Component {
1094                id: 2,
1095                h: 1,
1096                v: 1,
1097                tq: 4,
1098                dc_table: 4,
1099                ac_table: 4,
1100                _dc_pred: 0,
1101            },
1102        ];
1103        build_frame_header(&mut buf, 5, 100, 150, &components);
1104        assert_eq!(
1105            buf,
1106            [5, 0, 150, 0, 100, 2, 1, (1 << 4) | 1, 5, 2, (1 << 4) | 1, 4]
1107        );
1108    }
1109
1110    #[test]
1111    fn test_build_scan_header() {
1112        let mut buf = vec![];
1113        let components = vec![
1114            Component {
1115                id: 1,
1116                h: 1,
1117                v: 1,
1118                tq: 5,
1119                dc_table: 5,
1120                ac_table: 5,
1121                _dc_pred: 0,
1122            },
1123            Component {
1124                id: 2,
1125                h: 1,
1126                v: 1,
1127                tq: 4,
1128                dc_table: 4,
1129                ac_table: 4,
1130                _dc_pred: 0,
1131            },
1132        ];
1133        build_scan_header(&mut buf, &components);
1134        assert_eq!(buf, [2, 1, (5 << 4) | 5, 2, (4 << 4) | 4, 0, 63, 0]);
1135    }
1136
1137    #[test]
1138    fn test_build_huffman_segment() {
1139        let mut buf = vec![];
1140        build_huffman_segment(
1141            &mut buf,
1142            DCCLASS,
1143            LUMADESTINATION,
1144            &STD_LUMA_DC_CODE_LENGTHS,
1145            &STD_LUMA_DC_VALUES,
1146        );
1147        assert_eq!(
1148            buf,
1149            vec![
1150                0, 0, 1, 5, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
1151                10, 11
1152            ]
1153        );
1154    }
1155
1156    #[test]
1157    fn test_build_quantization_segment() {
1158        let mut buf = vec![];
1159        let qtable = [0u8; 64];
1160        build_quantization_segment(&mut buf, 8, 1, &qtable);
1161        let mut expected = vec![];
1162        expected.push(1);
1163        expected.extend_from_slice(&[0; 64]);
1164        assert_eq!(buf, expected);
1165    }
1166
1167    #[test]
1168    fn check_color_types() {
1169        const ALL: &[ColorType] = &[
1170            ColorType::L8,
1171            ColorType::L16,
1172            ColorType::La8,
1173            ColorType::Rgb8,
1174            ColorType::Rgba8,
1175            ColorType::La16,
1176            ColorType::Rgb16,
1177            ColorType::Rgba16,
1178            ColorType::Rgb32F,
1179            ColorType::Rgba32F,
1180        ];
1181
1182        for color in ALL {
1183            let image = DynamicImage::new(1, 1, *color);
1184
1185            image
1186                .write_to(&mut Cursor::new(vec![]), ImageFormat::Jpeg)
1187                .expect("supported or converted");
1188        }
1189    }
1190
1191    #[cfg(feature = "benchmarks")]
1192    #[bench]
1193    fn bench_jpeg_encoder_new(b: &mut Bencher) {
1194        b.iter(|| {
1195            let mut y = vec![];
1196            let _x = JpegEncoder::new(&mut y);
1197        });
1198    }
1199}
1200
1201// Tests regressions of `encode_image` against #1412, confusion about the subimage's position vs.
1202// dimensions. (We no longer have a position, four `u32` returns was confusing).
1203#[test]
1204fn sub_image_encoder_regression_1412() {
1205    let image = DynamicImage::new_rgb8(1280, 720);
1206    let subimg = crate::imageops::crop_imm(&image, 0, 358, 425, 361);
1207
1208    let mut encoded_crop = vec![];
1209    let mut encoder = JpegEncoder::new(&mut encoded_crop);
1210
1211    let result = encoder.encode_image(&*subimg);
1212    assert!(result.is_ok(), "Failed to encode subimage: {result:?}");
1213}