expand bilateral_filter() to color images

examples/bilateral.rs

@@ -0,0 +1,38 @@
+use std::env;
+
+use image::open;
+use imageproc::filter::{bilateral::GaussianEuclideanColorDistance, bilateral_filter};
+
+fn main() {
+    if env::args().len() != 2 {
+        panic!("Please enter an input image file path")
+    }
+
+    let input_path = env::args().nth(1).unwrap();
+
+    let dynamic_image =
+        open(&input_path).unwrap_or_else(|_| panic!("Could not load image at {:?}", input_path));
+
+    let image_grey = dynamic_image.to_luma8();
+    let image_color = dynamic_image.to_rgb8();
+
+    let radius = 16;
+    let sigma_color = 40.0;
+    let sigma_spatial = 40.0;
+
+    let bilateral_grey = bilateral_filter(
+        &image_grey,
+        radius,
+        sigma_spatial,
+        GaussianEuclideanColorDistance { sigma_squared: sigma_color },
+    );
+    let bilateral_color = bilateral_filter(
+        &image_color,
+        radius,
+        sigma_spatial,
+        GaussianEuclideanColorDistance { sigma_squared: sigma_color },
+    );
+
+    bilateral_grey.save("bilateral_grey.png").unwrap();
+    bilateral_color.save("bilateral_color.png").unwrap();
+}

src/filter/bilateral.rs

