DarkRoom/core/dr-gpu/tests/noise_reduction.rs

//! Noise reduction, end to end on a real device.
//!
//! `dr-pipeline`'s own tests assert what the operation *composes* — how many
//! passes, what radius, in what unit. None of them can tell whether the WGSL
//! compiles, whether the filter actually preserves an edge, or whether the
//! luminance and chroma halves stay out of each other's way once real floats
//! run through them. Those are questions only a GPU answers.
//!
//! # Reading the expected values
//!
//! Sources are uploaded through `DemosaicedImage::from_rgba8`, which flags
//! them non-linear, so the generated shader decodes sRGB before any operation
//! runs and the detail stage sees linear values. The last detail pass
//! re-encodes. So every assertion here decodes the readback back to linear
//! before comparing — comparing 8-bit code values directly would fold the
//! transfer function's varying slope into every tolerance.
//!
//! Almost everything is measured **against a baseline render of the same
//! image with the amount at zero**, rather than against an absolute
//! expectation. That is deliberate: it isolates what noise reduction did from
//! everything else the pipeline does to a pixel, and it stays correct if a
//! later change to the chain moves the values this stage is handed.

use dr_gpu::{AdjustPass, DemosaicedImage, GpuContext};
use dr_pipeline::detail::RenderScale;
use dr_pipeline::ops::noise_reduction::{CHROMA, ID, LUMINANCE};
use dr_pipeline::{Affects, EditGraph};
use dr_types::ColourSpace;

fn ctx() -> Option<GpuContext> {
    // CI runners and headless machines may have no usable adapter. Skip rather
    // than fail, exactly as the rest of this crate's device tests do.
    match pollster::block_on(GpuContext::new_headless()) {
        Ok(c) => Some(c),
        Err(e) => {
            eprintln!("skipping: no GPU adapter ({e})");
            None
        }
    }
}

fn graph_with(luminance: f32, chroma: f32) -> EditGraph {
    let mut graph = EditGraph::default_chain();
    graph.set_param(ID, LUMINANCE, luminance);
    graph.set_param(ID, CHROMA, chroma);
    graph
}

/// Render one graph and read the pixels back, at the scale the graph itself
/// works out — which is what a frontend does.
fn render(
    ctx: &GpuContext,
    pass: &mut AdjustPass,
    graph: &EditGraph,
    source: &DemosaicedImage,
    out: (u32, u32),
) -> Vec<u8> {
    let scale = graph.render_scale(source.size(), out);
    render_at(ctx, pass, graph, source, out, scale)
}

/// Render with an explicitly chosen [`RenderScale`].
///
/// Split out for one test only — the one that needs to compose the detail
/// stage at the *wrong* scale on purpose, to show that the conversion from
/// source pixels to render pixels is load-bearing rather than decorative.
fn render_at(
    _ctx: &GpuContext,
    pass: &mut AdjustPass,
    graph: &EditGraph,
    source: &DemosaicedImage,
    out: (u32, u32),
    scale: RenderScale,
) -> Vec<u8> {
    let shader = graph.compose_for(ColourSpace::Srgb);
    let detail = graph.compose_detail_for(scale, ColourSpace::Srgb);
    let key = graph.invalidation().through(Affects::Colour);
    pass.render_detailed(source, &shader, out.0, out.1, None, &detail, key)
        .expect("render");
    pass.export_pixels().expect("readback").0
}

fn srgb_decode(byte: u8) -> f32 {
    let e = byte as f32 / 255.0;
    if e <= 0.040_45 {
        e / 12.92
    } else {
        ((e + 0.055) / 1.055).powf(2.4)
    }
}

fn luminance(c: [f32; 3]) -> f32 {
    0.2126 * c[0] + 0.7152 * c[1] + 0.0722 * c[2]
}

/// One pixel of a readback, as linear RGB.
fn linear(pixels: &[u8], width: u32, x: u32, y: u32) -> [f32; 3] {
    let i = ((y * width + x) * 4) as usize;
    [
        srgb_decode(pixels[i]),
        srgb_decode(pixels[i + 1]),
        srgb_decode(pixels[i + 2]),
    ]
}

