neural_network_rs/plotter/
png_plotter.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5fn convert_to_256_color_scale(value: &f64) -> u8 {
6    (value * 255.0) as u8
7}
8
9pub fn plot_png(
10    name: &str,
11    dims: (usize, usize),
12    data: &Vec<Vec<f64>>,
13    color_type: png::ColorType,
14) -> Result<(), std::io::Error> {
15    let path = Path::new("images").join(name).with_extension("png");
16
17    if let Some(parent) = path.parent() {
18        std::fs::create_dir_all(parent)?;
19    }
20
21    let file = File::create(&path)?;
22    let ref mut w = BufWriter::new(file);
23
24    let width = dims.0 as u32;
25    let height = dims.1 as u32;
26
27    let mut encoder = png::Encoder::new(w, width, height);
28    encoder.set_color(color_type);
29
30    let mut writer = encoder.write_header()?;
31
32    let data_uint8 = data
33        .iter()
34        .flatten()
35        .map(convert_to_256_color_scale)
36        .collect::<Vec<u8>>();
37
38    writer.write_image_data(&data_uint8)?;
39
40    Ok(())
41}