iris/image/
mod.rs

1pub mod geometric;
2pub mod io;
3pub mod ops;
4
5pub use geometric::GeometricTransform;
6
7use burn::tensor::{Tensor, backend::Backend};
8
9/// Represents an image represented by a 3D Burn Tensor of shape [C, H, W] (Channels, Height, Width).
10/// The pixel values are floats, normally scaled to [0.0, 1.0].
11#[derive(Clone, Debug)]
12pub struct Image<B: Backend> {
13    /// The underlying 3D tensor of shape [C, H, W].
14    pub tensor: Tensor<B, 3>,
15}
16
17impl<B: Backend> Image<B> {
18    /// Creates a new Image from a 3D Burn Tensor.
19    pub fn new(tensor: Tensor<B, 3>) -> Self {
20        Self { tensor }
21    }
22
23    /// Loads an image from a file path using the default/given backend device.
24    pub fn open(
25        path: impl AsRef<std::path::Path>,
26        device: &B::Device,
27    ) -> crate::error::Result<Self> {
28        io::load_image(path, device)
29    }
30
31    /// Saves the image to a file path.
32    pub fn save(&self, path: impl AsRef<std::path::Path>) -> crate::error::Result<()> {
33        io::save_image(self, path)
34    }
35
36    /// Returns the shape [C, H, W] of the image tensor.
37    pub fn shape(&self) -> [usize; 3] {
38        let dims = self.tensor.dims();
39        [dims[0], dims[1], dims[2]]
40    }
41
42    /// Returns the number of channels (C) of the image.
43    pub fn channels(&self) -> usize {
44        self.tensor.dims()[0]
45    }
46
47    /// Returns the height (H) of the image.
48    pub fn height(&self) -> usize {
49        self.tensor.dims()[1]
50    }
51
52    /// Returns the width (W) of the image.
53    pub fn width(&self) -> usize {
54        self.tensor.dims()[2]
55    }
56
57    /// Performs template matching and returns the result map.
58    ///
59    /// The template slides over the source image and computes a match score at each position.
60    /// Returns a 2D tensor of shape `[H - th + 1, W - tw + 1]` where `(th, tw)` is the template size.
61    pub fn template_match(
62        &self,
63        template: &Image<B>,
64        method: crate::features::TemplateMatchMethod,
65    ) -> crate::error::Result<Tensor<B, 2>> {
66        use crate::features::template_match;
67        template_match(self, template, method)
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::test_helpers::{TestBackend, test_device};
75    use burn::backend::ndarray::NdArrayDevice;
76    use burn::tensor::TensorData;
77
78    fn get_test_device() -> NdArrayDevice {
79        test_device()
80    }
81
82    #[test]
83    fn test_image_creation() {
84        let device = get_test_device();
85        let data = TensorData::new(vec![0.5f32; 3 * 10 * 10], [3, 10, 10]);
86        let tensor = Tensor::<TestBackend, 3>::from_data(data, &device);
87        let img = Image::new(tensor);
88
89        assert_eq!(img.channels(), 3);
90        assert_eq!(img.height(), 10);
91        assert_eq!(img.width(), 10);
92        assert_eq!(img.shape(), [3, 10, 10]);
93    }
94
95    #[test]
96    fn test_image_ops() {
97        let device = get_test_device();
98        let data = TensorData::new(vec![0.5f32; 3 * 8 * 8], [3, 8, 8]);
99        let tensor = Tensor::<TestBackend, 3>::from_data(data, &device);
100        let img = Image::new(tensor);
101
102        // Crop
103        let cropped = img.crop(2, 2, 4, 4).unwrap();
104        assert_eq!(cropped.shape(), [3, 4, 4]);
105
106        // Flip
107        let flipped = img.flip(true, false).unwrap();
108        assert_eq!(flipped.shape(), [3, 8, 8]);
109
110        // Rotate
111        let rotated = img.rotate(90).unwrap();
112        assert_eq!(rotated.shape(), [3, 8, 8]); // width and height same in this case
113
114        // Grayscale
115        let gray = img.grayscale().unwrap();
116        assert_eq!(gray.channels(), 1);
117        assert_eq!(gray.height(), 8);
118        assert_eq!(gray.width(), 8);
119
120        // To RGB
121        let rgb = gray.to_rgb().unwrap();
122        assert_eq!(rgb.channels(), 3);
123    }
124}
iris/image/mod.rs

iris/image/
mod.rs
