diff --git a/src/filter/bilateral.rs b/src/filter/bilateral.rs index 901353b8..75d39e8f 100644 --- a/src/filter/bilateral.rs +++ b/src/filter/bilateral.rs @@ -1,6 +1,6 @@ //! Bilateral Filter and associated items. -use image::{GenericImage, Pixel}; +use image::{GenericImage, ImageBuffer, Luma, Pixel, Rgb, Rgba}; use num::cast::AsPrimitive; use crate::definitions::Image; @@ -56,6 +56,158 @@ where } } +/// A gaussian function of the euclidean distance between two pixel's colors, implemented using +/// a look up table. This gives substantial speed up for pixel's with a u8 subpixel. +/// +/// This implements [`ColorDistance`]. +pub struct LutGaussianEuclideanColorDistance { + lut: Vec, +} + +impl LutGaussianEuclideanColorDistance { + /// Creates a new [`GaussianEuclideanColorDistance`] using a given sigma value, which + /// must be positive. + /// + /// Internally, this is stored as sigma squared for performance. + /// + /// # Panics + /// + /// 1. If `sigma <= 0` + pub fn new(sigma: f32) -> Self { + assert!( + sigma > 0.0, + "GaussianEuclideanColorDistance sigma must be positive" + ); + let mut lut = Vec::with_capacity(512); + for diff in 0..512 { + let diff = (diff - 255) as f32; + let weight = fast_exp_negative(-0.5 * diff.powi(2) / sigma.powi(2)); + lut.push(weight); + } + Self { lut } + } +} + +impl ColorDistance> for LutGaussianEuclideanColorDistance { + #[inline(always)] + fn color_distance(&self, pixel1: &Luma, pixel2: &Luma) -> f32 { + let diff = (255 + pixel1.0[0] as usize) - pixel2.0[0] as usize; + self.lut[diff] + } +} + +impl ColorDistance> for LutGaussianEuclideanColorDistance { + #[inline(always)] + fn color_distance(&self, pixel1: &Rgb, pixel2: &Rgb) -> f32 { + let diffr = (255 + pixel1.0[0] as usize) - pixel2.0[0] as usize; + let diffg = (255 + pixel1.0[1] as usize) - pixel2.0[1] as usize; + let diffb = (255 + pixel1.0[2] as usize) - pixel2.0[2] as usize; + let lutr = self.lut[diffr]; + let lutg = self.lut[diffg]; + let lutb = self.lut[diffb]; + lutr * lutg * lutb + } +} + +impl ColorDistance> for LutGaussianEuclideanColorDistance { + #[inline(always)] + fn color_distance(&self, pixel1: &Rgba, pixel2: &Rgba) -> f32 { + let diffr = (255 + pixel1.0[0] as usize) - pixel2.0[0] as usize; + let diffg = (255 + pixel1.0[1] as usize) - pixel2.0[1] as usize; + let diffb = (255 + pixel1.0[2] as usize) - pixel2.0[2] as usize; + let diffa = (255 + pixel1.0[3] as usize) - pixel2.0[3] as usize; + let lutr = self.lut[diffr]; + let lutg = self.lut[diffg]; + let lutb = self.lut[diffb]; + let luta = self.lut[diffa]; + lutr * lutg * lutb * luta + } +} + +/// Loop-invariant context for a single `bilateral_filter` run. +/// +/// The per-pixel work is identical for interior and border pixels except for how the +/// window coordinates are computed, so it is expressed once in [`BilateralCtx::filter`] +/// and specialized at compile time via the `CLAMP` const generic. +struct BilateralCtx<'a, I, C> { + image: &'a I, + color_distance: &'a C, + spatial_distance_lookup: &'a [f32], + radius: u32, + radius_range: u32, + width: u32, + height: u32, +} + +impl BilateralCtx<'_, I, C> +where + I: GenericImage, + P: Pixel, + C: ColorDistance

, +

