1pub fn filter(
2 width: u32,
3 height: u32,
4 color: &[[f32; 4]],
5 albedo: &[[f32; 4]],
6 normal: &[[f32; 4]],
7 samples: u32,
8) -> Result<Vec<[f32; 4]>, String> {
9 let len = (width as usize)
10 .checked_mul(height as usize)
11 .ok_or("image dimensions overflow")?;
12 if len != color.len() || len != albedo.len() || len != normal.len() || samples == 0 {
13 return Err("invalid denoise inputs".into());
14 }
15 let mut source = color.to_vec();
16 let mut target = source.clone();
17 let kernel = [0.375, 0.25, 0.0625];
18 for step in [1i32, 2, 4] {
19 for y in 0..height as i32 {
20 for x in 0..width as i32 {
21 let index = y as usize * width as usize + x as usize;
22 let center = source[index];
23 let base = albedo[index];
24 let n = normal[index];
25 let luma = luminance(center);
26 let tolerance = 2.8 / (samples as f32).sqrt() * (luma + 0.05);
27 let mut sum = [0.0f32; 3];
28 let mut weight_sum = 0.0;
29 for dy in -2i32..=2 {
30 for dx in -2i32..=2 {
31 let xx = x + dx * step;
32 let yy = y + dy * step;
33 if xx < 0 || yy < 0 || xx >= width as i32 || yy >= height as i32 {
34 continue;
35 }
36 let other = yy as usize * width as usize + xx as usize;
37 let b = albedo[other];
38 let m = normal[other];
39 let base_gap = (0..3).map(|c| (base[c] - b[c]).abs()).sum::<f32>();
40 let normal_gap = (0..3).map(|c| (n[c] - m[c]).abs()).sum::<f32>();
41 let light_gap = (luminance(source[other]) - luma).abs();
42 let guide = (-base_gap * 8.0 - normal_gap * 4.0).exp();
43 let light = (-light_gap / tolerance.max(1e-5)).exp();
44 let weight = kernel[dx.unsigned_abs() as usize]
45 * kernel[dy.unsigned_abs() as usize]
46 * guide
47 * light;
48 for (c, value) in sum.iter_mut().enumerate() {
49 *value += source[other][c] * weight;
50 }
51 weight_sum += weight;
52 }
53 }
54 for (c, value) in sum.iter().enumerate() {
55 target[index][c] = *value / weight_sum.max(1e-8);
56 }
57 }
58 }
59 std::mem::swap(&mut source, &mut target);
60 }
61 Ok(source)
62}
63
64fn luminance(color: [f32; 4]) -> f32 {
65 color[0] * 0.2126 + color[1] * 0.7152 + color[2] * 0.0722
66}
67
68#[cfg(test)]
69mod tests {
70 use super::*;
71
72 #[test]
73 fn preserves_a_material_edge_and_reduces_grain() {
74 let width = 32;
75 let height = 16;
76 let mut color = Vec::new();
77 let mut albedo = Vec::new();
78 let mut normal = Vec::new();
79 for y in 0..height {
80 for x in 0..width {
81 let side = if x < width / 2 { 0.2 } else { 0.8 };
82 let grain = if (x + y) % 2 == 0 { 0.08 } else { -0.08 };
83 color.push([side + grain, side + grain, side + grain, 1.0]);
84 albedo.push([side, side, side, 1.0]);
85 normal.push([0.0, 0.0, 1.0, 1.0]);
86 }
87 }
88 let result = filter(width, height, &color, &albedo, &normal, 64).unwrap();
89 let left = result[(height / 2 * width + width / 2 - 1) as usize][0];
90 let right = result[(height / 2 * width + width / 2) as usize][0];
91 assert!(right - left > 0.5);
92 let index = (height / 2 * width + width / 4) as usize;
93 assert!((result[index][0] - 0.2).abs() < (color[index][0] - 0.2).abs());
94 assert_eq!(
95 result,
96 filter(width, height, &color, &albedo, &normal, 64).unwrap()
97 );
98 }
99}