Repair photosites beyond 8 sigma of every neighbour before the network

The app's hot-pixel pass takes gross defects only; at ISO 6400-25600 a 6D
frame keeps 1000-2000 photosites more than 8 sigma beyond all their
same-colour and adjacent neighbours, which the network turned into specks.
The same two tests with the threshold in the photosite's own sigma, plus
the factor of two that keeps a bright point of light (where 8 sigma is a
sliver of the signal). The next model is trained behind exactly this; on
an ISO 25600 frame the Rust and training code both repair 935.

core/dr-denoise/src/lib.rs

@@ -19,6 +19,7 @@
 pub mod noise;
 #[cfg(feature = "onnx")]
 pub mod onnx;
+pub mod repair;
 pub mod tile;
 
 use dr_decode::RawImage;
@@ -67,12 +68,23 @@ pub fn denoise(
     }
     let active = noise::active(raw);
     let (h, w) = (active.h, active.w);
+    // The active area laid out once, then the noise-aware repair the model
+    // was trained behind (see `repair`).
+    let mut mosaic: Vec<f32> = (0..h * w).map(|i| active.at(i / w, i % w)).collect();
+    let pattern = raw.cfa_pattern;
+    let repaired = repair::repair(&mut mosaic, h, w, repair::REPAIR_K, &|y, x, v| {
+        noise.sigma(pattern.colour_at(x as u32, y as u32) as usize, v)
+    });
+    log::info!(
+        "learned denoise: {repaired} photosites beyond {}σ of every neighbour repaired",
+        repair::REPAIR_K
+    );
     tile::run_tiled(
         net,
         h,
         w,
         raw.cfa_pattern,
-        &|y, x| active.at(y, x),
+        &|y, x| mosaic[y * w + x],
         &|c, v| noise.sigma(c, v),
         progress,
     )

core/dr-denoise/src/repair.rs

