image_proc.rs

use image::*;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use crate::{Error, CHAN};

pub struct RawImage {
    pub rows: usize,
    pub cols: usize,
    pub pixels: Vec<u8>,
}

fn read_bin_u32(f: &mut File) -> Result<u32, Error> {
    let mut buffer = [0; 4];
    f.read_exact(&mut buffer)?;
    Ok(u32::from_le_bytes(buffer))
}

pub fn read_raw<P: AsRef<Path>>(image: P) -> Result<RawImage, Error> {
    let mut file = File::open(image)?;
    let rows = read_bin_u32(&mut file)? as usize;
    let cols = read_bin_u32(&mut file)? as usize;
    let chans = read_bin_u32(&mut file)? as usize;

    assert!(
        chans == CHAN,
        "Channel size read from the binary file doesn't match the default value"
    );

    let size: usize = rows * cols * CHAN;
    let mut buffer: Vec<u8> = vec![0; size];
    // The file has pixels is a 3D array with dimensions height, width, channel
    file.read_exact(&mut buffer)?;

    let chunked_channels = buffer.chunks(CHAN).collect::<Vec<_>>();
    let chunked = chunked_channels.chunks(cols).collect::<Vec<_>>();
    let input: &[&[&[u8]]] = chunked.as_slice();

    // We need the image in a 3D array with dimensions channel, height, width
    let mut pixels: Vec<u8> = vec![0; size];
    let mut pixels_columns = pixels.chunks_mut(cols).collect::<Vec<_>>();
    let mut pixels_chunked = pixels_columns.chunks_mut(rows).collect::<Vec<_>>();
    let result: &mut [&mut [&mut [u8]]] = pixels_chunked.as_mut_slice();

    for row in 0..rows {
        for col in 0..cols {
            for chan in 0..CHAN {
                result[chan][row][col] = input[row][col][chan];
            }
        }
    }

    Ok(RawImage { rows, cols, pixels })
}

pub fn extern_image(rows: usize, cols: usize, image: &[u8]) -> RgbImage {
    let chunked_columns = image.chunks(cols).collect::<Vec<_>>();
    let chunked = chunked_columns.chunks(rows).collect::<Vec<_>>();
    let pixels: &[&[&[u8]]] = chunked.as_slice();

    RgbImage::from_fn(cols as u32, rows as u32, |c, r| {
        image::Rgb([
            pixels[0][r as usize][c as usize],
            pixels[1][r as usize][c as usize],
            pixels[2][r as usize][c as usize],
        ])
    })
}