iris/core/
utils.rs

1use crate::core::types::Point;
2use crate::error::{IrisError, Result};
3use crate::image::Image;
4use burn::tensor::{Tensor, TensorData, backend::Backend};
5
6impl<B: Backend> Image<B> {
7    /// Computes the element-wise sum of two images.
8    pub fn add(&self, other: &Self) -> Result<Self> {
9        if self.shape() != other.shape() {
10            return Err(IrisError::DimensionMismatch {
11                expected: self.shape().to_vec(),
12                actual: other.shape().to_vec(),
13            });
14        }
15        let added = self.tensor.clone().add(other.tensor.clone());
16        Ok(Image::new(added))
17    }
18
19    /// Computes the element-wise difference of two images.
20    pub fn subtract(&self, other: &Self) -> Result<Self> {
21        if self.shape() != other.shape() {
22            return Err(IrisError::DimensionMismatch {
23                expected: self.shape().to_vec(),
24                actual: other.shape().to_vec(),
25            });
26        }
27        let subbed = self.tensor.clone().sub(other.tensor.clone());
28        Ok(Image::new(subbed))
29    }
30
31    /// Computes the element-wise multiplication of two images.
32    pub fn multiply(&self, other: &Self) -> Result<Self> {
33        if self.shape() != other.shape() {
34            return Err(IrisError::DimensionMismatch {
35                expected: self.shape().to_vec(),
36                actual: other.shape().to_vec(),
37            });
38        }
39        let muled = self.tensor.clone().mul(other.tensor.clone());
40        Ok(Image::new(muled))
41    }
42
43    /// Computes the element-wise division of two images.
44    pub fn divide(&self, other: &Self) -> Result<Self> {
45        if self.shape() != other.shape() {
46            return Err(IrisError::DimensionMismatch {
47                expected: self.shape().to_vec(),
48                actual: other.shape().to_vec(),
49            });
50        }
51        let dived = self.tensor.clone().div(other.tensor.clone());
52        Ok(Image::new(dived))
53    }
54
55    /// Computes the absolute difference between two images.
56    pub fn absdiff(&self, other: &Self) -> Result<Self> {
57        if self.shape() != other.shape() {
58            return Err(IrisError::DimensionMismatch {
59                expected: self.shape().to_vec(),
60                actual: other.shape().to_vec(),
61            });
62        }
63        let diff = self.tensor.clone().sub(other.tensor.clone()).abs();
64        Ok(Image::new(diff))
65    }
66
67    /// Computes bitwise AND of two images at pixel/byte level (0..255).
68    pub fn bitwise_and(&self, other: &Self) -> Result<Self> {
69        self.bitwise_op(other, |a, b| a & b)
70    }
71
72    /// Computes bitwise OR of two images.
73    pub fn bitwise_or(&self, other: &Self) -> Result<Self> {
74        self.bitwise_op(other, |a, b| a | b)
75    }
76
77    /// Computes bitwise XOR of two images.
78    pub fn bitwise_xor(&self, other: &Self) -> Result<Self> {
79        self.bitwise_op(other, |a, b| a ^ b)
80    }
81
82    /// Computes bitwise NOT of the image.
83    pub fn bitwise_not(&self) -> Result<Self> {
84        let dims = self.tensor.dims();
85        let c = dims[0];
86        let h = dims[1];
87        let w = dims[2];
88
89        let device = self.tensor.device();
90        let tensor_data = self.tensor.clone().into_data();
91        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
92        let mut out_vals = vec![0.0f32; c * h * w];
93
94        {
95            use rayon::prelude::*;
96            out_vals.par_iter_mut().enumerate().for_each(|(i, val)| {
97                let pixel_val = (flat_vals[i].clamp(0.0, 1.0) * 255.0) as u8;
98                *val = f32::from(!pixel_val) / 255.0;
99            });
100        }
101
102        let new_data = TensorData::new(out_vals, [c, h, w]);
103        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
104        Ok(Image::new(new_tensor))
105    }
106
107    /// Returns the average value of all elements in the image.
108    pub fn mean(&self) -> Result<Vec<f64>> {
109        let dims = self.tensor.dims();
110        let c = dims[0];
111        let h = dims[1];
112        let w = dims[2];
113
114        let tensor_data = self.tensor.clone().into_data();
115        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
116        let mut channel_means = vec![0.0; c];
117
118        for ch in 0..c {
119            let mut sum = 0.0;
120            for y in 0..h {
121                for x in 0..w {
122                    sum += f64::from(flat_vals[ch * h * w + y * w + x]);
123                }
124            }
125            channel_means[ch] = sum / ((h * w) as f64);
126        }
127
128        Ok(channel_means)
129    }
130
131    /// Computes the mean and standard deviation of image elements channel-wise.
132    pub fn mean_std_dev(&self) -> Result<(Vec<f64>, Vec<f64>)> {
133        let dims = self.tensor.dims();
134        let c = dims[0];
135        let h = dims[1];
136        let w = dims[2];
137
138        let tensor_data = self.tensor.clone().into_data();
139        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
140
141        let mut channel_means = vec![0.0; c];
142        let mut channel_stddevs = vec![0.0; c];
143
144        let n = (h * w) as f64;
145
146        for ch in 0..c {
147            let mut sum = 0.0;
148            for y in 0..h {
149                for x in 0..w {
150                    sum += f64::from(flat_vals[ch * h * w + y * w + x]);
151                }
152            }
153            let mean = sum / n;
154            channel_means[ch] = mean;
155
156            let mut sq_sum = 0.0;
157            for y in 0..h {
158                for x in 0..w {
159                    let diff = f64::from(flat_vals[ch * h * w + y * w + x]) - mean;
160                    sq_sum += diff * diff;
161                }
162            }
163            channel_stddevs[ch] = (sq_sum / n).sqrt();
164        }
165
166        Ok((channel_means, channel_stddevs))
167    }
168
169    /// Finds global minimum and maximum values and their coordinate locations in a single-channel image.
170    pub fn min_max_loc(&self) -> Result<(f64, f64, Point<usize>, Point<usize>)> {
171        let dims = self.tensor.dims();
172        let c = dims[0];
173        let h = dims[1];
174        let w = dims[2];
175
176        if c != 1 {
177            return Err(IrisError::InvalidParameter(
178                "min_max_loc requires a single-channel image".into(),
179            ));
180        }
181
182        let tensor_data = self.tensor.clone().into_data();
183        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
184
185        let mut min_val = f64::MAX;
186        let mut max_val = f64::MIN;
187        let mut min_loc = Point::new(0, 0);
188        let mut max_loc = Point::new(0, 0);
189
190        for y in 0..h {
191            for x in 0..w {
192                let val = f64::from(flat_vals[y * w + x]);
193                if val < min_val {
194                    min_val = val;
195                    min_loc = Point::new(x, y);
196                }
197                if val > max_val {
198                    max_val = val;
199                    max_loc = Point::new(x, y);
200                }
201            }
202        }
203
204        Ok((min_val, max_val, min_loc, max_loc))
205    }
206
207    /// Counts non-zero pixels (value > 0.0).
208    pub fn count_non_zero(&self) -> Result<usize> {
209        let tensor_data = self.tensor.clone().into_data();
210        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
211        let count = flat_vals.iter().filter(|&&x| x > 0.0).count();
212        Ok(count)
213    }
214
215    /// Produces a binary mask where pixels within the range [low, high] are set to 1.0.
216    /// For multi-channel images, all channels must be in range for the pixel to be marked.
217    pub fn in_range(&self, low: &[f32], high: &[f32]) -> Result<Self> {
218        let dims = self.tensor.dims();
219        let c = dims[0];
220        let h = dims[1];
221        let w = dims[2];
222
223        if low.len() != c || high.len() != c {
224            return Err(IrisError::InvalidParameter(format!(
225                "low/high length ({}/{}) must match channels ({c})",
226                low.len(),
227                high.len(),
228            )));
229        }
230
231        let device = self.tensor.device();
232        let tensor_data = self.tensor.clone().into_data();
233        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
234        let mut out_vals = vec![0.0f32; h * w];
235
236        let pixels = h * w;
237        for i in 0..pixels {
238            let mut in_range = true;
239            for ch in 0..c {
240                let val = flat_vals[ch * pixels + i];
241                if val < low[ch] || val > high[ch] {
242                    in_range = false;
243                    break;
244                }
245            }
246            out_vals[i] = if in_range { 1.0 } else { 0.0 };
247        }
248
249        let new_data = TensorData::new(out_vals, [1, h, w]);
250        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
251        Ok(Image::new(new_tensor))
252    }
253
254    /// Normalizes the image to the given range [min_val, max_val].
255    pub fn normalize(&self, min_val: f32, max_val: f32) -> Result<Self> {
256        let dims = self.tensor.dims();
257        let c = dims[0];
258        let h = dims[1];
259        let w = dims[2];
260
261        let device = self.tensor.device();
262        let tensor_data = self.tensor.clone().into_data();
263        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
264        let mut out_vals = vec![0.0f32; c * h * w];
265
266        for ch in 0..c {
267            let mut ch_min = f32::MAX;
268            let mut ch_max = f32::MIN;
269            let pixels = h * w;
270            for i in 0..pixels {
271                let v = flat_vals[ch * pixels + i];
272                if v < ch_min {
273                    ch_min = v;
274                }
275                if v > ch_max {
276                    ch_max = v;
277                }
278            }
279            let range = ch_max - ch_min;
280            for i in 0..pixels {
281                let v = flat_vals[ch * pixels + i];
282                out_vals[ch * pixels + i] = if range.abs() < 1e-10 {
283                    min_val
284                } else {
285                    min_val + (v - ch_min) / range * (max_val - min_val)
286                };
287            }
288        }
289
290        let new_data = TensorData::new(out_vals, [c, h, w]);
291        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
292        Ok(Image::new(new_tensor))
293    }
294
295    // Helper for pixel-wise logical operations
296    fn bitwise_op(&self, other: &Self, op: impl Fn(u8, u8) -> u8 + Sync + Send) -> Result<Self> {
297        if self.shape() != other.shape() {
298            return Err(IrisError::DimensionMismatch {
299                expected: self.shape().to_vec(),
300                actual: other.shape().to_vec(),
301            });
302        }
303        let dims = self.tensor.dims();
304        let c = dims[0];
305        let h = dims[1];
306        let w = dims[2];
307
308        let device = self.tensor.device();
309        let data_self = self.tensor.clone().into_data();
310        let data_other = other.tensor.clone().into_data();
311
312        let vals_self: Vec<f32> = data_self.iter::<f32>().collect();
313        let vals_other: Vec<f32> = data_other.iter::<f32>().collect();
314        let mut out_vals = vec![0.0f32; c * h * w];
315
316        {
317            use rayon::prelude::*;
318            out_vals.par_iter_mut().enumerate().for_each(|(i, val)| {
319                let b1 = (vals_self[i].clamp(0.0, 1.0) * 255.0) as u8;
320                let b2 = (vals_other[i].clamp(0.0, 1.0) * 255.0) as u8;
321                *val = f32::from(op(b1, b2)) / 255.0;
322            });
323        }
324
325        let new_data = TensorData::new(out_vals, [c, h, w]);
326        let new_tensor = Tensor::<B, 3>::from_data(new_data, &device);
327        Ok(Image::new(new_tensor))
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use crate::test_helpers::{TestBackend, test_device};
335    use burn::backend::ndarray::NdArrayDevice;
336
337    fn get_test_device() -> NdArrayDevice {
338        test_device()
339    }
340
341    #[test]
342    fn test_image_math_and_bitwise() {
343        let device = get_test_device();
344        let data1 = TensorData::new(vec![0.5f32; 3 * 4 * 4], [3, 4, 4]);
345        let data2 = TensorData::new(vec![0.2f32; 3 * 4 * 4], [3, 4, 4]);
346
347        let img1 = Image::new(Tensor::<TestBackend, 3>::from_data(data1, &device));
348        let img2 = Image::new(Tensor::<TestBackend, 3>::from_data(data2, &device));
349
350        // Math
351        let added = img1.add(&img2).unwrap();
352        assert_eq!(added.shape(), [3, 4, 4]);
353
354        let subbed = img1.subtract(&img2).unwrap();
355        assert_eq!(subbed.shape(), [3, 4, 4]);
356
357        let absdiff = img1.absdiff(&img2).unwrap();
358        assert_eq!(absdiff.shape(), [3, 4, 4]);
359
360        // Bitwise
361        let bit_and = img1.bitwise_and(&img2).unwrap();
362        assert_eq!(bit_and.shape(), [3, 4, 4]);
363
364        let bit_not = img1.bitwise_not().unwrap();
365        assert_eq!(bit_not.shape(), [3, 4, 4]);
366
367        // Stats
368        let mean_vals = img1.mean().unwrap();
369        assert!((mean_vals[0] - 0.5).abs() < 1e-4);
370
371        let count = img1.count_non_zero().unwrap();
372        assert_eq!(count, 3 * 4 * 4);
373    }
374
375    #[test]
376    fn test_in_range() {
377        let device = get_test_device();
378        let data = vec![0.1, 0.5, 0.9, 0.3, 0.6, 0.2, 0.7, 0.8, 0.4];
379        let img = Image::new(Tensor::<TestBackend, 3>::from_data(
380            TensorData::new(data, [3, 1, 3]),
381            &device,
382        ));
383
384        let mask = img.in_range(&[0.2, 0.2, 0.2], &[0.8, 0.8, 0.8]).unwrap();
385        assert_eq!(mask.shape(), [1, 1, 3]);
386        let vals: Vec<f32> = mask.tensor.into_data().iter::<f32>().collect();
387        // pixel 0: 0.1<0.2 -> 0.0; pixel 1: 0.5 in range -> 1.0; pixel 2: 0.9>0.8 -> 0.0
388        assert!((vals[0]).abs() < 1e-5);
389        assert!((vals[1] - 1.0).abs() < 1e-5);
390        assert!((vals[2]).abs() < 1e-5);
391    }
392
393    #[test]
394    fn test_normalize() {
395        let device = get_test_device();
396        let data = vec![0.2, 0.4, 0.6, 0.8];
397        let img = Image::new(Tensor::<TestBackend, 3>::from_data(
398            TensorData::new(data, [1, 1, 4]),
399            &device,
400        ));
401
402        let normalized = img.normalize(0.0, 1.0).unwrap();
403        assert_eq!(normalized.shape(), [1, 1, 4]);
404        let vals: Vec<f32> = normalized.tensor.into_data().iter::<f32>().collect();
405        assert!((vals[0]).abs() < 1e-5); // min => 0.0
406        assert!((vals[3] - 1.0).abs() < 1e-5); // max => 1.0
407    }
408
409    #[test]
410    fn test_in_range_invalid_length() {
411        let device = get_test_device();
412        let data = vec![0.5f32; 3 * 4 * 4];
413        let img = Image::new(Tensor::<TestBackend, 3>::from_data(
414            TensorData::new(data, [3, 4, 4]),
415            &device,
416        ));
417        assert!(img.in_range(&[0.0], &[0.5, 0.5, 0.5]).is_err());
418    }
419}
iris/core/utils.rs

iris/core/
utils.rs