@@ -0,0 +1,145 @@
+//! TRACES: FR-DEV-3g
+//! Hot and dead photosites, judged against the noise, before the network.
+//!
+//! The app's own pass (`dr_gpu::Demosaicer::repair_hot_pixels`) runs first and
+//! takes the gross defects. At high ISO it leaves thousands of photosites per
+//! 6D frame more than 8σ beyond every neighbour, which the network turns into
+//! specks. This second pass uses that pass's two tests with the threshold in
+//! units of the photosite's own σ from the noise model:
+//!
+//! - beyond every same-colour neighbour (two photosites away, the 3×3 of its
+//!   plane) by more than `k·σ`, and
+//! - beyond every adjacent photosite, whatever its colour, by more than
+//!   `k·σ` **and** by a factor of two — what keeps a real point of light,
+//!   which lights its neighbours through the lens and the anti-aliasing
+//!   filter. A margin in σ alone is not enough: on a bright star 8σ is a
+//!   sliver of the signal, and the star would be flattened.
+//!
+//! A hot one becomes its brightest same-colour neighbour, a dead one its
+//! darkest. The shipped model was trained on input repaired exactly so
+//! (darkroom-denoise `denoise/repair.py`, `--repair-k 8`): the threshold
+//! belongs to the model, and changes with it. Neighbours off the frame are
+//! the nearest photosite on it, as the training code reads them.
+
+/// The threshold the shipped model was trained with, in σ.
+pub const REPAIR_K: f32 = 8.0;
+
+/// Repair `mosaic` (`h×w`, row-major, normalised) in place; `sigma(y, x, v)`
+/// is the photosite's σ. Returns how many photosites changed.
+pub fn repair(
+    mosaic: &mut [f32],
+    h: usize,
+    w: usize,
+    k: f32,
+    sigma: &(dyn Fn(usize, usize, f32) -> f32 + Sync),
+) -> usize {
+    let copy = mosaic.to_vec();
+    let original = ©
+    let at = |y: isize, x: isize| {
+        let y = y.clamp(0, h as isize - 1) as usize;
+        let x = x.clamp(0, w as isize - 1) as usize;
+        original[y * w + x]
+    };
+    let threads = std::thread::available_parallelism().map_or(1, |n| n.get());
+    let rows_per = h.div_ceil(threads).max(1);
+    let mut counts = vec![0usize; h.div_ceil(rows_per)];
+    std::thread::scope(|scope| {
+        for ((chunk, rows), count) in mosaic
+            .chunks_mut(rows_per * w)
+            .enumerate()
+            .zip(counts.iter_mut())
+        {
+            let at = &at;
+            scope.spawn(move || {
+                for (i, row) in rows.chunks_mut(w).enumerate() {
+                    let y = chunk * rows_per + i;
+                    for (x, out) in row.iter_mut().enumerate() {
+                        let v = original[y * w + x];
+                        let (yi, xi) = (y as isize, x as isize);
+                        let (mut s_hi, mut s_lo) = (f32::MIN, f32::MAX);
+                        let (mut a_hi, mut a_lo) = (f32::MIN, f32::MAX);
+                        for dy in -1isize..=1 {
+                            for dx in -1isize..=1 {
+                                if dy == 0 && dx == 0 {
+                                    continue;
+                                }
+                                let s = at(yi + 2 * dy, xi + 2 * dx);
+                                s_hi = s_hi.max(s);
+                                s_lo = s_lo.min(s);
+                                let a = at(yi + dy, xi + dx);
+                                a_hi = a_hi.max(a);
+                                a_lo = a_lo.min(a);
+                            }
+                        }
+                        let t = k * sigma(y, x, v);
+                        if v - s_hi > t && v - a_hi > t && a_hi < 0.5 * v {
+                            *out = s_hi;
+                            *count += 1;
+                        } else if s_lo - v > t && a_lo - v > t && v < 0.5 * a_lo {
+                            *out = s_lo;
+                            *count += 1;
+                        }
+                    }
+                }
+            });
+        }
+    });
+    counts.iter().sum()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const N: usize = 16;
+
+    fn flat(level: f32) -> Vec<f32> {
+        vec![level; N * N]
+    }
+
+    fn run(m: &mut [f32]) -> usize {
+        repair(m, N, N, REPAIR_K, &|_, _, _| 0.01)
+    }
+
+    #[test]
+    fn a_hot_photosite_becomes_its_brightest_same_colour_neighbour() {
+        let mut m = flat(0.1);
+        m[8 * N + 8] = 0.5; // 40σ above everything around it
+        m[8 * N + 10] = 0.12; // a same-colour neighbour, a little brighter
+        assert_eq!(run(&mut m), 1);
+        assert_eq!(m[8 * N + 8], 0.12);
+    }
+
+    #[test]
+    fn a_dead_photosite_in_a_lit_area_is_repaired() {
+        let mut m = flat(0.5);
+        m[5 * N + 5] = 0.0;
+        assert_eq!(run(&mut m), 1);
+        assert_eq!(m[5 * N + 5], 0.5);
+    }
+
+    #[test]
+    fn a_point_of_real_light_is_kept() {
+        // Light through a lens lands on a patch: its adjacent photosites are
+        // lit too, so the second test refuses it.
+        let mut m = flat(0.1);
+        for dy in 0..3 {
+            for dx in 0..3 {
+                m[(7 + dy) * N + 7 + dx] = if (dy, dx) == (1, 1) { 0.9 } else { 0.6 };
+            }
+        }
+        let before = m.clone();
+        assert_eq!(run(&mut m), 0);
+        assert_eq!(m, before);
+    }
+
+    #[test]
+    fn noise_within_the_threshold_is_left_alone() {
+        let mut m: Vec<f32> = (0..N * N)
+            .map(|i| 0.1 + 0.005 * ((i * 7919 % 13) as f32 - 6.0) / 6.0)
+            .collect();
+        let before = m.clone();
+        assert_eq!(run(&mut m), 0);
+        assert_eq!(m, before);
+    }
+}