ndarray_vision/format/
netpbm.rs

1use crate::core::{normalise_pixel_value, Image, ImageBase, PixelBound, RGB};
2use crate::format::{Decoder, Encoder};
3use ndarray::Data;
4use num_traits::cast::{FromPrimitive, NumCast};
5use num_traits::{Num, NumAssignOps};
6use std::fmt::Display;
7use std::io::{Error, ErrorKind};
8
9#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
10enum EncodingType {
11    Binary,
12    Plaintext,
13}
14
15/// Encoder type for a PPM image.
16#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
17pub struct PpmEncoder {
18    encoding: EncodingType,
19}
20
21/// Decoder type for a PPM image.
22#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Default)]
23pub struct PpmDecoder;
24
25impl Default for PpmEncoder {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31/// Implements the encoder trait for the PpmEncoder.
32///
33/// The ColourModel type argument is locked to RGB - this prevents calling
34/// RGB::into::<RGB>() unnecessarily which is unavoidable until trait specialisation is
35/// stabilised.
36impl<T, U> Encoder<T, U, RGB> for PpmEncoder
37where
38    U: Data<Elem = T>,
39    T: Copy
40        + Clone
41        + Num
42        + NumAssignOps
43        + NumCast
44        + PartialOrd
45        + Display
46        + PixelBound
47        + FromPrimitive,
48{
49    fn encode(&self, image: &ImageBase<U, RGB>) -> Vec<u8> {
50        use EncodingType::*;
51        match self.encoding {
52            Plaintext => self.encode_plaintext(image),
53            Binary => self.encode_binary(image),
54        }
55    }
56}
57
58impl PpmEncoder {
59    /// Create a new PPM encoder or decoder
60    pub fn new() -> Self {
61        PpmEncoder {
62            encoding: EncodingType::Binary,
63        }
64    }
65
66    /// Creates a new PPM format to encode plain-text. This results in very large
67    /// file sizes so isn't recommended in general use
68    pub fn new_plaintext_encoder() -> Self {
69        PpmEncoder {
70            encoding: EncodingType::Plaintext,
71        }
72    }
73
74    /// Gets the maximum pixel value in the image across all channels. This is
75    /// used in the PPM header
76    fn get_max_value<T, U>(image: &ImageBase<U, RGB>) -> Option<u8>
77    where
78        U: Data<Elem = T>,
79        T: Copy + Clone + Num + NumAssignOps + NumCast + PartialOrd + Display + PixelBound,
80    {
81        image
82            .data
83            .iter()
84            .fold(T::zero(), |ref acc, x| if x > acc { *x } else { *acc })
85            .to_u8()
86    }
87
88    ///! Generate the header string for the image
89    fn generate_header(self, rows: usize, cols: usize, max_value: u8) -> String {
90        use EncodingType::*;
91        match self.encoding {
92            Plaintext => format!("P3\n{} {} {}\n", rows, cols, max_value),
93            Binary => format!("P6\n{} {} {}\n", rows, cols, max_value),
94        }
95    }
96
97    /// Encode the image into the binary PPM format (P6) returning the bytes
98    fn encode_binary<T, U>(self, image: &ImageBase<U, RGB>) -> Vec<u8>
99    where
100        U: Data<Elem = T>,
101        T: Copy + Clone + Num + NumAssignOps + NumCast + PartialOrd + Display + PixelBound,
102    {
103        let max_val = Self::get_max_value(image).unwrap_or(255);
104
105        let mut result = self
106            .generate_header(image.rows(), image.cols(), max_val)
107            .into_bytes();
108
109        result.reserve(result.len() + (image.rows() * image.cols() * 3));
110
111        for data in image.data.iter() {
112            let value = (normalise_pixel_value(*data) * 255.0f64) as u8;
113            result.push(value);
114        }
115        result
116    }
117
118    /// Encode the image into the plaintext PPM format (P3) returning the text as
119    /// an array of bytes
120    fn encode_plaintext<T, U>(self, image: &ImageBase<U, RGB>) -> Vec<u8>
121    where
122        U: Data<Elem = T>,
123        T: Copy + Clone + Num + NumAssignOps + NumCast + PartialOrd + Display + PixelBound,
124    {
125        let max_val = 255;
126
127        let mut result = self.generate_header(image.rows(), image.cols(), max_val);
128        // Not very accurate as a reserve, doesn't factor in max storage for
129        // a pixel or spaces. But somewhere between best and worst case
130        result.reserve(image.rows() * image.cols() * 5);
131
132        // There is a 70 character line length in PPM using another string to keep track
133        let mut temp = String::new();
134        let max_margin = 70 - 12;
135        temp.reserve(max_margin);
136
137        for data in image.data.iter() {
138            let value = (normalise_pixel_value(*data) * 255.0f64) as u8;
139            temp.push_str(&format!("{} ", value));
140            if temp.len() > max_margin {
141                result.push_str(&temp);
142                result.push('\n');
143                temp.clear();
144            }
145        }
146        if !temp.is_empty() {
147            result.push_str(&temp);
148        }
149        result.into_bytes()
150    }
151}
152
153/// Implements the decoder trait for the PpmDecoder.
154///
155/// The ColourModel type argument is locked to RGB - this prevents calling
156/// RGB::into::<RGB>() unnecessarily which is unavoidable until trait specialisation is
157/// stabilised.
158impl<T> Decoder<T, RGB> for PpmDecoder
159where
160    T: Copy
161        + Clone
162        + Num
163        + NumAssignOps
164        + NumCast
165        + PartialOrd
166        + Display
167        + PixelBound
168        + FromPrimitive,
169{
170    fn decode(&self, bytes: &[u8]) -> std::io::Result<Image<T, RGB>> {
171        if bytes.len() < 9 {
172            Err(Error::new(
173                ErrorKind::InvalidData,
174                "File is below minimum size of ppm",
175            ))
176        } else if bytes.starts_with(b"P3") {
177            Self::decode_plaintext(&bytes[2..])
178        } else if bytes.starts_with(b"P6") {
179            Self::decode_binary(&bytes[2..])
180        } else {
181            Err(Error::new(
182                ErrorKind::InvalidData,
183                "File is below minimum size of ppm",
184            ))
185        }
186    }
187}
188
189impl PpmDecoder {
190    /// Decodes a PPM header getting (rows, cols, maximum value) or returning
191    /// an io::Error if the header is malformed
192    fn decode_header(bytes: &[u8]) -> std::io::Result<(usize, usize, usize)> {
193        let err = || Error::new(ErrorKind::InvalidData, "Error in file header");
194        let mut keep = true;
195        let bytes = bytes
196            .iter()
197            .filter(|x| {
198                if *x == &b'#' {
199                    keep = false;
200                    false
201                } else if !keep {
202                    if *x == &b'\n' || *x == &b'\r' {
203                        keep = true;
204                    }
205                    false
206                } else {
207                    true
208                }
209            })
210            .cloned()
211            .collect::<Vec<_>>();
212
213        if let Ok(s) = String::from_utf8(bytes) {
214            let res = s
215                .split_whitespace()
216                .map(|x| x.parse::<usize>().unwrap_or(0))
217                .collect::<Vec<_>>();
218            if res.len() == 3 {
219                Ok((res[0], res[1], res[2]))
220            } else {
221                Err(err())
222            }
223        } else {
224            Err(err())
225        }
226    }
227
228    fn decode_binary<T>(bytes: &[u8]) -> std::io::Result<Image<T, RGB>>
229    where
230        T: Copy
231            + Clone
232            + Num
233            + NumAssignOps
234            + NumCast
235            + PartialOrd
236            + Display
237            + PixelBound
238            + FromPrimitive,
239    {
240        let err = || Error::new(ErrorKind::InvalidData, "Error in file encoding");
241        const WHITESPACE: &[u8] = b" \t\n\r";
242
243        let mut image_bytes = Vec::<T>::new();
244
245        let mut last_saw_whitespace = false;
246        let mut is_comment = false;
247        let mut val_count = 0;
248        let header_end = bytes
249            .iter()
250            .position(|&b| {
251                if b == b'#' {
252                    is_comment = true;
253                } else if is_comment {
254                    if b == b'\r' || b == b'\n' {
255                        is_comment = false;
256                    }
257                } else if last_saw_whitespace && !WHITESPACE.contains(&b) {
258                    val_count += 1;
259                    last_saw_whitespace = false;
260                } else if WHITESPACE.contains(&b) {
261                    last_saw_whitespace = true;
262                }
263                val_count == 3 && WHITESPACE.contains(&b)
264            })
265            .ok_or_else(err)?;
266
267        let (rows, cols, max_val) = Self::decode_header(&bytes[0..header_end])?;
268        for b in bytes.iter().skip(header_end + 1) {
269            let real_pixel = (*b as f64) * (255.0f64 / (max_val as f64));
270            image_bytes.push(T::from_u8(real_pixel as u8).unwrap_or_else(T::zero));
271        }
272
273        if image_bytes.is_empty() || image_bytes.len() != (rows * cols * 3) {
274            Err(err())
275        } else {
276            let image = Image::<T, RGB>::from_shape_data(rows, cols, image_bytes);
277            Ok(image)
278        }
279    }
280
281    fn decode_plaintext<T>(bytes: &[u8]) -> std::io::Result<Image<T, RGB>>
282    where
283        T: Copy
284            + Clone
285            + Num
286            + NumAssignOps
287            + NumCast
288            + PartialOrd
289            + Display
290            + PixelBound
291            + FromPrimitive,
292    {
293        let err = || Error::new(ErrorKind::InvalidData, "Error in file encoding");
294        // plaintext is easier than binary because the whole thing is a string
295        let data = String::from_utf8(bytes.to_vec()).map_err(|_| err())?;
296
297        let mut rows = -1;
298        let mut cols = -1;
299        let mut max_val = -1;
300        let mut image_bytes = Vec::<T>::new();
301        for line in data.lines().filter(|l| !l.starts_with('#')) {
302            for value in line.split_whitespace().take_while(|x| !x.starts_with('#')) {
303                let temp = value.parse::<isize>().map_err(|_| err())?;
304                if rows < 0 {
305                    rows = temp;
306                } else if cols < 0 {
307                    cols = temp;
308                    image_bytes.reserve((rows * cols * 3) as usize);
309                } else if max_val < 0 {
310                    max_val = temp;
311                } else {
312                    let real_pixel = (temp as f64) * (255.0f64 / (max_val as f64));
313                    image_bytes.push(T::from_f64(real_pixel).unwrap_or_else(T::zero));
314                }
315            }
316        }
317        if image_bytes.is_empty() || image_bytes.len() != ((rows * cols * 3) as usize) {
318            Err(err())
319        } else {
320            let image = Image::<T, RGB>::from_shape_data(rows as usize, cols as usize, image_bytes);
321            Ok(image)
322        }
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use crate::core::colour_models::*;
330    use ndarray::prelude::*;
331    use ndarray_rand::RandomExt;
332    use rand::distributions::Uniform;
333    use std::fs::remove_file;
334
335    #[test]
336    fn max_value_test() {
337        let full_range = "P3 1 1 255 0 255 0";
338        let clamped = "P3 1 1 1 0 1 0";
339
340        let decoder = PpmDecoder::default();
341        let full_image: Image<u8, RGB> = decoder.decode(full_range.as_bytes()).unwrap();
342        let clamp_image: Image<u8, RGB> = decoder.decode(clamped.as_bytes()).unwrap();
343
344        assert_eq!(full_image, clamp_image);
345        assert_eq!(full_image.pixel(0, 0), arr1(&[0, 255, 0]));
346    }
347
348    #[test]
349    fn encoding_consistency() {
350        let image_str = "P3 
351            3 3 255 
352            255 255 255  0 0 0  255 0 0 
353            0 255 0  0 0 255  255 255 0
354            0 255 255  127 127 127  0 0 0";
355
356        let decoder = PpmDecoder::default();
357        let image: Image<u8, RGB> = decoder.decode(image_str.as_bytes()).unwrap();
358
359        let encoder = PpmEncoder::new();
360        let image_bytes = encoder.encode(&image);
361
362        let restored: Image<u8, RGB> = decoder.decode(&image_bytes).unwrap();
363
364        assert_eq!(image, restored);
365
366        let encoder = PpmEncoder::new_plaintext_encoder();
367        let image_bytes = encoder.encode(&image);
368        let restored: Image<u8, RGB> = decoder.decode(&image_bytes).unwrap();
369
370        assert_eq!(image, restored);
371    }
372
373    #[test]
374    fn binary_comments() {
375        let image_str = "P3 
376            3 3 255 
377            255 255 255  0 0 0  255 0 0 
378            0 255 0  0 0 255  255 255 0
379            0 255 255  127 127 127  0 0 0";
380
381        let decoder = PpmDecoder::default();
382        let image: Image<u8, RGB> = decoder.decode(image_str.as_bytes()).unwrap();
383
384        let encoder = PpmEncoder::new();
385        let mut image_bytes = encoder.encode(&image);
386        let comment = b"# This is a comment\n";
387        for i in 0..comment.len() {
388            image_bytes.insert(2 + i, comment[i]);
389        }
390        let restored: Image<u8, RGB> = decoder.decode(&image_bytes).unwrap();
391
392        assert_eq!(image, restored);
393    }
394
395    #[test]
396    fn binary_file_save() {
397        let mut image = Image::<u8, RGB>::new(480, 640);
398        let new_data = Array3::<u8>::random(image.data.dim(), Uniform::new(0, 255));
399        image.data = new_data;
400
401        let bin_encoder = PpmEncoder::new();
402
403        let filename = "bintest.ppm";
404
405        bin_encoder.encode_file(&image, filename).unwrap();
406
407        let decoder = PpmDecoder::default();
408        let new_image = decoder.decode_file(filename).unwrap();
409        let _ = remove_file(filename);
410
411        image_compare(&new_image, &image);
412    }
413
414    #[test]
415    fn plaintext_file_save() {
416        let mut image = Image::<u8, RGB>::new(480, 640);
417        let new_data = Array3::<u8>::random(image.data.dim(), Uniform::new(0, 255));
418        image.data = new_data;
419
420        let bin_encoder = PpmEncoder::new_plaintext_encoder();
421        let filename = "texttest.ppm";
422
423        bin_encoder.encode_file(&image, filename).unwrap();
424
425        let decoder = PpmDecoder::default();
426        let new_image = decoder.decode_file(filename).unwrap();
427        let _ = remove_file(filename);
428
429        image_compare(&new_image, &image);
430    }
431
432    fn image_compare<C>(actual: &Image<u8, C>, expected: &Image<u8, C>)
433    where
434        C: ColourModel,
435    {
436        assert_eq!(actual.data.shape(), expected.data.shape());
437
438        for (act, exp) in actual.data.iter().zip(expected.data.iter()) {
439            let delta = (*act as i16 - *exp as i16).abs();
440            // An error of 1 is acceptable on any value due to rounding
441            assert!(delta < 2);
442        }
443    }
444}
ndarray_vision/format/netpbm.rs

ndarray_vision/format/
netpbm.rs