/// A pixel split the way the operation itself splits it: a luminance, and a
/// colour difference whose own luminance is zero.
fn split(pixels: &[u8], width: u32, x: u32, y: u32) -> (f32, [f32; 3]) {
    let c = linear(pixels, width, x, y);
    let y0 = luminance(c);
    (y0, [c[0] - y0, c[1] - y0, c[2] - y0])
}

/// Upload an image built from a per-pixel closure.
fn upload(ctx: &GpuContext, size: u32, f: impl Fn(u32, u32) -> [u8; 3]) -> DemosaicedImage {
    let data: Vec<u8> = (0..size * size)
        .flat_map(|i| {
            let (x, y) = (i % size, i / size);
            let [r, g, b] = f(x, y);
            [r, g, b, 255]
        })
        .collect();
    DemosaicedImage::from_rgba8(ctx, &data, size, size).expect("upload")
}

/// A vertical step edge of a chosen height, centred on the frame.
fn step_edge(ctx: &GpuContext, size: u32, low: u8, high: u8) -> DemosaicedImage {
    upload(ctx, size, move |x, _| {
        let v = if x < size / 2 { low } else { high };
        [v, v, v]
    })
}

#[test]
fn a_difference_below_the_threshold_is_averaged_and_one_above_it_is_not() {
    // The defining property, and the reason this is a bilateral rather than a
    // Gaussian. Both images are step edges and the filter is identical; the
    // only thing that differs is how tall the step is relative to the noise
    // threshold. A Gaussian would soften both by exactly the same amount, and
    // that indiscriminate softening is what "denoised" pictures look like.
    let Some(ctx) = ctx() else { return };
    const SIZE: u32 = 64;
    let mid = SIZE / 2;
    let row = SIZE / 2;

    let measure = |source: &DemosaicedImage| -> f32 {
        let mut off = AdjustPass::new(&ctx);
        let plain = render(&ctx, &mut off, &graph_with(0.0, 0.0), source, (SIZE, SIZE));
        let mut on = AdjustPass::new(&ctx);
        let denoised = render(&ctx, &mut on, &graph_with(100.0, 0.0), source, (SIZE, SIZE));
        // How far the pixel just inside the bright side moved, in linear
        // luminance. An averaging filter pulls it down towards the dark half;
        // an edge-preserving one leaves it where it was.
        let (before, _) = split(&plain, SIZE, mid, row);
        let (after, _) = split(&denoised, SIZE, mid, row);
        before - after
    };

    // A step of eight code values around mid-grey is about 0.028 in linear
    // luminance, against a threshold of roughly 0.035 at full amount: within
    // the range where the filter is meant to treat a difference as noise.
    //
    // Worked out on paper, since it cannot be run here: at amount 100 the
    // kernel is 3 render pixels and sigma_k is 0.075, so at y = 0.2159 the
    // range sigma is 0.075·sqrt(0.2159 + 0.0025) = 0.0350 and a neighbour
    // 0.0280 away is weighted exp(-0.320) = 0.726. Summing the 7×7 kernel's
    // spatial weights over the four bright columns and the three dark ones
    // gives a filtered luminance of 0.2076 against 0.2159 — a move of 0.0082,
    // which survives the 8-bit readback as about **0.0072**. The threshold
    // below is set well under that rather than at it: what would be a bug is
    // the filter declining to average at all.
    let quiet = measure(&step_edge(&ctx, SIZE, 120, 128));
    assert!(
        quiet > 0.004,
        "a difference below the threshold was left alone: moved {quiet}"
    );

    // Black to white is thirteen times the threshold — 1.0 against a sigma of
    // 0.075·sqrt(1.0025) = 0.0751 — so a neighbour across it is weighted
    // exp(-88), which is zero in any arithmetic. The edge has to survive
    // intact; one that softens here is a halo in every high-contrast picture.
    let loud = measure(&step_edge(&ctx, SIZE, 0, 255));
    assert!(
        loud.abs() < 0.004,
        "an edge far above the threshold was smoothed: moved {loud}"
    );
    assert!(
        quiet > loud.abs() * 3.0,
        "the filter did not distinguish noise from an edge: {quiet} vs {loud}"
    );
}

