ndarray_vision/processing/
filter.rs1use 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
10pub trait MedianFilterExt {
13 type Output;
14 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}