::Subpixel: 'static, + f32: From + AsPrimitive, +{ + /// Filter a single output pixel. + /// + /// When `CLAMP` is `true` the window coordinates are clamped into the image bounds + /// (used for border pixels); when `false` the window is assumed fully in-bounds + /// (used for the interior). Because `CLAMP` is a const generic, each variant + /// monomorphizes to straight-line code with the unused branch eliminated. + #[inline(always)] + fn filter(&self, x: u32, y: u32) -> P { + const MAX_CHANNELS: usize = 4; + + debug_assert!(self.image.in_bounds(x, y)); + // Safety: callers only pass (x, y) within the image bounds. + let center_pixel = unsafe { self.image.unsafe_get_pixel(x, y) }; + + let mut channel_sums = [0f32; MAX_CHANNELS]; + let mut weight_sum = 0f32; + + // Both branches walk the window in the same (w_y, w_x) order that + // `spatial_distance_lookup` was built in, so a sequential iterator gives the + // right weight without the per-pixel `window_len * w_y + w_x` index math and + // bounds check. + let mut spatial = self.spatial_distance_lookup.iter(); + for w_y in 0..self.radius_range { + for w_x in 0..self.radius_range { + let (window_x, window_y) = if CLAMP { + ( + (x + w_x).saturating_sub(self.radius).min(self.width - 1), + (y + w_y).saturating_sub(self.radius).min(self.height - 1), + ) + } else { + ((x + w_x) - self.radius, (y + w_y) - self.radius) + }; + + debug_assert!(self.image.in_bounds(window_x, window_y)); + // Safety: for CLAMP=true the coords are clamped in-bounds; for CLAMP=false + // the interior loop ranges guarantee the whole window is in-bounds. + let window_pixel = unsafe { self.image.unsafe_get_pixel(window_x, window_y) }; + + let spatial_weight = spatial.next().unwrap(); + let color_weight = self + .color_distance + .color_distance(¢er_pixel, &window_pixel); + let weight = spatial_weight * color_weight; + + weight_sum += weight; + for (i, c) in window_pixel.channels().iter().enumerate() { + channel_sums[i] += weight * f32::from(*c); + } + } + } + + let mut out_pixel = center_pixel; + let num_channels = P::CHANNEL_COUNT as usize; + let out_channels = out_pixel.channels_mut(); + for i in 0..num_channels { + out_channels[i] = (channel_sums[i] / weight_sum).as_(); + } + out_pixel + } +} + /// Denoise an 8-bit image while preserving edges using bilateral filtering. /// /// # Arguments @@ -124,7 +276,9 @@ where assert_ne!(image.height(), 0); assert!(spatial_sigma > 0.0, "spatial_sigma must be positive"); - let radius = i16::from(radius); + let (width, height) = image.dimensions(); + let radius = (radius as u32).min(width).min(height); + let radius = radius as i16; let spatial_sigma_squared = spatial_sigma.powi(2); let mut spatial_distance_lookup = @@ -138,54 +292,69 @@ where } } - let (width, height) = image.dimensions(); - let window_len = 2 * radius + 1; - - 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 radius = radius as u32; + let radius_range = 2 * radius + 1; - let mut channel_sums = [0f32; MAX_CHANNELS]; - let mut weight_sum = 0f32; + let mut out_image = ImageBuffer::new(width, height); - for w_y in -radius..=radius { - for w_x in -radius..=radius { - // 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 ctx = BilateralCtx { + image, + color_distance: &color_distance, + spatial_distance_lookup: &spatial_distance_lookup, + radius, + radius_range, + width, + height, + }; - let (window_y, window_x) = (window_y as u32, window_x as u32); + // Top edge + for y in 0..radius { + for x in 0..width { + let val = ctx.filter::(x, y); + unsafe { + out_image.unsafe_put_pixel(x, y, val); + } + } + } - 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) }; + // Middle Rows + for y in radius..(height - radius) { + // Left Edge + for x in 0..radius { + let val = ctx.filter::(x, y); + unsafe { + out_image.unsafe_put_pixel(x, y, val); + } + } - let spatial_weight = spatial_distance_lookup - [(window_len * (w_y + radius) + (w_x + radius)) as usize]; - let color_weight = color_distance.color_distance(¢er_pixel, &window_pixel); - let weight = spatial_weight * color_weight; + // Middle, unclamped section + for x in radius..(width - radius) { + let val = ctx.filter::(x, y); + unsafe { + out_image.unsafe_put_pixel(x, y, val); + } + } - weight_sum += weight; - for (i, c) in window_pixel.channels().iter().enumerate() { - channel_sums[i] += weight * f32::from(*c); - } + // Right edge + for x in (width - radius)..width { + let val = ctx.filter::(x, y); + unsafe { + out_image.unsafe_put_pixel(x, y, val); } } + } - let mut out_pixel = center_pixel; - let num_channels = P::CHANNEL_COUNT as usize; - let out_channels = out_pixel.channels_mut(); - for i in 0..num_channels { - out_channels[i] = (channel_sums[i] / weight_sum).as_(); + // Bottom edge + for y in (height - radius)..height { + for x in 0..width { + let val = ctx.filter::(x, y); + unsafe { + out_image.unsafe_put_pixel(x, y, val); + } } - out_pixel - }; + } - Image::from_fn(width, height, bilateral_pixel_filter) + out_image } /// Un-normalized Gaussian Weight