Skip to main content

ndarray_ndimage/filters/
median.rs

1use ndarray::{s, ArrayRef3, 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(mask: &ArrayRef3<bool>) -> Mask
10{
11    let range = |i, max| {
12        if i == 0 {
13            0..2
14        } else if i == max {
15            max - 1..max + 1
16        } else {
17            i - 1..i + 2
18        }
19    };
20
21    let (width, height, depth) = dim_minus(mask, 1);
22    let ranges_x: Vec<_> = (0..=width).map(|x| range(x, width)).collect();
23    let ranges_y: Vec<_> = (0..=height).map(|y| range(y, height)).collect();
24    let ranges_z: Vec<_> = (0..=depth).map(|z| range(z, depth)).collect();
25
26    // `from_shape_fn` is strangely much slower here
27    let mut new_mask = array_like(mask, mask.dim(), false);
28    Zip::indexed(&mut new_mask).for_each(|idx, new_mask| {
29        let r_x = &ranges_x[idx.0];
30        let r_y = &ranges_y[idx.1];
31        let r_z = &ranges_z[idx.2];
32
33        // For binary images, the median filter can be replaced with a simple majority vote
34        let nb_required = ((r_x.len() * r_y.len() * r_z.len()) as u8 - 1) / 2;
35        *new_mask = mask
36            .slice(s![r_x.clone(), r_y.clone(), r_z.clone()])
37            .iter()
38            .fold(0, |acc, &m| acc + m as u8)
39            > nb_required;
40    });
41    new_mask
42}