@@ -0,0 +1,286 @@
+//! Bilateral Filter and associated items.
+
+use image::{GenericImage, Pixel};
+use itertools::Itertools;
+use num::cast::AsPrimitive;
+
+use crate::definitions::Image;
+
+/// A trait which provides a distance metric between two pixels based on their colors.
+///
+/// This trait is used with the `bilateral_filter()` function.
+pub trait ColorDistance<P> {
+    /// Returns a distance measure between the two pixels based on their colors
+    fn color_distance(&self, pixel1: &P, pixel2: &P) -> f32;
+}
+
+/// A gaussian function of the euclidean distance between two pixel's colors.
+///
+/// This implements [`ColorDistance`].
+pub struct GaussianEuclideanColorDistance {
+    /// The square of the standard deviation of the gaussian distribution used on
+    /// the euclidean distance between pixels' colors.
+    pub sigma_squared: f32,
+}
+
+impl<P> ColorDistance<P> for GaussianEuclideanColorDistance
+where
+    P: Pixel,
+    f32: From<P::Subpixel>,
+{
+    fn color_distance(&self, pixel1: &P, pixel2: &P) -> f32 {
+        let euclidean_distance_squared = pixel1
+            .channels()
+            .iter()
+            .zip(pixel2.channels().iter())
+            .map(|(c1, c2)| (f32::from(*c1) - f32::from(*c2)).powi(2))
+            .sum::<f32>();
+
+        gaussian_weight(euclidean_distance_squared, self.sigma_squared)
+    }
+}
+
+/// Denoise an 8-bit image while preserving edges using bilateral filtering.
+///
+/// # Arguments
+///
+/// * `image` - Image to be filtered.
+/// * `window_size` - Window size for filtering.
+/// * `spatial_sigma` - Standard deviation for euclidean spatial distance. A larger value results in
+///     averaging of pixels with larger spatial distances.
+/// * `color_distance` - A type which implements [`ColorDistance`]. This defines the metric used to
+///     define how different two pixels are based on their colors. Common examples may include simple
+///     absolute difference for greyscale pixels or cartesian distance in the CIE-Lab color space
+///     \[1\].
+///
+/// This filter averages pixels based on their spatial distance as well as their color
+/// distance. Spatial distance is measured by the Gaussian function of the Euclidean distance
+/// between two pixels with the user-specified standard deviation (`sigma_spatial`).
+///
+/// # References
+///
+///   \[1\] C. Tomasi and R. Manduchi. "Bilateral Filtering for Gray and Color
+///        Images." IEEE International Conference on Computer Vision (1998)
+///        839-846. DOI: 10.1109/ICCV.1998.710815
+///
+/// # Panics
+///
+/// 1. If `image.width() > i32::MAX as u32`
+/// 2. If `image.height() > i32::MAX as u32`.
+/// 3. If `image.width() == 0`
+/// 4. If `image.height() == 0`
+///
+/// # Examples
+///
+/// ```
+/// use imageproc::filter::bilateral::{bilateral_filter, GaussianEuclideanColorDistance};
+/// use imageproc::utils::gray_bench_image;
+///
+/// let image = gray_bench_image(50, 50);
+///
+/// let filtered = bilateral_filter(&image, 2, 3., GaussianEuclideanColorDistance { sigma: 10.0 });
+/// ```
+#[must_use = "the function does not modify the original image"]
+pub fn bilateral_filter<I, P, C>(
+    image: &I,
+    radius: u8,
+    spatial_sigma: f32,
+    color_distance: C,
+) -> Image<P>
+where
+    I: GenericImage<Pixel = P>,
+    P: Pixel,
+    C: ColorDistance<P>,
+    <P as image::Pixel>::Subpixel: 'static,
+    f32: From<P::Subpixel> + AsPrimitive<P::Subpixel>,
+{
+    assert!(!image.width() > i32::MAX as u32);
+    assert!(!image.height() > i32::MAX as u32);
+    assert_ne!(image.width(), 0);
+    assert_ne!(image.height(), 0);
+
+    let radius = i16::from(radius);
+
+    let window_range = -radius..=radius;
+    let spatial_distance_lookup = window_range
+        .clone()
+        .cartesian_product(window_range.clone())
+        .map(|(w_y, w_x)| {
+            //The gaussian of the euclidean spatial distance
+            gaussian_weight(
+                <f32 as From<i16>>::from(w_x).powi(2) + <f32 as From<i16>>::from(w_y).powi(2),
+                spatial_sigma.powi(2),
+            )
+        })
+        .collect_vec();
+
+    let (width, height) = image.dimensions();
+    let bilateral_pixel_filter = |x, y| {
+        debug_assert!(image.in_bounds(x, y));
+        // Safety: `Image::from_fn` yields `col` in [0, width) and `row` in [0, height).
+        let center_pixel = unsafe { image.unsafe_get_pixel(x, y) };
+
+        let window_len = 2 * radius + 1;
+        let weights_and_values = window_range
+            .clone()
+            .cartesian_product(window_range.clone())
+            .map(|(w_y, w_x)| {
+                // these casts will always be correct due to asserts made at the beginning of the
+                // function about the image width/height
+                //
+                // the subtraction will also never overflow due to the `is_empty()` assert
+                let window_y = (i32::from(w_y) + (y as i32)).clamp(0, (height as i32) - 1);
+                let window_x = (i32::from(w_x) + (x as i32)).clamp(0, (width as i32) - 1);
+
+                let (window_y, window_x) = (window_y as u32, window_x as u32);
+
+                debug_assert!(image.in_bounds(window_x, window_y));
+                // Safety: we clamped `window_x` and `window_y` to be in bounds.
+                let window_pixel = unsafe { image.unsafe_get_pixel(window_x, window_y) };
+
+                let spatial_distance = spatial_distance_lookup
+                    [(window_len * (w_y + radius) + (w_x + radius)) as usize];
+                let color_distance = color_distance.color_distance(&center_pixel, &window_pixel);
+                let weight = spatial_distance * color_distance;
+
+                (weight, window_pixel)
+            });
+
+        weighted_average(weights_and_values)
+    };
+
+    Image::from_fn(width, height, bilateral_pixel_filter)
+}
+
+fn weighted_average<P>(weights_and_values: impl Iterator<Item = (f32, P)>) -> P
+where
+    P: Pixel,
+    <P as image::Pixel>::Subpixel: 'static,
+    f32: From<P::Subpixel> + AsPrimitive<P::Subpixel>,
+{
+    let (weights_sum, weighted_channel_sums) = weights_and_values
+        .map(|(w, v)| {
+            (
+                w,
+                v.channels().iter().map(|s| w * f32::from(*s)).collect_vec(),
+            )
+        })
+        .reduce(|(w1, channels1), (w2, channels2)| {
+            (
+                w1 + w2,
+                channels1
+                    .into_iter()
+                    .zip_eq(channels2)
+                    .map(|(c1, c2)| c1 + c2)
+                    .collect_vec(),
+            )
+        })
+        .expect("cannot find a weighted average given no weights and values");
+
+    let channel_averages = weighted_channel_sums.iter().map(|x| x / weights_sum);
+
+    *P::from_slice(
+        &channel_averages
+            .map(<f32 as AsPrimitive<P::Subpixel>>::as_)
+            .collect_vec(),
+    )
+}
+
+/// Un-normalized Gaussian Weight
+fn gaussian_weight(x: f32, sigma_squared: f32) -> f32 {
+    (-0.5 * x.powi(2) / sigma_squared).exp()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_bilateral_filter_greyscale() {
+        let image = gray_image!(
+            1, 2, 3;
+            4, 5, 6;
+            7, 8, 9);
+        let expect = gray_image!(
+            2, 3, 4;
+            4, 5, 6;
+            6, 7, 8);
+        let actual = bilateral_filter(
+            &image,
+            1,
+            3.,
+            GaussianEuclideanColorDistance { sigma_squared: 10. },
+        );
+        assert_pixels_eq!(expect, actual);
+    }
+}
+
+#[cfg(not(miri))]
+#[cfg(test)]
+mod proptests {
+    use super::*;
+    use crate::proptest_utils::arbitrary_image;
+    use image::Luma;
+    use image::Rgb;
+    use proptest::prelude::*;
+
+    proptest! {
+        #[test]
+        fn proptest_bilateral_filter_greyscale(
+            img in arbitrary_image::<Luma<u8>>(1..40, 1..40),
+            radius in 0..10u8,
+            sigma_color in any::<f32>(),
+            sigma_spatial in any::<f32>(),
+        ) {
+            let out = bilateral_filter(&img, radius, sigma_spatial, GaussianEuclideanColorDistance { sigma_squared: sigma_color });
+            assert_eq!(out.dimensions(), img.dimensions());
+        }
+
+        #[test]
+        fn proptest_bilateral_filter_rgb(
+            img in arbitrary_image::<Rgb<u8>>(1..40, 1..40),
+            window_size in 0..25u8,
+            sigma_color in any::<f32>(),
+            sigma_spatial in any::<f32>(),
+        ) {
+            let out = bilateral_filter(&img, window_size, sigma_spatial, GaussianEuclideanColorDistance { sigma_squared: sigma_color });
+            assert_eq!(out.dimensions(), img.dimensions());
+        }
+    }
+}
+
+#[cfg(not(miri))]
+#[cfg(test)]
+mod benches {
+    use super::*;
+    use crate::utils::{gray_bench_image, rgb_bench_image};
+    use test::{black_box, Bencher};
+
+    #[bench]
+    fn bench_bilateral_filter_greyscale(b: &mut Bencher) {
+        let image = gray_bench_image(100, 100);
+        b.iter(|| {
+            let filtered = bilateral_filter(
+                &image,
+                5,
+                3.,
+                GaussianEuclideanColorDistance { sigma_squared: 10. },
+            );
+            black_box(filtered);
+        });
+    }
+
+    #[bench]
+    fn bench_bilateral_filter_rgb(b: &mut Bencher) {
+        let image = rgb_bench_image(100, 100);
+        b.iter(|| {
+            let filtered = bilateral_filter(
+                &image,
+                5,
+                3.,
+                GaussianEuclideanColorDistance { sigma_squared: 10. },
+            );
+            black_box(filtered);
+        });
+    }
+}

src/filter/median.rs

@@ -323,15 +323,8 @@ struct HistSet {
 
 impl HistSet {
     fn new(num_channels: u8, expected_count: u32) -> HistSet {
-        // Can't use vec![[0u32; 256], num_channels as usize]
-        // because arrays of length > 32 aren't cloneable.
-        let mut data = Vec::with_capacity(num_channels as usize);
-        for _ in 0..num_channels {
-            data.push([0u32; 256]);
-        }
-
         HistSet {
-            data,
+            data: vec![[0u32; 256]; num_channels.into()],
             expected_count,
         }
     }