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
21static SOF0: u8 = 0xC0;
24static DHT: u8 = 0xC4;
26static SOI: u8 = 0xD8;
28static EOI: u8 = 0xD9;
30static SOS: u8 = 0xDA;
32static DQT: u8 = 0xDB;
34static APP0: u8 = 0xE0;
36static APP1: u8 = 0xE1;
37static APP2: u8 = 0xE2;
38
39#[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#[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
66static 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
79static 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
91static 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
113static 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#[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
157static EXIF_HEADER: [u8; 6] = [0x45, 0x78, 0x69, 0x66, 0x00, 0x00];
160
161#[derive(Copy, Clone)]
163struct Component {
164 id: u8,
166
167 h: u8,
169
170 v: u8,
172
173 tq: u8,
175
176 dc_table: u8,
178
179 ac_table: u8,
181
182 _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 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 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
293pub enum PixelDensityUnit {
294 PixelAspectRatio,
297
298 Inches,
300
301 Centimeters,
303}
304
305#[derive(Clone, Copy, Debug, Eq, PartialEq)]
315pub struct PixelDensity {
316 pub density: (u16, u16),
318 pub unit: PixelDensityUnit,
320}
321
322impl PixelDensity {
323 #[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 fn default() -> Self {
338 PixelDensity {
339 density: (1, 1),
340 unit: PixelDensityUnit::PixelAspectRatio,
341 }
342 }
343}
344
345pub 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 pub fn new(w: W) -> JpegEncoder<W> {
366 JpegEncoder::new_with_quality(w, 75)
367 }
368
369 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 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 pub fn set_pixel_density(&mut self, pixel_density: PixelDensity) {
440 self.pixel_density = pixel_density;
441 }
442
443 #[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 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 self.write_icc_profile_chunks()?;
525
526 build_frame_header(
527 &mut buf,
528 8,
529 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 transform::fdct(&yblock, &mut dct_yblock);
617
618 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 copy_blocks_ycbcr(image, x, y, &mut yblock, &mut cb_block, &mut cr_block);
650
651 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 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 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 const C_YR: i32 = 19595; const C_YG: i32 = 38469; const C_YB: i32 = 7471; const Y_ROUNDING: i32 = (1 << 15) - 1; const C_UR: i32 = 11059; const C_UG: i32 = 21709; const C_UB: i32 = 32768; const UV_BIAS_ROUNDING: i32 = (128 * (1 << 16)) + ((1 << 15) - 1); const C_VR: i32 = C_UB; const C_VG: i32 = 27439; const C_VB: i32 = 5329; 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#[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 let img = [255u8, 0, 0];
973
974 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 {
985 let decoded = decode(&encoded_img);
986 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 let img = [255u8, 0, 0, 255];
999
1000 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 {
1011 let decoded = decode(&encoded_img);
1012 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, 1, 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, ]
1044 );
1045 }
1046
1047 #[test]
1048 fn test_image_too_large() {
1049 let img = [0; 65_536];
1052 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#[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}