/// A fine chroma pattern: red and blue swung in opposite directions by the
/// same number of code values, green held.
///
/// This is what chroma noise looks like to the filter — a colour difference
/// alternating over a few pixels — and it is the one pattern that can tell the
/// two halves of this operation apart.
///
/// It is not a *pure* colour pattern, and the tests must not assume it is.
/// Rec. 709 weights red at 0.2126 and blue at 0.0722, so swinging one up and
/// the other down by equal amounts moves lightness by about a seventh of the
/// swing. That residue is real, it is not the chroma filter's to remove, and
/// [`chroma_r`] is what keeps it out of the measurements.
fn chroma_pattern(ctx: &GpuContext, size: u32, half_period: u32, swing: i32) -> DemosaicedImage {
    upload(ctx, size, move |x, _| {
        let on = (x / half_period).is_multiple_of(2);
        let d = if on { swing } else { -swing };
        [(128 + d) as u8, 128, (128 - d) as u8]
    })
}

/// How much of a known alternating pattern survived, as the correlation of a
/// chosen measurement of the middle row against the pattern's own sign.
///
/// A matched filter rather than a peak-to-peak reading. The readback is eight
/// bits, and the modulation these tests work with is only a handful of code
/// values — deliberately, because a larger one would read as a real colour
/// boundary and the filter would refuse to touch it. Correlating over a whole
/// number of periods averages the quantisation down instead of letting it set
/// the noise floor of the measurement.
///
/// `margin` covers a whole number of periods too, so the window is unbiased by
/// the row's mean, and it keeps the measurement clear of the borders where a
/// clamped kernel legitimately behaves differently.
///
/// `sample` is what to measure. Passing [`chroma_r`] rather than the raw red
/// channel matters more than it looks: a colour square wave built from 8-bit
/// code values carries a *luminance* square wave under it — Rec. 709 does not
/// weight red and blue equally, so swinging one up and the other down moves
/// lightness too — and that component is not the chroma filter's to remove.
/// Left in the measurement it is a constant floor under every reading, which
/// compresses every ratio this file asserts towards one and would leave the
/// tests unable to tell a correct kernel from one twice the size.
fn modulation(
    pixels: &[u8],
    width: u32,
    half_period: u32,
    sample: impl Fn([f32; 3]) -> f32,
) -> f32 {
    let row = width / 2;
    let margin = half_period * 4;
    let mut total = 0.0;
    let mut count = 0.0;
    for x in margin..(width - margin) {
        let value = sample(linear(pixels, width, x, row));
        let sign = if (x / half_period).is_multiple_of(2) {
            1.0
        } else {
            -1.0
        };
        total += value * sign;
        count += 1.0;
    }
    total / count
}

/// The red component of the colour difference — red with its lightness taken
/// out, which is the quantity the chroma passes actually filter.
fn chroma_r(c: [f32; 3]) -> f32 {
    c[0] - luminance(c)
}

