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::ExtendedColorType;
11use crate::error::{
12 ImageError, ImageResult, ParameterError, ParameterErrorKind, UnsupportedError,
13 UnsupportedErrorKind,
14};
15use crate::image::{ImageEncoder, ImageFormat};
16
17use byteorder_lite::{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: ExtendedColorType,
148 ) -> ImageResult<()>
149 where
150 S: Into<FlatSamples<'s>>,
151 {
152 let image = image.into();
153 match self.header {
154 HeaderStrategy::Dynamic => self.write_dynamic_header(image, width, height, color),
155 HeaderStrategy::Subtype(subtype) => {
156 self.write_subtyped_header(subtype, image, width, height, color)
157 }
158 HeaderStrategy::Chosen(ref header) => {
159 Self::write_with_header(&mut self.writer, header, image, width, height, color)
160 }
161 }
162 }
163
164 /// Choose any valid pnm format that the image can be expressed in and write its header.
165 ///
166 /// Returns how the body should be written if successful.
167 fn write_dynamic_header(
168 &mut self,
169 image: FlatSamples,
170 width: u32,
171 height: u32,
172 color: ExtendedColorType,
173 ) -> ImageResult<()> {
174 let depth = u32::from(color.channel_count());
175 let (maxval, tupltype) = match color {
176 ExtendedColorType::L1 => (1, ArbitraryTuplType::BlackAndWhite),
177 ExtendedColorType::L8 => (0xff, ArbitraryTuplType::Grayscale),
178 ExtendedColorType::L16 => (0xffff, ArbitraryTuplType::Grayscale),
179 ExtendedColorType::La1 => (1, ArbitraryTuplType::BlackAndWhiteAlpha),
180 ExtendedColorType::La8 => (0xff, ArbitraryTuplType::GrayscaleAlpha),
181 ExtendedColorType::La16 => (0xffff, ArbitraryTuplType::GrayscaleAlpha),
182 ExtendedColorType::Rgb8 => (0xff, ArbitraryTuplType::RGB),
183 ExtendedColorType::Rgb16 => (0xffff, ArbitraryTuplType::RGB),
184 ExtendedColorType::Rgba8 => (0xff, ArbitraryTuplType::RGBAlpha),
185 ExtendedColorType::Rgba16 => (0xffff, ArbitraryTuplType::RGBAlpha),
186 _ => {
187 return Err(ImageError::Unsupported(
188 UnsupportedError::from_format_and_kind(
189 ImageFormat::Pnm.into(),
190 UnsupportedErrorKind::Color(color),
191 ),
192 ))
193 }
194 };
195
196 let header = PnmHeader {
197 decoded: HeaderRecord::Arbitrary(ArbitraryHeader {
198 width,
199 height,
200 depth,
201 maxval,
202 tupltype: Some(tupltype),
203 }),
204 encoded: None,
205 };
206
207 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
208 }
209
210 /// Try to encode the image with the chosen format, give its corresponding pixel encoding type.
211 fn write_subtyped_header(
212 &mut self,
213 subtype: PnmSubtype,
214 image: FlatSamples,
215 width: u32,
216 height: u32,
217 color: ExtendedColorType,
218 ) -> ImageResult<()> {
219 let header = match (subtype, color) {
220 (PnmSubtype::ArbitraryMap, color) => {
221 return self.write_dynamic_header(image, width, height, color)
222 }
223 (PnmSubtype::Pixmap(encoding), ExtendedColorType::Rgb8) => PnmHeader {
224 decoded: HeaderRecord::Pixmap(PixmapHeader {
225 encoding,
226 width,
227 height,
228 maxval: 255,
229 }),
230 encoded: None,
231 },
232 (PnmSubtype::Graymap(encoding), ExtendedColorType::L8) => PnmHeader {
233 decoded: HeaderRecord::Graymap(GraymapHeader {
234 encoding,
235 width,
236 height,
237 maxwhite: 255,
238 }),
239 encoded: None,
240 },
241 (PnmSubtype::Bitmap(encoding), ExtendedColorType::L8 | ExtendedColorType::L1) => {
242 PnmHeader {
243 decoded: HeaderRecord::Bitmap(BitmapHeader {
244 encoding,
245 height,
246 width,
247 }),
248 encoded: None,
249 }
250 }
251 (_, _) => {
252 return Err(ImageError::Parameter(ParameterError::from_kind(
253 ParameterErrorKind::Generic(
254 "Color type can not be represented in the chosen format".to_owned(),
255 ),
256 )));
257 }
258 };
259
260 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
261 }
262
263 /// Try to encode the image with the chosen header, checking if values are correct.
264 ///
265 /// Returns how the body should be written if successful.
266 fn write_with_header(
267 writer: &mut dyn Write,
268 header: &PnmHeader,
269 image: FlatSamples,
270 width: u32,
271 height: u32,
272 color: ExtendedColorType,
273 ) -> ImageResult<()> {
274 let unchecked = UncheckedHeader { header };
275
276 unchecked
277 .check_header_dimensions(width, height)?
278 .check_header_color(color)?
279 .check_sample_values(image)?
280 .write_header(writer)?
281 .write_image(writer)
282 }
283}
284
285impl<W: Write> ImageEncoder for PnmEncoder<W> {
286 #[track_caller]
287 fn write_image(
288 mut self,
289 buf: &[u8],
290 width: u32,
291 height: u32,
292 color_type: ExtendedColorType,
293 ) -> ImageResult<()> {
294 let expected_buffer_len: u64 = color_type.buffer_size(width, height);
295 assert_eq!(
296 expected_buffer_len,
297 buf.len() as u64,
298 "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
299 buf.len(),
300 );
301
302 self.encode(image:buf, width, height, color_type)
303 }
304}
305
306impl<'a> CheckedImageBuffer<'a> {
307 fn check(
308 image: FlatSamples<'a>,
309 width: u32,
310 height: u32,
311 color: ExtendedColorType,
312 ) -> ImageResult<CheckedImageBuffer<'a>> {
313 let components = color.channel_count() as usize;
314 let uwidth = width as usize;
315 let uheight = height as usize;
316 let expected_len = components
317 .checked_mul(uwidth)
318 .and_then(|v| v.checked_mul(uheight));
319 if Some(image.len()) != expected_len {
320 // Image buffer does not correspond to size and colour.
321 return Err(ImageError::Parameter(ParameterError::from_kind(
322 ParameterErrorKind::DimensionMismatch,
323 )));
324 }
325 Ok(CheckedImageBuffer {
326 _image: image,
327 _width: width,
328 _height: height,
329 _color: color,
330 })
331 }
332}
333
334impl<'a> UncheckedHeader<'a> {
335 fn check_header_dimensions(
336 self,
337 width: u32,
338 height: u32,
339 ) -> ImageResult<CheckedDimensions<'a>> {
340 if self.header.width() != width || self.header.height() != height {
341 // Chosen header does not match Image dimensions.
342 return Err(ImageError::Parameter(ParameterError::from_kind(
343 ParameterErrorKind::DimensionMismatch,
344 )));
345 }
346
347 Ok(CheckedDimensions {
348 unchecked: self,
349 width,
350 height,
351 })
352 }
353}
354
355impl<'a> CheckedDimensions<'a> {
356 // Check color compatibility with the header. This will only error when we are certain that
357 // the combination is bogus (e.g. combining Pixmap and Palette) but allows uncertain
358 // combinations (basically a ArbitraryTuplType::Custom with any color of fitting depth).
359 fn check_header_color(self, color: ExtendedColorType) -> ImageResult<CheckedHeaderColor<'a>> {
360 let components = u32::from(color.channel_count());
361
362 match *self.unchecked.header {
363 PnmHeader {
364 decoded: HeaderRecord::Bitmap(_),
365 ..
366 } => match color {
367 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
368 _ => {
369 return Err(ImageError::Parameter(ParameterError::from_kind(
370 ParameterErrorKind::Generic(
371 "PBM format only support luma color types".to_owned(),
372 ),
373 )))
374 }
375 },
376 PnmHeader {
377 decoded: HeaderRecord::Graymap(_),
378 ..
379 } => match color {
380 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
381 _ => {
382 return Err(ImageError::Parameter(ParameterError::from_kind(
383 ParameterErrorKind::Generic(
384 "PGM format only support luma color types".to_owned(),
385 ),
386 )))
387 }
388 },
389 PnmHeader {
390 decoded: HeaderRecord::Pixmap(_),
391 ..
392 } => match color {
393 ExtendedColorType::Rgb8 => (),
394 _ => {
395 return Err(ImageError::Parameter(ParameterError::from_kind(
396 ParameterErrorKind::Generic(
397 "PPM format only support ExtendedColorType::Rgb8".to_owned(),
398 ),
399 )))
400 }
401 },
402 PnmHeader {
403 decoded:
404 HeaderRecord::Arbitrary(ArbitraryHeader {
405 depth,
406 ref tupltype,
407 ..
408 }),
409 ..
410 } => match (tupltype, color) {
411 (&Some(ArbitraryTuplType::BlackAndWhite), ExtendedColorType::L1) => (),
412 (&Some(ArbitraryTuplType::BlackAndWhiteAlpha), ExtendedColorType::La8) => (),
413
414 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L1) => (),
415 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L8) => (),
416 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L16) => (),
417 (&Some(ArbitraryTuplType::GrayscaleAlpha), ExtendedColorType::La8) => (),
418
419 (&Some(ArbitraryTuplType::RGB), ExtendedColorType::Rgb8) => (),
420 (&Some(ArbitraryTuplType::RGBAlpha), ExtendedColorType::Rgba8) => (),
421
422 (&None, _) if depth == components => (),
423 (&Some(ArbitraryTuplType::Custom(_)), _) if depth == components => (),
424 _ if depth != components => {
425 return Err(ImageError::Parameter(ParameterError::from_kind(
426 ParameterErrorKind::Generic(format!(
427 "Depth mismatch: header {depth} vs. color {components}"
428 )),
429 )))
430 }
431 _ => {
432 return Err(ImageError::Parameter(ParameterError::from_kind(
433 ParameterErrorKind::Generic(
434 "Invalid color type for selected PAM color type".to_owned(),
435 ),
436 )))
437 }
438 },
439 }
440
441 Ok(CheckedHeaderColor {
442 dimensions: self,
443 color,
444 })
445 }
446}
447
448impl<'a> CheckedHeaderColor<'a> {
449 fn check_sample_values(self, image: FlatSamples<'a>) -> ImageResult<CheckedHeader<'a>> {
450 let header_maxval = match self.dimensions.unchecked.header.decoded {
451 HeaderRecord::Bitmap(_) => 1,
452 HeaderRecord::Graymap(GraymapHeader { maxwhite, .. }) => maxwhite,
453 HeaderRecord::Pixmap(PixmapHeader { maxval, .. }) => maxval,
454 HeaderRecord::Arbitrary(ArbitraryHeader { maxval, .. }) => maxval,
455 };
456
457 // We trust the image color bit count to be correct at least.
458 let max_sample = match self.color {
459 ExtendedColorType::Unknown(n) if n <= 16 => (1 << n) - 1,
460 ExtendedColorType::L1 => 1,
461 ExtendedColorType::L8
462 | ExtendedColorType::La8
463 | ExtendedColorType::Rgb8
464 | ExtendedColorType::Rgba8
465 | ExtendedColorType::Bgr8
466 | ExtendedColorType::Bgra8 => 0xff,
467 ExtendedColorType::L16
468 | ExtendedColorType::La16
469 | ExtendedColorType::Rgb16
470 | ExtendedColorType::Rgba16 => 0xffff,
471 _ => {
472 // Unsupported target color type.
473 return Err(ImageError::Unsupported(
474 UnsupportedError::from_format_and_kind(
475 ImageFormat::Pnm.into(),
476 UnsupportedErrorKind::Color(self.color),
477 ),
478 ));
479 }
480 };
481
482 // Avoid the performance heavy check if possible, e.g. if the header has been chosen by us.
483 if header_maxval < max_sample && !image.all_smaller(header_maxval) {
484 // Sample value greater than allowed for chosen header.
485 return Err(ImageError::Unsupported(
486 UnsupportedError::from_format_and_kind(
487 ImageFormat::Pnm.into(),
488 UnsupportedErrorKind::GenericFeature(
489 "Sample value greater than allowed for chosen header".to_owned(),
490 ),
491 ),
492 ));
493 }
494
495 let encoding = image.encoding_for(&self.dimensions.unchecked.header.decoded);
496
497 let image = CheckedImageBuffer::check(
498 image,
499 self.dimensions.width,
500 self.dimensions.height,
501 self.color,
502 )?;
503
504 Ok(CheckedHeader {
505 color: self,
506 encoding,
507 _image: image,
508 })
509 }
510}
511
512impl<'a> CheckedHeader<'a> {
513 fn write_header(self, writer: &mut dyn Write) -> ImageResult<TupleEncoding<'a>> {
514 self.header().write(writer)?;
515 Ok(self.encoding)
516 }
517
518 fn header(&self) -> &PnmHeader {
519 self.color.dimensions.unchecked.header
520 }
521}
522
523struct SampleWriter<'a>(&'a mut dyn Write);
524
525impl SampleWriter<'_> {
526 fn write_samples_ascii<V>(self, samples: V) -> io::Result<()>
527 where
528 V: Iterator,
529 V::Item: fmt::Display,
530 {
531 let mut auto_break_writer = AutoBreak::new(self.0, 70);
532 for value in samples {
533 write!(auto_break_writer, "{value} ")?;
534 }
535 auto_break_writer.flush()
536 }
537
538 fn write_pbm_bits<V>(self, samples: &[V], width: u32) -> io::Result<()>
539 /* Default gives 0 for all primitives. TODO: replace this with `Zeroable` once it hits stable */
540 where
541 V: Default + Eq + Copy,
542 {
543 // The length of an encoded scanline
544 let line_width = (width - 1) / 8 + 1;
545
546 // We'll be writing single bytes, so buffer
547 let mut line_buffer = Vec::with_capacity(line_width as usize);
548
549 for line in samples.chunks(width as usize) {
550 for byte_bits in line.chunks(8) {
551 let mut byte = 0u8;
552 for i in 0..8 {
553 // Black pixels are encoded as 1s
554 if let Some(&v) = byte_bits.get(i) {
555 if v == V::default() {
556 byte |= 1u8 << (7 - i);
557 }
558 }
559 }
560 line_buffer.push(byte);
561 }
562 self.0.write_all(line_buffer.as_slice())?;
563 line_buffer.clear();
564 }
565
566 self.0.flush()
567 }
568}
569
570impl<'a> FlatSamples<'a> {
571 fn len(&self) -> usize {
572 match *self {
573 FlatSamples::U8(arr) => arr.len(),
574 FlatSamples::U16(arr) => arr.len(),
575 }
576 }
577
578 fn all_smaller(&self, max_val: u32) -> bool {
579 match *self {
580 FlatSamples::U8(arr) => arr.iter().any(|&val| u32::from(val) > max_val),
581 FlatSamples::U16(arr) => arr.iter().any(|&val| u32::from(val) > max_val),
582 }
583 }
584
585 fn encoding_for(&self, header: &HeaderRecord) -> TupleEncoding<'a> {
586 match *header {
587 HeaderRecord::Bitmap(BitmapHeader {
588 encoding: SampleEncoding::Binary,
589 width,
590 ..
591 }) => TupleEncoding::PbmBits {
592 samples: *self,
593 width,
594 },
595
596 HeaderRecord::Bitmap(BitmapHeader {
597 encoding: SampleEncoding::Ascii,
598 ..
599 }) => TupleEncoding::Ascii { samples: *self },
600
601 HeaderRecord::Arbitrary(_) => TupleEncoding::Bytes { samples: *self },
602
603 HeaderRecord::Graymap(GraymapHeader {
604 encoding: SampleEncoding::Ascii,
605 ..
606 })
607 | HeaderRecord::Pixmap(PixmapHeader {
608 encoding: SampleEncoding::Ascii,
609 ..
610 }) => TupleEncoding::Ascii { samples: *self },
611
612 HeaderRecord::Graymap(GraymapHeader {
613 encoding: SampleEncoding::Binary,
614 ..
615 })
616 | HeaderRecord::Pixmap(PixmapHeader {
617 encoding: SampleEncoding::Binary,
618 ..
619 }) => TupleEncoding::Bytes { samples: *self },
620 }
621 }
622}
623
624impl<'a> From<&'a [u8]> for FlatSamples<'a> {
625 fn from(samples: &'a [u8]) -> Self {
626 FlatSamples::U8(samples)
627 }
628}
629
630impl<'a> From<&'a [u16]> for FlatSamples<'a> {
631 fn from(samples: &'a [u16]) -> Self {
632 FlatSamples::U16(samples)
633 }
634}
635
636impl TupleEncoding<'_> {
637 fn write_image(&self, writer: &mut dyn Write) -> ImageResult<()> {
638 match *self {
639 TupleEncoding::PbmBits {
640 samples: FlatSamples::U8(samples),
641 width,
642 } => SampleWriter(writer)
643 .write_pbm_bits(samples, width)
644 .map_err(ImageError::IoError),
645 TupleEncoding::PbmBits {
646 samples: FlatSamples::U16(samples),
647 width,
648 } => SampleWriter(writer)
649 .write_pbm_bits(samples, width)
650 .map_err(ImageError::IoError),
651
652 TupleEncoding::Bytes {
653 samples: FlatSamples::U8(samples),
654 } => writer.write_all(samples).map_err(ImageError::IoError),
655 TupleEncoding::Bytes {
656 samples: FlatSamples::U16(samples),
657 } => samples.iter().try_for_each(|&sample| {
658 writer
659 .write_u16::<BigEndian>(sample)
660 .map_err(ImageError::IoError)
661 }),
662
663 TupleEncoding::Ascii {
664 samples: FlatSamples::U8(samples),
665 } => SampleWriter(writer)
666 .write_samples_ascii(samples.iter())
667 .map_err(ImageError::IoError),
668 TupleEncoding::Ascii {
669 samples: FlatSamples::U16(samples),
670 } => SampleWriter(writer)
671 .write_samples_ascii(samples.iter())
672 .map_err(ImageError::IoError),
673 }
674 }
675}
676