1 | use crate::error::{Error, Result, UnsupportedFeature}; |
2 | use crate::huffman::{fill_default_mjpeg_tables, HuffmanDecoder, HuffmanTable}; |
3 | use crate::marker::Marker; |
4 | use crate::parser::{ |
5 | parse_app, parse_com, parse_dht, parse_dqt, parse_dri, parse_sof, parse_sos, |
6 | AdobeColorTransform, AppData, CodingProcess, Component, Dimensions, EntropyCoding, FrameInfo, |
7 | IccChunk, ScanInfo, |
8 | }; |
9 | use crate::read_u8; |
10 | use crate::upsampler::Upsampler; |
11 | use crate::worker::{compute_image_parallel, PreferWorkerKind, RowData, Worker, WorkerScope}; |
12 | use alloc::borrow::ToOwned; |
13 | use alloc::sync::Arc; |
14 | use alloc::vec::Vec; |
15 | use alloc::{format, vec}; |
16 | use core::cmp; |
17 | use core::mem; |
18 | use core::ops::Range; |
19 | use std::io::Read; |
20 | |
21 | pub const MAX_COMPONENTS: usize = 4; |
22 | |
23 | mod lossless; |
24 | use self::lossless::compute_image_lossless; |
25 | |
26 | #[rustfmt::skip] |
27 | static UNZIGZAG: [u8; 64] = [ |
28 | 0, 1, 8, 16, 9, 2, 3, 10, |
29 | 17, 24, 32, 25, 18, 11, 4, 5, |
30 | 12, 19, 26, 33, 40, 48, 41, 34, |
31 | 27, 20, 13, 6, 7, 14, 21, 28, |
32 | 35, 42, 49, 56, 57, 50, 43, 36, |
33 | 29, 22, 15, 23, 30, 37, 44, 51, |
34 | 58, 59, 52, 45, 38, 31, 39, 46, |
35 | 53, 60, 61, 54, 47, 55, 62, 63, |
36 | ]; |
37 | |
38 | /// An enumeration over combinations of color spaces and bit depths a pixel can have. |
39 | #[derive (Clone, Copy, Debug, PartialEq)] |
40 | pub enum PixelFormat { |
41 | /// Luminance (grayscale), 8 bits |
42 | L8, |
43 | /// Luminance (grayscale), 16 bits |
44 | L16, |
45 | /// RGB, 8 bits per channel |
46 | RGB24, |
47 | /// CMYK, 8 bits per channel |
48 | CMYK32, |
49 | } |
50 | |
51 | impl PixelFormat { |
52 | /// Determine the size in bytes of each pixel in this format |
53 | pub fn pixel_bytes(&self) -> usize { |
54 | match self { |
55 | PixelFormat::L8 => 1, |
56 | PixelFormat::L16 => 2, |
57 | PixelFormat::RGB24 => 3, |
58 | PixelFormat::CMYK32 => 4, |
59 | } |
60 | } |
61 | } |
62 | |
63 | /// Represents metadata of an image. |
64 | #[derive (Clone, Copy, Debug, PartialEq)] |
65 | pub struct ImageInfo { |
66 | /// The width of the image, in pixels. |
67 | pub width: u16, |
68 | /// The height of the image, in pixels. |
69 | pub height: u16, |
70 | /// The pixel format of the image. |
71 | pub pixel_format: PixelFormat, |
72 | /// The coding process of the image. |
73 | pub coding_process: CodingProcess, |
74 | } |
75 | |
76 | /// Describes the colour transform to apply before binary data is returned |
77 | #[derive (Debug, Clone, Copy, PartialEq, Eq, Hash)] |
78 | #[non_exhaustive ] |
79 | pub enum ColorTransform { |
80 | /// No transform should be applied and the data is returned as-is. |
81 | None, |
82 | /// Unknown colour transformation |
83 | Unknown, |
84 | /// Grayscale transform should be applied (expects 1 channel) |
85 | Grayscale, |
86 | /// RGB transform should be applied. |
87 | RGB, |
88 | /// YCbCr transform should be applied. |
89 | YCbCr, |
90 | /// CMYK transform should be applied. |
91 | CMYK, |
92 | /// YCCK transform should be applied. |
93 | YCCK, |
94 | /// big gamut Y/Cb/Cr, bg-sYCC |
95 | JcsBgYcc, |
96 | /// big gamut red/green/blue, bg-sRGB |
97 | JcsBgRgb, |
98 | } |
99 | |
100 | /// JPEG decoder |
101 | pub struct Decoder<R> { |
102 | reader: R, |
103 | |
104 | frame: Option<FrameInfo>, |
105 | dc_huffman_tables: Vec<Option<HuffmanTable>>, |
106 | ac_huffman_tables: Vec<Option<HuffmanTable>>, |
107 | quantization_tables: [Option<Arc<[u16; 64]>>; 4], |
108 | |
109 | restart_interval: u16, |
110 | |
111 | adobe_color_transform: Option<AdobeColorTransform>, |
112 | color_transform: Option<ColorTransform>, |
113 | |
114 | is_jfif: bool, |
115 | is_mjpeg: bool, |
116 | |
117 | icc_markers: Vec<IccChunk>, |
118 | |
119 | exif_data: Option<Vec<u8>>, |
120 | xmp_data: Option<Vec<u8>>, |
121 | psir_data: Option<Vec<u8>>, |
122 | |
123 | // Used for progressive JPEGs. |
124 | coefficients: Vec<Vec<i16>>, |
125 | // Bitmask of which coefficients has been completely decoded. |
126 | coefficients_finished: [u64; MAX_COMPONENTS], |
127 | |
128 | // Maximum allowed size of decoded image buffer |
129 | decoding_buffer_size_limit: usize, |
130 | } |
131 | |
132 | impl<R: Read> Decoder<R> { |
133 | /// Creates a new `Decoder` using the reader `reader`. |
134 | pub fn new(reader: R) -> Decoder<R> { |
135 | Decoder { |
136 | reader, |
137 | frame: None, |
138 | dc_huffman_tables: vec![None, None, None, None], |
139 | ac_huffman_tables: vec![None, None, None, None], |
140 | quantization_tables: [None, None, None, None], |
141 | restart_interval: 0, |
142 | adobe_color_transform: None, |
143 | color_transform: None, |
144 | is_jfif: false, |
145 | is_mjpeg: false, |
146 | icc_markers: Vec::new(), |
147 | exif_data: None, |
148 | xmp_data: None, |
149 | psir_data: None, |
150 | coefficients: Vec::new(), |
151 | coefficients_finished: [0; MAX_COMPONENTS], |
152 | decoding_buffer_size_limit: usize::MAX, |
153 | } |
154 | } |
155 | |
156 | /// Colour transform to use when decoding the image. App segments relating to colour transforms |
157 | /// will be ignored. |
158 | pub fn set_color_transform(&mut self, transform: ColorTransform) { |
159 | self.color_transform = Some(transform); |
160 | } |
161 | |
162 | /// Set maximum buffer size allowed for decoded images |
163 | pub fn set_max_decoding_buffer_size(&mut self, max: usize) { |
164 | self.decoding_buffer_size_limit = max; |
165 | } |
166 | |
167 | /// Returns metadata about the image. |
168 | /// |
169 | /// The returned value will be `None` until a call to either `read_info` or `decode` has |
170 | /// returned `Ok`. |
171 | pub fn info(&self) -> Option<ImageInfo> { |
172 | match self.frame { |
173 | Some(ref frame) => { |
174 | let pixel_format = match frame.components.len() { |
175 | 1 => match frame.precision { |
176 | 2..=8 => PixelFormat::L8, |
177 | 9..=16 => PixelFormat::L16, |
178 | _ => panic!(), |
179 | }, |
180 | 3 => PixelFormat::RGB24, |
181 | 4 => PixelFormat::CMYK32, |
182 | _ => panic!(), |
183 | }; |
184 | |
185 | Some(ImageInfo { |
186 | width: frame.output_size.width, |
187 | height: frame.output_size.height, |
188 | pixel_format, |
189 | coding_process: frame.coding_process, |
190 | }) |
191 | } |
192 | None => None, |
193 | } |
194 | } |
195 | |
196 | /// Returns raw exif data, starting at the TIFF header, if the image contains any. |
197 | /// |
198 | /// The returned value will be `None` until a call to `decode` has returned `Ok`. |
199 | pub fn exif_data(&self) -> Option<&[u8]> { |
200 | self.exif_data.as_deref() |
201 | } |
202 | |
203 | /// Returns the raw XMP packet if there is any. |
204 | /// |
205 | /// The returned value will be `None` until a call to `decode` has returned `Ok`. |
206 | pub fn xmp_data(&self) -> Option<&[u8]> { |
207 | self.xmp_data.as_deref() |
208 | } |
209 | |
210 | /// Returns the embeded icc profile if the image contains one. |
211 | pub fn icc_profile(&self) -> Option<Vec<u8>> { |
212 | let mut marker_present: [Option<&IccChunk>; 256] = [None; 256]; |
213 | let num_markers = self.icc_markers.len(); |
214 | if num_markers == 0 || num_markers >= 255 { |
215 | return None; |
216 | } |
217 | // check the validity of the markers |
218 | for chunk in &self.icc_markers { |
219 | if usize::from(chunk.num_markers) != num_markers { |
220 | // all the lengths must match |
221 | return None; |
222 | } |
223 | if chunk.seq_no == 0 { |
224 | return None; |
225 | } |
226 | if marker_present[usize::from(chunk.seq_no)].is_some() { |
227 | // duplicate seq_no |
228 | return None; |
229 | } else { |
230 | marker_present[usize::from(chunk.seq_no)] = Some(chunk); |
231 | } |
232 | } |
233 | |
234 | // assemble them together by seq_no failing if any are missing |
235 | let mut data = Vec::new(); |
236 | // seq_no's start at 1 |
237 | for &chunk in marker_present.get(1..=num_markers)? { |
238 | data.extend_from_slice(&chunk?.data); |
239 | } |
240 | Some(data) |
241 | } |
242 | |
243 | /// Heuristic to avoid starting thread, synchronization if we expect a small amount of |
244 | /// parallelism to be utilized. |
245 | fn select_worker(frame: &FrameInfo, worker_preference: PreferWorkerKind) -> PreferWorkerKind { |
246 | const PARALLELISM_THRESHOLD: u64 = 128 * 128; |
247 | |
248 | match worker_preference { |
249 | PreferWorkerKind::Immediate => PreferWorkerKind::Immediate, |
250 | PreferWorkerKind::Multithreaded => { |
251 | let width: u64 = frame.output_size.width.into(); |
252 | let height: u64 = frame.output_size.width.into(); |
253 | if width * height > PARALLELISM_THRESHOLD { |
254 | PreferWorkerKind::Multithreaded |
255 | } else { |
256 | PreferWorkerKind::Immediate |
257 | } |
258 | } |
259 | } |
260 | } |
261 | |
262 | /// Tries to read metadata from the image without decoding it. |
263 | /// |
264 | /// If successful, the metadata can be obtained using the `info` method. |
265 | pub fn read_info(&mut self) -> Result<()> { |
266 | WorkerScope::with(|worker| self.decode_internal(true, worker)).map(|_| ()) |
267 | } |
268 | |
269 | /// Configure the decoder to scale the image during decoding. |
270 | /// |
271 | /// This efficiently scales the image by the smallest supported scale |
272 | /// factor that produces an image larger than or equal to the requested |
273 | /// size in at least one axis. The currently implemented scale factors |
274 | /// are 1/8, 1/4, 1/2 and 1. |
275 | /// |
276 | /// To generate a thumbnail of an exact size, pass the desired size and |
277 | /// then scale to the final size using a traditional resampling algorithm. |
278 | pub fn scale(&mut self, requested_width: u16, requested_height: u16) -> Result<(u16, u16)> { |
279 | self.read_info()?; |
280 | let frame = self.frame.as_mut().unwrap(); |
281 | let idct_size = crate::idct::choose_idct_size( |
282 | frame.image_size, |
283 | Dimensions { |
284 | width: requested_width, |
285 | height: requested_height, |
286 | }, |
287 | ); |
288 | frame.update_idct_size(idct_size)?; |
289 | Ok((frame.output_size.width, frame.output_size.height)) |
290 | } |
291 | |
292 | /// Decodes the image and returns the decoded pixels if successful. |
293 | pub fn decode(&mut self) -> Result<Vec<u8>> { |
294 | WorkerScope::with(|worker| self.decode_internal(false, worker)) |
295 | } |
296 | |
297 | fn decode_internal( |
298 | &mut self, |
299 | stop_after_metadata: bool, |
300 | worker_scope: &WorkerScope, |
301 | ) -> Result<Vec<u8>> { |
302 | if stop_after_metadata && self.frame.is_some() { |
303 | // The metadata has already been read. |
304 | return Ok(Vec::new()); |
305 | } else if self.frame.is_none() |
306 | && (read_u8(&mut self.reader)? != 0xFF |
307 | || Marker::from_u8(read_u8(&mut self.reader)?) != Some(Marker::SOI)) |
308 | { |
309 | return Err(Error::Format( |
310 | "first two bytes are not an SOI marker" .to_owned(), |
311 | )); |
312 | } |
313 | |
314 | let mut previous_marker = Marker::SOI; |
315 | let mut pending_marker = None; |
316 | let mut scans_processed = 0; |
317 | let mut planes = vec![ |
318 | Vec::<u8>::new(); |
319 | self.frame |
320 | .as_ref() |
321 | .map_or(0, |frame| frame.components.len()) |
322 | ]; |
323 | let mut planes_u16 = vec![ |
324 | Vec::<u16>::new(); |
325 | self.frame |
326 | .as_ref() |
327 | .map_or(0, |frame| frame.components.len()) |
328 | ]; |
329 | |
330 | loop { |
331 | let marker = match pending_marker.take() { |
332 | Some(m) => m, |
333 | None => self.read_marker()?, |
334 | }; |
335 | |
336 | match marker { |
337 | // Frame header |
338 | Marker::SOF(..) => { |
339 | // Section 4.10 |
340 | // "An image contains only one frame in the cases of sequential and |
341 | // progressive coding processes; an image contains multiple frames for the |
342 | // hierarchical mode." |
343 | if self.frame.is_some() { |
344 | return Err(Error::Unsupported(UnsupportedFeature::Hierarchical)); |
345 | } |
346 | |
347 | let frame = parse_sof(&mut self.reader, marker)?; |
348 | let component_count = frame.components.len(); |
349 | |
350 | if frame.is_differential { |
351 | return Err(Error::Unsupported(UnsupportedFeature::Hierarchical)); |
352 | } |
353 | if frame.entropy_coding == EntropyCoding::Arithmetic { |
354 | return Err(Error::Unsupported( |
355 | UnsupportedFeature::ArithmeticEntropyCoding, |
356 | )); |
357 | } |
358 | if frame.precision != 8 && frame.coding_process != CodingProcess::Lossless { |
359 | return Err(Error::Unsupported(UnsupportedFeature::SamplePrecision( |
360 | frame.precision, |
361 | ))); |
362 | } |
363 | if !(2..=16).contains(&frame.precision) { |
364 | return Err(Error::Unsupported(UnsupportedFeature::SamplePrecision( |
365 | frame.precision, |
366 | ))); |
367 | } |
368 | if component_count != 1 && component_count != 3 && component_count != 4 { |
369 | return Err(Error::Unsupported(UnsupportedFeature::ComponentCount( |
370 | component_count as u8, |
371 | ))); |
372 | } |
373 | |
374 | // Make sure we support the subsampling ratios used. |
375 | let _ = Upsampler::new( |
376 | &frame.components, |
377 | frame.image_size.width, |
378 | frame.image_size.height, |
379 | )?; |
380 | |
381 | self.frame = Some(frame); |
382 | |
383 | if stop_after_metadata { |
384 | return Ok(Vec::new()); |
385 | } |
386 | |
387 | planes = vec![Vec::new(); component_count]; |
388 | planes_u16 = vec![Vec::new(); component_count]; |
389 | } |
390 | |
391 | // Scan header |
392 | Marker::SOS => { |
393 | if self.frame.is_none() { |
394 | return Err(Error::Format("scan encountered before frame" .to_owned())); |
395 | } |
396 | |
397 | let frame = self.frame.clone().unwrap(); |
398 | let scan = parse_sos(&mut self.reader, &frame)?; |
399 | |
400 | if frame.coding_process == CodingProcess::DctProgressive |
401 | && self.coefficients.is_empty() |
402 | { |
403 | self.coefficients = frame |
404 | .components |
405 | .iter() |
406 | .map(|c| { |
407 | let block_count = |
408 | c.block_size.width as usize * c.block_size.height as usize; |
409 | vec![0; block_count * 64] |
410 | }) |
411 | .collect(); |
412 | } |
413 | |
414 | if frame.coding_process == CodingProcess::Lossless { |
415 | let (marker, data) = self.decode_scan_lossless(&frame, &scan)?; |
416 | |
417 | for (i, plane) in data |
418 | .into_iter() |
419 | .enumerate() |
420 | .filter(|(_, plane)| !plane.is_empty()) |
421 | { |
422 | planes_u16[i] = plane; |
423 | } |
424 | pending_marker = marker; |
425 | } else { |
426 | // This was previously buggy, so let's explain the log here a bit. When a |
427 | // progressive frame is encoded then the coefficients (DC, AC) of each |
428 | // component (=color plane) can be split amongst scans. In particular it can |
429 | // happen or at least occurs in the wild that a scan contains coefficient 0 of |
430 | // all components. If now one but not all components had all other coefficients |
431 | // delivered in previous scans then such a scan contains all components but |
432 | // completes only some of them! (This is technically NOT permitted for all |
433 | // other coefficients as the standard dictates that scans with coefficients |
434 | // other than the 0th must only contain ONE component so we would either |
435 | // complete it or not. We may want to detect and error in case more component |
436 | // are part of a scan than allowed.) What a weird edge case. |
437 | // |
438 | // But this means we track precisely which components get completed here. |
439 | let mut finished = [false; MAX_COMPONENTS]; |
440 | |
441 | if scan.successive_approximation_low == 0 { |
442 | for (&i, component_finished) in |
443 | scan.component_indices.iter().zip(&mut finished) |
444 | { |
445 | if self.coefficients_finished[i] == !0 { |
446 | continue; |
447 | } |
448 | for j in scan.spectral_selection.clone() { |
449 | self.coefficients_finished[i] |= 1 << j; |
450 | } |
451 | if self.coefficients_finished[i] == !0 { |
452 | *component_finished = true; |
453 | } |
454 | } |
455 | } |
456 | |
457 | let preference = |
458 | Self::select_worker(&frame, PreferWorkerKind::Multithreaded); |
459 | |
460 | let (marker, data) = worker_scope |
461 | .get_or_init_worker(preference, |worker| { |
462 | self.decode_scan(&frame, &scan, worker, &finished) |
463 | })?; |
464 | |
465 | if let Some(data) = data { |
466 | for (i, plane) in data |
467 | .into_iter() |
468 | .enumerate() |
469 | .filter(|(_, plane)| !plane.is_empty()) |
470 | { |
471 | if self.coefficients_finished[i] == !0 { |
472 | planes[i] = plane; |
473 | } |
474 | } |
475 | } |
476 | |
477 | pending_marker = marker; |
478 | } |
479 | |
480 | scans_processed += 1; |
481 | } |
482 | |
483 | // Table-specification and miscellaneous markers |
484 | // Quantization table-specification |
485 | Marker::DQT => { |
486 | let tables = parse_dqt(&mut self.reader)?; |
487 | |
488 | for (i, &table) in tables.iter().enumerate() { |
489 | if let Some(table) = table { |
490 | let mut unzigzagged_table = [0u16; 64]; |
491 | |
492 | for j in 0..64 { |
493 | unzigzagged_table[UNZIGZAG[j] as usize] = table[j]; |
494 | } |
495 | |
496 | self.quantization_tables[i] = Some(Arc::new(unzigzagged_table)); |
497 | } |
498 | } |
499 | } |
500 | // Huffman table-specification |
501 | Marker::DHT => { |
502 | let is_baseline = self.frame.as_ref().map(|frame| frame.is_baseline); |
503 | let (dc_tables, ac_tables) = parse_dht(&mut self.reader, is_baseline)?; |
504 | |
505 | let current_dc_tables = mem::take(&mut self.dc_huffman_tables); |
506 | self.dc_huffman_tables = dc_tables |
507 | .into_iter() |
508 | .zip(current_dc_tables) |
509 | .map(|(a, b)| a.or(b)) |
510 | .collect(); |
511 | |
512 | let current_ac_tables = mem::take(&mut self.ac_huffman_tables); |
513 | self.ac_huffman_tables = ac_tables |
514 | .into_iter() |
515 | .zip(current_ac_tables) |
516 | .map(|(a, b)| a.or(b)) |
517 | .collect(); |
518 | } |
519 | // Arithmetic conditioning table-specification |
520 | Marker::DAC => { |
521 | return Err(Error::Unsupported( |
522 | UnsupportedFeature::ArithmeticEntropyCoding, |
523 | )) |
524 | } |
525 | // Restart interval definition |
526 | Marker::DRI => self.restart_interval = parse_dri(&mut self.reader)?, |
527 | // Comment |
528 | Marker::COM => { |
529 | let _comment = parse_com(&mut self.reader)?; |
530 | } |
531 | // Application data |
532 | Marker::APP(..) => { |
533 | if let Some(data) = parse_app(&mut self.reader, marker)? { |
534 | match data { |
535 | AppData::Adobe(color_transform) => { |
536 | self.adobe_color_transform = Some(color_transform) |
537 | } |
538 | AppData::Jfif => { |
539 | // From the JFIF spec: |
540 | // "The APP0 marker is used to identify a JPEG FIF file. |
541 | // The JPEG FIF APP0 marker is mandatory right after the SOI marker." |
542 | // Some JPEGs in the wild does not follow this though, so we allow |
543 | // JFIF headers anywhere APP0 markers are allowed. |
544 | /* |
545 | if previous_marker != Marker::SOI { |
546 | return Err(Error::Format("the JFIF APP0 marker must come right after the SOI marker".to_owned())); |
547 | } |
548 | */ |
549 | |
550 | self.is_jfif = true; |
551 | } |
552 | AppData::Avi1 => self.is_mjpeg = true, |
553 | AppData::Icc(icc) => self.icc_markers.push(icc), |
554 | AppData::Exif(data) => self.exif_data = Some(data), |
555 | AppData::Xmp(data) => self.xmp_data = Some(data), |
556 | AppData::Psir(data) => self.psir_data = Some(data), |
557 | } |
558 | } |
559 | } |
560 | // Restart |
561 | Marker::RST(..) => { |
562 | // Some encoders emit a final RST marker after entropy-coded data, which |
563 | // decode_scan does not take care of. So if we encounter one, we ignore it. |
564 | if previous_marker != Marker::SOS { |
565 | return Err(Error::Format( |
566 | "RST found outside of entropy-coded data" .to_owned(), |
567 | )); |
568 | } |
569 | } |
570 | |
571 | // Define number of lines |
572 | Marker::DNL => { |
573 | // Section B.2.1 |
574 | // "If a DNL segment (see B.2.5) is present, it shall immediately follow the first scan." |
575 | if previous_marker != Marker::SOS || scans_processed != 1 { |
576 | return Err(Error::Format( |
577 | "DNL is only allowed immediately after the first scan" .to_owned(), |
578 | )); |
579 | } |
580 | |
581 | return Err(Error::Unsupported(UnsupportedFeature::DNL)); |
582 | } |
583 | |
584 | // Hierarchical mode markers |
585 | Marker::DHP | Marker::EXP => { |
586 | return Err(Error::Unsupported(UnsupportedFeature::Hierarchical)) |
587 | } |
588 | |
589 | // End of image |
590 | Marker::EOI => break, |
591 | |
592 | _ => { |
593 | return Err(Error::Format(format!( |
594 | " {:?} marker found where not allowed" , |
595 | marker |
596 | ))) |
597 | } |
598 | } |
599 | |
600 | previous_marker = marker; |
601 | } |
602 | |
603 | if self.frame.is_none() { |
604 | return Err(Error::Format( |
605 | "end of image encountered before frame" .to_owned(), |
606 | )); |
607 | } |
608 | |
609 | let frame = self.frame.as_ref().unwrap(); |
610 | let preference = Self::select_worker(frame, PreferWorkerKind::Multithreaded); |
611 | |
612 | worker_scope.get_or_init_worker(preference, |worker| { |
613 | self.decode_planes(worker, planes, planes_u16) |
614 | }) |
615 | } |
616 | |
617 | fn decode_planes( |
618 | &mut self, |
619 | worker: &mut dyn Worker, |
620 | mut planes: Vec<Vec<u8>>, |
621 | planes_u16: Vec<Vec<u16>>, |
622 | ) -> Result<Vec<u8>> { |
623 | if self.frame.is_none() { |
624 | return Err(Error::Format( |
625 | "end of image encountered before frame" .to_owned(), |
626 | )); |
627 | } |
628 | |
629 | let frame = self.frame.as_ref().unwrap(); |
630 | |
631 | if frame |
632 | .components |
633 | .len() |
634 | .checked_mul(frame.output_size.width.into()) |
635 | .and_then(|m| m.checked_mul(frame.output_size.height.into())) |
636 | .map_or(true, |m| self.decoding_buffer_size_limit < m) |
637 | { |
638 | return Err(Error::Format( |
639 | "size of decoded image exceeds maximum allowed size" .to_owned(), |
640 | )); |
641 | } |
642 | |
643 | // If we're decoding a progressive jpeg and a component is unfinished, render what we've got |
644 | if frame.coding_process == CodingProcess::DctProgressive |
645 | && self.coefficients.len() == frame.components.len() |
646 | { |
647 | for (i, component) in frame.components.iter().enumerate() { |
648 | // Only dealing with unfinished components |
649 | if self.coefficients_finished[i] == !0 { |
650 | continue; |
651 | } |
652 | |
653 | let quantization_table = |
654 | match self.quantization_tables[component.quantization_table_index].clone() { |
655 | Some(quantization_table) => quantization_table, |
656 | None => continue, |
657 | }; |
658 | |
659 | // Get the worker prepared |
660 | let row_data = RowData { |
661 | index: i, |
662 | component: component.clone(), |
663 | quantization_table, |
664 | }; |
665 | worker.start(row_data)?; |
666 | |
667 | // Send the rows over to the worker and collect the result |
668 | let coefficients_per_mcu_row = usize::from(component.block_size.width) |
669 | * usize::from(component.vertical_sampling_factor) |
670 | * 64; |
671 | |
672 | let mut tasks = (0..frame.mcu_size.height).map(|mcu_y| { |
673 | let offset = usize::from(mcu_y) * coefficients_per_mcu_row; |
674 | let row_coefficients = |
675 | self.coefficients[i][offset..offset + coefficients_per_mcu_row].to_vec(); |
676 | (i, row_coefficients) |
677 | }); |
678 | |
679 | // FIXME: additional potential work stealing opportunities for rayon case if we |
680 | // also internally can parallelize over components. |
681 | worker.append_rows(&mut tasks)?; |
682 | planes[i] = worker.get_result(i)?; |
683 | } |
684 | } |
685 | |
686 | if frame.coding_process == CodingProcess::Lossless { |
687 | compute_image_lossless(frame, planes_u16) |
688 | } else { |
689 | compute_image( |
690 | &frame.components, |
691 | planes, |
692 | frame.output_size, |
693 | self.determine_color_transform(), |
694 | ) |
695 | } |
696 | } |
697 | |
698 | fn determine_color_transform(&self) -> ColorTransform { |
699 | if let Some(color_transform) = self.color_transform { |
700 | return color_transform; |
701 | } |
702 | |
703 | let frame = self.frame.as_ref().unwrap(); |
704 | |
705 | if frame.components.len() == 1 { |
706 | return ColorTransform::Grayscale; |
707 | } |
708 | |
709 | // Using logic for determining colour as described here: https://entropymine.wordpress.com/2018/10/22/how-is-a-jpeg-images-color-type-determined/ |
710 | |
711 | if frame.components.len() == 3 { |
712 | match ( |
713 | frame.components[0].identifier, |
714 | frame.components[1].identifier, |
715 | frame.components[2].identifier, |
716 | ) { |
717 | (1, 2, 3) => { |
718 | return ColorTransform::YCbCr; |
719 | } |
720 | (1, 34, 35) => { |
721 | return ColorTransform::JcsBgYcc; |
722 | } |
723 | (82, 71, 66) => { |
724 | return ColorTransform::RGB; |
725 | } |
726 | (114, 103, 98) => { |
727 | return ColorTransform::JcsBgRgb; |
728 | } |
729 | _ => {} |
730 | } |
731 | |
732 | if self.is_jfif { |
733 | return ColorTransform::YCbCr; |
734 | } |
735 | } |
736 | |
737 | if let Some(colour_transform) = self.adobe_color_transform { |
738 | match colour_transform { |
739 | AdobeColorTransform::Unknown => { |
740 | if frame.components.len() == 3 { |
741 | return ColorTransform::RGB; |
742 | } else if frame.components.len() == 4 { |
743 | return ColorTransform::CMYK; |
744 | } |
745 | } |
746 | AdobeColorTransform::YCbCr => { |
747 | return ColorTransform::YCbCr; |
748 | } |
749 | AdobeColorTransform::YCCK => { |
750 | return ColorTransform::YCCK; |
751 | } |
752 | } |
753 | } else if frame.components.len() == 4 { |
754 | return ColorTransform::CMYK; |
755 | } |
756 | |
757 | if frame.components.len() == 4 { |
758 | ColorTransform::YCCK |
759 | } else if frame.components.len() == 3 { |
760 | ColorTransform::YCbCr |
761 | } else { |
762 | ColorTransform::Unknown |
763 | } |
764 | } |
765 | |
766 | fn read_marker(&mut self) -> Result<Marker> { |
767 | loop { |
768 | // This should be an error as the JPEG spec doesn't allow extraneous data between marker segments. |
769 | // libjpeg allows this though and there are images in the wild utilising it, so we are |
770 | // forced to support this behavior. |
771 | // Sony Ericsson P990i is an example of a device which produce this sort of JPEGs. |
772 | while read_u8(&mut self.reader)? != 0xFF {} |
773 | |
774 | // Section B.1.1.2 |
775 | // All markers are assigned two-byte codes: an X’FF’ byte followed by a |
776 | // byte which is not equal to 0 or X’FF’ (see Table B.1). Any marker may |
777 | // optionally be preceded by any number of fill bytes, which are bytes |
778 | // assigned code X’FF’. |
779 | let mut byte = read_u8(&mut self.reader)?; |
780 | |
781 | // Section B.1.1.2 |
782 | // "Any marker may optionally be preceded by any number of fill bytes, which are bytes assigned code X’FF’." |
783 | while byte == 0xFF { |
784 | byte = read_u8(&mut self.reader)?; |
785 | } |
786 | |
787 | if byte != 0x00 && byte != 0xFF { |
788 | return Ok(Marker::from_u8(byte).unwrap()); |
789 | } |
790 | } |
791 | } |
792 | |
793 | #[allow (clippy::type_complexity)] |
794 | fn decode_scan( |
795 | &mut self, |
796 | frame: &FrameInfo, |
797 | scan: &ScanInfo, |
798 | worker: &mut dyn Worker, |
799 | finished: &[bool; MAX_COMPONENTS], |
800 | ) -> Result<(Option<Marker>, Option<Vec<Vec<u8>>>)> { |
801 | assert!(scan.component_indices.len() <= MAX_COMPONENTS); |
802 | |
803 | let components: Vec<Component> = scan |
804 | .component_indices |
805 | .iter() |
806 | .map(|&i| frame.components[i].clone()) |
807 | .collect(); |
808 | |
809 | // Verify that all required quantization tables has been set. |
810 | if components |
811 | .iter() |
812 | .any(|component| self.quantization_tables[component.quantization_table_index].is_none()) |
813 | { |
814 | return Err(Error::Format("use of unset quantization table" .to_owned())); |
815 | } |
816 | |
817 | if self.is_mjpeg { |
818 | fill_default_mjpeg_tables( |
819 | scan, |
820 | &mut self.dc_huffman_tables, |
821 | &mut self.ac_huffman_tables, |
822 | ); |
823 | } |
824 | |
825 | // Verify that all required huffman tables has been set. |
826 | if scan.spectral_selection.start == 0 |
827 | && scan |
828 | .dc_table_indices |
829 | .iter() |
830 | .any(|&i| self.dc_huffman_tables[i].is_none()) |
831 | { |
832 | return Err(Error::Format( |
833 | "scan makes use of unset dc huffman table" .to_owned(), |
834 | )); |
835 | } |
836 | if scan.spectral_selection.end > 1 |
837 | && scan |
838 | .ac_table_indices |
839 | .iter() |
840 | .any(|&i| self.ac_huffman_tables[i].is_none()) |
841 | { |
842 | return Err(Error::Format( |
843 | "scan makes use of unset ac huffman table" .to_owned(), |
844 | )); |
845 | } |
846 | |
847 | // Prepare the worker thread for the work to come. |
848 | for (i, component) in components.iter().enumerate() { |
849 | if finished[i] { |
850 | let row_data = RowData { |
851 | index: i, |
852 | component: component.clone(), |
853 | quantization_table: self.quantization_tables |
854 | [component.quantization_table_index] |
855 | .clone() |
856 | .unwrap(), |
857 | }; |
858 | |
859 | worker.start(row_data)?; |
860 | } |
861 | } |
862 | |
863 | let is_progressive = frame.coding_process == CodingProcess::DctProgressive; |
864 | let is_interleaved = components.len() > 1; |
865 | let mut dummy_block = [0i16; 64]; |
866 | let mut huffman = HuffmanDecoder::new(); |
867 | let mut dc_predictors = [0i16; MAX_COMPONENTS]; |
868 | let mut mcus_left_until_restart = self.restart_interval; |
869 | let mut expected_rst_num = 0; |
870 | let mut eob_run = 0; |
871 | let mut mcu_row_coefficients = vec![vec![]; components.len()]; |
872 | |
873 | if !is_progressive { |
874 | for (i, component) in components.iter().enumerate().filter(|&(i, _)| finished[i]) { |
875 | let coefficients_per_mcu_row = component.block_size.width as usize |
876 | * component.vertical_sampling_factor as usize |
877 | * 64; |
878 | mcu_row_coefficients[i] = vec![0i16; coefficients_per_mcu_row]; |
879 | } |
880 | } |
881 | |
882 | // 4.8.2 |
883 | // When reading from the stream, if the data is non-interleaved then an MCU consists of |
884 | // exactly one block (effectively a 1x1 sample). |
885 | let (mcu_horizontal_samples, mcu_vertical_samples) = if is_interleaved { |
886 | let horizontal = components |
887 | .iter() |
888 | .map(|component| component.horizontal_sampling_factor as u16) |
889 | .collect::<Vec<_>>(); |
890 | let vertical = components |
891 | .iter() |
892 | .map(|component| component.vertical_sampling_factor as u16) |
893 | .collect::<Vec<_>>(); |
894 | (horizontal, vertical) |
895 | } else { |
896 | (vec![1], vec![1]) |
897 | }; |
898 | |
899 | // This also affects how many MCU values we read from stream. If it's a non-interleaved stream, |
900 | // the MCUs will be exactly the block count. |
901 | let (max_mcu_x, max_mcu_y) = if is_interleaved { |
902 | (frame.mcu_size.width, frame.mcu_size.height) |
903 | } else { |
904 | ( |
905 | components[0].block_size.width, |
906 | components[0].block_size.height, |
907 | ) |
908 | }; |
909 | |
910 | for mcu_y in 0..max_mcu_y { |
911 | if mcu_y * 8 >= frame.image_size.height { |
912 | break; |
913 | } |
914 | |
915 | for mcu_x in 0..max_mcu_x { |
916 | if mcu_x * 8 >= frame.image_size.width { |
917 | break; |
918 | } |
919 | |
920 | if self.restart_interval > 0 { |
921 | if mcus_left_until_restart == 0 { |
922 | match huffman.take_marker(&mut self.reader)? { |
923 | Some(Marker::RST(n)) => { |
924 | if n != expected_rst_num { |
925 | return Err(Error::Format(format!( |
926 | "found RST {} where RST {} was expected" , |
927 | n, expected_rst_num |
928 | ))); |
929 | } |
930 | |
931 | huffman.reset(); |
932 | // Section F.2.1.3.1 |
933 | dc_predictors = [0i16; MAX_COMPONENTS]; |
934 | // Section G.1.2.2 |
935 | eob_run = 0; |
936 | |
937 | expected_rst_num = (expected_rst_num + 1) % 8; |
938 | mcus_left_until_restart = self.restart_interval; |
939 | } |
940 | Some(marker) => { |
941 | return Err(Error::Format(format!( |
942 | "found marker {:?} inside scan where RST {} was expected" , |
943 | marker, expected_rst_num |
944 | ))) |
945 | } |
946 | None => { |
947 | return Err(Error::Format(format!( |
948 | "no marker found where RST {} was expected" , |
949 | expected_rst_num |
950 | ))) |
951 | } |
952 | } |
953 | } |
954 | |
955 | mcus_left_until_restart -= 1; |
956 | } |
957 | |
958 | for (i, component) in components.iter().enumerate() { |
959 | for v_pos in 0..mcu_vertical_samples[i] { |
960 | for h_pos in 0..mcu_horizontal_samples[i] { |
961 | let coefficients = if is_progressive { |
962 | let block_y = (mcu_y * mcu_vertical_samples[i] + v_pos) as usize; |
963 | let block_x = (mcu_x * mcu_horizontal_samples[i] + h_pos) as usize; |
964 | let block_offset = |
965 | (block_y * component.block_size.width as usize + block_x) * 64; |
966 | &mut self.coefficients[scan.component_indices[i]] |
967 | [block_offset..block_offset + 64] |
968 | } else if finished[i] { |
969 | // Because the worker thread operates in batches as if we were always interleaved, we |
970 | // need to distinguish between a single-shot buffer and one that's currently in process |
971 | // (for a non-interleaved) stream |
972 | let mcu_batch_current_row = if is_interleaved { |
973 | 0 |
974 | } else { |
975 | mcu_y % component.vertical_sampling_factor as u16 |
976 | }; |
977 | |
978 | let block_y = (mcu_batch_current_row * mcu_vertical_samples[i] |
979 | + v_pos) as usize; |
980 | let block_x = (mcu_x * mcu_horizontal_samples[i] + h_pos) as usize; |
981 | let block_offset = |
982 | (block_y * component.block_size.width as usize + block_x) * 64; |
983 | &mut mcu_row_coefficients[i][block_offset..block_offset + 64] |
984 | } else { |
985 | &mut dummy_block[..64] |
986 | } |
987 | .try_into() |
988 | .unwrap(); |
989 | |
990 | if scan.successive_approximation_high == 0 { |
991 | decode_block( |
992 | &mut self.reader, |
993 | coefficients, |
994 | &mut huffman, |
995 | self.dc_huffman_tables[scan.dc_table_indices[i]].as_ref(), |
996 | self.ac_huffman_tables[scan.ac_table_indices[i]].as_ref(), |
997 | scan.spectral_selection.clone(), |
998 | scan.successive_approximation_low, |
999 | &mut eob_run, |
1000 | &mut dc_predictors[i], |
1001 | )?; |
1002 | } else { |
1003 | decode_block_successive_approximation( |
1004 | &mut self.reader, |
1005 | coefficients, |
1006 | &mut huffman, |
1007 | self.ac_huffman_tables[scan.ac_table_indices[i]].as_ref(), |
1008 | scan.spectral_selection.clone(), |
1009 | scan.successive_approximation_low, |
1010 | &mut eob_run, |
1011 | )?; |
1012 | } |
1013 | } |
1014 | } |
1015 | } |
1016 | } |
1017 | |
1018 | // Send the coefficients from this MCU row to the worker thread for dequantization and idct. |
1019 | for (i, component) in components.iter().enumerate() { |
1020 | if finished[i] { |
1021 | // In the event of non-interleaved streams, if we're still building the buffer out, |
1022 | // keep going; don't send it yet. We also need to ensure we don't skip over the last |
1023 | // row(s) of the image. |
1024 | if !is_interleaved |
1025 | && (mcu_y + 1) * 8 < frame.image_size.height |
1026 | && (mcu_y + 1) % component.vertical_sampling_factor as u16 > 0 |
1027 | { |
1028 | continue; |
1029 | } |
1030 | |
1031 | let coefficients_per_mcu_row = component.block_size.width as usize |
1032 | * component.vertical_sampling_factor as usize |
1033 | * 64; |
1034 | |
1035 | let row_coefficients = if is_progressive { |
1036 | // Because non-interleaved streams will have multiple MCU rows concatenated together, |
1037 | // the row for calculating the offset is different. |
1038 | let worker_mcu_y = if is_interleaved { |
1039 | mcu_y |
1040 | } else { |
1041 | // Explicitly doing floor-division here |
1042 | mcu_y / component.vertical_sampling_factor as u16 |
1043 | }; |
1044 | |
1045 | let offset = worker_mcu_y as usize * coefficients_per_mcu_row; |
1046 | self.coefficients[scan.component_indices[i]] |
1047 | [offset..offset + coefficients_per_mcu_row] |
1048 | .to_vec() |
1049 | } else { |
1050 | mem::replace( |
1051 | &mut mcu_row_coefficients[i], |
1052 | vec![0i16; coefficients_per_mcu_row], |
1053 | ) |
1054 | }; |
1055 | |
1056 | // FIXME: additional potential work stealing opportunities for rayon case if we |
1057 | // also internally can parallelize over components. |
1058 | worker.append_row((i, row_coefficients))?; |
1059 | } |
1060 | } |
1061 | } |
1062 | |
1063 | let mut marker = huffman.take_marker(&mut self.reader)?; |
1064 | while let Some(Marker::RST(_)) = marker { |
1065 | marker = self.read_marker().ok(); |
1066 | } |
1067 | |
1068 | if finished.iter().any(|&c| c) { |
1069 | // Retrieve all the data from the worker thread. |
1070 | let mut data = vec![Vec::new(); frame.components.len()]; |
1071 | |
1072 | for (i, &component_index) in scan.component_indices.iter().enumerate() { |
1073 | if finished[i] { |
1074 | data[component_index] = worker.get_result(i)?; |
1075 | } |
1076 | } |
1077 | |
1078 | Ok((marker, Some(data))) |
1079 | } else { |
1080 | Ok((marker, None)) |
1081 | } |
1082 | } |
1083 | } |
1084 | |
1085 | #[allow (clippy::too_many_arguments)] |
1086 | fn decode_block<R: Read>( |
1087 | reader: &mut R, |
1088 | coefficients: &mut [i16; 64], |
1089 | huffman: &mut HuffmanDecoder, |
1090 | dc_table: Option<&HuffmanTable>, |
1091 | ac_table: Option<&HuffmanTable>, |
1092 | spectral_selection: Range<u8>, |
1093 | successive_approximation_low: u8, |
1094 | eob_run: &mut u16, |
1095 | dc_predictor: &mut i16, |
1096 | ) -> Result<()> { |
1097 | debug_assert_eq!(coefficients.len(), 64); |
1098 | |
1099 | if spectral_selection.start == 0 { |
1100 | // Section F.2.2.1 |
1101 | // Figure F.12 |
1102 | let value = huffman.decode(reader, dc_table.unwrap())?; |
1103 | let diff = match value { |
1104 | 0 => 0, |
1105 | 1..=11 => huffman.receive_extend(reader, value)?, |
1106 | _ => { |
1107 | // Section F.1.2.1.1 |
1108 | // Table F.1 |
1109 | return Err(Error::Format( |
1110 | "invalid DC difference magnitude category" .to_owned(), |
1111 | )); |
1112 | } |
1113 | }; |
1114 | |
1115 | // Malicious JPEG files can cause this add to overflow, therefore we use wrapping_add. |
1116 | // One example of such a file is tests/crashtest/images/dc-predictor-overflow.jpg |
1117 | *dc_predictor = dc_predictor.wrapping_add(diff); |
1118 | coefficients[0] = *dc_predictor << successive_approximation_low; |
1119 | } |
1120 | |
1121 | let mut index = cmp::max(spectral_selection.start, 1); |
1122 | |
1123 | if index < spectral_selection.end && *eob_run > 0 { |
1124 | *eob_run -= 1; |
1125 | return Ok(()); |
1126 | } |
1127 | |
1128 | // Section F.1.2.2.1 |
1129 | while index < spectral_selection.end { |
1130 | if let Some((value, run)) = huffman.decode_fast_ac(reader, ac_table.unwrap())? { |
1131 | index += run; |
1132 | |
1133 | if index >= spectral_selection.end { |
1134 | break; |
1135 | } |
1136 | |
1137 | coefficients[UNZIGZAG[index as usize] as usize] = value << successive_approximation_low; |
1138 | index += 1; |
1139 | } else { |
1140 | let byte = huffman.decode(reader, ac_table.unwrap())?; |
1141 | let r = byte >> 4; |
1142 | let s = byte & 0x0f; |
1143 | |
1144 | if s == 0 { |
1145 | match r { |
1146 | 15 => index += 16, // Run length of 16 zero coefficients. |
1147 | _ => { |
1148 | *eob_run = (1 << r) - 1; |
1149 | |
1150 | if r > 0 { |
1151 | *eob_run += huffman.get_bits(reader, r)?; |
1152 | } |
1153 | |
1154 | break; |
1155 | } |
1156 | } |
1157 | } else { |
1158 | index += r; |
1159 | |
1160 | if index >= spectral_selection.end { |
1161 | break; |
1162 | } |
1163 | |
1164 | coefficients[UNZIGZAG[index as usize] as usize] = |
1165 | huffman.receive_extend(reader, s)? << successive_approximation_low; |
1166 | index += 1; |
1167 | } |
1168 | } |
1169 | } |
1170 | |
1171 | Ok(()) |
1172 | } |
1173 | |
1174 | fn decode_block_successive_approximation<R: Read>( |
1175 | reader: &mut R, |
1176 | coefficients: &mut [i16; 64], |
1177 | huffman: &mut HuffmanDecoder, |
1178 | ac_table: Option<&HuffmanTable>, |
1179 | spectral_selection: Range<u8>, |
1180 | successive_approximation_low: u8, |
1181 | eob_run: &mut u16, |
1182 | ) -> Result<()> { |
1183 | debug_assert_eq!(coefficients.len(), 64); |
1184 | |
1185 | let bit = 1 << successive_approximation_low; |
1186 | |
1187 | if spectral_selection.start == 0 { |
1188 | // Section G.1.2.1 |
1189 | |
1190 | if huffman.get_bits(reader, 1)? == 1 { |
1191 | coefficients[0] |= bit; |
1192 | } |
1193 | } else { |
1194 | // Section G.1.2.3 |
1195 | |
1196 | if *eob_run > 0 { |
1197 | *eob_run -= 1; |
1198 | refine_non_zeroes(reader, coefficients, huffman, spectral_selection, 64, bit)?; |
1199 | return Ok(()); |
1200 | } |
1201 | |
1202 | let mut index = spectral_selection.start; |
1203 | |
1204 | while index < spectral_selection.end { |
1205 | let byte = huffman.decode(reader, ac_table.unwrap())?; |
1206 | let r = byte >> 4; |
1207 | let s = byte & 0x0f; |
1208 | |
1209 | let mut zero_run_length = r; |
1210 | let mut value = 0; |
1211 | |
1212 | match s { |
1213 | 0 => { |
1214 | match r { |
1215 | 15 => { |
1216 | // Run length of 16 zero coefficients. |
1217 | // We don't need to do anything special here, zero_run_length is 15 |
1218 | // and then value (which is zero) gets written, resulting in 16 |
1219 | // zero coefficients. |
1220 | } |
1221 | _ => { |
1222 | *eob_run = (1 << r) - 1; |
1223 | |
1224 | if r > 0 { |
1225 | *eob_run += huffman.get_bits(reader, r)?; |
1226 | } |
1227 | |
1228 | // Force end of block. |
1229 | zero_run_length = 64; |
1230 | } |
1231 | } |
1232 | } |
1233 | 1 => { |
1234 | if huffman.get_bits(reader, 1)? == 1 { |
1235 | value = bit; |
1236 | } else { |
1237 | value = -bit; |
1238 | } |
1239 | } |
1240 | _ => return Err(Error::Format("unexpected huffman code" .to_owned())), |
1241 | } |
1242 | |
1243 | let range = Range { |
1244 | start: index, |
1245 | end: spectral_selection.end, |
1246 | }; |
1247 | index = refine_non_zeroes(reader, coefficients, huffman, range, zero_run_length, bit)?; |
1248 | |
1249 | if value != 0 { |
1250 | coefficients[UNZIGZAG[index as usize] as usize] = value; |
1251 | } |
1252 | |
1253 | index += 1; |
1254 | } |
1255 | } |
1256 | |
1257 | Ok(()) |
1258 | } |
1259 | |
1260 | fn refine_non_zeroes<R: Read>( |
1261 | reader: &mut R, |
1262 | coefficients: &mut [i16; 64], |
1263 | huffman: &mut HuffmanDecoder, |
1264 | range: Range<u8>, |
1265 | zrl: u8, |
1266 | bit: i16, |
1267 | ) -> Result<u8> { |
1268 | debug_assert_eq!(coefficients.len(), 64); |
1269 | |
1270 | let last = range.end - 1; |
1271 | let mut zero_run_length = zrl; |
1272 | |
1273 | for i in range { |
1274 | let index = UNZIGZAG[i as usize] as usize; |
1275 | |
1276 | let coefficient = &mut coefficients[index]; |
1277 | |
1278 | if *coefficient == 0 { |
1279 | if zero_run_length == 0 { |
1280 | return Ok(i); |
1281 | } |
1282 | |
1283 | zero_run_length -= 1; |
1284 | } else if huffman.get_bits(reader, 1)? == 1 && *coefficient & bit == 0 { |
1285 | if *coefficient > 0 { |
1286 | *coefficient = coefficient |
1287 | .checked_add(bit) |
1288 | .ok_or_else(|| Error::Format("Coefficient overflow" .to_owned()))?; |
1289 | } else { |
1290 | *coefficient = coefficient |
1291 | .checked_sub(bit) |
1292 | .ok_or_else(|| Error::Format("Coefficient overflow" .to_owned()))?; |
1293 | } |
1294 | } |
1295 | } |
1296 | |
1297 | Ok(last) |
1298 | } |
1299 | |
1300 | fn compute_image( |
1301 | components: &[Component], |
1302 | mut data: Vec<Vec<u8>>, |
1303 | output_size: Dimensions, |
1304 | color_transform: ColorTransform, |
1305 | ) -> Result<Vec<u8>> { |
1306 | if data.is_empty() || data.iter().any(Vec::is_empty) { |
1307 | return Err(Error::Format("not all components have data" .to_owned())); |
1308 | } |
1309 | |
1310 | if components.len() == 1 { |
1311 | let component = &components[0]; |
1312 | let mut decoded: Vec<u8> = data.remove(0); |
1313 | |
1314 | let width = component.size.width as usize; |
1315 | let height = component.size.height as usize; |
1316 | let size = width * height; |
1317 | let line_stride = component.block_size.width as usize * component.dct_scale; |
1318 | |
1319 | // if the image width is a multiple of the block size, |
1320 | // then we don't have to move bytes in the decoded data |
1321 | if usize::from(output_size.width) != line_stride { |
1322 | // The first line already starts at index 0, so we need to move only lines 1..height |
1323 | // We move from the top down because all lines are being moved backwards. |
1324 | for y in 1..height { |
1325 | let destination_idx = y * width; |
1326 | let source_idx = y * line_stride; |
1327 | let end = source_idx + width; |
1328 | decoded.copy_within(source_idx..end, destination_idx); |
1329 | } |
1330 | } |
1331 | decoded.resize(size, 0); |
1332 | Ok(decoded) |
1333 | } else { |
1334 | compute_image_parallel(components, data, output_size, color_transform) |
1335 | } |
1336 | } |
1337 | |
1338 | #[allow (clippy::type_complexity)] |
1339 | pub(crate) fn choose_color_convert_func( |
1340 | component_count: usize, |
1341 | color_transform: ColorTransform, |
1342 | ) -> Result<fn(&[Vec<u8>], &mut [u8])> { |
1343 | match component_count { |
1344 | 3 => match color_transform { |
1345 | ColorTransform::None => Ok(color_no_convert), |
1346 | ColorTransform::Grayscale => Err(Error::Format( |
1347 | "Invalid number of channels (3) for Grayscale data" .to_string(), |
1348 | )), |
1349 | ColorTransform::RGB => Ok(color_convert_line_rgb), |
1350 | ColorTransform::YCbCr => Ok(color_convert_line_ycbcr), |
1351 | ColorTransform::CMYK => Err(Error::Format( |
1352 | "Invalid number of channels (3) for CMYK data" .to_string(), |
1353 | )), |
1354 | ColorTransform::YCCK => Err(Error::Format( |
1355 | "Invalid number of channels (3) for YCCK data" .to_string(), |
1356 | )), |
1357 | ColorTransform::JcsBgYcc => Err(Error::Unsupported( |
1358 | UnsupportedFeature::ColorTransform(ColorTransform::JcsBgYcc), |
1359 | )), |
1360 | ColorTransform::JcsBgRgb => Err(Error::Unsupported( |
1361 | UnsupportedFeature::ColorTransform(ColorTransform::JcsBgRgb), |
1362 | )), |
1363 | ColorTransform::Unknown => Err(Error::Format("Unknown colour transform" .to_string())), |
1364 | }, |
1365 | 4 => match color_transform { |
1366 | ColorTransform::None => Ok(color_no_convert), |
1367 | ColorTransform::Grayscale => Err(Error::Format( |
1368 | "Invalid number of channels (4) for Grayscale data" .to_string(), |
1369 | )), |
1370 | ColorTransform::RGB => Err(Error::Format( |
1371 | "Invalid number of channels (4) for RGB data" .to_string(), |
1372 | )), |
1373 | ColorTransform::YCbCr => Err(Error::Format( |
1374 | "Invalid number of channels (4) for YCbCr data" .to_string(), |
1375 | )), |
1376 | ColorTransform::CMYK => Ok(color_convert_line_cmyk), |
1377 | ColorTransform::YCCK => Ok(color_convert_line_ycck), |
1378 | |
1379 | ColorTransform::JcsBgYcc => Err(Error::Unsupported( |
1380 | UnsupportedFeature::ColorTransform(ColorTransform::JcsBgYcc), |
1381 | )), |
1382 | ColorTransform::JcsBgRgb => Err(Error::Unsupported( |
1383 | UnsupportedFeature::ColorTransform(ColorTransform::JcsBgRgb), |
1384 | )), |
1385 | ColorTransform::Unknown => Err(Error::Format("Unknown colour transform" .to_string())), |
1386 | }, |
1387 | _ => panic!(), |
1388 | } |
1389 | } |
1390 | |
1391 | fn color_convert_line_rgb(data: &[Vec<u8>], output: &mut [u8]) { |
1392 | assert!(data.len() == 3, "wrong number of components for rgb" ); |
1393 | let [r: &Vec, g: &Vec, b: &Vec]: &[Vec<u8>; 3] = data.try_into().unwrap(); |
1394 | for (((chunk: &mut [u8], r: &u8), g: &u8), b: &u8) in outputimpl Iterator |
1395 | .chunks_exact_mut(chunk_size:3) |
1396 | .zip(r.iter()) |
1397 | .zip(g.iter()) |
1398 | .zip(b.iter()) |
1399 | { |
1400 | chunk[0] = *r; |
1401 | chunk[1] = *g; |
1402 | chunk[2] = *b; |
1403 | } |
1404 | } |
1405 | |
1406 | fn color_convert_line_ycbcr(data: &[Vec<u8>], output: &mut [u8]) { |
1407 | assert!(data.len() == 3, "wrong number of components for ycbcr" ); |
1408 | let [y, cb, cr]: &[_; 3] = data.try_into().unwrap(); |
1409 | |
1410 | #[cfg (not(feature = "platform_independent" ))] |
1411 | let arch_specific_pixels = { |
1412 | if let Some(ycbcr) = crate::arch::get_color_convert_line_ycbcr() { |
1413 | #[allow (unsafe_code)] |
1414 | unsafe { |
1415 | ycbcr(y, cb, cr, output) |
1416 | } |
1417 | } else { |
1418 | 0 |
1419 | } |
1420 | }; |
1421 | |
1422 | #[cfg (feature = "platform_independent" )] |
1423 | let arch_specific_pixels = 0; |
1424 | |
1425 | for (((chunk, y), cb), cr) in output |
1426 | .chunks_exact_mut(3) |
1427 | .zip(y.iter()) |
1428 | .zip(cb.iter()) |
1429 | .zip(cr.iter()) |
1430 | .skip(arch_specific_pixels) |
1431 | { |
1432 | let (r, g, b) = ycbcr_to_rgb(*y, *cb, *cr); |
1433 | chunk[0] = r; |
1434 | chunk[1] = g; |
1435 | chunk[2] = b; |
1436 | } |
1437 | } |
1438 | |
1439 | fn color_convert_line_ycck(data: &[Vec<u8>], output: &mut [u8]) { |
1440 | assert!(data.len() == 4, "wrong number of components for ycck" ); |
1441 | let [c: &Vec, m: &Vec, y: &Vec, k: &Vec]: &[Vec<u8>; 4] = data.try_into().unwrap(); |
1442 | |
1443 | for ((((chunk: &mut [u8], c: &u8), m: &u8), y: &u8), k: &u8) in outputimpl Iterator |
1444 | .chunks_exact_mut(chunk_size:4) |
1445 | .zip(c.iter()) |
1446 | .zip(m.iter()) |
1447 | .zip(y.iter()) |
1448 | .zip(k.iter()) |
1449 | { |
1450 | let (r: u8, g: u8, b: u8) = ycbcr_to_rgb(*c, *m, *y); |
1451 | chunk[0] = r; |
1452 | chunk[1] = g; |
1453 | chunk[2] = b; |
1454 | chunk[3] = 255 - *k; |
1455 | } |
1456 | } |
1457 | |
1458 | fn color_convert_line_cmyk(data: &[Vec<u8>], output: &mut [u8]) { |
1459 | assert!(data.len() == 4, "wrong number of components for cmyk" ); |
1460 | let [c: &Vec, m: &Vec, y: &Vec, k: &Vec]: &[Vec<u8>; 4] = data.try_into().unwrap(); |
1461 | |
1462 | for ((((chunk: &mut [u8], c: &u8), m: &u8), y: &u8), k: &u8) in outputimpl Iterator |
1463 | .chunks_exact_mut(chunk_size:4) |
1464 | .zip(c.iter()) |
1465 | .zip(m.iter()) |
1466 | .zip(y.iter()) |
1467 | .zip(k.iter()) |
1468 | { |
1469 | chunk[0] = 255 - c; |
1470 | chunk[1] = 255 - m; |
1471 | chunk[2] = 255 - y; |
1472 | chunk[3] = 255 - k; |
1473 | } |
1474 | } |
1475 | |
1476 | fn color_no_convert(data: &[Vec<u8>], output: &mut [u8]) { |
1477 | let mut output_iter: IterMut<'_, u8> = output.iter_mut(); |
1478 | |
1479 | for pixel: &Vec in data { |
1480 | for d: &u8 in pixel { |
1481 | *(output_iter.next().unwrap()) = *d; |
1482 | } |
1483 | } |
1484 | } |
1485 | |
1486 | const FIXED_POINT_OFFSET: i32 = 20; |
1487 | const HALF: i32 = (1 << FIXED_POINT_OFFSET) / 2; |
1488 | |
1489 | // ITU-R BT.601 |
1490 | // Based on libjpeg-turbo's jdcolext.c |
1491 | fn ycbcr_to_rgb(y: u8, cb: u8, cr: u8) -> (u8, u8, u8) { |
1492 | let y: i32 = y as i32 * (1 << FIXED_POINT_OFFSET) + HALF; |
1493 | let cb: i32 = cb as i32 - 128; |
1494 | let cr: i32 = cr as i32 - 128; |
1495 | |
1496 | let r: u8 = clamp_fixed_point(y + stbi_f2f(1.40200) * cr); |
1497 | let g: u8 = clamp_fixed_point(y - stbi_f2f(0.34414) * cb - stbi_f2f(0.71414) * cr); |
1498 | let b: u8 = clamp_fixed_point(y + stbi_f2f(1.77200) * cb); |
1499 | (r, g, b) |
1500 | } |
1501 | |
1502 | fn stbi_f2f(x: f32) -> i32 { |
1503 | (x * ((1 << FIXED_POINT_OFFSET) as f32) + 0.5) as i32 |
1504 | } |
1505 | |
1506 | fn clamp_fixed_point(value: i32) -> u8 { |
1507 | (value >> FIXED_POINT_OFFSET).min(255).max(0) as u8 |
1508 | } |
1509 | |