#[test]
fn chroma_noise_reduction_never_moves_lightness() {
    // Half of the claim the two-slider design rests on. The split is into a
    // luminance and a colour difference whose own luminance is zero, so the
    // chroma passes reconstruct with the lightness this pixel arrived with —
    // exactly, not approximately.
    let Some(ctx) = ctx() else { return };
    const SIZE: u32 = 64;

    // The strongest available form of the assertion, on an image that has no
    // colour to filter: every colour difference is zero, so the filter is the
    // identity and the output must be the *same bytes*. A reconstruction that
    // used a filtered luminance instead of this pixel's own would soften the
    // step and show up here immediately.
    let grey = step_edge(&ctx, SIZE, 90, 110);
    let mut off = AdjustPass::new(&ctx);
    let plain = render(&ctx, &mut off, &graph_with(0.0, 0.0), &grey, (SIZE, SIZE));
    let mut on = AdjustPass::new(&ctx);
    let denoised = render(&ctx, &mut on, &graph_with(0.0, 100.0), &grey, (SIZE, SIZE));
    assert_eq!(
        plain, denoised,
        "chroma noise reduction altered an image with no colour in it"
    );

    // And on an image that does have colour to filter, where the two halves
    // could actually interfere: the colour modulation must fall while the
    // luminance modulation under it stays where it was.
    //
    // On paper, at amount 100 over a half-period of 4: the kernel is 12 render
    // pixels, both chroma thresholds are 0.16·sqrt(y + floor) ≈ 0.076 and
    // 0.20·… ≈ 0.095, so an opposite-coloured neighbour is weighted 0.410, and
    // summing the separable kernel over the eight phases leaves about **0.42**
    // of the chroma amplitude — 0.0158 of 0.0377. The luminance amplitude must
    // not move at all: every colour difference the pass averages has zero
    // luminance by construction, so their weighted mean does too.
    //
    // Both tolerances are wide of those numbers because both readings pass
    // through an eight-bit readback twice over; the failure they guard against
    // is not a drift of a few percent but a collapse, which is what a leak
    // between the two components would be.
    let source = chroma_pattern(&ctx, SIZE, 4, 12);
    let mut off = AdjustPass::new(&ctx);
    let plain = render(&ctx, &mut off, &graph_with(0.0, 0.0), &source, (SIZE, SIZE));
    let mut on = AdjustPass::new(&ctx);
    let chroma = render(
        &ctx,
        &mut on,
        &graph_with(0.0, 100.0),
        &source,
        (SIZE, SIZE),
    );

    let colour_before = modulation(&plain, SIZE, 4, chroma_r);
    let colour_after = modulation(&chroma, SIZE, 4, chroma_r);
    assert!(
        colour_after < colour_before * 0.7,
        "chroma noise reduction did not reduce the colour swing: \
         {colour_after} of {colour_before}"
    );

    let light_before = modulation(&plain, SIZE, 4, luminance);
    let light_after = modulation(&chroma, SIZE, 4, luminance);
    assert!(
        (light_after - light_before).abs() < light_before.abs() * 0.3,
        "chroma noise reduction moved lightness: {light_after} was {light_before}"
    );
}

#[test]
fn luminance_noise_reduction_never_moves_colour() {
    // The other half. The luminance pass adds the *change* in lightness back
    // to the colour it was given, so the colour difference passes through
    // untouched however hard the luminance is filtered. Written the obvious
    // way instead — filtering the three channels and calling it a luminance
    // filter — the colour would desaturate as the amount rose, and the chroma
    // slider would stop meaning anything on its own.
    let Some(ctx) = ctx() else { return };
    const SIZE: u32 = 64;
    let source = chroma_pattern(&ctx, SIZE, 4, 12);

    let mut off = AdjustPass::new(&ctx);
    let plain = render(&ctx, &mut off, &graph_with(0.0, 0.0), &source, (SIZE, SIZE));
    let mut on = AdjustPass::new(&ctx);
    let luma = render(
        &ctx,
        &mut on,
        &graph_with(100.0, 0.0),
        &source,
        (SIZE, SIZE),
    );

    // The colour difference — not the raw channel, which follows lightness.
    let row = SIZE / 2;
    let interior = 16..(SIZE - 16);
    let mut worst = 0.0f32;
    for x in interior {
        let (_, before) = split(&plain, SIZE, x, row);
        let (_, after) = split(&luma, SIZE, x, row);
        worst = worst.max((after[0] - before[0]).abs());
    }
    // A code value at this brightness, doubled for the two readbacks the
    // comparison passes through. The leak this guards against would be a
    // sizeable fraction of the pattern's own 0.037 swing, not a rounding.
    let quantum = srgb_decode(129) - srgb_decode(128);
    assert!(
        worst < quantum * 3.0,
        "luminance noise reduction moved colour by {worst} \
         (one code value is {quantum})"
    );
}

