WIP: noise reduction

Checkpoint committed by the coordinator, not by the authoring agent: the
session hit its API limit mid-task and left this work uncommitted. Committed
so it survives, NOT because it is finished - expect failing tests and
half-applied changes. The agent resumes from here.
This commit is contained in:
2026-08-22 19:01:18 +02:00
parent c963dafd09
commit b4e55b47c1
6 changed files with 1469 additions and 5 deletions
+471
View File
@@ -0,0 +1,471 @@
//! 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.034 at full amount: within
// the range where the filter is meant to treat a difference as noise.
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 thirty times the threshold. It has to survive intact
// — an edge 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 on a constant grey: colour that alternates every
/// `half_period` pixels with the lightness very nearly fixed.
///
/// This is what chroma noise looks like to the filter — a colour difference
/// with almost no luminance difference under it — and it is the one pattern
/// that can tell the two halves of this operation apart.
fn chroma_pattern(ctx: &GpuContext, size: u32, half_period: u32, swing: i32) -> DemosaicedImage {
upload(ctx, size, move |x, _| {
let on = (x / half_period) % 2 == 0;
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 one
/// linear channel 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.
fn modulation(pixels: &[u8], width: u32, half_period: u32, channel: usize) -> 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 sample = match channel {
usize::MAX => luminance(linear(pixels, width, x, row)),
c => linear(pixels, width, x, row)[c],
};
let sign = if (x / half_period) % 2 == 0 { 1.0 } else { -1.0 };
total += sample * sign;
count += 1.0;
}
total / count
}
/// Correlate against luminance rather than a channel.
const AS_LUMINANCE: usize = usize::MAX;
#[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. The tolerance is wide
// because both readings pass through an eight-bit readback twice over; the
// failure it guards 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, 0);
let colour_after = modulation(&chroma, SIZE, 4, 0);
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, AS_LUMINANCE);
let light_after = modulation(&chroma, SIZE, 4, AS_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 proxy sees the same pattern at half the period, with no
// resampling of its own to confuse the measurement.
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, 0);
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, 0);
// Both must be doing something: two flat images would agree perfectly and
// prove nothing.
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, 0)
};
assert!(
export_amp < untouched * 0.8,
"the denoiser did nothing: {export_amp} of {untouched}"
);
assert!(
(proxy_amp - export_amp).abs() < export_amp * 0.2,
"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, 0);
assert!(
wrong_amp < export_amp * 0.7,
"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"
);
}