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

retrofire_core/util/
pnm.rs