#[test]
fn the_same_edit_denoises_the_same_at_two_resolutions() {
    // TRACES: FR-DSP-1 — the thing this operation is most likely to get wrong.
    //
    // A radius is stored in *source* pixels and converted to render pixels at
    // every render, because noise is made by photosites. Get that conversion
    // wrong and the develop view and the exported file are different
    // photographs: tune the slider on a half-size proxy and the export is
    // denoised at half the strength, or twice it.
    //
    // The subject is a chroma square wave with a period that is a power of two
    // and aligned to the frame, so halving the render resolution decimates it
    // exactly. The fused pass loads the nearest source pixel when the framing
    // is unrotated, so output column `x` reads source column `2x` — always the
    // same half of an eight-wide block as `2x + 1` — and the proxy sees the
    // same two colours at half the period, with no resampling of its own to
    // confuse the measurement.
    //
    // # The expected numbers
    //
    // Worked out on paper, because the tolerances below are meaningless
    // without knowing what they are tolerances *around*.
    //
    // The two colours are (134, 128, 122) and (122, 128, 134), which decode to
    // linear (0.2384, 0.2159, 0.1946) and its mirror. Their colour differences
    // are ±(0.0193, -0.0033, -0.0245), so the pattern's chroma amplitude —
    // what `modulation` with `chroma_r` reads — is 0.0188 before filtering.
    //
    // At amount 60 both chroma thresholds are 0.12·sqrt(y + floor) ≈ 0.0565,
    // so a neighbour of the opposite colour is weighted
    // `exp(-(dl²/2σ_g² + |dc|²/2σ_c²)) ≈ exp(-0.624) ≈ 0.536`: attenuated, but
    // far from rejected, which is the regime where the kernel's *width* is
    // what decides the answer. That is the point — a test where the range
    // weights dominated would pass whatever the radius conversion did.
    //
    // Summing the separable kernel's spatial weights over the eight phases of
    // the pattern gives a mean surviving fraction of **0.521** at export (a
    // radius of 8 render pixels over a half-period of 8) and **0.521** on the
    // proxy (4 over 4) — the two arrangements are the same filter sampled at
    // two rates, and they agree to three decimal places. So both amplitudes
    // land at about 0.0098.
    //
    // The control lands at **0.305**, about 0.0057: a kernel of 8 render
    // pixels over a half-period of 4 is twice as wide in the terms that
    // matter. The tolerances are set wide of those numbers rather than tight
    // to them, because the readback is eight bits and each of the handful of
    // distinct output values carries up to half a code value of rounding — a
    // floor of a few percent on any of these readings that no amount of
    // averaging over a larger frame removes, since the error is periodic
    // rather than random.
    let Some(ctx) = ctx() else { return };
    const SOURCE: u32 = 256;
    const HALF_PERIOD: u32 = 8; // in source pixels
                                // Six code values of swing. Small on purpose: the colour difference has to
                                // land near the filter's threshold, because a larger one is a colour
                                // boundary and the whole point of a bilateral is that it refuses to cross
                                // those. There would be nothing to measure at either resolution.
    let source = chroma_pattern(&ctx, SOURCE, HALF_PERIOD, 6);

    // Sixty percent is an eight-source-pixel radius, which halves to exactly
    // four render pixels on a half-size proxy — so the rounding to an integer
    // kernel is not what this test is measuring.
    let graph = graph_with(0.0, 60.0);

    let mut export_pass = AdjustPass::new(&ctx);
    let export = render(&ctx, &mut export_pass, &graph, &source, (SOURCE, SOURCE));
    let export_amp = modulation(&export, SOURCE, HALF_PERIOD, chroma_r);

    let proxy_size = SOURCE / 2;
    let mut proxy_pass = AdjustPass::new(&ctx);
    let proxy = render(
        &ctx,
        &mut proxy_pass,
        &graph,
        &source,
        (proxy_size, proxy_size),
    );
    let proxy_amp = modulation(&proxy, proxy_size, HALF_PERIOD / 2, chroma_r);

    // Both must be doing something: two flat images would agree perfectly and
    // prove nothing. About 0.0098 of 0.0188, by the derivation above.
    let untouched = {
        let mut pass = AdjustPass::new(&ctx);
        let plain = render(
            &ctx,
            &mut pass,
            &graph_with(0.0, 0.0),
            &source,
            (SOURCE, SOURCE),
        );
        modulation(&plain, SOURCE, HALF_PERIOD, chroma_r)
    };
    assert!(
        export_amp < untouched * 0.8,
        "the denoiser did nothing: {export_amp} of {untouched}"
    );

    assert!(
        (proxy_amp - export_amp).abs() < export_amp * 0.25,
        "the same edit left {proxy_amp} of the pattern on the proxy and \
         {export_amp} on the export"
    );

    // The control, and the reason the tolerance above means something. Compose
    // the detail stage as though the proxy were a full-resolution render —
    // which is exactly the bug of storing a radius in render pixels — and the
    // kernel is twice as wide in source terms. If the conversion were not
    // load-bearing, this would land in the same place as the other two.
    let mut wrong_pass = AdjustPass::new(&ctx);
    let wrong = render_at(
        &ctx,
        &mut wrong_pass,
        &graph,
        &source,
        (proxy_size, proxy_size),
        RenderScale::full((proxy_size, proxy_size)),
    );
    let wrong_amp = modulation(&wrong, proxy_size, HALF_PERIOD / 2, chroma_r);
    assert!(
        wrong_amp < export_amp * 0.8,
        "an unconverted radius was indistinguishable from a converted one: \
         {wrong_amp} against {export_amp}"
    );
}

