iris-cv 0.0.0

A fast computer vision library framework in Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use crate::error::Result;
use crate::image::Image;
use burn::tensor::{Tensor, TensorData, backend::Backend};
use std::path::Path;

/// Loads an image from the specified file path.
/// The image is automatically converted to RGB/Grayscale and represented as a [C, H, W] float tensor.
pub fn load_image<B: Backend>(path: impl AsRef<Path>, device: &B::Device) -> Result<Image<B>> {
    let img = image::open(path)?;
    let img = img.to_rgb8();
    let (width, height) = img.dimensions();
    let w = width as usize;
    let h = height as usize;
    let c = 3usize;

    let mut flat_data = vec![0.0f32; c * h * w];
    let pixels = img.as_flat_samples();
    let slice = pixels.as_slice();

    for y in 0..h {
        for x in 0..w {
            let idx = (y * w + x) * 3;
            // Map HWC (file layout) to CHW (standard deep learning tensor layout)
            flat_data[y * w + x] = f32::from(slice[idx]) / 255.0; // R
            flat_data[h * w + y * w + x] = f32::from(slice[idx + 1]) / 255.0; // G
            flat_data[2 * h * w + y * w + x] = f32::from(slice[idx + 2]) / 255.0; // B
        }
    }

    let tensor_data = TensorData::new(flat_data, [c, h, w]);
    let tensor = Tensor::<B, 3>::from_data(tensor_data, device);
    Ok(Image::new(tensor))
}

/// Saves the image tensor to the specified file path.
/// Assumes image tensor layout is [C, H, W] with float values in [0.0, 1.0].
pub fn save_image<B: Backend>(image: &Image<B>, path: impl AsRef<Path>) -> Result<()> {
    let dims = image.tensor.dims();
    let c = dims[0];
    let h = dims[1];
    let w = dims[2];

    // Read the tensor data back to the host CPU
    // Under Burn 0.21.0, into_data() fetches the tensor data synchronously on the CPU.
    let tensor_data = image.tensor.clone().into_data();
    let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();

    let mut img_buf = image::ImageBuffer::new(w as u32, h as u32);

    for y in 0..h {
        for x in 0..w {
            let r_val = flat_vals[y * w + x];
            let g_val = if c > 1 {
                flat_vals[h * w + y * w + x]
            } else {
                r_val
            };
            let b_val = if c > 2 {
                flat_vals[2 * h * w + y * w + x]
            } else {
                r_val
            };

            let r = (r_val.clamp(0.0, 1.0) * 255.0) as u8;
            let g = (g_val.clamp(0.0, 1.0) * 255.0) as u8;
            let b = (b_val.clamp(0.0, 1.0) * 255.0) as u8;

            img_buf.put_pixel(x as u32, y as u32, image::Rgb([r, g, b]));
        }
    }

    img_buf.save(path)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::{TestBackend, test_device};

    #[test]
    fn test_image_io() {
        let device = test_device();
        let flat_data = vec![0.5f32; 3 * 8 * 8];
        let tensor =
            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 8, 8]), &device);
        let img = Image::new(tensor);

        let temp_path = "temp_test_io.png";
        save_image(&img, temp_path).unwrap();

        let loaded = load_image::<TestBackend>(temp_path, &device).unwrap();
        assert_eq!(loaded.shape(), [3, 8, 8]);

        let _ = std::fs::remove_file(temp_path);
    }
}