ndarray_vision/processing/
filter.rs

1use crate::core::{ColourModel, Image, ImageBase};
2use ndarray::prelude::*;
3use ndarray::{Data, IntoDimension, OwnedRepr, Zip};
4use ndarray_stats::interpolate::*;
5use ndarray_stats::Quantile1dExt;
6use noisy_float::types::n64;
7use num_traits::{FromPrimitive, Num, ToPrimitive};
8use std::marker::PhantomData;
9
10/// Median filter, given a region to move over the image, each pixel is given
11/// the median value of itself and it's neighbours
12pub trait MedianFilterExt {
13    type Output;
14    /// Run the median filter given the region. Median is assumed to be calculated
15    /// independently for each channel.
16    fn median_filter<E>(&self, region: E) -> Self::Output
17    where
18        E: IntoDimension<Dim = Ix2>;
19}
20
21impl<T, U> MedianFilterExt for ArrayBase<U, Ix3>
22where
23    U: Data<Elem = T>,
24    T: Copy + Clone + FromPrimitive + ToPrimitive + Num + Ord,
25{
26    type Output = ArrayBase<OwnedRepr<T>, Ix3>;
27
28    fn median_filter<E>(&self, region: E) -> Self::Output
29    where
30        E: IntoDimension<Dim = Ix2>,
31    {
32        let shape = region.into_dimension();
33        let r_offset = shape[0] / 2;
34        let c_offset = shape[1] / 2;
35        let region = (shape[0], shape[1], 1);
36        let mut result = Array3::<T>::zeros(self.dim());
37        Zip::indexed(self.windows(region)).for_each(|(i, j, k), window| {
38            let mut flat_window = Array::from_iter(window.iter()).mapv(|x| *x);
39            if let Ok(v) = flat_window.quantile_mut(n64(0.5f64), &Linear {}) {
40                if let Some(r) = result.get_mut([i + r_offset, j + c_offset, k]) {
41                    *r = v;
42                }
43            }
44        });
45        result
46    }
47}
48
49impl<T, U, C> MedianFilterExt for ImageBase<U, C>
50where
51    U: Data<Elem = T>,
52    T: Copy + Clone + FromPrimitive + ToPrimitive + Num + Ord,
53    C: ColourModel,
54{
55    type Output = Image<T, C>;
56
57    fn median_filter<E>(&self, region: E) -> Self::Output
58    where
59        E: IntoDimension<Dim = Ix2>,
60    {
61        let data = self.data.median_filter(region);
62        Image {
63            data,
64            model: PhantomData,
65        }
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use crate::core::colour_models::{Gray, RGB};
73
74    #[test]
75    fn simple_median() {
76        let mut pixels = Vec::<u8>::new();
77        for i in 0..9 {
78            pixels.extend_from_slice(&[i, i + 1, i + 2]);
79        }
80        let image = Image::<_, RGB>::from_shape_data(3, 3, pixels);
81
82        let image = image.median_filter((3, 3));
83
84        let mut expected = Image::<u8, RGB>::new(3, 3);
85        expected.pixel_mut(1, 1).assign(&arr1(&[4, 5, 6]));
86
87        assert_eq!(image, expected);
88    }
89
90    #[test]
91    fn row_median() {
92        let pixels = vec![1, 2, 3, 4, 5, 6, 7];
93        let image = Image::<_, Gray>::from_shape_data(7, 1, pixels);
94        let image = image.median_filter((3, 1));
95
96        let expected_pixels = vec![0, 2, 3, 4, 5, 6, 0];
97        let expected = Image::<_, Gray>::from_shape_data(7, 1, expected_pixels);
98
99        assert_eq!(image, expected);
100    }
101}
ndarray_vision/processing/filter.rs

ndarray_vision/processing/
filter.rs