#[test]
fn each_amount_costs_only_the_dispatches_it_needs() {
    // The cost story, which is invisible in the picture and therefore has to
    // be asserted on a counter. Luminance is one exact two-dimensional pass;
    // chroma is two, because at its radius the exact form is quadratic and
    // unaffordable. An edit using neither must pay for neither — and must
    // produce pixels identical to a chain that has no denoiser in it at all.
    let Some(ctx) = ctx() else { return };
    const SIZE: u32 = 48;
    let source = step_edge(&ctx, SIZE, 40, 200);

    for (luminance, chroma, expected) in [(60.0, 0.0, 1), (0.0, 60.0, 2), (60.0, 60.0, 3)] {
        let mut pass = AdjustPass::new(&ctx);
        render(
            &ctx,
            &mut pass,
            &graph_with(luminance, chroma),
            &source,
            (SIZE, SIZE),
        );
        assert_eq!(
            pass.detail_dispatches(),
            expected,
            "luminance {luminance}, chroma {chroma}"
        );
        assert_eq!(pass.colour_dispatches(), 1);
    }

    let mut neutral = AdjustPass::new(&ctx);
    let a = render(
        &ctx,
        &mut neutral,
        &graph_with(0.0, 0.0),
        &source,
        (SIZE, SIZE),
    );
    assert_eq!(neutral.detail_dispatches(), 0);
    assert_eq!(neutral.detail_allocations(), 0, "nothing was allocated");

    // Byte-identical, not merely close: an operation at its defaults must not
    // touch the image, and a stage that ran and wrote back the same values
    // would still have quantised twice.
    let mut absent = AdjustPass::new(&ctx);
    let b = render(
        &ctx,
        &mut absent,
        &EditGraph::default_chain(),
        &source,
        (SIZE, SIZE),
    );
    assert_eq!(a, b, "a neutral denoiser changed the picture");
}

#[test]
fn dragging_either_slider_recompiles_nothing_and_reallocates_nothing() {
    // TRACES: FR-DEV-3d. Both of these are ruinous per frame and invisible in
    // the output, which is why they need a counter rather than an eye. A
    // radius rides in a uniform buffer, so moving a slider re-runs the detail
    // dispatches against the pipelines already compiled — and does not re-run
    // the colour pass at all, since nothing it depends on moved.
    let Some(ctx) = ctx() else { return };
    const SIZE: u32 = 48;
    let source = step_edge(&ctx, SIZE, 40, 200);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = graph_with(50.0, 50.0);
    render(&ctx, &mut pass, &graph, &source, (SIZE, SIZE));
    let pipelines = pass.cached_detail_pipelines();
    let allocations = pass.detail_allocations();
    assert_eq!(pipelines, 3, "one per pass: luminance, then two for chroma");

    for amount in [55.0, 60.0, 65.0, 70.0] {
        graph.set_param(ID, LUMINANCE, amount);
        graph.set_param(ID, CHROMA, amount);
        render(&ctx, &mut pass, &graph, &source, (SIZE, SIZE));
    }
    assert_eq!(
        pass.cached_detail_pipelines(),
        pipelines,
        "an amount is a uniform, not a shader"
    );
    assert_eq!(
        pass.detail_allocations(),
        allocations,
        "a steady viewport must allocate nothing"
    );
    assert_eq!(
        pass.colour_dispatches(),
        1,
        "the fused colour pass re-ran for a change it does not depend on"
    );
}