DarkRoom/core/dr-denoise/src/repair.rs

//! 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);
    }
}