//! Raw upload, black/white normalisation, and Bayer demosaic. //! //! The first real pipeline stage (ARCH §5.2). It takes CFA sensor data from //! `dr-decode`, uploads it once, and produces a linear scene-referred //! RGBA16Float texture in *camera* colour space. Everything downstream — the //! camera matrix, white balance, the tone operations — works on that texture //! and never sees the CFA pattern. //! //! Uploaded once per image, not per frame. Moving a slider re-runs the adjust //! pass over this texture; it does not re-demosaic, which is what keeps the //! interaction budget (NFR-P9) reachable on a 24 MP file. use dr_decode::{CfaPattern, RawImage}; use wgpu::util::DeviceExt; use crate::{GpuContext, GpuError}; /// Uniform block for the demosaic pass. Layout must match `demosaic.wgsl`. #[repr(C)] #[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)] struct DemosaicParams { width: u32, height: u32, crop_x: u32, crop_y: u32, stride: u32, pattern: u32, _pad0: u32, _pad1: u32, black: [f32; 4], inv_range: [f32; 4], } /// A demosaiced image living on the GPU. /// /// Linear, scene-referred, camera colour space, RGBA16Float. This is the /// input every adjustment operates on, and the reason the ops need no /// knowledge of sensors or CFA patterns. pub struct DemosaicedImage { texture: wgpu::Texture, view: wgpu::TextureView, width: u32, height: u32, /// Carried through for the camera→sRGB transform in the adjust pass. color_matrix: [f32; 9], /// As-shot white balance, the neutral starting point for the WB control. as_shot_wb: [f32; 3], } impl DemosaicedImage { pub const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba16Float; pub fn texture(&self) -> &wgpu::Texture { &self.texture } pub fn view(&self) -> &wgpu::TextureView { &self.view } pub fn size(&self) -> (u32, u32) { (self.width, self.height) } /// Camera RGB → linear sRGB, row-major. Identity where the body is /// uncalibrated, so the image renders uncalibrated rather than black. pub fn color_matrix(&self) -> [f32; 9] { self.color_matrix } /// As-shot white balance multipliers, green-normalised. /// /// The white balance control is expressed *relative* to these, so its /// neutral position reproduces what the camera chose. pub fn as_shot_wb(&self) -> [f32; 3] { self.as_shot_wb } } /// Runs the demosaic pass. Holds the pipeline so repeated images reuse it. pub struct Demosaicer { ctx: GpuContext, pipeline: wgpu::ComputePipeline, bind_group_layout: wgpu::BindGroupLayout, } impl Demosaicer { pub fn new(ctx: &GpuContext) -> Result { let shader = ctx .device .create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("demosaic"), source: wgpu::ShaderSource::Wgsl(include_str!("shaders/demosaic.wgsl").into()), }); let bind_group_layout = ctx.device .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("demosaic-bgl"), entries: &[ // Raw samples, packed two u16 per u32. wgpu::BindGroupLayoutEntry { binding: 0, visibility: wgpu::ShaderStages::COMPUTE, ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, count: None, }, wgpu::BindGroupLayoutEntry { binding: 1, visibility: wgpu::ShaderStages::COMPUTE, ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None, }, count: None, }, wgpu::BindGroupLayoutEntry { binding: 2, visibility: wgpu::ShaderStages::COMPUTE, ty: wgpu::BindingType::StorageTexture { access: wgpu::StorageTextureAccess::WriteOnly, format: DemosaicedImage::FORMAT, view_dimension: wgpu::TextureViewDimension::D2, }, count: None, }, ], }); let layout = ctx .device .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("demosaic-layout"), bind_group_layouts: &[&bind_group_layout], push_constant_ranges: &[], }); let pipeline = ctx .device .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("demosaic-pipeline"), layout: Some(&layout), module: &shader, entry_point: Some("main"), compilation_options: Default::default(), cache: None, }); Ok(Self { ctx: ctx.clone(), pipeline, bind_group_layout, }) } /// Upload and demosaic one image. pub fn run(&self, raw: &RawImage) -> Result { if raw.cfa_pattern.is_xtrans() { // Better to say so than to render a maze of artefacts that reads // as a corrupt file (FR-RAW-5). return Err(GpuError::UnsupportedCfa( "X-Trans demosaic not yet implemented".into(), )); } let pattern = match raw.cfa_pattern { CfaPattern::Rggb => 0u32, CfaPattern::Bggr => 1, CfaPattern::Grbg => 2, CfaPattern::Gbrg => 3, other => { return Err(GpuError::UnsupportedCfa(format!("{other:?}"))); } }; let (width, height) = (raw.crop.width.max(1), raw.crop.height.max(1)); let limits = self.ctx.device.limits(); if width > limits.max_texture_dimension_2d || height > limits.max_texture_dimension_2d { return Err(GpuError::TooLarge(format!( "{width}×{height} exceeds the device limit of {}", limits.max_texture_dimension_2d ))); } // Pack the u16 samples two per u32. WGSL has no u16 storage type, so // unpacking happens in the shader. let packed = pack_samples(&raw.data); let raw_buf = self .ctx .device .create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("raw-samples"), contents: bytemuck::cast_slice(&packed), usage: wgpu::BufferUsages::STORAGE, }); let params = DemosaicParams { width, height, crop_x: raw.crop.x, crop_y: raw.crop.y, stride: raw.width, pattern, _pad0: 0, _pad1: 0, black: black_per_cell(raw), inv_range: inv_range_per_cell(raw), }; let params_buf = self .ctx .device .create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("demosaic-params"), contents: bytemuck::bytes_of(¶ms), usage: wgpu::BufferUsages::UNIFORM, }); let texture = self.ctx.device.create_texture(&wgpu::TextureDescriptor { label: Some("demosaiced"), size: wgpu::Extent3d { width, height, depth_or_array_layers: 1, }, mip_level_count: 1, sample_count: 1, dimension: wgpu::TextureDimension::D2, format: DemosaicedImage::FORMAT, // STORAGE to write here, TEXTURE_BINDING so the adjust pass can // sample it. COPY_SRC only for tests. usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_SRC, view_formats: &[], }); let view = texture.create_view(&Default::default()); let bind_group = self .ctx .device .create_bind_group(&wgpu::BindGroupDescriptor { label: Some("demosaic-bg"), layout: &self.bind_group_layout, entries: &[ wgpu::BindGroupEntry { binding: 0, resource: raw_buf.as_entire_binding(), }, wgpu::BindGroupEntry { binding: 1, resource: params_buf.as_entire_binding(), }, wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::TextureView(&view), }, ], }); let mut enc = self .ctx .device .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("demosaic-encoder"), }); { let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("demosaic-pass"), timestamp_writes: None, }); pass.set_pipeline(&self.pipeline); pass.set_bind_group(0, &bind_group, &[]); pass.dispatch_workgroups(width.div_ceil(8), height.div_ceil(8), 1); } self.ctx.queue.submit(Some(enc.finish())); Ok(DemosaicedImage { texture, view, width, height, // Identity where the body is uncalibrated: the image renders with // no colour transform rather than not at all. color_matrix: raw.color_matrix.unwrap_or(IDENTITY_3X3), as_shot_wb: [raw.wb_coeffs[0], raw.wb_coeffs[1], raw.wb_coeffs[2]], }) } } const IDENTITY_3X3: [f32; 9] = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]; /// Pack u16 samples two per u32, little-endian within the word. /// /// WGSL has no 16-bit storage type without an optional feature, so the shader /// unpacks. An odd sample count pads with a zero, which is never addressed: /// the shader indexes by pixel, not by word. fn pack_samples(data: &[u16]) -> Vec { let mut out = Vec::with_capacity(data.len().div_ceil(2)); let mut chunks = data.chunks_exact(2); for pair in &mut chunks { out.push(u32::from(pair[0]) | (u32::from(pair[1]) << 16)); } if let Some(&last) = chunks.remainder().first() { out.push(u32::from(last)); } out } /// Black level per CFA cell position, indexed `(y & 1) * 2 + (x & 1)`. /// /// `dr-decode` reports four levels in CFA order, which is already this /// layout. Bodies reporting a single level get it broadcast. fn black_per_cell(raw: &RawImage) -> [f32; 4] { let b = raw.black_level; if b[1] == 0 && b[2] == 0 && b[3] == 0 { return [f32::from(b[0]); 4]; } [ f32::from(b[0]), f32::from(b[1]), f32::from(b[2]), f32::from(b[3]), ] } /// Reciprocal of the usable range per cell, so the shader avoids a division. /// /// A white level at or below black would divide by zero; such a file is /// malformed, and falling back to full scale renders something inspectable /// rather than a NaN texture. fn inv_range_per_cell(raw: &RawImage) -> [f32; 4] { let white = f32::from(raw.white_level); let black = black_per_cell(raw); let mut out = [0.0f32; 4]; for (i, &b) in black.iter().enumerate() { let range = white - b; out[i] = if range > 1.0 { 1.0 / range } else { 1.0 / 65535.0 }; } out } #[cfg(test)] mod tests { use super::*; use dr_decode::CropRect; fn raw_for(black: [u16; 4], white: u16) -> RawImage { RawImage { width: 4, height: 4, data: vec![0; 16], cfa_pattern: CfaPattern::Rggb, black_level: black, white_level: white, wb_coeffs: [1.0, 1.0, 1.0, 1.0], color_matrix: None, crop: CropRect { x: 0, y: 0, width: 4, height: 4, }, } } #[test] fn samples_pack_two_per_word() { let packed = pack_samples(&[0x1234, 0xABCD]); assert_eq!(packed, vec![0xABCD_1234]); } #[test] fn an_odd_sample_count_does_not_lose_the_last_value() { // A sensor with an odd sample count would otherwise drop its final // photosite, or worse, read past the buffer. let packed = pack_samples(&[0x0001, 0x0002, 0x0003]); assert_eq!(packed.len(), 2); assert_eq!(packed[1] & 0xFFFF, 3); } #[test] fn packing_preserves_every_sample() { let data: Vec = (0..64).map(|i| i * 1000).collect(); let packed = pack_samples(&data); for (i, &expected) in data.iter().enumerate() { let word = packed[i / 2]; let got = if i % 2 == 0 { word & 0xFFFF } else { word >> 16 }; assert_eq!(got as u16, expected, "sample {i}"); } } #[test] fn a_single_black_level_is_broadcast_to_every_cell() { // Many bodies report one level rather than four; treating the absent // three as zero would leave three quarters of the image lifted. let raw = raw_for([512, 0, 0, 0], 16383); assert_eq!(black_per_cell(&raw), [512.0; 4]); } #[test] fn per_cell_black_levels_are_kept_distinct() { let raw = raw_for([2047, 2048, 2048, 2049], 15070); assert_eq!(black_per_cell(&raw), [2047.0, 2048.0, 2048.0, 2049.0]); } #[test] fn normalisation_maps_white_to_one() { let raw = raw_for([2048, 2048, 2048, 2048], 15070); let inv = inv_range_per_cell(&raw); let normalised = (15070.0 - 2048.0) * inv[0]; assert!( (normalised - 1.0).abs() < 1e-5, "the white level must land on 1.0, got {normalised}" ); } #[test] fn a_degenerate_range_does_not_divide_by_zero() { // A malformed file reporting white <= black must not produce NaN // across the whole texture. let raw = raw_for([5000, 5000, 5000, 5000], 4000); let inv = inv_range_per_cell(&raw); assert!(inv.iter().all(|v| v.is_finite() && *v > 0.0)); } // ---- GPU tests ----------------------------------------------------- // // These exercise the shader itself. The CPU tests above cover the // parameter maths; only running the pass proves the kernels and the CFA // indexing are right. fn ctx() -> Option { match pollster::block_on(GpuContext::new_headless()) { Ok(c) => Some(c), Err(e) => { eprintln!("skipping: no GPU adapter ({e})"); None } } } /// Build a synthetic CFA image of a uniform colour. /// /// Each photosite carries its own channel's value, which is what a /// sensor looking at a flat patch would record. A correct demosaic must /// return that colour at every pixel. fn flat_cfa(pattern: CfaPattern, size: u32, rgb: [u16; 3], black: u16, white: u16) -> RawImage { let mut data = vec![0u16; (size * size) as usize]; for y in 0..size { for x in 0..size { let c = pattern.colour_at(x, y) as usize; data[(y * size + x) as usize] = black + rgb[c]; } } RawImage { width: size, height: size, data, cfa_pattern: pattern, black_level: [black; 4], white_level: white, wb_coeffs: [1.0, 1.0, 1.0, 1.0], color_matrix: None, crop: CropRect { x: 0, y: 0, width: size, height: size, }, } } /// Read back the demosaiced texture as f32 RGBA. fn read_rgba(ctx: &GpuContext, img: &DemosaicedImage) -> Vec<[f32; 4]> { let (w, h) = img.size(); let unpadded = w * 8; // RGBA16Float = 8 bytes per pixel let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT; let padded = unpadded.div_ceil(align) * align; let buf = ctx.device.create_buffer(&wgpu::BufferDescriptor { label: Some("demosaic-readback"), 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( wgpu::ImageCopyTexture { texture: img.texture(), mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All, }, wgpu::ImageCopyBuffer { buffer: &buf, layout: wgpu::ImageDataLayout { 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(..); let (tx, rx) = std::sync::mpsc::channel(); slice.map_async(wgpu::MapMode::Read, move |r| { let _ = tx.send(r); }); ctx.device.poll(wgpu::Maintain::Wait); rx.recv().expect("map").expect("map ok"); let data = slice.get_mapped_range(); let mut out = Vec::with_capacity((w * h) as usize); for y in 0..h { let row = (y * padded) as usize; for x in 0..w { let px = row + (x * 8) as usize; let mut c = [0.0f32; 4]; for (i, slot) in c.iter_mut().enumerate() { let o = px + i * 2; let bits = u16::from_le_bytes([data[o], data[o + 1]]); *slot = half_to_f32(bits); } out.push(c); } } drop(data); buf.unmap(); out } /// Decode an IEEE 754 binary16 value. fn half_to_f32(bits: u16) -> f32 { let sign = f32::from_bits(u32::from(bits & 0x8000) << 16); let exp = (bits >> 10) & 0x1F; let mant = bits & 0x03FF; let magnitude = match exp { 0 => f32::from(mant) * 2f32.powi(-24), 0x1F => { if mant == 0 { f32::INFINITY } else { f32::NAN } } _ => (1.0 + f32::from(mant) / 1024.0) * 2f32.powi(i32::from(exp) - 15), }; magnitude.copysign(sign) } #[test] fn a_flat_patch_demosaics_to_that_colour() { // The fundamental correctness property: a sensor looking at a uniform // colour must reconstruct that colour everywhere. Any error in the // kernels, the CFA indexing, or the normalisation breaks it. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let black = 2048u16; let white = 16383u16; let range = f32::from(white - black); // A colour with three clearly distinct channels, so a swap is loud. let rgb = [8000u16, 4000, 2000]; let expected = [ f32::from(rgb[0]) / range, f32::from(rgb[1]) / range, f32::from(rgb[2]) / range, ]; let raw = flat_cfa(CfaPattern::Rggb, 32, rgb, black, white); let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); // Interior pixels only: the border reflects, and a 2-pixel margin is // where that shows. let (w, _) = img.size(); for y in 2..30u32 { for x in 2..30u32 { let p = px[(y * w + x) as usize]; for (ch, &want) in expected.iter().enumerate() { assert!( (p[ch] - want).abs() < 0.01, "pixel ({x},{y}) channel {ch}: got {}, want {want}", p[ch] ); } } } } #[test] fn every_bayer_layout_reconstructs_the_same_colour() { // The packed CFA constants in the shader are easy to get wrong — two // of the four were wrong on the first attempt, and a wrong one swaps // red and blue. Each layout describes the same scene, so each must // produce the same output. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let (black, white) = (0u16, 16383u16); let rgb = [9000u16, 5000, 1500]; let expected = [ f32::from(rgb[0]) / f32::from(white), f32::from(rgb[1]) / f32::from(white), f32::from(rgb[2]) / f32::from(white), ]; for pattern in [ CfaPattern::Rggb, CfaPattern::Bggr, CfaPattern::Grbg, CfaPattern::Gbrg, ] { let raw = flat_cfa(pattern, 32, rgb, black, white); let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, _) = img.size(); let p = px[(16 * w + 16) as usize]; for (ch, &want) in expected.iter().enumerate() { assert!( (p[ch] - want).abs() < 0.01, "{pattern:?} channel {ch}: got {}, want {want} — \ a wrong CFA constant swaps channels", p[ch] ); } } } #[test] fn an_odd_crop_origin_still_reconstructs_correctly() { // The case CropRect::shifts_cfa_phase exists for. Cropping to an odd // origin re-phases the pattern; if the shader is given the unshifted // one, red and blue swap. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let (black, white) = (0u16, 16383u16); let rgb = [9000u16, 5000, 1500]; let mut raw = flat_cfa(CfaPattern::Rggb, 34, rgb, black, white); // Crop one photosite in on both axes, as a body with an odd active // area would. The visible top-left is now green-on-a-red-row. raw.crop = CropRect { x: 1, y: 1, width: 32, height: 32, }; let (dx, dy) = raw.crop.shifts_cfa_phase(); assert!(dx && dy, "the fixture must actually shift the phase"); raw.cfa_pattern = raw.cfa_pattern.shifted(dx, dy); let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, _) = img.size(); let p = px[(16 * w + 16) as usize]; let expected = [ f32::from(rgb[0]) / f32::from(white), f32::from(rgb[1]) / f32::from(white), f32::from(rgb[2]) / f32::from(white), ]; for (ch, &want) in expected.iter().enumerate() { assert!( (p[ch] - want).abs() < 0.01, "channel {ch}: got {}, want {want} — the crop origin \ re-phases the CFA and the shader must see the shifted pattern", p[ch] ); } } #[test] fn output_is_free_of_nan_and_negatives() { // f16 NaN propagates silently through every later stage; a negative // value breaks the ratio-based operations downstream. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); // A high-contrast checkerboard, which is where the gradient // correction overshoots hardest. let size = 32u32; let mut data = vec![0u16; (size * size) as usize]; for y in 0..size { for x in 0..size { data[(y * size + x) as usize] = if (x / 2 + y / 2) % 2 == 0 { 16000 } else { 40 }; } } let raw = RawImage { width: size, height: size, data, cfa_pattern: CfaPattern::Rggb, black_level: [32; 4], white_level: 16383, wb_coeffs: [1.0, 1.0, 1.0, 1.0], color_matrix: None, crop: CropRect { x: 0, y: 0, width: size, height: size, }, }; let img = d.run(&raw).expect("demosaic"); for (i, p) in read_rgba(&ctx, &img).iter().enumerate() { for (ch, v) in p.iter().take(3).enumerate() { assert!(v.is_finite(), "pixel {i} channel {ch} is {v}"); assert!(*v >= 0.0, "pixel {i} channel {ch} is negative: {v}"); } } } #[test] fn xtrans_is_refused_rather_than_rendered_wrong() { let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let mut raw = flat_cfa(CfaPattern::Rggb, 8, [100, 100, 100], 0, 1000); raw.cfa_pattern = CfaPattern::XTrans; assert!( matches!(d.run(&raw), Err(GpuError::UnsupportedCfa(_))), "an X-Trans file must report the gap, not render artefacts" ); } }