ndarray_ndimage/filters/
median.rs

1use ndarray::{s, ArrayBase, Data, Ix3, Zip};
2
3use crate::{array_like, dim_minus, Mask};
4
5/// Binary median filter.
6///
7/// A 3x3 structuring element (`Kernel3d::Full`) is used except on the borders, where a smaller
8/// structuring element is used.
9pub fn median_filter<S>(mask: &ArrayBase<S, Ix3>) -> Mask
10where
11    S: Data<Elem = bool>,
12{
13    let range = |i, max| {
14        if i == 0 {
15            0..2
16        } else if i == max {
17            max - 1..max + 1
18        } else {
19            i - 1..i + 2
20        }
21    };
22
23    let (width, height, depth) = dim_minus(mask, 1);
24    let ranges_x: Vec<_> = (0..=width).map(|x| range(x, width)).collect();
25    let ranges_y: Vec<_> = (0..=height).map(|y| range(y, height)).collect();
26    let ranges_z: Vec<_> = (0..=depth).map(|z| range(z, depth)).collect();
27
28    // `from_shape_fn` is strangely much slower here
29    let mut new_mask = array_like(mask, mask.dim(), false);
30    Zip::indexed(&mut new_mask).for_each(|idx, new_mask| {
31        let r_x = &ranges_x[idx.0];
32        let r_y = &ranges_y[idx.1];
33        let r_z = &ranges_z[idx.2];
34
35        // For binary images, the median filter can be replaced with a simple majority vote
36        let nb_required = ((r_x.len() * r_y.len() * r_z.len()) as u8 - 1) / 2;
37        *new_mask = mask
38            .slice(s![r_x.clone(), r_y.clone(), r_z.clone()])
39            .iter()
40            .fold(0, |acc, &m| acc + m as u8)
41            > nb_required;
42    });
43    new_mask
44}