1//! Encoding of PNM Images
2use std::fmt;
3use std::io;
4
5use std::io::Write;
6
7use super::AutoBreak;
8use super::{ArbitraryHeader, ArbitraryTuplType, BitmapHeader, GraymapHeader, PixmapHeader};
9use super::{HeaderRecord, PnmHeader, PnmSubtype, SampleEncoding};
10use crate::color::{ColorType, ExtendedColorType};
11use crate::error::{
12 ImageError, ImageResult, ParameterError, ParameterErrorKind, UnsupportedError,
13 UnsupportedErrorKind,
14};
15use crate::image::{ImageEncoder, ImageFormat};
16
17use byteorder::{BigEndian, WriteBytesExt};
18
19enum HeaderStrategy {
20 Dynamic,
21 Subtype(PnmSubtype),
22 Chosen(PnmHeader),
23}
24
25#[derive(Clone, Copy)]
26pub enum FlatSamples<'a> {
27 U8(&'a [u8]),
28 U16(&'a [u16]),
29}
30
31/// Encodes images to any of the `pnm` image formats.
32pub struct PnmEncoder<W: Write> {
33 writer: W,
34 header: HeaderStrategy,
35}
36
37/// Encapsulate the checking system in the type system. Non of the fields are actually accessed
38/// but requiring them forces us to validly construct the struct anyways.
39struct CheckedImageBuffer<'a> {
40 _image: FlatSamples<'a>,
41 _width: u32,
42 _height: u32,
43 _color: ExtendedColorType,
44}
45
46// Check the header against the buffer. Each struct produces the next after a check.
47struct UncheckedHeader<'a> {
48 header: &'a PnmHeader,
49}
50
51struct CheckedDimensions<'a> {
52 unchecked: UncheckedHeader<'a>,
53 width: u32,
54 height: u32,
55}
56
57struct CheckedHeaderColor<'a> {
58 dimensions: CheckedDimensions<'a>,
59 color: ExtendedColorType,
60}
61
62struct CheckedHeader<'a> {
63 color: CheckedHeaderColor<'a>,
64 encoding: TupleEncoding<'a>,
65 _image: CheckedImageBuffer<'a>,
66}
67
68enum TupleEncoding<'a> {
69 PbmBits {
70 samples: FlatSamples<'a>,
71 width: u32,
72 },
73 Ascii {
74 samples: FlatSamples<'a>,
75 },
76 Bytes {
77 samples: FlatSamples<'a>,
78 },
79}
80
81impl<W: Write> PnmEncoder<W> {
82 /// Create new PnmEncoder from the `writer`.
83 ///
84 /// The encoded images will have some `pnm` format. If more control over the image type is
85 /// required, use either one of `with_subtype` or `with_header`. For more information on the
86 /// behaviour, see `with_dynamic_header`.
87 pub fn new(writer: W) -> Self {
88 PnmEncoder {
89 writer,
90 header: HeaderStrategy::Dynamic,
91 }
92 }
93
94 /// Encode a specific pnm subtype image.
95 ///
96 /// The magic number and encoding type will be chosen as provided while the rest of the header
97 /// data will be generated dynamically. Trying to encode incompatible images (e.g. encoding an
98 /// RGB image as Graymap) will result in an error.
99 ///
100 /// This will overwrite the effect of earlier calls to `with_header` and `with_dynamic_header`.
101 pub fn with_subtype(self, subtype: PnmSubtype) -> Self {
102 PnmEncoder {
103 writer: self.writer,
104 header: HeaderStrategy::Subtype(subtype),
105 }
106 }
107
108 /// Enforce the use of a chosen header.
109 ///
110 /// While this option gives the most control over the actual written data, the encoding process
111 /// will error in case the header data and image parameters do not agree. It is the users
112 /// obligation to ensure that the width and height are set accordingly, for example.
113 ///
114 /// Choose this option if you want a lossless decoding/encoding round trip.
115 ///
116 /// This will overwrite the effect of earlier calls to `with_subtype` and `with_dynamic_header`.
117 pub fn with_header(self, header: PnmHeader) -> Self {
118 PnmEncoder {
119 writer: self.writer,
120 header: HeaderStrategy::Chosen(header),
121 }
122 }
123
124 /// Create the header dynamically for each image.
125 ///
126 /// This is the default option upon creation of the encoder. With this, most images should be
127 /// encodable but the specific format chosen is out of the users control. The pnm subtype is
128 /// chosen arbitrarily by the library.
129 ///
130 /// This will overwrite the effect of earlier calls to `with_subtype` and `with_header`.
131 pub fn with_dynamic_header(self) -> Self {
132 PnmEncoder {
133 writer: self.writer,
134 header: HeaderStrategy::Dynamic,
135 }
136 }
137
138 /// Encode an image whose samples are represented as `u8`.
139 ///
140 /// Some `pnm` subtypes are incompatible with some color options, a chosen header most
141 /// certainly with any deviation from the original decoded image.
142 pub fn encode<'s, S>(
143 &mut self,
144 image: S,
145 width: u32,
146 height: u32,
147 color: ColorType,
148 ) -> ImageResult<()>
149 where
150 S: Into<FlatSamples<'s>>,
151 {
152 let image = image.into();
153 match self.header {
154 HeaderStrategy::Dynamic => {
155 self.write_dynamic_header(image, width, height, color.into())
156 }
157 HeaderStrategy::Subtype(subtype) => {
158 self.write_subtyped_header(subtype, image, width, height, color.into())
159 }
160 HeaderStrategy::Chosen(ref header) => Self::write_with_header(
161 &mut self.writer,
162 header,
163 image,
164 width,
165 height,
166 color.into(),
167 ),
168 }
169 }
170
171 /// Choose any valid pnm format that the image can be expressed in and write its header.
172 ///
173 /// Returns how the body should be written if successful.
174 fn write_dynamic_header(
175 &mut self,
176 image: FlatSamples,
177 width: u32,
178 height: u32,
179 color: ExtendedColorType,
180 ) -> ImageResult<()> {
181 let depth = u32::from(color.channel_count());
182 let (maxval, tupltype) = match color {
183 ExtendedColorType::L1 => (1, ArbitraryTuplType::BlackAndWhite),
184 ExtendedColorType::L8 => (0xff, ArbitraryTuplType::Grayscale),
185 ExtendedColorType::L16 => (0xffff, ArbitraryTuplType::Grayscale),
186 ExtendedColorType::La1 => (1, ArbitraryTuplType::BlackAndWhiteAlpha),
187 ExtendedColorType::La8 => (0xff, ArbitraryTuplType::GrayscaleAlpha),
188 ExtendedColorType::La16 => (0xffff, ArbitraryTuplType::GrayscaleAlpha),
189 ExtendedColorType::Rgb8 => (0xff, ArbitraryTuplType::RGB),
190 ExtendedColorType::Rgb16 => (0xffff, ArbitraryTuplType::RGB),
191 ExtendedColorType::Rgba8 => (0xff, ArbitraryTuplType::RGBAlpha),
192 ExtendedColorType::Rgba16 => (0xffff, ArbitraryTuplType::RGBAlpha),
193 _ => {
194 return Err(ImageError::Unsupported(
195 UnsupportedError::from_format_and_kind(
196 ImageFormat::Pnm.into(),
197 UnsupportedErrorKind::Color(color),
198 ),
199 ))
200 }
201 };
202
203 let header = PnmHeader {
204 decoded: HeaderRecord::Arbitrary(ArbitraryHeader {
205 width,
206 height,
207 depth,
208 maxval,
209 tupltype: Some(tupltype),
210 }),
211 encoded: None,
212 };
213
214 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
215 }
216
217 /// Try to encode the image with the chosen format, give its corresponding pixel encoding type.
218 fn write_subtyped_header(
219 &mut self,
220 subtype: PnmSubtype,
221 image: FlatSamples,
222 width: u32,
223 height: u32,
224 color: ExtendedColorType,
225 ) -> ImageResult<()> {
226 let header = match (subtype, color) {
227 (PnmSubtype::ArbitraryMap, color) => {
228 return self.write_dynamic_header(image, width, height, color)
229 }
230 (PnmSubtype::Pixmap(encoding), ExtendedColorType::Rgb8) => PnmHeader {
231 decoded: HeaderRecord::Pixmap(PixmapHeader {
232 encoding,
233 width,
234 height,
235 maxval: 255,
236 }),
237 encoded: None,
238 },
239 (PnmSubtype::Graymap(encoding), ExtendedColorType::L8) => PnmHeader {
240 decoded: HeaderRecord::Graymap(GraymapHeader {
241 encoding,
242 width,
243 height,
244 maxwhite: 255,
245 }),
246 encoded: None,
247 },
248 (PnmSubtype::Bitmap(encoding), ExtendedColorType::L8)
249 | (PnmSubtype::Bitmap(encoding), ExtendedColorType::L1) => PnmHeader {
250 decoded: HeaderRecord::Bitmap(BitmapHeader {
251 encoding,
252 width,
253 height,
254 }),
255 encoded: None,
256 },
257 (_, _) => {
258 return Err(ImageError::Parameter(ParameterError::from_kind(
259 ParameterErrorKind::Generic(
260 "Color type can not be represented in the chosen format".to_owned(),
261 ),
262 )));
263 }
264 };
265
266 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
267 }
268
269 /// Try to encode the image with the chosen header, checking if values are correct.
270 ///
271 /// Returns how the body should be written if successful.
272 fn write_with_header(
273 writer: &mut dyn Write,
274 header: &PnmHeader,
275 image: FlatSamples,
276 width: u32,
277 height: u32,
278 color: ExtendedColorType,
279 ) -> ImageResult<()> {
280 let unchecked = UncheckedHeader { header };
281
282 unchecked
283 .check_header_dimensions(width, height)?
284 .check_header_color(color)?
285 .check_sample_values(image)?
286 .write_header(writer)?
287 .write_image(writer)
288 }
289}
290
291impl<W: Write> ImageEncoder for PnmEncoder<W> {
292 #[track_caller]
293 fn write_image(
294 mut self,
295 buf: &[u8],
296 width: u32,
297 height: u32,
298 color_type: ColorType,
299 ) -> ImageResult<()> {
300 let expected_buffer_len: u64 =
301 (width as u64 * height as u64).saturating_mul(color_type.bytes_per_pixel() as u64);
302 assert_eq!(
303 expected_buffer_len,
304 buf.len() as u64,
305 "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
306 buf.len(),
307 );
308
309 self.encode(image:buf, width, height, color_type)
310 }
311}
312
313impl<'a> CheckedImageBuffer<'a> {
314 fn check(
315 image: FlatSamples<'a>,
316 width: u32,
317 height: u32,
318 color: ExtendedColorType,
319 ) -> ImageResult<CheckedImageBuffer<'a>> {
320 let components = color.channel_count() as usize;
321 let uwidth = width as usize;
322 let uheight = height as usize;
323 let expected_len = components
324 .checked_mul(uwidth)
325 .and_then(|v| v.checked_mul(uheight));
326 if Some(image.len()) != expected_len {
327 // Image buffer does not correspond to size and colour.
328 return Err(ImageError::Parameter(ParameterError::from_kind(
329 ParameterErrorKind::DimensionMismatch,
330 )));
331 }
332 Ok(CheckedImageBuffer {
333 _image: image,
334 _width: width,
335 _height: height,
336 _color: color,
337 })
338 }
339}
340
341impl<'a> UncheckedHeader<'a> {
342 fn check_header_dimensions(
343 self,
344 width: u32,
345 height: u32,
346 ) -> ImageResult<CheckedDimensions<'a>> {
347 if self.header.width() != width || self.header.height() != height {
348 // Chosen header does not match Image dimensions.
349 return Err(ImageError::Parameter(ParameterError::from_kind(
350 ParameterErrorKind::DimensionMismatch,
351 )));
352 }
353
354 Ok(CheckedDimensions {
355 unchecked: self,
356 width,
357 height,
358 })
359 }
360}
361
362impl<'a> CheckedDimensions<'a> {
363 // Check color compatibility with the header. This will only error when we are certain that
364 // the combination is bogus (e.g. combining Pixmap and Palette) but allows uncertain
365 // combinations (basically a ArbitraryTuplType::Custom with any color of fitting depth).
366 fn check_header_color(self, color: ExtendedColorType) -> ImageResult<CheckedHeaderColor<'a>> {
367 let components = u32::from(color.channel_count());
368
369 match *self.unchecked.header {
370 PnmHeader {
371 decoded: HeaderRecord::Bitmap(_),
372 ..
373 } => match color {
374 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
375 _ => {
376 return Err(ImageError::Parameter(ParameterError::from_kind(
377 ParameterErrorKind::Generic(
378 "PBM format only support luma color types".to_owned(),
379 ),
380 )))
381 }
382 },
383 PnmHeader {
384 decoded: HeaderRecord::Graymap(_),
385 ..
386 } => match color {
387 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
388 _ => {
389 return Err(ImageError::Parameter(ParameterError::from_kind(
390 ParameterErrorKind::Generic(
391 "PGM format only support luma color types".to_owned(),
392 ),
393 )))
394 }
395 },
396 PnmHeader {
397 decoded: HeaderRecord::Pixmap(_),
398 ..
399 } => match color {
400 ExtendedColorType::Rgb8 => (),
401 _ => {
402 return Err(ImageError::Parameter(ParameterError::from_kind(
403 ParameterErrorKind::Generic(
404 "PPM format only support ExtendedColorType::Rgb8".to_owned(),
405 ),
406 )))
407 }
408 },
409 PnmHeader {
410 decoded:
411 HeaderRecord::Arbitrary(ArbitraryHeader {
412 depth,
413 ref tupltype,
414 ..
415 }),
416 ..
417 } => match (tupltype, color) {
418 (&Some(ArbitraryTuplType::BlackAndWhite), ExtendedColorType::L1) => (),
419 (&Some(ArbitraryTuplType::BlackAndWhiteAlpha), ExtendedColorType::La8) => (),
420
421 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L1) => (),
422 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L8) => (),
423 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L16) => (),
424 (&Some(ArbitraryTuplType::GrayscaleAlpha), ExtendedColorType::La8) => (),
425
426 (&Some(ArbitraryTuplType::RGB), ExtendedColorType::Rgb8) => (),
427 (&Some(ArbitraryTuplType::RGBAlpha), ExtendedColorType::Rgba8) => (),
428
429 (&None, _) if depth == components => (),
430 (&Some(ArbitraryTuplType::Custom(_)), _) if depth == components => (),
431 _ if depth != components => {
432 return Err(ImageError::Parameter(ParameterError::from_kind(
433 ParameterErrorKind::Generic(format!(
434 "Depth mismatch: header {} vs. color {}",
435 depth, components
436 )),
437 )))
438 }
439 _ => {
440 return Err(ImageError::Parameter(ParameterError::from_kind(
441 ParameterErrorKind::Generic(
442 "Invalid color type for selected PAM color type".to_owned(),
443 ),
444 )))
445 }
446 },
447 }
448
449 Ok(CheckedHeaderColor {
450 dimensions: self,
451 color,
452 })
453 }
454}
455
456impl<'a> CheckedHeaderColor<'a> {
457 fn check_sample_values(self, image: FlatSamples<'a>) -> ImageResult<CheckedHeader<'a>> {
458 let header_maxval = match self.dimensions.unchecked.header.decoded {
459 HeaderRecord::Bitmap(_) => 1,
460 HeaderRecord::Graymap(GraymapHeader { maxwhite, .. }) => maxwhite,
461 HeaderRecord::Pixmap(PixmapHeader { maxval, .. }) => maxval,
462 HeaderRecord::Arbitrary(ArbitraryHeader { maxval, .. }) => maxval,
463 };
464
465 // We trust the image color bit count to be correct at least.
466 let max_sample = match self.color {
467 ExtendedColorType::Unknown(n) if n <= 16 => (1 << n) - 1,
468 ExtendedColorType::L1 => 1,
469 ExtendedColorType::L8
470 | ExtendedColorType::La8
471 | ExtendedColorType::Rgb8
472 | ExtendedColorType::Rgba8
473 | ExtendedColorType::Bgr8
474 | ExtendedColorType::Bgra8 => 0xff,
475 ExtendedColorType::L16
476 | ExtendedColorType::La16
477 | ExtendedColorType::Rgb16
478 | ExtendedColorType::Rgba16 => 0xffff,
479 _ => {
480 // Unsupported target color type.
481 return Err(ImageError::Unsupported(
482 UnsupportedError::from_format_and_kind(
483 ImageFormat::Pnm.into(),
484 UnsupportedErrorKind::Color(self.color),
485 ),
486 ));
487 }
488 };
489
490 // Avoid the performance heavy check if possible, e.g. if the header has been chosen by us.
491 if header_maxval < max_sample && !image.all_smaller(header_maxval) {
492 // Sample value greater than allowed for chosen header.
493 return Err(ImageError::Unsupported(
494 UnsupportedError::from_format_and_kind(
495 ImageFormat::Pnm.into(),
496 UnsupportedErrorKind::GenericFeature(
497 "Sample value greater than allowed for chosen header".to_owned(),
498 ),
499 ),
500 ));
501 }
502
503 let encoding = image.encoding_for(&self.dimensions.unchecked.header.decoded);
504
505 let image = CheckedImageBuffer::check(
506 image,
507 self.dimensions.width,
508 self.dimensions.height,
509 self.color,
510 )?;
511
512 Ok(CheckedHeader {
513 color: self,
514 encoding,
515 _image: image,
516 })
517 }
518}
519
520impl<'a> CheckedHeader<'a> {
521 fn write_header(self, writer: &mut dyn Write) -> ImageResult<TupleEncoding<'a>> {
522 self.header().write(writer)?;
523 Ok(self.encoding)
524 }
525
526 fn header(&self) -> &PnmHeader {
527 self.color.dimensions.unchecked.header
528 }
529}
530
531struct SampleWriter<'a>(&'a mut dyn Write);
532
533impl<'a> SampleWriter<'a> {
534 fn write_samples_ascii<V>(self, samples: V) -> io::Result<()>
535 where
536 V: Iterator,
537 V::Item: fmt::Display,
538 {
539 let mut auto_break_writer = AutoBreak::new(self.0, 70);
540 for value in samples {
541 write!(auto_break_writer, "{} ", value)?;
542 }
543 auto_break_writer.flush()
544 }
545
546 fn write_pbm_bits<V>(self, samples: &[V], width: u32) -> io::Result<()>
547 /* Default gives 0 for all primitives. TODO: replace this with `Zeroable` once it hits stable */
548 where
549 V: Default + Eq + Copy,
550 {
551 // The length of an encoded scanline
552 let line_width = (width - 1) / 8 + 1;
553
554 // We'll be writing single bytes, so buffer
555 let mut line_buffer = Vec::with_capacity(line_width as usize);
556
557 for line in samples.chunks(width as usize) {
558 for byte_bits in line.chunks(8) {
559 let mut byte = 0u8;
560 for i in 0..8 {
561 // Black pixels are encoded as 1s
562 if let Some(&v) = byte_bits.get(i) {
563 if v == V::default() {
564 byte |= 1u8 << (7 - i)
565 }
566 }
567 }
568 line_buffer.push(byte)
569 }
570 self.0.write_all(line_buffer.as_slice())?;
571 line_buffer.clear();
572 }
573
574 self.0.flush()
575 }
576}
577
578impl<'a> FlatSamples<'a> {
579 fn len(&self) -> usize {
580 match *self {
581 FlatSamples::U8(arr) => arr.len(),
582 FlatSamples::U16(arr) => arr.len(),
583 }
584 }
585
586 fn all_smaller(&self, max_val: u32) -> bool {
587 match *self {
588 FlatSamples::U8(arr) => arr.iter().any(|&val| u32::from(val) > max_val),
589 FlatSamples::U16(arr) => arr.iter().any(|&val| u32::from(val) > max_val),
590 }
591 }
592
593 fn encoding_for(&self, header: &HeaderRecord) -> TupleEncoding<'a> {
594 match *header {
595 HeaderRecord::Bitmap(BitmapHeader {
596 encoding: SampleEncoding::Binary,
597 width,
598 ..
599 }) => TupleEncoding::PbmBits {
600 samples: *self,
601 width,
602 },
603
604 HeaderRecord::Bitmap(BitmapHeader {
605 encoding: SampleEncoding::Ascii,
606 ..
607 }) => TupleEncoding::Ascii { samples: *self },
608
609 HeaderRecord::Arbitrary(_) => TupleEncoding::Bytes { samples: *self },
610
611 HeaderRecord::Graymap(GraymapHeader {
612 encoding: SampleEncoding::Ascii,
613 ..
614 })
615 | HeaderRecord::Pixmap(PixmapHeader {
616 encoding: SampleEncoding::Ascii,
617 ..
618 }) => TupleEncoding::Ascii { samples: *self },
619
620 HeaderRecord::Graymap(GraymapHeader {
621 encoding: SampleEncoding::Binary,
622 ..
623 })
624 | HeaderRecord::Pixmap(PixmapHeader {
625 encoding: SampleEncoding::Binary,
626 ..
627 }) => TupleEncoding::Bytes { samples: *self },
628 }
629 }
630}
631
632impl<'a> From<&'a [u8]> for FlatSamples<'a> {
633 fn from(samples: &'a [u8]) -> Self {
634 FlatSamples::U8(samples)
635 }
636}
637
638impl<'a> From<&'a [u16]> for FlatSamples<'a> {
639 fn from(samples: &'a [u16]) -> Self {
640 FlatSamples::U16(samples)
641 }
642}
643
644impl<'a> TupleEncoding<'a> {
645 fn write_image(&self, writer: &mut dyn Write) -> ImageResult<()> {
646 match *self {
647 TupleEncoding::PbmBits {
648 samples: FlatSamples::U8(samples),
649 width,
650 } => SampleWriter(writer)
651 .write_pbm_bits(samples, width)
652 .map_err(ImageError::IoError),
653 TupleEncoding::PbmBits {
654 samples: FlatSamples::U16(samples),
655 width,
656 } => SampleWriter(writer)
657 .write_pbm_bits(samples, width)
658 .map_err(ImageError::IoError),
659
660 TupleEncoding::Bytes {
661 samples: FlatSamples::U8(samples),
662 } => writer.write_all(samples).map_err(ImageError::IoError),
663 TupleEncoding::Bytes {
664 samples: FlatSamples::U16(samples),
665 } => samples.iter().try_for_each(|&sample| {
666 writer
667 .write_u16::<BigEndian>(sample)
668 .map_err(ImageError::IoError)
669 }),
670
671 TupleEncoding::Ascii {
672 samples: FlatSamples::U8(samples),
673 } => SampleWriter(writer)
674 .write_samples_ascii(samples.iter())
675 .map_err(ImageError::IoError),
676 TupleEncoding::Ascii {
677 samples: FlatSamples::U16(samples),
678 } => SampleWriter(writer)
679 .write_samples_ascii(samples.iter())
680 .map_err(ImageError::IoError),
681 }
682 }
683}
684