Skip to main content

ai_image/codecs/pnm/
header.rs

1use alloc::{string::String, vec::Vec};
2use core::fmt;
3use no_std_io::io;
4
5/// The kind of encoding used to store sample values
6#[derive(Clone, Copy, PartialEq, Eq, Debug)]
7pub enum SampleEncoding {
8    /// Samples are unsigned binary integers in big endian
9    Binary,
10
11    /// Samples are encoded as decimal ascii strings separated by whitespace
12    Ascii,
13}
14
15/// Denotes the category of the magic number
16#[derive(Clone, Copy, PartialEq, Eq, Debug)]
17pub enum PnmSubtype {
18    /// Magic numbers P1 and P4
19    Bitmap(SampleEncoding),
20
21    /// Magic numbers P2 and P5
22    Graymap(SampleEncoding),
23
24    /// Magic numbers P3 and P6
25    Pixmap(SampleEncoding),
26
27    /// Magic number P7
28    ArbitraryMap,
29}
30
31/// Stores the complete header data of a file.
32///
33/// Internally, provides mechanisms for lossless reencoding. After reading a file with the decoder
34/// it is possible to recover the header and construct an encoder. Using the encoder on the just
35/// loaded image should result in a byte copy of the original file (for single image pnms without
36/// additional trailing data).
37#[derive(Clone)]
38pub struct PnmHeader {
39    pub(crate) decoded: HeaderRecord,
40    pub(crate) encoded: Option<Vec<u8>>,
41}
42
43#[derive(Clone)]
44pub(crate) enum HeaderRecord {
45    Bitmap(BitmapHeader),
46    Graymap(GraymapHeader),
47    Pixmap(PixmapHeader),
48    Arbitrary(ArbitraryHeader),
49}
50
51/// Header produced by a `pbm` file ("Portable Bit Map")
52#[derive(Clone, Copy, Debug)]
53pub struct BitmapHeader {
54    /// Binary or Ascii encoded file
55    pub encoding: SampleEncoding,
56
57    /// Height of the image file
58    pub height: u32,
59
60    /// Width of the image file
61    pub width: u32,
62}
63
64/// Header produced by a `pgm` file ("Portable Gray Map")
65#[derive(Clone, Copy, Debug)]
66pub struct GraymapHeader {
67    /// Binary or Ascii encoded file
68    pub encoding: SampleEncoding,
69
70    /// Height of the image file
71    pub height: u32,
72
73    /// Width of the image file
74    pub width: u32,
75
76    /// Maximum sample value within the image
77    pub maxwhite: u32,
78}
79
80/// Header produced by a `ppm` file ("Portable Pixel Map")
81#[derive(Clone, Copy, Debug)]
82pub struct PixmapHeader {
83    /// Binary or Ascii encoded file
84    pub encoding: SampleEncoding,
85
86    /// Height of the image file
87    pub height: u32,
88
89    /// Width of the image file
90    pub width: u32,
91
92    /// Maximum sample value within the image
93    pub maxval: u32,
94}
95
96/// Header produced by a `pam` file ("Portable Arbitrary Map")
97#[derive(Clone, Debug)]
98pub struct ArbitraryHeader {
99    /// Height of the image file
100    pub height: u32,
101
102    /// Width of the image file
103    pub width: u32,
104
105    /// Number of color channels
106    pub depth: u32,
107
108    /// Maximum sample value within the image
109    pub maxval: u32,
110
111    /// Color interpretation of image pixels
112    pub tupltype: Option<ArbitraryTuplType>,
113}
114
115/// Standardized tuple type specifiers in the header of a `pam`.
116#[derive(Clone, Debug)]
117pub enum ArbitraryTuplType {
118    /// Pixels are either black (0) or white (1)
119    BlackAndWhite,
120
121    /// Pixels are either black (0) or white (1) and a second alpha channel
122    BlackAndWhiteAlpha,
123
124    /// Pixels represent the amount of white
125    Grayscale,
126
127    /// Grayscale with an additional alpha channel
128    GrayscaleAlpha,
129
130    /// Three channels: Red, Green, Blue
131    RGB,
132
133    /// Four channels: Red, Green, Blue, Alpha
134    RGBAlpha,
135
136    /// An image format which is not standardized
137    Custom(String),
138}
139
140impl ArbitraryTuplType {
141    pub(crate) fn name(&self) -> &str {
142        match self {
143            ArbitraryTuplType::BlackAndWhite => "BLACKANDWHITE",
144            ArbitraryTuplType::BlackAndWhiteAlpha => "BLACKANDWHITE_ALPHA",
145            ArbitraryTuplType::Grayscale => "GRAYSCALE",
146            ArbitraryTuplType::GrayscaleAlpha => "GRAYSCALE_ALPHA",
147            ArbitraryTuplType::RGB => "RGB",
148            ArbitraryTuplType::RGBAlpha => "RGB_ALPHA",
149            ArbitraryTuplType::Custom(custom) => custom,
150        }
151    }
152}
153
154impl PnmSubtype {
155    /// Get the two magic constant bytes corresponding to this format subtype.
156    #[must_use]
157    pub fn magic_constant(self) -> &'static [u8; 2] {
158        match self {
159            PnmSubtype::Bitmap(SampleEncoding::Ascii) => b"P1",
160            PnmSubtype::Graymap(SampleEncoding::Ascii) => b"P2",
161            PnmSubtype::Pixmap(SampleEncoding::Ascii) => b"P3",
162            PnmSubtype::Bitmap(SampleEncoding::Binary) => b"P4",
163            PnmSubtype::Graymap(SampleEncoding::Binary) => b"P5",
164            PnmSubtype::Pixmap(SampleEncoding::Binary) => b"P6",
165            PnmSubtype::ArbitraryMap => b"P7",
166        }
167    }
168
169    /// Whether samples are stored as binary or as decimal ascii
170    #[must_use]
171    pub fn sample_encoding(self) -> SampleEncoding {
172        match self {
173            PnmSubtype::ArbitraryMap => SampleEncoding::Binary,
174            PnmSubtype::Bitmap(enc) => enc,
175            PnmSubtype::Graymap(enc) => enc,
176            PnmSubtype::Pixmap(enc) => enc,
177        }
178    }
179}
180
181impl PnmHeader {
182    /// Retrieve the format subtype from which the header was created.
183    #[must_use]
184    pub fn subtype(&self) -> PnmSubtype {
185        match self.decoded {
186            HeaderRecord::Bitmap(BitmapHeader { encoding, .. }) => PnmSubtype::Bitmap(encoding),
187            HeaderRecord::Graymap(GraymapHeader { encoding, .. }) => PnmSubtype::Graymap(encoding),
188            HeaderRecord::Pixmap(PixmapHeader { encoding, .. }) => PnmSubtype::Pixmap(encoding),
189            HeaderRecord::Arbitrary(ArbitraryHeader { .. }) => PnmSubtype::ArbitraryMap,
190        }
191    }
192
193    /// The width of the image this header is for.
194    #[must_use]
195    pub fn width(&self) -> u32 {
196        match self.decoded {
197            HeaderRecord::Bitmap(BitmapHeader { width, .. }) => width,
198            HeaderRecord::Graymap(GraymapHeader { width, .. }) => width,
199            HeaderRecord::Pixmap(PixmapHeader { width, .. }) => width,
200            HeaderRecord::Arbitrary(ArbitraryHeader { width, .. }) => width,
201        }
202    }
203
204    /// The height of the image this header is for.
205    #[must_use]
206    pub fn height(&self) -> u32 {
207        match self.decoded {
208            HeaderRecord::Bitmap(BitmapHeader { height, .. }) => height,
209            HeaderRecord::Graymap(GraymapHeader { height, .. }) => height,
210            HeaderRecord::Pixmap(PixmapHeader { height, .. }) => height,
211            HeaderRecord::Arbitrary(ArbitraryHeader { height, .. }) => height,
212        }
213    }
214
215    /// The biggest value a sample can have. In other words, the colour resolution.
216    #[must_use]
217    pub fn maximal_sample(&self) -> u32 {
218        match self.decoded {
219            HeaderRecord::Bitmap(BitmapHeader { .. }) => 1,
220            HeaderRecord::Graymap(GraymapHeader { maxwhite, .. }) => maxwhite,
221            HeaderRecord::Pixmap(PixmapHeader { maxval, .. }) => maxval,
222            HeaderRecord::Arbitrary(ArbitraryHeader { maxval, .. }) => maxval,
223        }
224    }
225
226    /// Retrieve the underlying bitmap header if any
227    #[must_use]
228    pub fn as_bitmap(&self) -> Option<&BitmapHeader> {
229        match self.decoded {
230            HeaderRecord::Bitmap(ref bitmap) => Some(bitmap),
231            _ => None,
232        }
233    }
234
235    /// Retrieve the underlying graymap header if any
236    #[must_use]
237    pub fn as_graymap(&self) -> Option<&GraymapHeader> {
238        match self.decoded {
239            HeaderRecord::Graymap(ref graymap) => Some(graymap),
240            _ => None,
241        }
242    }
243
244    /// Retrieve the underlying pixmap header if any
245    #[must_use]
246    pub fn as_pixmap(&self) -> Option<&PixmapHeader> {
247        match self.decoded {
248            HeaderRecord::Pixmap(ref pixmap) => Some(pixmap),
249            _ => None,
250        }
251    }
252
253    /// Retrieve the underlying arbitrary header if any
254    #[must_use]
255    pub fn as_arbitrary(&self) -> Option<&ArbitraryHeader> {
256        match self.decoded {
257            HeaderRecord::Arbitrary(ref arbitrary) => Some(arbitrary),
258            _ => None,
259        }
260    }
261
262    /// Write the header back into a binary stream
263    pub fn write(&self, writer: &mut dyn io::Write) -> io::Result<()> {
264        writer.write_all(self.subtype().magic_constant())?;
265        match *self {
266            PnmHeader {
267                encoded: Some(ref content),
268                ..
269            } => writer.write_all(content),
270            PnmHeader {
271                decoded:
272                    HeaderRecord::Bitmap(BitmapHeader {
273                        encoding: _encoding,
274                        width,
275                        height,
276                    }),
277                ..
278            } => writeln!(writer, "\n{width} {height}"),
279            PnmHeader {
280                decoded:
281                    HeaderRecord::Graymap(GraymapHeader {
282                        encoding: _encoding,
283                        width,
284                        height,
285                        maxwhite,
286                    }),
287                ..
288            } => writeln!(writer, "\n{width} {height} {maxwhite}"),
289            PnmHeader {
290                decoded:
291                    HeaderRecord::Pixmap(PixmapHeader {
292                        encoding: _encoding,
293                        width,
294                        height,
295                        maxval,
296                    }),
297                ..
298            } => writeln!(writer, "\n{width} {height} {maxval}"),
299            PnmHeader {
300                decoded:
301                    HeaderRecord::Arbitrary(ArbitraryHeader {
302                        width,
303                        height,
304                        depth,
305                        maxval,
306                        ref tupltype,
307                    }),
308                ..
309            } => {
310                struct TupltypeWriter<'a>(&'a Option<ArbitraryTuplType>);
311                impl fmt::Display for TupltypeWriter<'_> {
312                    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
313                        match self.0 {
314                            Some(tt) => writeln!(f, "TUPLTYPE {}", tt.name()),
315                            None => Ok(()),
316                        }
317                    }
318                }
319
320                writeln!(
321                    writer,
322                    "\nWIDTH {}\nHEIGHT {}\nDEPTH {}\nMAXVAL {}\n{}ENDHDR",
323                    width,
324                    height,
325                    depth,
326                    maxval,
327                    TupltypeWriter(tupltype)
328                )
329            }
330        }
331    }
332}
333
334impl From<BitmapHeader> for PnmHeader {
335    fn from(header: BitmapHeader) -> Self {
336        PnmHeader {
337            decoded: HeaderRecord::Bitmap(header),
338            encoded: None,
339        }
340    }
341}
342
343impl From<GraymapHeader> for PnmHeader {
344    fn from(header: GraymapHeader) -> Self {
345        PnmHeader {
346            decoded: HeaderRecord::Graymap(header),
347            encoded: None,
348        }
349    }
350}
351
352impl From<PixmapHeader> for PnmHeader {
353    fn from(header: PixmapHeader) -> Self {
354        PnmHeader {
355            decoded: HeaderRecord::Pixmap(header),
356            encoded: None,
357        }
358    }
359}
360
361impl From<ArbitraryHeader> for PnmHeader {
362    fn from(header: ArbitraryHeader) -> Self {
363        PnmHeader {
364            decoded: HeaderRecord::Arbitrary(header),
365            encoded: None,
366        }
367    }
368}