iris/image/
io.rs

1use crate::error::Result;
2use crate::image::Image;
3use burn::tensor::{Tensor, TensorData, backend::Backend};
4use std::path::Path;
5
6/// Loads an image from the specified file path.
7/// The image is automatically converted to RGB/Grayscale and represented as a [C, H, W] float tensor.
8pub fn load_image<B: Backend>(path: impl AsRef<Path>, device: &B::Device) -> Result<Image<B>> {
9    let img = image::open(path)?;
10    let img = img.to_rgb8();
11    let (width, height) = img.dimensions();
12    let w = width as usize;
13    let h = height as usize;
14    let c = 3usize;
15
16    let mut flat_data = vec![0.0f32; c * h * w];
17    let pixels = img.as_flat_samples();
18    let slice = pixels.as_slice();
19
20    for y in 0..h {
21        for x in 0..w {
22            let idx = (y * w + x) * 3;
23            // Map HWC (file layout) to CHW (standard deep learning tensor layout)
24            flat_data[y * w + x] = f32::from(slice[idx]) / 255.0; // R
25            flat_data[h * w + y * w + x] = f32::from(slice[idx + 1]) / 255.0; // G
26            flat_data[2 * h * w + y * w + x] = f32::from(slice[idx + 2]) / 255.0; // B
27        }
28    }
29
30    let tensor_data = TensorData::new(flat_data, [c, h, w]);
31    let tensor = Tensor::<B, 3>::from_data(tensor_data, device);
32    Ok(Image::new(tensor))
33}
34
35/// Saves the image tensor to the specified file path.
36/// Assumes image tensor layout is [C, H, W] with float values in [0.0, 1.0].
37pub fn save_image<B: Backend>(image: &Image<B>, path: impl AsRef<Path>) -> Result<()> {
38    let dims = image.tensor.dims();
39    let c = dims[0];
40    let h = dims[1];
41    let w = dims[2];
42
43    // Read the tensor data back to the host CPU
44    // Under Burn 0.21.0, into_data() fetches the tensor data synchronously on the CPU.
45    let tensor_data = image.tensor.clone().into_data();
46    let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
47
48    let mut img_buf = image::ImageBuffer::new(w as u32, h as u32);
49
50    for y in 0..h {
51        for x in 0..w {
52            let r_val = flat_vals[y * w + x];
53            let g_val = if c > 1 {
54                flat_vals[h * w + y * w + x]
55            } else {
56                r_val
57            };
58            let b_val = if c > 2 {
59                flat_vals[2 * h * w + y * w + x]
60            } else {
61                r_val
62            };
63
64            let r = (r_val.clamp(0.0, 1.0) * 255.0) as u8;
65            let g = (g_val.clamp(0.0, 1.0) * 255.0) as u8;
66            let b = (b_val.clamp(0.0, 1.0) * 255.0) as u8;
67
68            img_buf.put_pixel(x as u32, y as u32, image::Rgb([r, g, b]));
69        }
70    }
71
72    img_buf.save(path)?;
73    Ok(())
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use crate::test_helpers::{TestBackend, test_device};
80
81    #[test]
82    fn test_image_io() {
83        let device = test_device();
84        let flat_data = vec![0.5f32; 3 * 8 * 8];
85        let tensor =
86            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 8, 8]), &device);
87        let img = Image::new(tensor);
88
89        let temp_path = "temp_test_io.png";
90        save_image(&img, temp_path).unwrap();
91
92        let loaded = load_image::<TestBackend>(temp_path, &device).unwrap();
93        assert_eq!(loaded.shape(), [3, 8, 8]);
94
95        let _ = std::fs::remove_file(temp_path);
96    }
97}
iris/image/io.rs

iris/image/
io.rs
