Skip to content

Commit 9cb856e

Browse files
youssefelzedyKeavonTrueDoctor
authored
Feature: Add median filter option to blur node (#3196)
* Add median filter option to blur node: - Add median parameter to blur function for noise reduction - Implement median_filter_algorithm with efficient quickselect - Support gamma space calculations for median filtering - Preserve edges while removing noise, complementing existing blur options Feature in Issue: #912 * Optimize median filter algorithm by reusing buffers and improving pixel neighborhood collection * Add median filter option to blur node: - Add median parameter to blur function for noise reduction - Implement median_filter_algorithm with efficient quickselect - Support gamma space calculations for median filtering - Preserve edges while removing noise, complementing existing blur options Feature in Issue: #912 * Optimize median filter algorithm by reusing buffers and improving pixel neighborhood collection * Improve median filter's NaN handling by using total_cmp for safe comparisons * Add median filter node for noise reduction - Create dedicated median_filter node separate from blur functionality - Implement median_filter_algorithm with efficient quickselect - Support gamma space calculations for consistency - Use safe NaN handling with f32::total_cmp to prevent panics - Optimize performance with memory reuse and O(n) median selection * Refactor median filter to remove gamma handling and simplify algorithm --------- Co-authored-by: Keavon Chambers <keavon@keavon.com> Co-authored-by: Dennis Kobert <dennis@kobert.dev>
1 parent e88db02 commit 9cb856e

1 file changed

Lines changed: 84 additions & 1 deletion

File tree

node-graph/nodes/raster/src/filter.rs

Lines changed: 84 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use raster_types::Image;
66
use raster_types::{Bitmap, BitmapMut};
77
use raster_types::{CPU, Raster};
88

9-
/// Blurs the image with a Gaussian or blur kernel filter.
9+
/// Blurs the image with a Gaussian or box blur kernel filter.
1010
#[node_macro::node(category("Raster: Filter"))]
1111
async fn blur(
1212
_: impl Ctx,
@@ -42,6 +42,36 @@ async fn blur(
4242
.collect()
4343
}
4444

45+
/// Applies a median filter to reduce noise while preserving edges.
46+
#[node_macro::node(category("Raster: Filter"))]
47+
async fn median_filter(
48+
_: impl Ctx,
49+
/// The image to be filtered.
50+
image_frame: Table<Raster<CPU>>,
51+
/// The radius of the filter kernel. Larger values remove more noise but may blur fine details.
52+
#[range((0., 50.))]
53+
#[hard_min(0.)]
54+
radius: PixelLength,
55+
) -> Table<Raster<CPU>> {
56+
image_frame
57+
.into_iter()
58+
.map(|mut row| {
59+
let image = row.element.clone();
60+
61+
// Apply median filter
62+
let filtered_image = if radius < 0.5 {
63+
// Minimum filter radius
64+
image.clone()
65+
} else {
66+
Raster::new_cpu(median_filter_algorithm(image.into_data(), radius as u32))
67+
};
68+
69+
row.element = filtered_image;
70+
row
71+
})
72+
.collect()
73+
}
74+
4575
// 1D gaussian kernel
4676
fn gaussian_kernel(radius: f64) -> Vec<f64> {
4777
// Given radius, compute the size of the kernel that's approximately three times the radius
@@ -179,3 +209,56 @@ fn box_blur_algorithm(mut original_buffer: Image<Color>, radius: f64, gamma: boo
179209

180210
y_axis
181211
}
212+
213+
fn median_filter_algorithm(original_buffer: Image<Color>, radius: u32) -> Image<Color> {
214+
let (width, height) = original_buffer.dimensions();
215+
let mut output = Image::new(width, height, Color::TRANSPARENT);
216+
217+
// Pre-allocate and reuse buffers outside the loops to avoid repeated allocations.
218+
let window_capacity = ((2 * radius + 1).pow(2)) as usize;
219+
let mut r_vals: Vec<f32> = Vec::with_capacity(window_capacity);
220+
let mut g_vals: Vec<f32> = Vec::with_capacity(window_capacity);
221+
let mut b_vals: Vec<f32> = Vec::with_capacity(window_capacity);
222+
let mut a_vals: Vec<f32> = Vec::with_capacity(window_capacity);
223+
224+
for y in 0..height {
225+
for x in 0..width {
226+
r_vals.clear();
227+
g_vals.clear();
228+
b_vals.clear();
229+
a_vals.clear();
230+
231+
// Use saturating_add to avoid potential overflow in extreme cases
232+
let y_max = y.saturating_add(radius).min(height - 1);
233+
let x_max = x.saturating_add(radius).min(width - 1);
234+
235+
for ny in y.saturating_sub(radius)..=y_max {
236+
for nx in x.saturating_sub(radius)..=x_max {
237+
if let Some(px) = original_buffer.get_pixel(nx, ny) {
238+
r_vals.push(px.r());
239+
g_vals.push(px.g());
240+
b_vals.push(px.b());
241+
a_vals.push(px.a());
242+
}
243+
}
244+
}
245+
246+
let r = median_quickselect(&mut r_vals);
247+
let g = median_quickselect(&mut g_vals);
248+
let b = median_quickselect(&mut b_vals);
249+
let a = median_quickselect(&mut a_vals);
250+
251+
output.set_pixel(x, y, Color::from_rgbaf32_unchecked(r, g, b, a));
252+
}
253+
}
254+
255+
output
256+
}
257+
/// Finds the median of a slice using quickselect for O(n) average case performance.
258+
/// This is more efficient than sorting the entire slice which would be O(n log n).
259+
fn median_quickselect(values: &mut [f32]) -> f32 {
260+
let mid: usize = values.len() / 2;
261+
// nth_unstable is like quickselect: average O(n)
262+
// Use total_cmp for safe NaN handling instead of partial_cmp().unwrap()
263+
*values.select_nth_unstable_by(mid, |a, b| a.total_cmp(b)).1
264+
}

0 commit comments

Comments
 (0)