diff --git a/crates/algo/src/denoise.rs b/crates/algo/src/denoise.rs --- a/crates/algo/src/denoise.rs +++ b/crates/algo/src/denoise.rs @@ -28,28 +28,57 @@ pub(crate) fn reach() -> f32 { (2 * MAX_LUMA_R + 2 * MAX_CHROMA_R + 6) as f32 } -/// Guided filter of `p` with guide `g` (both `w × h`), box radius `r`. -fn guided(g: &[f32], p: &[f32], w: usize, h: usize, r: usize, eps: f32) -> Vec { - let len = w * h; - // Interleave the four box-filtered moments to run one pass. - let mut m = vec![0.0f32; len * 4]; - for i in 0..len { - m[i * 4] = g[i]; - m[i * 4 + 1] = p[i]; - m[i * 4 + 2] = g[i] * g[i]; - m[i * 4 + 3] = g[i] * p[i]; - } - box_blur_n(&mut m, w, h, 4, r); - let mut ab = vec![0.0f32; len * 2]; - for i in 0..len { - let (mg, mp, mgg, mgp) = (m[i * 4], m[i * 4 + 1], m[i * 4 + 2], m[i * 4 + 3]); - let var = (mgg - mg * mg).max(0.0); - let a = (mgp - mg * mp) / (var + eps); - ab[i * 2] = a; - ab[i * 2 + 1] = mp - a * mg; +/// Shared first and second moments of a guide (two filtered channels). +fn guide_moments(g: &[f32], w: usize, h: usize, r: usize) -> Vec { + let mut moments = Vec::with_capacity(g.len() * 2); + for &v in g { + moments.extend_from_slice(&[v, v * v]); + } + box_blur_n(&mut moments, w, h, 2, r); + moments +} + +/// Self-guided luminance: mean(p) = mean(g), mean(g*p) = mean(g*g). +fn guided_luma(g: &[f32], w: usize, h: usize, r: usize, eps: f32) -> Vec { + let mut ab = guide_moments(g, w, h, r); + for pair in ab.as_chunks_mut::<2>().0.iter_mut() { + let mean = pair[0]; + let covariance = pair[1] - mean * mean; + // Keep the original unclamped numerator, including roundoff below zero. + let a = covariance / (covariance.max(0.0) + eps); + pair[0] = a; + pair[1] = mean - a * mean; } box_blur_n(&mut ab, w, h, 2, r); - (0..len).map(|i| ab[i * 2] * g[i] + ab[i * 2 + 1]).collect() + g.iter().zip(ab.as_chunks::<2>().0.iter()).map(|(&v, pair)| pair[0] * v + pair[1]).collect() +} + +/// Chroma planes share the cleaned guide's mean and regularised variance. +/// Two guide channels plus two moments and two coefficients per chroma plane: +/// ten filtered channels, versus twelve for two independent guided filters. +fn guided_chroma(g: &[f32], planes: [&mut [f32]; 2], w: usize, h: usize, r: usize, eps: f32, amount: f32) { + let mut stats = guide_moments(g, w, h, r); + for pair in stats.as_chunks_mut::<2>().0.iter_mut() { + pair[1] = (pair[1] - pair[0] * pair[0]).max(0.0) + eps; + } + for plane in planes { + let mut ab = Vec::with_capacity(g.len() * 2); + for (&guide, &p) in g.iter().zip(plane.iter()) { + ab.extend_from_slice(&[p, guide * p]); + } + box_blur_n(&mut ab, w, h, 2, r); + for (pair, guide) in ab.as_chunks_mut::<2>().0.iter_mut().zip(stats.as_chunks::<2>().0.iter()) { + let mean = pair[0]; + let a = (pair[1] - guide[0] * mean) / guide[1]; + pair[0] = a; + pair[1] = mean - a * guide[0]; + } + box_blur_n(&mut ab, w, h, 2, r); + for ((v, &guide), pair) in plane.iter_mut().zip(g).zip(ab.as_chunks::<2>().0.iter()) { + let filtered = pair[0] * guide + pair[1]; + *v += (filtered - *v) * amount; + } + } } pub(crate) fn reduce_noise(src: &Image, out: Rect, ctx: &Ctx, spec: &DenoiseSpec) -> Vec { @@ -125,7 +154,7 @@ pub(crate) fn reduce_noise(src: &Image, out: Rect, ctx: &Ctx, spec: &DenoiseSpec let r = (1.0 + spec.strength.clamp(0.0, 10.0) * 0.4).round() as usize; let keep = spec.preserve.clamp(0.0, 100.0) / 100.0; let eps = (0.004 + 0.05 * s).powi(2) * (1.05 - keep); - let d = guided(&yv, &yv, ww, wh, r.min(MAX_LUMA_R), eps); + let d = guided_luma(&yv, ww, wh, r.min(MAX_LUMA_R), eps); let k = s.sqrt(); for (v, dv) in yv.iter_mut().zip(d) { *v += (dv - *v) * k; @@ -135,12 +164,7 @@ pub(crate) fn reduce_noise(src: &Image, out: Rect, ctx: &Ctx, spec: &DenoiseSpec if col > 0.0 && model != Model::Gray { let r = (1.0 + col * 8.0).round() as usize; let eps = (0.01 + 0.05 * col).powi(2); - for plane in [&mut c1, &mut c2] { - let d = guided(&yv, plane, ww, wh, r.min(MAX_CHROMA_R), eps); - for (v, dv) in plane.iter_mut().zip(d) { - *v += (dv - *v) * col; - } - } + guided_chroma(&yv, [&mut c1, &mut c2], ww, wh, r.min(MAX_CHROMA_R), eps, col); } // Sharpen Details: unsharp mask on luma. if sh > 0.0 { @@ -178,3 +202,106 @@ pub(crate) fn reduce_noise(src: &Image, out: Rect, ctx: &Ctx, spec: &DenoiseSpec } res } + +#[cfg(test)] +mod tests { + use super::*; + + /// Guided filter of `p` with guide `g` (both `w × h`), box radius `r`. + fn reference_guided(g: &[f32], p: &[f32], w: usize, h: usize, r: usize, eps: f32) -> Vec { + let len = w * h; + // Interleave the four box-filtered moments to run one pass. + let mut m = vec![0.0f32; len * 4]; + for i in 0..len { + m[i * 4] = g[i]; + m[i * 4 + 1] = p[i]; + m[i * 4 + 2] = g[i] * g[i]; + m[i * 4 + 3] = g[i] * p[i]; + } + box_blur_n(&mut m, w, h, 4, r); + let mut ab = vec![0.0f32; len * 2]; + for i in 0..len { + let (mg, mp, mgg, mgp) = (m[i * 4], m[i * 4 + 1], m[i * 4 + 2], m[i * 4 + 3]); + let var = (mgg - mg * mg).max(0.0); + let a = (mgp - mg * mp) / (var + eps); + ab[i * 2] = a; + ab[i * 2 + 1] = mp - a * mg; + } + box_blur_n(&mut ab, w, h, 2, r); + (0..len).map(|i| ab[i * 2] * g[i] + ab[i * 2 + 1]).collect() + } + + #[test] + fn shared_statistics_match_independent_filters_bit_for_bit() { + for (w, h) in [(1, 1), (1, 19), (23, 1), (7, 5), (33, 17)] { + let len = w * h; + for flat in [false, true] { + let g: Vec = (0..len).map(|i| if flat { 0.1 } else { ((i * 37 % 101) as f32 - 20.0) / 63.0 }).collect(); + let p: Vec = (0..len).map(|i| (i * 13 % 71) as f32 / 70.0 - 0.5).collect(); + let q: Vec = (0..len).map(|i| (i * 29 % 83) as f32 / 41.0 - 1.0).collect(); + for r in [0, 1, MAX_LUMA_R, MAX_CHROMA_R, 25] { + for eps in [0.000001, 0.0036] { + let expected = reference_guided(&g, &g, w, h, r, eps); + let actual = guided_luma(&g, w, h, r, eps); + assert_bits(&actual, &expected); + for amount in [0.01, 0.45, 1.0] { + let mut actual_p = p.clone(); + let mut actual_q = q.clone(); + guided_chroma(&g, [&mut actual_p, &mut actual_q], w, h, r, eps, amount); + for (original, actual) in [(&p, &actual_p), (&q, &actual_q)] { + let filtered = reference_guided(&g, original, w, h, r, eps); + let expected: Vec = original.iter().zip(filtered).map(|(&v, d)| v + (d - v) * amount).collect(); + assert_bits(actual, &expected); + } + } + } + } + } + } + } + + #[test] + fn reduce_noise_guides_both_chroma_planes_by_the_cleaned_luma() { + // Caller wiring: Lab keeps its planes as they are, so each output plane is compared with + // the independent filters run on the planes in order (luma by itself, chroma by the + // cleaned luma). + let (w, h) = (19, 13); + let rect = Rect::new(0, 0, w as i32, h as i32); + let mut src = Image::new(rect, 3); + for (i, px) in src.data.as_chunks_mut::<3>().0.iter_mut().enumerate() { + px[0] = (i * 37 % 101) as f32 / 100.0; + px[1] = (i * 13 % 71) as f32 / 70.0 - 0.5; + px[2] = (i * 29 % 83) as f32 / 41.0 - 1.0; + } + let ctx = Ctx { bounds: rect, mode: ColorMode::Lab, alpha: false }; + let spec = DenoiseSpec { strength: 6.0, preserve: 30.0, color: 50.0, sharpen: 0.0, jpeg: false }; + let out = reduce_noise(&src, rect, &ctx, &spec); + + let plane = |c: usize| -> Vec { src.data.as_chunks::<3>().0.iter().map(|px| px[c]).collect() }; + let (luma, chroma) = (plane(0), [plane(1), plane(2)]); + let s = spec.strength / 10.0; + let eps = (0.004 + 0.05 * s).powi(2) * (1.05 - spec.preserve / 100.0); + let smooth = reference_guided(&luma, &luma, w, h, (1.0 + spec.strength * 0.4).round() as usize, eps); + let cleaned: Vec = luma.iter().zip(&smooth).map(|(&v, &d)| v + (d - v) * s.sqrt()).collect(); + let col = spec.color / 100.0; + let chroma_r = (1.0 + col * 8.0).round() as usize; + let chroma_eps = (0.01 + 0.05 * col).powi(2); + let mut expected = vec![cleaned]; + for p in &chroma { + let filtered = reference_guided(&expected[0], p, w, h, chroma_r, chroma_eps); + expected.push(p.iter().zip(filtered).map(|(&v, f)| v + (f - v) * col).collect()); + } + for (c, expected) in expected.iter().enumerate() { + let actual: Vec = out.as_chunks::<3>().0.iter().map(|px| px[c]).collect(); + assert_bits(&actual, expected); + assert!(actual.iter().zip(&plane(c)).any(|(a, o)| (a - o).abs() > 1e-3), "plane {c} was filtered"); + } + } + + fn assert_bits(actual: &[f32], expected: &[f32]) { + assert_eq!(actual.len(), expected.len()); + for (i, (actual, expected)) in actual.iter().zip(expected).enumerate() { + assert_eq!(actual.to_bits(), expected.to_bits(), "pixel {i}: {actual} != {expected}"); + } + } +}