iris/noise/
mod.rs

1use crate::error::{IrisError, Result};
2use crate::image::Image;
3use burn::tensor::{Tensor, TensorData, backend::Backend};
4
5impl<B: Backend> Image<B> {
6    /// Adds Gaussian noise with the given mean and standard deviation.
7    pub fn add_gaussian_noise(&self, mean: f32, std_dev: f32) -> Result<Self> {
8        let dims = self.tensor.dims();
9        let c = dims[0];
10        let h = dims[1];
11        let w = dims[2];
12
13        let device = self.tensor.device();
14        let tensor_data = self.tensor.clone().into_data();
15        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
16        let mut out_vals = vec![0.0f32; c * h * w];
17
18        // Box-Muller transform for Gaussian random numbers
19        let mut seed: u64 = 0x1234_5678_9ABC_DEF0;
20        let mut next_gaussian: Option<f32> = None;
21
22        for i in 0..(c * h * w) {
23            let gaussian = if let Some(g) = next_gaussian.take() {
24                g
25            } else {
26                // Marsaglia polar method
27                loop {
28                    let u1 = {
29                        seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
30                        ((seed >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
31                    };
32                    let u2 = {
33                        seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
34                        ((seed >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
35                    };
36                    let s = u1 * u1 + u2 * u2;
37                    if s > 0.0 && s < 1.0 {
38                        let factor = (-2.0 * s.ln() / s).sqrt();
39                        let g1 = u1 * factor;
40                        let g2 = u2 * factor;
41                        next_gaussian = Some(g2);
42                        break g1;
43                    }
44                }
45            };
46
47            let noise = mean + std_dev * gaussian;
48            out_vals[i] = (flat_vals[i] + noise).clamp(0.0, 1.0);
49        }
50
51        let new_data = TensorData::new(out_vals, [c, h, w]);
52        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
53        Ok(Image::new(new_tensor))
54    }
55
56    /// Adds salt-and-pepper (impulse) noise with the given probability.
57    /// `amount` is the fraction of pixels to corrupt (0.0 to 1.0).
58    pub fn add_salt_pepper_noise(&self, amount: f32) -> Result<Self> {
59        if !(0.0..=1.0).contains(&amount) {
60            return Err(IrisError::InvalidParameter(
61                "amount must be in [0.0, 1.0]".into(),
62            ));
63        }
64
65        let dims = self.tensor.dims();
66        let c = dims[0];
67        let h = dims[1];
68        let w = dims[2];
69
70        let device = self.tensor.device();
71        let tensor_data = self.tensor.clone().into_data();
72        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
73        let mut out_vals = flat_vals.clone();
74
75        let total_pixels = h * w;
76        let num_noise = (total_pixels as f32 * amount) as usize;
77
78        let mut seed: u64 = 0xABCD_EF01_2345_6789;
79
80        for _ in 0..num_noise {
81            let py = {
82                seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
83                ((seed >> 33) as usize) % h
84            };
85            let px = {
86                seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
87                ((seed >> 33) as usize) % w
88            };
89            let is_salt = {
90                seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
91                (seed >> 33) & 1 == 0
92            };
93
94            for ch in 0..c {
95                out_vals[ch * total_pixels + py * w + px] = if is_salt { 1.0 } else { 0.0 };
96            }
97        }
98
99        let new_data = TensorData::new(out_vals, [c, h, w]);
100        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
101        Ok(Image::new(new_tensor))
102    }
103
104    /// Adds speckle (multiplicative) noise: pixel = pixel + pixel * noise.
105    pub fn add_speckle_noise(&self, std_dev: f32) -> Result<Self> {
106        let dims = self.tensor.dims();
107        let c = dims[0];
108        let h = dims[1];
109        let w = dims[2];
110
111        let device = self.tensor.device();
112        let tensor_data = self.tensor.clone().into_data();
113        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
114        let mut out_vals = vec![0.0f32; c * h * w];
115
116        let mut seed: u64 = 0x1111_2222_3333_4444;
117        let mut next_gaussian: Option<f32> = None;
118
119        for i in 0..(c * h * w) {
120            let gaussian = if let Some(g) = next_gaussian.take() {
121                g
122            } else {
123                loop {
124                    let u1 = {
125                        seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
126                        ((seed >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
127                    };
128                    let u2 = {
129                        seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
130                        ((seed >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
131                    };
132                    let s = u1 * u1 + u2 * u2;
133                    if s > 0.0 && s < 1.0 {
134                        let factor = (-2.0 * s.ln() / s).sqrt();
135                        let g1 = u1 * factor;
136                        let g2 = u2 * factor;
137                        next_gaussian = Some(g2);
138                        break g1;
139                    }
140                }
141            };
142
143            let noise = flat_vals[i] * std_dev * gaussian;
144            out_vals[i] = (flat_vals[i] + noise).clamp(0.0, 1.0);
145        }
146
147        let new_data = TensorData::new(out_vals, [c, h, w]);
148        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
149        Ok(Image::new(new_tensor))
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use crate::test_helpers::{TestBackend, test_device};
157
158    #[test]
159    fn test_gaussian_noise() {
160        let device = test_device();
161        let data = TensorData::new(vec![0.5f32; 3 * 8 * 8], [3, 8, 8]);
162        let img = Image::new(Tensor::<TestBackend, 3>::from_data(data, &device));
163        let noisy = img.add_gaussian_noise(0.0, 0.05).unwrap();
164        assert_eq!(noisy.shape(), [3, 8, 8]);
165    }
166
167    #[test]
168    fn test_salt_pepper_noise() {
169        let device = test_device();
170        let data = TensorData::new(vec![0.5f32; 3 * 8 * 8], [3, 8, 8]);
171        let img = Image::new(Tensor::<TestBackend, 3>::from_data(data, &device));
172        let noisy = img.add_salt_pepper_noise(0.1).unwrap();
173        assert_eq!(noisy.shape(), [3, 8, 8]);
174    }
175
176    #[test]
177    fn test_speckle_noise() {
178        let device = test_device();
179        let data = TensorData::new(vec![0.5f32; 3 * 8 * 8], [3, 8, 8]);
180        let img = Image::new(Tensor::<TestBackend, 3>::from_data(data, &device));
181        let noisy = img.add_speckle_noise(0.1).unwrap();
182        assert_eq!(noisy.shape(), [3, 8, 8]);
183    }
184
185    #[test]
186    fn test_noise_invalid_amount() {
187        let device = test_device();
188        let data = TensorData::new(vec![0.5f32; 3 * 8 * 8], [3, 8, 8]);
189        let img = Image::new(Tensor::<TestBackend, 3>::from_data(data, &device));
190        assert!(img.add_salt_pepper_noise(1.5).is_err());
191    }
192}
iris/noise/mod.rs

iris/noise/
mod.rs
