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

//! 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,
        baseline_exposure: 0.0,
        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}"
            );
        }
    }
}