retrofire_core/util/
pnm.rs

1//! PNM, also known as NetPBM, file format support.
2//!
3//! PNM is a venerable family of extremely simple image formats, each
4//! consisting of a simple textual header followed by either textual or
5//! binary pixel data.
6//!
7//! ```text
8//! Type  | Txt | Bin | Pixel format
9//! ------+-----+-----+-----------------
10//! PBM   | P1  | P4  | 1 bpp monochrome
11//! PGM   | P2  | P5  | 8 bpp grayscale
12//! PPM   | P3  | P6  | 3x8 bpp RGB
13//! ```
14
15use alloc::{string::String, vec::Vec};
16use core::{
17    fmt::{self, Debug, Display, Formatter},
18    num::{IntErrorKind, ParseIntError},
19    str::FromStr,
20};
21#[cfg(feature = "std")]
22use std::{
23    fs::File,
24    io::{self, BufReader, BufWriter, Read, Write},
25    path::Path,
26};
27
28use crate::math::{Color3, rgb};
29
30use super::{Dims, buf::Buf2};
31#[cfg(feature = "std")]
32use super::{
33    buf::AsSlice2,
34    pixfmt::{IntoPixel, Rgb888},
35};
36
37use Error::*;
38use Format::*;
39
40/// The header of a PNM image.
41#[derive(Copy, Clone, Debug, Eq, PartialEq)]
42struct Header {
43    format: Format,
44    dims: Dims,
45    #[allow(unused)]
46    // TODO Currently not used
47    max: u16,
48}
49
50/// The format of a PNM image.
51#[derive(Copy, Clone, Debug, Eq, PartialEq)]
52#[allow(unused)]
53#[repr(u16)]
54enum Format {
55    /// 1-bit monochrome image, text encoding.
56    TextBitmap = magic(b"P1"),
57    /// Grayscale image, text encoding.
58    TextGraymap = magic(b"P2"),
59    /// RGB image, text encoding.
60    TextPixmap = magic(b"P3"),
61    /// 1-bit monochrome image, packed binary encoding.
62    BinaryBitmap = magic(b"P4"),
63    /// Grayscale image, binary encoding. 1 byte per pixel.
64    BinaryGraymap = magic(b"P5"),
65    /// RGB image, binary encoding. 3 bytes per pixel.
66    BinaryPixmap = magic(b"P6"),
67}
68
69const fn magic(bytes: &[u8; 2]) -> u16 {
70    u16::from_be_bytes(*bytes)
71}
72
73impl Display for Format {
74    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
75        write!(f, "P{}", *self as u8 as char)
76    }
77}
78
79impl TryFrom<[u8; 2]> for Format {
80    type Error = Error;
81    fn try_from(magic: [u8; 2]) -> Result<Self> {
82        Ok(match &magic {
83            b"P2" => TextGraymap,
84            b"P3" => TextPixmap,
85            b"P4" => BinaryBitmap,
86            b"P5" => BinaryGraymap,
87            b"P6" => BinaryPixmap,
88            other => Err(Unsupported(*other))?,
89        })
90    }
91}
92
93// Error during loading or decoding a PNM file.
94#[derive(Debug, Eq, PartialEq)]
95pub enum Error {
96    /// An I/O error occurred.
97    #[cfg(feature = "std")]
98    Io(io::ErrorKind),
99    /// Unsupported magic number.
100    Unsupported([u8; 2]),
101    /// Unexpected end of input while decoding.
102    UnexpectedEnd,
103    /// Invalid numeric value encountered.
104    InvalidNumber,
105}
106
107/// Result of loading or decoding a PNM file.
108pub type Result<T> = core::result::Result<T, Error>;
109
110impl core::error::Error for Error {}
111
112impl Display for Error {
113    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
114        match *self {
115            #[cfg(feature = "std")]
116            Io(kind) => write!(f, "i/o error {kind}"),
117            Unsupported([c, d]) => {
118                write!(f, "unsupported magic number {}{}", c as char, d as char)
119            }
120            UnexpectedEnd => write!(f, "unexpected end of input"),
121            InvalidNumber => write!(f, "invalid numeric value"),
122        }
123    }
124}
125
126impl From<ParseIntError> for Error {
127    fn from(e: ParseIntError) -> Self {
128        if *e.kind() == IntErrorKind::Empty {
129            UnexpectedEnd
130        } else {
131            InvalidNumber
132        }
133    }
134}
135
136#[cfg(feature = "std")]
137impl From<io::Error> for Error {
138    fn from(e: io::Error) -> Self {
139        Io(e.kind())
140    }
141}
142
143impl Header {
144    /// Attempts to parse a PNM header from an iterator.
145    ///
146    /// Currently supported formats are P2, P3, P4, P5, and P6.
147    fn parse(input: impl IntoIterator<Item = u8>) -> Result<Self> {
148        let mut it = input.into_iter();
149        let magic = [
150            it.next().ok_or(UnexpectedEnd)?,
151            it.next().ok_or(UnexpectedEnd)?,
152        ];
153        let format = magic.try_into()?;
154        let dims = (parse_num(&mut it)?, parse_num(&mut it)?);
155        let max: u16 = match &format {
156            TextBitmap | BinaryBitmap => 1,
157            _ => parse_num(&mut it)?,
158        };
159        Ok(Self { format, dims, max })
160    }
161    /// Writes `self` to `dest` as a valid PNM header,
162    /// including a trailing newline.
163    #[cfg(feature = "std")]
164    fn write(&self, mut dest: impl Write) -> io::Result<()> {
165        let Self { format, dims: (w, h), max } = *self;
166        let max: &dyn Display = match format {
167            TextBitmap | BinaryBitmap => &"",
168            _ => &max,
169        };
170        writeln!(dest, "{format} {w} {h} {max}")
171    }
172}
173
174/// Loads a PNM image from a path into a buffer.
175///
176/// Currently supported formats are P2, P3, P4, P5, and P6.
177/// Read more about the formats in the [module docs][self].
178///
179/// # Errors
180/// Returns [`pnm::Error`][Error] in case of an I/O error or invalid PNM image.
181#[cfg(feature = "std")]
182pub fn load_pnm(path: impl AsRef<Path>) -> Result<Buf2<Color3>> {
183    let r = &mut BufReader::new(File::open(path)?);
184    read_pnm(r)
185}
186
187/// Reads a PNM image into a buffer.
188///
189/// Currently supported PNM formats are P2, P3, P4, P5, and P6.
190/// Read more about the formats in the [module docs][self].
191///
192/// # Errors
193/// Returns [`pnm::Error`][Error] in case of an I/O error or invalid PNM image.
194#[cfg(feature = "std")]
195pub fn read_pnm(input: impl Read) -> Result<Buf2<Color3>> {
196    let input = BufReader::new(input);
197    parse_pnm(input.bytes().map_while(io::Result::ok))
198}
199
200/// Attempts to decode a PNM image from an iterator of bytes.
201///
202/// Currently supported PNM formats are P2, P3, P4, P5, and P6.
203/// Read more about the formats in the [module docs][self].
204///
205/// # Errors
206/// Returns [`Error`] in case of an invalid or unrecognized PNM image.
207pub fn parse_pnm(input: impl IntoIterator<Item = u8>) -> Result<Buf2<Color3>> {
208    let mut it = input.into_iter();
209    let h = Header::parse(&mut it)?;
210
211    let count = h.dims.0 * h.dims.1;
212    let data: Vec<Color3> = match h.format {
213        BinaryPixmap => {
214            let mut col = [0u8; 3];
215            it.zip((0..3).cycle())
216                .flat_map(|(c, i)| {
217                    col[i] = c;
218                    (i == 2).then(|| col.into())
219                })
220                .take(count as usize)
221                .collect()
222        }
223        BinaryGraymap => it //
224            .map(|c| rgb(c, c, c))
225            .collect(),
226        BinaryBitmap => it
227            .flat_map(|byte| (0..8).rev().map(move |i| (byte >> i) & 1))
228            .map(|bit| {
229                // Conventionally in PBM 0 is white, 1 is black
230                let ch = (1 - bit) * 0xFF;
231                rgb(ch, ch, ch)
232            })
233            .collect(),
234        TextPixmap => {
235            let mut col = [0u8; 3];
236            (0..3)
237                .cycle()
238                .flat_map(|i| {
239                    col[i] = match parse_num(&mut it) {
240                        Ok(c) => c,
241                        Err(e) => return Some(Err(e)),
242                    };
243                    (i == 2).then(|| Ok(col.into()))
244                })
245                .take(count as usize)
246                .collect::<Result<Vec<_>>>()?
247        }
248        TextGraymap => (0..count)
249            .map(|_| {
250                let val = parse_num(&mut it)?;
251                Ok(rgb(val, val, val))
252            })
253            .collect::<Result<Vec<_>>>()?,
254        _ => return Err(Unsupported((h.format as u16).to_be_bytes())),
255    };
256
257    if data.len() < count as usize {
258        Err(UnexpectedEnd)
259    } else {
260        Ok(Buf2::new_from(h.dims, data))
261    }
262}
263
264/// Writes an image to a file in PPM format, P6 sub-format
265/// (binary 8-bits-per-channel RGB).
266///
267/// Caution: This function overwrites the file if it already exists.
268/// Use [`write_ppm`] for more control over file creation.
269///
270/// # Errors
271/// Returns [`std::io::Error`] if an error occurs while writing.
272#[cfg(feature = "std")]
273pub fn save_ppm<T>(
274    path: impl AsRef<Path>,
275    data: impl AsSlice2<T>,
276) -> io::Result<()>
277where
278    T: IntoPixel<[u8; 3], Rgb888> + Copy,
279{
280    let out = BufWriter::new(File::create(path)?);
281    write_ppm(out, data)
282}
283
284/// Writes an image to `out` in PPM format, P6 sub-format
285/// (binary 8-bits-per-channel RGB).
286///
287/// # Errors
288/// Returns [`std::io::Error`] if an error occurs while writing.
289#[cfg(feature = "std")]
290pub fn write_ppm<T>(
291    mut out: impl Write,
292    data: impl AsSlice2<T>,
293) -> io::Result<()>
294where
295    T: IntoPixel<[u8; 3], Rgb888> + Copy,
296{
297    let slice = data.as_slice2();
298    Header {
299        format: BinaryPixmap,
300        dims: slice.dims(),
301        max: 255,
302    }
303    .write(&mut out)?;
304
305    // Appease the borrow checker
306    slice
307        .rows()
308        .flatten()
309        .map(|c| c.into_pixel())
310        .try_for_each(|rgb| out.write_all(&rgb[..]))
311}
312
313/// Parses a numeric value from `src`, skipping whitespace and comments.
314fn parse_num<T>(src: impl IntoIterator<Item = u8>) -> Result<T>
315where
316    T: FromStr,
317    Error: From<T::Err>,
318{
319    let mut whitespace_or_comment = {
320        let mut in_comment = false;
321        move |b: &u8| match *b {
322            b'#' => {
323                in_comment = true;
324                true
325            }
326            b'\n' => {
327                in_comment = false;
328                true
329            }
330            _ => in_comment || b.is_ascii_whitespace(),
331        }
332    };
333    let str = src
334        .into_iter()
335        .skip_while(whitespace_or_comment)
336        .take_while(|b| !whitespace_or_comment(b))
337        .map(char::from)
338        .collect::<String>();
339    Ok(str.parse()?)
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    #[test]
347    fn parse_value_int() {
348        assert_eq!(parse_num(*b"123"), Ok(123));
349        assert_eq!(parse_num(*b"12345"), Ok(12345));
350    }
351
352    #[test]
353    fn parse_num_empty() {
354        assert_eq!(parse_num::<i32>(*b""), Err(UnexpectedEnd));
355    }
356
357    #[test]
358    fn parse_num_with_whitespace() {
359        assert_eq!(parse_num(*b" \n\n   42 "), Ok(42));
360    }
361
362    #[test]
363    fn parse_num_with_comment_before() {
364        assert_eq!(parse_num(*b"# this is a comment\n42"), Ok(42));
365    }
366
367    #[test]
368    fn parse_num_with_comment_after() {
369        assert_eq!(parse_num(*b"42#this is a comment"), Ok(42));
370    }
371
372    #[test]
373    fn parse_header_whitespace() {
374        assert_eq!(
375            Header::parse(*b"P6 123\t \n\r321      255 "),
376            Ok(Header {
377                format: BinaryPixmap,
378                dims: (123, 321),
379                max: 255,
380            })
381        );
382    }
383
384    #[test]
385    fn parse_header_comment() {
386        assert_eq!(
387            Header::parse(*b"P6 # foo 42\n 123\n#bar\n#baz\n321 255 "),
388            Ok(Header {
389                format: BinaryPixmap,
390                dims: (123, 321),
391                max: 255,
392            })
393        );
394    }
395
396    #[test]
397    fn parse_header_p2() {
398        assert_eq!(
399            Header::parse(*b"P2 123 456 789"),
400            Ok(Header {
401                format: TextGraymap,
402                dims: (123, 456),
403                max: 789,
404            })
405        );
406    }
407
408    #[test]
409    fn parse_header_p3() {
410        assert_eq!(
411            Header::parse(*b"P3 123 456 789"),
412            Ok(Header {
413                format: TextPixmap,
414                dims: (123, 456),
415                max: 789,
416            })
417        );
418    }
419
420    #[test]
421    fn parse_header_p4() {
422        assert_eq!(
423            Header::parse(*b"P4 123 456 "),
424            Ok(Header {
425                format: BinaryBitmap,
426                dims: (123, 456),
427                max: 1,
428            })
429        );
430    }
431
432    #[test]
433    fn parse_header_p5() {
434        assert_eq!(
435            Header::parse(*b"P5 123 456 789 "),
436            Ok(Header {
437                format: BinaryGraymap,
438                dims: (123, 456),
439                max: 789,
440            })
441        );
442    }
443
444    #[test]
445    fn parse_header_p6() {
446        assert_eq!(
447            Header::parse(*b"P6 123 456 789 "),
448            Ok(Header {
449                format: BinaryPixmap,
450                dims: (123, 456),
451                max: 789,
452            })
453        );
454    }
455
456    #[test]
457    fn parse_header_unsupported_magic() {
458        let res = Header::parse(*b"P7 1 1 1 ");
459        assert_eq!(res, Err(Unsupported(*b"P7")));
460    }
461
462    #[test]
463    fn parse_header_invalid_magic() {
464        let res = Header::parse(*b"FOO");
465        assert_eq!(res, Err(Unsupported(*b"FO")));
466    }
467
468    #[test]
469    fn parse_header_invalid_dims() {
470        assert_eq!(Header::parse(*b"P5 abc 1 1 "), Err(InvalidNumber));
471        assert_eq!(Header::parse(*b"P5 1 1 "), Err(UnexpectedEnd));
472        assert_eq!(Header::parse(*b"P6 1 -1 1 "), Err(InvalidNumber));
473    }
474
475    #[test]
476    fn parse_pnm_truncated() {
477        let data = *b"P3 2 2 256 \n 0 0 0   123 0 42   0 64 128";
478        assert_eq!(parse_pnm(data).err(), Some(UnexpectedEnd));
479    }
480
481    #[cfg(feature = "std")]
482    #[test]
483    fn write_header_p1() {
484        let mut out = Vec::new();
485        let hdr = Header {
486            format: TextBitmap,
487            dims: (123, 456),
488            max: 1,
489        };
490        hdr.write(&mut out).unwrap();
491        assert_eq!(&out, b"P1 123 456 \n");
492    }
493
494    #[cfg(feature = "std")]
495    #[test]
496    fn write_header_p6() {
497        let mut out = Vec::new();
498        let hdr = Header {
499            format: BinaryPixmap,
500            dims: (123, 456),
501            max: 789,
502        };
503        hdr.write(&mut out).unwrap();
504        assert_eq!(&out, b"P6 123 456 789\n");
505    }
506
507    #[test]
508    fn read_pnm_p2() {
509        let data = *b"P2 2 2 128 \n 12 34 56 78";
510
511        let buf = parse_pnm(data).unwrap();
512
513        assert_eq!(buf.width(), 2);
514        assert_eq!(buf.height(), 2);
515
516        assert_eq!(buf[[0, 0]], rgb(12, 12, 12));
517        assert_eq!(buf[[1, 0]], rgb(34, 34, 34));
518        assert_eq!(buf[[0, 1]], rgb(56, 56, 56));
519        assert_eq!(buf[[1, 1]], rgb(78, 78, 78));
520    }
521
522    #[test]
523    fn read_pnm_p3() {
524        let data = *b"P3 2 2 256 \n 0 0 0   123 0 42   0 64 128   255 255 255";
525
526        let buf = parse_pnm(data).unwrap();
527
528        assert_eq!(buf.dims(), (2, 2));
529
530        assert_eq!(buf[[0, 0]], rgb(0, 0, 0));
531        assert_eq!(buf[[1, 0]], rgb(123, 0, 42));
532        assert_eq!(buf[[0, 1]], rgb(0, 64, 128));
533        assert_eq!(buf[[1, 1]], rgb(255, 255, 255));
534    }
535
536    #[test]
537    fn read_pnm_p4() {
538        // 0x69 == 0b0110_1001
539        let buf = parse_pnm(*b"P4 4 2\n\x69").unwrap();
540
541        assert_eq!(buf.dims(), (4, 2));
542
543        let b = rgb(0u8, 0, 0);
544        let w = rgb(0xFFu8, 0xFF, 0xFF);
545
546        assert_eq!(buf[0usize], [w, b, b, w]);
547        assert_eq!(buf[1usize], [b, w, w, b]);
548    }
549
550    #[test]
551    fn read_pnm_p5() {
552        let buf = parse_pnm(*b"P5 2 2 255\n\x01\x23\x45\x67").unwrap();
553
554        assert_eq!(buf.dims(), (2, 2));
555
556        assert_eq!(buf[0usize], [rgb(0x01, 0x01, 0x01), rgb(0x23, 0x23, 0x23)]);
557        assert_eq!(buf[1usize], [rgb(0x45, 0x45, 0x45), rgb(0x67, 0x67, 0x67)]);
558    }
559
560    #[test]
561    fn read_pnm_p6() {
562        let buf = parse_pnm(
563            *b"P6 2 2 255\n\
564            \x01\x12\x23\
565            \x34\x45\x56\
566            \x67\x78\x89\
567            \x9A\xAB\xBC",
568        )
569        .unwrap();
570
571        assert_eq!(buf.dims(), (2, 2));
572
573        assert_eq!(buf[0usize], [rgb(0x01, 0x12, 0x23), rgb(0x34, 0x45, 0x56)]);
574        assert_eq!(buf[1usize], [rgb(0x67, 0x78, 0x89), rgb(0x9A, 0xAB, 0xBC)]);
575    }
576
577    #[cfg(feature = "std")]
578    #[test]
579    fn write_ppm() {
580        use alloc::vec;
581        let buf = vec![
582            rgb(0xFF, 0, 0),
583            rgb(0, 0xFF, 0),
584            rgb(0, 0, 0xFF),
585            rgb(0xFF, 0xFF, 0),
586        ];
587
588        let mut out = vec![];
589        super::write_ppm(&mut out, Buf2::new_from((2, 2), buf)).unwrap();
590
591        assert_eq!(
592            &out,
593            b"P6 2 2 255\n\
594              \xFF\x00\x00\
595              \x00\xFF\x00\
596              \x00\x00\xFF\
597              \xFF\xFF\x00"
598        );
599    }
600}
retrofire_core/util/pnm.rs

retrofire_core/util/
pnm.rs
