rusty_vision/codec/decoders/
png.rs

1use flate2::read::ZlibDecoder;
2
3use std::fs::File;
4use std::io::{self, Read};
5use std::vec;
6
7use crate::color::ColorSpace;
8use crate::error::Error;
9use crate::geometry::Shape;
10use crate::image::Image;
11
12pub fn decode(file: &mut File) -> Result<Image, Error> {
13    let mut signature = [0; 8];
14    file.read_exact(&mut signature)?;
15
16    if &signature != b"\x89PNG\r\n\x1a\n" {
17        return Err(Error::ImageDecodeError(std::io::Error::new(
18            io::ErrorKind::InvalidData,
19            "Invalid PNG Signature",
20        )));
21    }
22
23    fn read_chunk<R: Read>(reader: &mut R) -> io::Result<([u8; 4], Vec<u8>)> {
24        let mut length_bytes = [0; 4];
25        reader.read_exact(&mut length_bytes)?;
26
27        let length = u32::from_be_bytes(length_bytes);
28
29        let mut chunk_type = [0; 4];
30        reader.read_exact(&mut chunk_type)?;
31
32        let mut data = vec![0; length as usize];
33        reader.read_exact(&mut data)?;
34
35        let mut crc_bytes = [0; 4];
36        reader.read_exact(&mut crc_bytes)?;
37
38        Ok((chunk_type, data))
39    }
40
41    let mut width = 0;
42    let mut height = 0;
43    let mut colortype = 0;
44    let mut palette = Vec::new();
45    let mut image_data = Vec::new();
46    'outer: loop {
47        let (chunk_type, chunk_data) = read_chunk(file).unwrap();
48        dbg!(String::from_utf8_lossy(&chunk_type));
49        match &chunk_type {
50            b"IHDR" => {
51                width = u32::from_be_bytes([
52                    chunk_data[0],
53                    chunk_data[1],
54                    chunk_data[2],
55                    chunk_data[3],
56                ]);
57                height = u32::from_be_bytes([
58                    chunk_data[4],
59                    chunk_data[5],
60                    chunk_data[6],
61                    chunk_data[7],
62                ]);
63
64                println!("Bit Depth: {}", chunk_data[8]);
65                colortype = chunk_data[9];
66
67                dbg!(chunk_data);
68            }
69            b"PLTE" => {
70                dbg!(&chunk_data);
71                palette = chunk_data;
72            }
73            b"IDAT" => {
74                image_data.extend(chunk_data);
75            }
76            b"IEND" => {
77                break 'outer;
78            }
79            value => {
80                panic!(
81                    "Block type {} decoding has not been implemented",
82                    String::from_utf8_lossy(value)
83                );
84            }
85        }
86    }
87
88    println!("{width:?}, {height:?}");
89    dbg!(palette.len());
90
91    let mut zlib_decoder = ZlibDecoder::new(&image_data[..]);
92    let mut decompressed = Vec::new();
93    zlib_decoder.read_to_end(&mut decompressed)?;
94
95    println!("Decompressed size:{}", decompressed.len());
96
97    // let bytes_per_pixel = colortype as usize; // RGB
98
99    let row_size = width as usize;
100    let mut data = Vec::with_capacity(width as usize * height as usize * 3);
101
102    dbg!(data.capacity());
103
104    for y in 0..height as usize {
105        // let filter_type = decompressed[y * (row_size + 1)];
106
107        let start = y * (row_size + 1) + 1;
108        let end = start + row_size;
109
110        match colortype {
111            3 => {
112                for &index in &decompressed[start..end] {
113                    // dbg!(index);
114                    let p_index = index as usize * 3;
115                    data.extend_from_slice(&palette[p_index..p_index + 3]);
116                }
117            }
118            _ => {
119                data.extend_from_slice(&decompressed[start..end]);
120            }
121        }
122    }
123
124    // TODO: Change to implicit ndim
125    let shape = Shape::new(width as usize, height as usize, Some(4));
126
127    Ok(Image::from_data(data, shape, ColorSpace::RGBA))
128}