Upload the learned demosaic's result, and blend grain back into it

DemosaicedImage::from_rgb_f32 takes the network's linear camera RGB and
stands it beside the classical source of the same photograph: the matrix,
profile tables and as-shot balance are that source's, the id is new, so
nothing downstream can tell which demosaic ran and every cache keyed on the
source sees a new one.

GrainBlend is the denoise's live control. It returns only the brightness of
the noise the network removed, taken after the as-shot balance and handed
back divided by it, so the grain is neutral in the finished picture; colour
speckle and demosaic false colour stay out. It writes a new source rather
than adding a term to the adjust shader: the blend depends on two images and
one number, a 20 MP pass is milliseconds, and a fresh source id is all the
adjust pass's caches need. The test reads it back: at 0 the network's
result, at 1 the same white-balanced step in every channel.

core/dr-gpu/tests/grain.rs

@@ -0,0 +1,145 @@
+//! TRACES: FR-DEV-3g
+//! The grain blend, read back off the device.
+
+use dr_decode::{CfaPattern, CropRect, RawImage};
+use dr_gpu::{DemosaicedImage, Demosaicer, GpuContext, GrainBlend};
+
+const W: u32 = 16;
+const H: u32 = 8;
+
+fn ctx() -> Option<GpuContext> {
+    pollster::block_on(GpuContext::new_headless()).ok()
+}
+
+/// A photograph to stand the uploads beside: its as-shot balance is what
+/// the grain is made neutral under.
+fn like(ctx: &GpuContext) -> DemosaicedImage {
+    let raw = RawImage {
+        width: W,
+        height: H,
+        data: vec![400; (W * H) as usize],
+        cfa_pattern: CfaPattern::Rggb,
+        black_level: [0; 4],
+        white_level: 4095,
+        wb_coeffs: [2.0, 1.0, 1.5, 1.0],
+        color_matrix: Some([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]),
+        samples_per_pixel: 1,
+        profile: None,
+        profile_tables: None,
+        make: String::new(),
+        model: String::new(),
+        crop: CropRect {
+            x: 0,
+            y: 0,
+            width: W,
+            height: H,
+        },
+    };
+    Demosaicer::new(ctx).unwrap().run(&raw).unwrap()
+}
+
+fn read(ctx: &GpuContext, img: &DemosaicedImage) -> Vec<[f32; 4]> {
+    let (w, h) = (img.texture().width(), img.texture().height());
+    let padded =
+        (w * 8).div_ceil(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT) * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
+    let buf = ctx.device.create_buffer(&wgpu::BufferDescriptor {
+        label: None,
+        size: (padded * h) as u64,
+        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+        mapped_at_creation: false,
+    });
+    let mut enc = ctx.device.create_command_encoder(&Default::default());
+    enc.copy_texture_to_buffer(
+        img.texture().as_image_copy(),
+        wgpu::TexelCopyBufferInfo {
+            buffer: &buf,
+            layout: wgpu::TexelCopyBufferLayout {
+                offset: 0,
+                bytes_per_row: Some(padded),
+                rows_per_image: Some(h),
+            },
+        },
+        wgpu::Extent3d {
+            width: w,
+            height: h,
+            depth_or_array_layers: 1,
+        },
+    );
+    ctx.queue.submit(Some(enc.finish()));
+    let slice = buf.slice(..);
+    slice.map_async(wgpu::MapMode::Read, |_| {});
+    ctx.device
+        .poll(wgpu::PollType::wait_indefinitely())
+        .unwrap();
+    let bytes = slice.get_mapped_range();
+    let mut out = Vec::new();
+    for y in 0..h as usize {
+        let row: &[u16] =
+            bytemuck::cast_slice(&bytes[y * padded as usize..y * padded as usize + w as usize * 8]);
+        for t in row.chunks(4) {
+            out.push([0, 1, 2, 3].map(|c| half_to_f32(t[c])));
+        }
+    }
+    out
+}
+
+fn half_to_f32(h: u16) -> f32 {
+    let s = if h & 0x8000 != 0 { -1.0 } else { 1.0 };
+    let e = ((h >> 10) & 0x1f) as i32;
+    let m = (h & 0x3ff) as f32;
+    if e == 0 {
+        s * m * 2f32.powi(-24)
+    } else {
+        s * (1.0 + m / 1024.0) * 2f32.powi(e - 15)
+    }
+}
+
+#[test]
+fn grain_returns_only_neutral_brightness() {
+    let Some(ctx) = ctx() else {
+        eprintln!("no GPU adapter; skipping");
+        return;
+    };
+    let base = like(&ctx);
+    let n = (W * H) as usize;
+    let d: Vec<f32> = (0..n).flat_map(|_| [0.20, 0.30, 0.10]).collect();
+    // The classical result: the same colour plus noise, coloured noise too.
+    let c: Vec<f32> = (0..n)
+        .flat_map(|i| {
+            let a = ((i * 37) % 11) as f32 / 110.0 - 0.05;
+            let b = ((i * 53) % 7) as f32 / 140.0 - 0.025;
+            [0.20 + a, 0.30 + b, 0.10 - a]
+        })
+        .collect();
+    let denoised = DemosaicedImage::from_rgb_f32(&ctx, &base, W, H, &d).unwrap();
+    let classical = DemosaicedImage::from_rgb_f32(&ctx, &base, W, H, &c).unwrap();
+    let blend = GrainBlend::new(&ctx);
+    let wb = [2.0f32, 1.0, 1.5];
+
+    let none = read(&ctx, &blend.blend(&denoised, &classical, 0.0).unwrap());
+    for p in &none {
+        for ch in 0..3 {
+            assert!(
+                (p[ch] - d[ch]).abs() < 1e-3,
+                "grain 0 must be the network's result: {p:?}"
+            );
+        }
+    }
+
+    let all = read(&ctx, &blend.blend(&denoised, &classical, 1.0).unwrap());
+    for (i, p) in all.iter().enumerate() {
+        let want_dy: f32 = [0.2126f32, 0.7152, 0.0722]
+            .iter()
+            .enumerate()
+            .map(|(ch, k)| k * wb[ch] * (c[i * 3 + ch] - d[ch]))
+            .sum();
+        // After white balance every channel moved by the same amount.
+        for ch in 0..3 {
+            let moved = wb[ch] * (p[ch] - d[ch]);
+            assert!(
+                (moved - want_dy).abs() < 2e-3,
+                "pixel {i} channel {ch}: moved {moved}, want {want_dy}"
+            );
+        }
+    }
+}