iris/edges/
gradients.rs

1use crate::error::Result;
2use crate::image::Image;
3use burn::tensor::{Tensor, TensorData, backend::Backend};
4
5impl<B: Backend> Image<B> {
6    /// Applies Scharr operator to find horizontal and vertical gradients.
7    pub fn scharr(&self) -> Result<Self> {
8        let gray = self.grayscale()?;
9        let dims = gray.tensor.dims();
10        let h = dims[1];
11        let w = dims[2];
12
13        let device = gray.tensor.device();
14        let tensor_data = gray.tensor.clone().into_data();
15        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
16        let mut out_vals = vec![0.0f32; h * w];
17
18        let kx = [[-3.0, 0.0, 3.0], [-10.0, 0.0, 10.0], [-3.0, 0.0, 3.0]];
19        let ky = [[-3.0, -10.0, -3.0], [0.0, 0.0, 0.0], [3.0, 10.0, 3.0]];
20
21        {
22            use rayon::prelude::*;
23            out_vals
24                .par_chunks_exact_mut(w)
25                .enumerate()
26                .skip(1)
27                .take(h - 2)
28                .for_each(|(y, row)| {
29                    for x in 1..(w - 1) {
30                        let mut gx = 0.0f32;
31                        let mut gy = 0.0f32;
32
33                        for dy in -1..=1 {
34                            for dx in -1..=1 {
35                                let val = flat_vals
36                                    [(y as isize + dy) as usize * w + (x as isize + dx) as usize];
37                                gx += val * kx[(dy + 1) as usize][(dx + 1) as usize] as f32;
38                                gy += val * ky[(dy + 1) as usize][(dx + 1) as usize] as f32;
39                            }
40                        }
41                        row[x] = (gx * gx + gy * gy).sqrt();
42                    }
43                });
44        }
45
46        let new_data = TensorData::new(out_vals, [1, h, w]);
47        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
48        Ok(Image::new(new_tensor))
49    }
50
51    /// Computes the Laplacian of an image.
52    pub fn laplacian(&self) -> Result<Self> {
53        let gray = self.grayscale()?;
54        let dims = gray.tensor.dims();
55        let h = dims[1];
56        let w = dims[2];
57
58        let device = gray.tensor.device();
59        let tensor_data = gray.tensor.clone().into_data();
60        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
61        let mut out_vals = vec![0.0f32; h * w];
62
63        // Standard 3x3 Laplacian kernel
64        let kernel = [[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]];
65
66        {
67            use rayon::prelude::*;
68            out_vals
69                .par_chunks_exact_mut(w)
70                .enumerate()
71                .skip(1)
72                .take(h - 2)
73                .for_each(|(y, row)| {
74                    for x in 1..(w - 1) {
75                        let mut sum = 0.0f32;
76                        for dy in -1..=1 {
77                            for dx in -1..=1 {
78                                let val = flat_vals
79                                    [(y as isize + dy) as usize * w + (x as isize + dx) as usize];
80                                sum += val * kernel[(dy + 1) as usize][(dx + 1) as usize] as f32;
81                            }
82                        }
83                        row[x] = sum.abs();
84                    }
85                });
86        }
87
88        let new_data = TensorData::new(out_vals, [1, h, w]);
89        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
90        Ok(Image::new(new_tensor))
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::test_helpers::{TestBackend, test_device};
98
99    #[test]
100    fn test_gradients() {
101        let device = test_device();
102        let flat_data = vec![0.5f32; 3 * 8 * 8];
103        let tensor_data = TensorData::new(flat_data, [3, 8, 8]);
104        let img = Image::new(Tensor::<TestBackend, 3>::from_data(tensor_data, &device));
105
106        let scharr_img = img.scharr().unwrap();
107        assert_eq!(scharr_img.shape(), [1, 8, 8]);
108
109        let laplace_img = img.laplacian().unwrap();
110        assert_eq!(laplace_img.shape(), [1, 8, 8]);
111    }
112}
iris/edges/gradients.rs

iris/edges/
gradients.rs
