pint/
decoder.rs

1use inflate::inflate_bytes_zlib;
2use std::str;
3use std::{fs::File, io::Read};
4
5mod tests;
6use crate::types::RGB;
7
8#[derive(Default)]
9struct PngChunk {
10    chunk_type: String,
11    data_len: usize,
12    chunk_len: usize,
13    data: Vec<u8>,
14    crc: Vec<u8>,
15}
16
17#[derive(PartialEq, Debug)]
18enum ColorType {
19    TrueColorRGB,
20    Indexed,
21}
22impl ColorType {
23    fn num_channels(&self) -> usize {
24        match self {
25            ColorType::TrueColorRGB => 3,
26            ColorType::Indexed => 1,
27        }
28    }
29}
30struct IHDRData {
31    width: u32,
32    height: u32,
33    bit_depth: u8,
34    color_type: ColorType,
35}
36
37impl Default for IHDRData {
38    fn default() -> IHDRData {
39        IHDRData {
40            width: 0,
41            height: 0,
42            bit_depth: 0,
43            color_type: ColorType::TrueColorRGB,
44        }
45    }
46}
47
48fn to_hex_string(bytes: Vec<u8>) -> Vec<String> {
49    bytes.iter().map(|b| format!("{:02X}", b)).collect()
50}
51
52pub fn check_valid_png(file: &mut File) {
53    const VALID_PNG: &str = "89504E470D0A1A0A";
54
55    let first_bytes: &mut Vec<u8> = &mut vec![0; 8];
56    let _ = File::read(file, first_bytes);
57    if to_hex_string(first_bytes.to_vec()).join("") != VALID_PNG {
58        eprintln!("pint: given file is not a valid png");
59        std::process::exit(0);
60    }
61}
62
63fn bytes_to_int(bytes_arr: &[u8]) -> u32 {
64    let mut dst = [0u8; 4];
65    dst.clone_from_slice(bytes_arr);
66    u32::from_be_bytes(dst)
67}
68fn parse_png_chunks(buf: &[u8]) -> Option<PngChunk> {
69    let mut result: PngChunk = PngChunk::default();
70    let chunk_types = ["IHDR", "PLTE", "IDAT", "IEND"];
71
72    result.data_len = bytes_to_int(&buf[0..4]) as usize;
73    result.chunk_len = result.data_len + 12;
74    result.chunk_type = match str::from_utf8(&buf[4..8]) {
75        Ok(s) => s.to_string(),
76        Err(_) => return None,
77    };
78
79    if !chunk_types.contains(&result.chunk_type.as_str()) {
80        return None;
81    }
82
83    result.data = buf[8..8 + result.data_len].to_vec();
84    result.crc = buf[8 + result.data_len..8 + result.data_len].to_vec();
85
86    Some(result)
87}
88
89fn parse_ihdr(data: Vec<u8>) -> IHDRData {
90    let result = IHDRData {
91        width: bytes_to_int(&data[0..4]),
92        height: bytes_to_int(&data[4..8]),
93        bit_depth: data[8],
94        color_type: match data[9] {
95            2 => ColorType::TrueColorRGB,
96            3 => ColorType::Indexed,
97            _ => panic!("Image Color-type is not Indexed or TrueColorRGB"),
98        },
99    };
100    assert!(result.bit_depth == 8, "images need bit-depth of 8");
101    result
102}
103
104#[derive(Clone, Copy, Debug)]
105enum RGBorU8 {
106    Rgb(RGB),
107    U8(u8),
108}
109
110fn get_current_pixel(
111    current_row: &[u8],
112    line_pos: usize,
113    plte: &Option<Vec<RGB>>,
114    color_type: &ColorType,
115) -> RGBorU8 {
116    match color_type {
117        ColorType::TrueColorRGB => RGBorU8::Rgb(RGB(
118            current_row[line_pos],
119            current_row[line_pos + 1],
120            current_row[line_pos + 2],
121        )),
122        ColorType::Indexed => RGBorU8::U8(current_row[line_pos]),
123    }
124}
125
126fn none(current_pixel: u8, prev_pixel: u8, up_pixel: u8, diag_pixel: u8) -> u8 {
127    current_pixel
128}
129fn sub_filter(current_pixel: u8, prev_pixel: u8, up_pixel: u8, diag_pixel: u8) -> u8 {
130    // can safely cast back to u8 because % 256 <= 255
131    ((current_pixel as u16 + prev_pixel as u16) % 256_u16) as u8
132}
133fn up_filter(current_pixel: u8, prev_pixel: u8, up_pixel: u8, diag_pixel: u8) -> u8 {
134    ((current_pixel as u16 + up_pixel as u16) % 256_u16) as u8
135}
136fn avg_filter(current_pixel: u8, prev_pixel: u8, up_pixel: u8, diag_pixel: u8) -> u8 {
137    ((current_pixel as i16 + (up_pixel + prev_pixel) as i16 / 2) % 256) as u8
138}
139fn paeth_filter(current_pixel: u8, prev_pixel: u8, up_pixel: u8, diag_pixel: u8) -> u8 {
140    let p = prev_pixel as i16 + up_pixel as i16 - diag_pixel as i16;
141    let p_left_pixel = (p - prev_pixel as i16).abs();
142    let p_up_pixel = (p - up_pixel as i16).abs();
143    let p_diag_pixel = (p - diag_pixel as i16).abs();
144
145    let mut prediction = 0;
146    if p_left_pixel <= p_up_pixel && p_left_pixel <= p_diag_pixel {
147        prediction = prev_pixel;
148    } else if p_up_pixel <= p_diag_pixel {
149        prediction = up_pixel;
150    } else {
151        prediction = diag_pixel;
152    }
153    ((current_pixel as i16 + prediction as i16) % 256) as u8
154}
155
156fn do_filter(
157    i: usize,
158    current: RGBorU8,
159    prev: &mut RGBorU8,
160    up: RGBorU8,
161    diag: RGBorU8,
162    plte: &Option<Vec<RGB>>,
163) -> RGB {
164    let filter = [none, sub_filter, up_filter, avg_filter, paeth_filter];
165    assert!(i < filter.len(), "No such filter exists");
166
167    match current {
168        RGBorU8::Rgb(c) => {
169            let p = match prev {
170                RGBorU8::Rgb(p) => p,
171                _ => panic!("RGB values dont store u8"),
172            };
173
174            let u = match up {
175                RGBorU8::Rgb(u) => u,
176                _ => panic!("RGB values dont store u8"),
177            };
178            let d = match diag {
179                RGBorU8::Rgb(d) => d,
180                _ => panic!("RGB values dont store u8"),
181            };
182            let applied = RGB(
183                filter[i](c.0, p.0, u.0, d.0),
184                filter[i](c.1, p.1, u.1, d.1),
185                filter[i](c.2, p.2, u.2, d.2),
186            );
187            *prev = RGBorU8::Rgb(applied);
188            applied
189        }
190        RGBorU8::U8(c) => {
191            let p = match prev {
192                RGBorU8::U8(p) => *p,
193                _ => panic!("Indexed values dont store RGB"),
194            };
195            let u = match up {
196                RGBorU8::U8(u) => u,
197                _ => panic!("U8 values dont store RGB"),
198            };
199            let d = match diag {
200                RGBorU8::U8(d) => d,
201                _ => panic!("U8 values dont store RGB"),
202            };
203
204            let applied = filter[i](c, p, u, d);
205            *prev = RGBorU8::U8(applied);
206            (*plte).as_ref().expect("plte exists on indexed color-type")[applied as usize]
207        }
208    }
209}
210
211fn get_diag_pixel(
212    prev_row: &[RGBorU8],
213    line_pos: usize,
214    default_val: RGBorU8,
215    color_type: &ColorType,
216) -> RGBorU8 {
217    if (line_pos as i16 / color_type.num_channels() as i16) > 0 {
218        prev_row[(line_pos / color_type.num_channels()) - 1]
219    } else {
220        default_val
221    }
222}
223
224fn apply_filter(
225    current_row: &[u8],
226    prev_row: &mut Vec<RGBorU8>,
227    filter_index: usize,
228    color_type: &ColorType,
229    plte: &Option<Vec<RGB>>,
230    default_val: RGBorU8,
231) -> Vec<RGB> {
232    let mut line_pos = 0;
233    let mut prev_pixel = default_val;
234    let mut result: Vec<RGB> = Vec::new();
235    let mut current_row_applied = Vec::new();
236
237    while line_pos < current_row.len() - (color_type.num_channels() - 1) {
238        let pixel = get_current_pixel(current_row, line_pos, plte, color_type);
239        let up = prev_row[line_pos / color_type.num_channels()];
240        let diag = get_diag_pixel(prev_row, line_pos, default_val, color_type);
241
242        result.push(do_filter(
243            filter_index,
244            pixel,
245            &mut prev_pixel,
246            up,
247            diag,
248            plte,
249        ));
250        current_row_applied.push(prev_pixel);
251
252        line_pos += color_type.num_channels();
253    }
254    *prev_row = current_row_applied;
255
256    result
257}
258
259fn parse_data(
260    image_data: Vec<Vec<u8>>,
261    meta_data: IHDRData,
262    plte: Option<Vec<RGB>>,
263    default_val: RGBorU8,
264) -> Vec<Vec<RGB>> {
265    let mut inflated = Vec::new();
266    for i in image_data {
267        inflated.append(&mut inflate_bytes_zlib(&i as &[u8]).unwrap());
268    }
269    let mut rgb_img: Vec<Vec<RGB>> = Vec::new();
270    let mut j = 0;
271    let byte_width = meta_data.width as usize * meta_data.color_type.num_channels();
272    let mut prev_row: Vec<RGBorU8> = vec![default_val; meta_data.width as usize];
273
274    let mut i = 0;
275    while j + byte_width < inflated.len() {
276        i += 1;
277        let filter = inflated[j] as usize;
278        let current_row = &inflated[j + 1..j + byte_width + 1]; // + 1 for the line-filter
279
280        rgb_img.push(apply_filter(
281            current_row,
282            &mut prev_row,
283            filter,
284            &meta_data.color_type,
285            &plte,
286            default_val,
287        ));
288
289        j += byte_width + 1;
290    }
291
292    rgb_img
293}
294
295fn parse_plte(data: Vec<u8>) -> Vec<RGB> {
296    assert!(data.len() % 3 == 0, "data should be splittable in triplets");
297
298    let mut result: Vec<RGB> = Vec::new();
299    for i in (0..data.len()).step_by(3) {
300        result.push(RGB(data[i], data[i + 1], data[i + 2]))
301    }
302    result
303}
304
305pub fn decode_png(mut file: File) -> Vec<Vec<RGB>> {
306    let file_byte_size = File::metadata(&file).unwrap().len();
307    let buf: &mut [u8] = &mut vec![0; (file_byte_size - 8) as usize]; // cut of beginning identifier sequence
308    let result = File::read(&mut file, buf);
309    let mut i: u32 = 0;
310    let mut meta_data: IHDRData = IHDRData::default();
311    let mut plte: Option<Vec<RGB>> = None;
312    let mut data: Vec<Vec<u8>> = Vec::new();
313    let mut rgb_img: Vec<Vec<RGB>> = Vec::new();
314
315    assert!(file_byte_size - 8 > 0);
316    while i < (file_byte_size - 8) as u32 {
317        let chunk = match parse_png_chunks(&buf[i as usize..]) {
318            Some(c) => c,
319            None => {
320                i += 1;
321                continue;
322            }
323        };
324
325        match chunk.chunk_type.as_str() {
326            "IHDR" => meta_data = parse_ihdr(chunk.data),
327            "PLTE" => plte = Some(parse_plte(chunk.data)),
328            "IDAT" => {
329                if meta_data.color_type == ColorType::Indexed && plte == None {
330                    // if color is indexed then it needs a palette
331                    panic!("pint: couldn't find PLTE chunk in image.");
332                }
333                data.push(chunk.data);
334            }
335            "IEND" => {
336                let default_val = match meta_data.color_type {
337                    ColorType::TrueColorRGB => RGBorU8::Rgb(RGB(0, 0, 0)),
338                    ColorType::Indexed => RGBorU8::U8(0),
339                };
340                rgb_img = parse_data(data, meta_data, plte, default_val);
341
342                break;
343            }
344            _ => {
345                i += 1;
346                continue;
347            }
348        }
349        i += chunk.chunk_len as u32;
350    }
351    rgb_img
352}
353pub fn infer_codel_size(rgb_img: &Vec<Vec<RGB>>) -> i32 {
354    let mut min_size = i32::MAX;
355    let mut current_size = 1i32;
356    let mut current_color: RGB;
357
358    // go from left to right through img
359    for y in 0..rgb_img.len() {
360        let mut prev_color: RGB = rgb_img[y][0];
361        for x in 1..rgb_img[0].len() {
362            current_color = rgb_img[y][x];
363            if prev_color == current_color {
364                current_size += 1;
365            } else if current_size < min_size {
366                min_size = current_size;
367                current_size = 1;
368            }
369            prev_color = current_color;
370        }
371    }
372    // go from top to bottom through img
373    for x in 0..rgb_img[0].len() {
374        let mut prev_color: RGB = rgb_img[0][x];
375        for y in 1..rgb_img.len() {
376            current_color = rgb_img[y][x];
377            if prev_color == current_color {
378                current_size += 1;
379            } else if current_size < min_size {
380                min_size = current_size;
381                current_size = 1;
382            }
383            prev_color = current_color;
384        }
385    }
386    if min_size < 1 {
387        eprintln!("warning: inferred codel-size less than 1 => defaults to 1");
388        1
389    } else if rgb_img[0].len() as i32 % min_size != 0 || rgb_img.len() as i32 % min_size != 0 {
390        eprintln!("warning: inferred codel-size doesnt fit image-dimensions => defaults to 1");
391        1
392    } else {
393        min_size
394    }
395}
pint/decoder.rs

pint/
decoder.rs
