//! Raw upload, black/white normalisation, and demosaic — Bayer and X-Trans. //! //! 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], } /// TRACES: FR-RAW-5 /// Uniform block for the X-Trans pass. Layout must match `xtrans.wgsl`. /// /// Separate from [`DemosaicParams`] rather than a superset of it: the two /// passes disagree about what a black level even is — four positional values /// on a 2×2 cell, one sensor-wide value on a 6×6 tile — and a shared block /// would have to carry both and let each shader pick. #[repr(C)] #[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)] struct XTransParams { width: u32, height: u32, crop_x: u32, crop_y: u32, stride: u32, black: f32, inv_range: f32, _pad0: u32, /// As-shot white balance gains, green-normalised. The demosaic /// interpolates in balanced space and undoes them before writing. wb: [f32; 4], inv_wb: [f32; 4], /// The 6×6 tile at this sensor's phase, two bits per photosite. tile: [u32; 4], } /// Uniform block for the hot-pixel repair. Layout must match /// `hot_pixels.wgsl`. /// /// One block for both colour filter arrays: the repair asks only "which /// photosites share this one's colour", and a 6×6 tile answers that for a /// Bayer cell as well as for X-Trans. #[repr(C)] #[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)] struct HotPixelParams { crop_x: u32, crop_y: u32, width: u32, height: u32, stride: u32, words: u32, row_invocations: u32, samples: u32, black: [f32; 4], inv_range: [f32; 4], tile: [u32; 4], } /// The repair's workgroup width. Must match `@workgroup_size` in /// `hot_pixels.wgsl`. const HOT_PIXEL_GROUP: u32 = 64; /// A demosaiced image living on the GPU. /// /// RGBA16Float, scene-referred, camera colour space. This is the input every /// adjustment operates on, and the reason the ops need no knowledge of sensors /// or CFA patterns. /// /// **Two producers, not one.** [`Demosaicer::run`] builds it from CFA sensor /// data; [`DemosaicedImage::from_rgba8`] builds it from an already-processed /// RGB image such as a JPEG. Nothing about the type is CFA-specific — it is /// simply "an image on the GPU, ready to adjust" — which is what lets develop /// mode work on a JPEG without the edit graph or any operation knowing that /// the source was not a RAW file. /// /// The one thing that does differ is the transfer function: sensor data is /// linear, a JPEG is gamma-encoded. That difference is carried by /// [`Self::is_non_linear`] and resolved once, in the generated shader's /// prologue, rather than being defended against by every operation. 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], /// TRACES: FR-DEV-3e /// The camera profile's tables, carried through with the matrix for the /// adjust pass to upload (D20). `None` for a JPEG and for a raw with no /// profile. profile_tables: Option>, /// As-shot white balance, the neutral starting point for the WB control. as_shot_wb: [f32; 3], /// Whether the texture holds gamma-encoded rather than linear values. non_linear: bool, /// Which upload this is, unique for the life of the process. See /// [`Self::id`]. id: u64, /// TRACES: FR-DSP-2 | NFR-RES-2 /// The whole frame's size in pixels — what [`Self::size`] reports. /// The texture's own size when it holds the whole frame at full /// resolution, which is every photograph that fits in one. frame: (u32, u32), /// Which part of the frame the texture holds, as origin and extent in /// normalised frame coordinates. `[0, 0, 1, 1]` for the whole frame, /// reduced or not. See [`Self::window_uniforms`]. window: [f32; 4], } /// The window of a texture that holds the whole frame. const WHOLE_FRAME: [f32; 4] = [0.0, 0.0, 1.0, 1.0]; /// The next [`DemosaicedImage::id`]. fn next_image_id() -> u64 { static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) } 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 } /// The size of the photograph this stands for, in its own pixels. /// /// **Not necessarily the texture's.** For a photograph larger than one /// texture this is a reduced copy of it or a window cut from it, and /// everything that sizes a render, a crop or a kernel has to go on /// measuring the photograph. What indexes the texture's texels asks /// [`Self::texture_size`] instead. pub fn size(&self) -> (u32, u32) { self.frame } /// The texture's own size in texels. pub fn texture_size(&self) -> (u32, u32) { (self.width, self.height) } /// TRACES: FR-DSP-2 | NFR-RES-2 /// The source window uniforms the fused shader reads, in the order /// `dr_pipeline::SOURCE_WINDOW_UNIFORM_FIELDS` declares them. /// /// The second `vec4` is zero for a texture that holds the whole frame at /// full resolution, so the shader measures the texture itself exactly as /// it did before windows existed. pub fn window_uniforms(&self) -> [f32; dr_pipeline::SOURCE_WINDOW_UNIFORM_FIELDS] { let [x, y, w, h] = self.window; let (fw, fh) = if self.is_whole() { (0.0, 0.0) } else { (self.frame.0 as f32, self.frame.1 as f32) }; [x, y, w, h, fw, fh, 0.0, 0.0] } /// Whether the texture is the whole frame at full resolution. pub fn is_whole(&self) -> bool { self.window == WHOLE_FRAME && self.frame == (self.width, self.height) } /// The window this texture holds, as origin and extent in normalised /// frame coordinates. pub fn window(&self) -> [f32; 4] { self.window } /// Which texture this is, as a number that is never reused. /// /// For a cache that has to know it is still looking at the same pixels /// (`AdjustPass`'s sample cache) without holding the texture alive to find /// out: keeping a handle would keep half a gigabyte of a closed photograph /// on the device, and comparing addresses would mistake a new upload for an /// old one the moment the allocator reused the slot. A texture here is /// never written after it is built, so the same id is the same pixels. pub(crate) fn id(&self) -> u64 { self.id } /// 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 } /// TRACES: FR-DEV-3e /// The camera profile's tables this source renders through, if any. pub fn profile_tables(&self) -> Option<&std::sync::Arc> { self.profile_tables.as_ref() } /// 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 } /// The largest edge this device can hold in one texture. /// /// Exposed because it is a *hardware* limit the caller has to plan around, /// not a failure to report after the fact: a 13728×8928 film scan exceeds /// the common 8192 limit, and the only way to develop it at all is to fit /// it first. 8192 is still four times a 4K display's long edge, so nothing /// visible is lost. pub fn max_dimension(ctx: &GpuContext) -> u32 { ctx.device.limits().max_texture_dimension_2d } /// Whether the texture is gamma-encoded rather than linear. /// /// True for a source that arrived already display-encoded — a JPEG. The /// adjust pass forwards this to the shader, which linearises before any /// operation runs, so the ops themselves always see linear colour. pub fn is_non_linear(&self) -> bool { self.non_linear } /// Build one from an already-processed RGBA8 image, skipping demosaic. /// /// The path a JPEG takes into develop mode (FR-RAW-4). There is no CFA to /// interpolate and no sensor to normalise: the pixels are uploaded as they /// arrived, gamma encoding intact, and flagged so the shader linearises /// them. /// /// The two sensor-derived transforms are deliberately neutral rather than /// absent: /// /// - **Colour matrix identity** — a JPEG is already in sRGB primaries, so /// there is no camera space to convert out of. Applying a real camera /// matrix here would be a second, unwanted colour transform. /// - **White balance neutral** — the camera applied its own before writing /// the file, and it cannot be undone from the encoded pixels. The WB /// control still works; its neutral position is simply "as the camera /// left it" rather than "as the sensor recorded it". /// /// `rgba` must be tightly packed, 4 bytes per pixel, `width * height` /// pixels, and must fit [`Self::max_dimension`] — a film scan can easily /// exceed it, so callers downscale first rather than being refused here. pub fn from_rgba8( ctx: &GpuContext, rgba: &[u8], width: u32, height: u32, ) -> Result { let (width, height) = (width.max(1), height.max(1)); let limits = 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 ))); } let expected = (width as usize) * (height as usize) * 4; if rgba.len() < expected { return Err(GpuError::TooLarge(format!( "{} bytes is short of the {expected} needed for {width}×{height}", rgba.len() ))); } // The staging texture is Rgba8Unorm because that is what the bytes // are; the pass below converts into the Rgba16Float the rest of the // pipeline expects. Writing f16 on the CPU instead would cost a // full-image conversion before the upload rather than after it. let half: Vec = rgba[..expected] .iter() .map(|&b| f32_to_f16_bits(f32::from(b) / 255.0)) .collect(); let texture = ctx.device.create_texture_with_data( &ctx.queue, &wgpu::TextureDescriptor { label: Some("jpeg-source"), size: wgpu::Extent3d { width, height, depth_or_array_layers: 1, }, mip_level_count: 1, sample_count: 1, dimension: wgpu::TextureDimension::D2, format: Self::FORMAT, // No STORAGE_BINDING: nothing writes to this one. The adjust // pass samples it, and tests copy from it. usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_SRC, view_formats: &[], }, wgpu::util::TextureDataOrder::LayerMajor, bytemuck::cast_slice(&half), ); let view = texture.create_view(&Default::default()); Ok(Self { texture, view, width, height, color_matrix: IDENTITY_3X3, profile_tables: None, as_shot_wb: [1.0, 1.0, 1.0], // **The identity, and this is the whole reason the field is here // rather than resolved further down.** A JPEG has already been // rendered by the camera; the view transform skips a source // flagged non-linear, since rendering the rendering would crush // the shadows and flatten the highlights of an image that was // already finished. non_linear: true, id: next_image_id(), frame: (width, height), window: WHOLE_FRAME, }) } } impl DemosaicedImage { /// TRACES: FR-MRG-3 /// A source that is already RGB in camera space: a linear DNG, which is /// what a merge writes. No demosaic; the samples are normalised by the /// file's black and white levels exactly as the demosaic kernel would /// normalise a photosite, and everything else — the matrix, the /// balance, the view transform — is carried through as for a CFA /// file, because the composite is developed as one photograph from the /// body that took its sources. pub fn from_linear_rgb16(ctx: &GpuContext, raw: &RawImage) -> Result { let (width, height) = (raw.crop.width.max(1), raw.crop.height.max(1)); Self::linear_rgb16_window(ctx, raw, [0, 0, width, height], 1) } /// TRACES: FR-DSP-2 | NFR-RES-2 /// Part of a linear DNG, or a reduced copy of it, for a photograph too /// large to hold in one texture. /// /// `region` is `[x, y, width, height]` in pixels of the frame (the /// file's crop), clamped to it. `reduce` averages `reduce × reduce` /// blocks into one texel — a box filter, which is what a reduced copy /// that is only ever displayed smaller than itself needs, and which keeps /// the samples in scene-linear light where an average means something. /// /// The texture then knows where it sits ([`Self::window`]) and how large /// the photograph is ([`Self::size`]), and the fused shader maps each /// output pixel's position in the *photograph* into it. So a crop, a /// rotation or a mask drawn on the reduced copy lands on the same pixels /// of a full-resolution window, and an export in tiles is the same /// picture as one that fitted. /// /// Refused only if the result itself does not fit the device. pub fn linear_rgb16_window( ctx: &GpuContext, raw: &RawImage, region: [u32; 4], reduce: u32, ) -> Result { let frame = (raw.crop.width.max(1), raw.crop.height.max(1)); let k = reduce.max(1); let x0 = region[0].min(frame.0 - 1); let y0 = region[1].min(frame.1 - 1); let rw = region[2].clamp(1, frame.0 - x0); let rh = region[3].clamp(1, frame.1 - y0); let (width, height) = (rw.div_ceil(k), rh.div_ceil(k)); let limits = 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 ))); } let stride = raw.width as usize * 3; let expected = raw.height as usize * stride; if raw.data.len() < expected { return Err(GpuError::TooLarge(format!( "{} samples is short of the {expected} a {}×{} RGB image needs", raw.data.len(), raw.width, raw.height ))); } let black = black_per_cell(raw); let inv = inv_range_per_cell(raw); // One output row per task: a 200-megapixel reduction is a second of // one core, and the rows are independent. let row_texels = width as usize * 4; let mut half = vec![0u16; row_texels * height as usize]; let fill_row = |ty: usize, out: &mut [u16]| { let sy0 = y0 as usize + ty * k as usize; let sy1 = (sy0 + k as usize).min((y0 + rh) as usize); for tx in 0..width as usize { let sx0 = x0 as usize + tx * k as usize; let sx1 = (sx0 + k as usize).min((x0 + rw) as usize); let mut acc = [0f32; 3]; for sy in sy0..sy1 { let row = (raw.crop.y as usize + sy) * stride + raw.crop.x as usize * 3; for sx in sx0..sx1 { let p = &raw.data[row + sx * 3..row + sx * 3 + 3]; for c in 0..3 { acc[c] += f32::from(p[c]); } } } let n = ((sy1 - sy0) * (sx1 - sx0)).max(1) as f32; let texel = &mut out[tx * 4..tx * 4 + 4]; for c in 0..3 { let v = (acc[c] / n - black[c]) * inv[c]; texel[c] = f32_to_f16_bits_unclamped(v); } texel[3] = f32_to_f16_bits(1.0); } }; let threads = std::thread::available_parallelism().map_or(1, |n| n.get()); let rows_per = (height as usize).div_ceil(threads).max(1); std::thread::scope(|scope| { for (chunk, rows) in half.chunks_mut(rows_per * row_texels).enumerate() { let fill_row = &fill_row; scope.spawn(move || { for (i, out) in rows.chunks_mut(row_texels).enumerate() { fill_row(chunk * rows_per + i, out); } }); } }); let texture = ctx.device.create_texture_with_data( &ctx.queue, &wgpu::TextureDescriptor { label: Some("linear-rgb-source"), size: wgpu::Extent3d { width, height, depth_or_array_layers: 1, }, mip_level_count: 1, sample_count: 1, dimension: wgpu::TextureDimension::D2, format: Self::FORMAT, usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_SRC, view_formats: &[], }, wgpu::util::TextureDataOrder::LayerMajor, bytemuck::cast_slice(&half), ); let view = texture.create_view(&Default::default()); // The extent is the texels' own, `width × k`, not the region's: the // last block of a reduction may run past the frame's edge, and // stretching it to fit would put every texel slightly off the // pixels it averaged. The shader's bounds test is on the frame, so // nothing past the edge is ever read. let window = [ x0 as f32 / frame.0 as f32, y0 as f32 / frame.1 as f32, (width * k) as f32 / frame.0 as f32, (height * k) as f32 / frame.1 as f32, ]; Ok(Self { texture, view, width, height, color_matrix: raw.color_matrix.unwrap_or(IDENTITY_3X3), profile_tables: raw.profile_tables.clone(), as_shot_wb: [raw.wb_coeffs[0], raw.wb_coeffs[1], raw.wb_coeffs[2]], non_linear: false, id: next_image_id(), frame, window, }) } } /// Convert an f32 to half-precision bits, the general case: sign, /// subnormals, round-to-nearest-even, saturation at the largest finite. /// /// `f32_to_f16_bits` below is the 8-bit special case and says why it can /// be; this one exists because a linear DNG is not that case. A 14-bit /// sensor's least significant step, normalised, is 6.1e-5 — right at f16's /// smallest normal (6.1e-5) — so the deepest shadows of a composite land /// in the subnormal range, and rounding them to zero would crush the /// shadows of exactly the file that was written to keep them. Values below /// zero (black subtraction on a noisy photosite) and above one (a highlight /// past the white level) are legitimate and kept. fn f32_to_f16_bits_unclamped(v: f32) -> u16 { let bits = v.to_bits(); let sign = ((bits >> 16) & 0x8000) as u16; let exp = ((bits >> 23) & 0xFF) as i32; let mant = bits & 0x7F_FFFF; if exp == 0xFF { // Infinity or NaN: a NaN sample is a decode fault; store the largest // finite rather than propagate it through a blend. return sign | 0x7BFF; } let e = exp - 127 + 15; if e >= 0x1F { return sign | 0x7BFF; } if e <= 0 { // Subnormal in f16 (or underflow). Shift the full mantissa with its // implicit bit right by the deficit, rounding to nearest even. if e < -10 { return sign; } let m = (mant | 0x80_0000) >> (1 - e); let shift = 13; let rounded = round_shift(m, shift); return sign | rounded as u16; } let rounded = round_shift(mant, 13); // Rounding can carry into the exponent; that is correct. sign | (((e as u32) << 10) + rounded) as u16 } /// `v >> shift`, rounded to nearest with ties to even. fn round_shift(v: u32, shift: u32) -> u32 { let half = 1u32 << (shift - 1); let mask = (1u32 << shift) - 1; let low = v & mask; let mut out = v >> shift; if low > half || (low == half && (out & 1) == 1) { out += 1; } out } /// Convert an f32 to IEEE 754 half-precision bits. /// /// Written out rather than pulled in as a dependency: the inputs here are /// `0.0..=1.0` from an 8-bit source, which is entirely inside the normal range /// of f16, so the subnormal and overflow cases a general converter must handle /// cannot arise. The clamp makes that assumption explicit rather than implicit. fn f32_to_f16_bits(v: f32) -> u16 { let v = v.clamp(0.0, 1.0); if v == 0.0 { return 0; } let bits = v.to_bits(); let exp = ((bits >> 23) & 0xFF) as i32 - 127 + 15; let mantissa = (bits >> 13) & 0x3FF; // v is in 0.0..=1.0, so the exponent cannot overflow f16's range; values // below f16's smallest normal round to zero rather than to a subnormal, // which at 8-bit source precision is a distinction without a difference. if exp <= 0 { return 0; } ((exp as u16) << 10) | mantissa as u16 } /// Runs the demosaic pass. Holds the pipelines so repeated images reuse them. /// /// Both CFA families are built up front rather than on first use. A Fujifilm /// file arriving mid-session would otherwise pay a shader compilation inside /// the interaction budget, and the compilation is the one part of this that /// can fail on a driver — better to learn that when the pipeline is created /// than when a photograph is opened. pub struct Demosaicer { ctx: GpuContext, pipeline: wgpu::ComputePipeline, xtrans_pipeline: wgpu::ComputePipeline, bind_group_layout: wgpu::BindGroupLayout, hot_pixel_pipeline: wgpu::ComputePipeline, hot_pixel_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 xtrans_shader = ctx .device .create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("xtrans-demosaic"), source: wgpu::ShaderSource::Wgsl(include_str!("shaders/xtrans.wgsl").into()), }); // One layout for both passes. They take the same three bindings — raw // samples, a uniform block, the output texture — and only the contents // of the uniform differ, so a second layout would be the same three // entries written twice. 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: &[Some(&bind_group_layout)], immediate_size: 0, }); 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, }); let xtrans_pipeline = ctx.device .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("xtrans-demosaic-pipeline"), layout: Some(&layout), module: &xtrans_shader, entry_point: Some("main"), compilation_options: Default::default(), cache: None, }); let (hot_pixel_pipeline, hot_pixel_layout) = hot_pixel_pipeline(ctx); Ok(Self { ctx: ctx.clone(), pipeline, xtrans_pipeline, bind_group_layout, hot_pixel_pipeline, hot_pixel_layout, }) } /// Upload and demosaic one image. /// /// Two kernels behind one entry point. The caller hands over a /// `RawImage`; which of the two CFA families it came off is this /// function's problem, not theirs. pub fn run(&self, raw: &RawImage) -> Result { if raw.samples_per_pixel == 3 { return DemosaicedImage::from_linear_rgb16(&self.ctx, raw); } 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 ))); } // Declared here and filled in one branch each, so the bytes handed to // the buffer outlive the `if` that chose them. let bayer_params; let xtrans_params; // Kept for the hot-pixel repair: finding the X-Trans phase reads the // whole frame on the CPU, and once per photograph is enough. let mut xtrans_tile = None; let (pipeline, params_bytes) = if raw.cfa_pattern.is_xtrans() { xtrans_params = xtrans_params_for(raw, width, height); xtrans_tile = Some(xtrans_params.tile); (&self.xtrans_pipeline, bytemuck::bytes_of(&xtrans_params)) } else { 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:?}"))); } }; bayer_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), }; (&self.pipeline, bytemuck::bytes_of(&bayer_params)) }; // 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, }); // TRACES: FR-RAW-3 // The mosaic the demosaic actually reads: the readout with its hot and // dead photosites repaired. A second buffer rather than in place, // because every photosite's verdict reads its neighbours' originals. let repaired = self.ctx.device.create_buffer(&wgpu::BufferDescriptor { label: Some("raw-repaired"), size: raw_buf.size(), usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false, }); let words = packed.len() as u32; let groups = words.div_ceil(HOT_PIXEL_GROUP).max(1); // A 24 MP readout is 190,000 workgroups, past the 65,535 one // dispatch dimension may hold, so the grid folds into rows. let groups_x = groups.min( self.ctx .device .limits() .max_compute_workgroups_per_dimension, ); let groups_y = groups.div_ceil(groups_x); let hot_params = hot_pixel_params( raw, (width, height), words, groups_x * HOT_PIXEL_GROUP, xtrans_tile, ); let hot_params_buf = self.ctx .device .create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("hot-pixel-params"), contents: bytemuck::bytes_of(&hot_params), usage: wgpu::BufferUsages::UNIFORM, }); let hot_bind_group = self .ctx .device .create_bind_group(&wgpu::BindGroupDescriptor { label: Some("hot-pixel-bg"), layout: &self.hot_pixel_layout, entries: &[ wgpu::BindGroupEntry { binding: 0, resource: raw_buf.as_entire_binding(), }, wgpu::BindGroupEntry { binding: 1, resource: hot_params_buf.as_entire_binding(), }, wgpu::BindGroupEntry { binding: 2, resource: repaired.as_entire_binding(), }, ], }); let params_buf = self .ctx .device .create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("demosaic-params"), contents: params_bytes, 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: repaired.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"), }); // Two passes in one submission. wgpu orders a storage write in one // pass before a read of the same buffer in the next, so the demosaic // sees every repair. { let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("hot-pixel-pass"), timestamp_writes: None, }); pass.set_pipeline(&self.hot_pixel_pipeline); pass.set_bind_group(0, &hot_bind_group, &[]); pass.dispatch_workgroups(groups_x, groups_y, 1); } { let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("demosaic-pass"), timestamp_writes: None, }); pass.set_pipeline(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), profile_tables: raw.profile_tables.clone(), as_shot_wb: [raw.wb_coeffs[0], raw.wb_coeffs[1], raw.wb_coeffs[2]], // Whatever the profile database had for this body (FR-DEV-3e), // resolved at decode because that is the only place the make and // model are known. // Sensor data is linear by construction — the demosaic shader // normalises against black and white levels and applies no // transfer function. non_linear: false, id: next_image_id(), frame: (width, height), window: WHOLE_FRAME, }) } } 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. fn inv_range_per_cell(raw: &RawImage) -> [f32; 4] { let white = f32::from(raw.white_level); let mut out = [0.0f32; 4]; for (slot, black) in out.iter_mut().zip(black_per_cell(raw)) { *slot = inv_range_above(white, black); } out } /// Reciprocal of the usable range above one black level. /// /// 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_above(white: f32, black: f32) -> f32 { let range = white - black; if range > 1.0 { 1.0 / range } else { 1.0 / 65535.0 } } // ---- X-Trans ---------------------------------------------------------- // // Fujifilm's 6×6 colour filter array (FR-RAW-5). Nothing in the Bayer path // generalises to it: there is no 2×2 cell, the black levels are not // positional, and the pattern's origin is a per-body fact rather than a // constant. /// TRACES: FR-RAW-5 /// The X-Trans tile, row-major from the sensor's own origin. 0=R, 1=G, 2=B. /// /// Transcribed from rawler's camera database — `color_pattern` in /// `data/cameras/fuji/x-t3.toml`, `"GGRGGBGGBGGRBRGRBGGGBGGRGGRGGBRBGBRG"` — /// and not derived by eye. It is the same tile every Fujifilm body uses; what /// differs between them is only where it starts, which is what /// [`detect_xtrans_phase`] is for. /// /// The structure worth knowing when reading the shader: twenty green, eight /// red, eight blue, with a red *and* a blue in every row and every column. /// That last property is the whole point of the design — no line of the /// sensor is blind to a colour, so there is no orientation along which the /// pattern aliases the way Bayer does. const XTRANS_TILE: [[u8; 6]; 6] = [ [1, 1, 0, 1, 1, 2], [1, 1, 2, 1, 1, 0], [2, 0, 1, 0, 2, 1], [1, 1, 2, 1, 1, 0], [1, 1, 0, 1, 1, 2], [0, 2, 1, 2, 0, 1], ]; /// Colour of the photosite at absolute sensor coordinates, for a given phase. fn xtrans_colour_at(phase: (u32, u32), x: u32, y: u32) -> u8 { XTRANS_TILE[((y + phase.1) % 6) as usize][((x + phase.0) % 6) as usize] } /// Pack the phase-rotated tile into the four words the shader indexes. /// /// Word `k` holds row `2k` in its low twelve bits and row `2k+1` in the next /// twelve, two bits per photosite; the fourth word is padding that keeps the /// uniform block's 16-byte alignment. Rotating on the CPU means the shader /// never has to know that a phase exists. fn pack_xtrans_tile(phase: (u32, u32)) -> [u32; 4] { let mut out = [0u32; 4]; for row in 0..6u32 { for col in 0..6u32 { let colour = u32::from(xtrans_colour_at(phase, col, row)); out[(row >> 1) as usize] |= colour << ((row & 1) * 12 + col * 2); } } out } /// As-shot white balance gains, green-normalised and bounded. /// /// Bounded because these reach a divisor in the shader: `dr-decode` already /// turns rawler's `NaN` fourth coefficient into 1.0, but a coefficient of /// 0.001 from a mis-parsed tag would survive that and turn one channel into /// a thousandfold amplifier. fn wb_gains(raw: &RawImage) -> [f32; 3] { let bounded = |v: f32| { if v.is_finite() { v.clamp(0.05, 20.0) } else { 1.0 } }; [ bounded(raw.wb_coeffs[0]), bounded(raw.wb_coeffs[1]), bounded(raw.wb_coeffs[2]), ] } /// The single black level and range the X-Trans pass normalises against. /// /// rawler reports Fujifilm black levels over the whole 6×6 tile, of which /// `dr-decode`'s four-element field keeps the first four. On every body /// examined those four are identical, so their mean *is* the level rather /// than an estimate of it — and averaging is what puts a body that does /// report distinct values in the middle rather than on whichever corner /// happened to be first. fn xtrans_levels(raw: &RawImage) -> (f32, f32) { let black = black_per_cell(raw).iter().sum::() / 4.0; (black, inv_range_above(f32::from(raw.white_level), black)) } /// TRACES: FR-RAW-5 /// Recover the 6×6 phase of the tile from the sensor data itself. /// /// **Why this is guesswork rather than a lookup.** rawler knows each body's /// pattern exactly — it is a 36-character string in the camera database — but /// `CfaPattern::XTrans` is a bare enum variant, so the phase is discarded /// before `dr-gpu` ever sees the file. It is not a constant that could simply /// be hard-coded: of the Fujifilm bodies in that database, the tile starts at /// four different origins, and choosing the wrong one mislabels every /// photosite on the sensor. Widening `dr-decode`'s type to carry the string /// is the real fix; until then the phase is read back out of the pixels. /// /// **How.** Average the sensor over each of the 36 positions in the tile. /// Photosites sharing a filter share a mean, so the correct phase is the one /// whose grouping of those 36 numbers into 20 green, 8 red and 8 blue has the /// least spread within each group. That criterion needs nothing from the /// scene and is decisive — except for one thing it cannot possibly see: /// translating the tile by three columns turns it into itself with red and /// blue exchanged, so the two labellings fit the data equally well. Red and /// blue are told apart by the as-shot white balance, on the argument that the /// camera's own gains should bring the three channel means towards each /// other, and only bring them together for the right assignment. /// /// That last step is an assumption about the scene, and a frame that is /// almost entirely one colour can defeat it. The failure is a red/blue swap, /// which is loud and obviously wrong rather than subtly wrong — and it goes /// away entirely once the pattern is plumbed through from the decoder. fn detect_xtrans_phase(raw: &RawImage) -> (u32, u32) { let mut sums = [[0.0f64; 6]; 6]; let mut counts = [[0.0f64; 6]; 6]; // Only the cropped area. The masked border a sensor readout carries is at // the black level in every position, and averaging it in flattens the very // differences this reads. let stride = raw.width as usize; let x0 = raw.crop.x as usize; let y0 = raw.crop.y as usize; let x1 = (x0 + raw.crop.width as usize).min(stride); let y1 = (y0 + raw.crop.height as usize).min(raw.height as usize); // A couple of hundred rows already give tens of thousands of samples per // tile position, and this runs over a 24 MP buffer. The step is kept // coprime with six so that skipping rows still visits all six rows of the // tile — a step of six would sample one row of it and nothing else. let mut step = y1.saturating_sub(y0) / 256; while step < 1 || step.is_multiple_of(2) || step.is_multiple_of(3) { step += 1; } for y in (y0..y1).step_by(step) { let row = y * stride; for x in x0..x1 { let Some(&value) = raw.data.get(row + x) else { continue; }; sums[y % 6][x % 6] += f64::from(value); counts[y % 6][x % 6] += 1.0; } } let mut means = [[0.0f64; 6]; 6]; let mut total = 0.0; for ((mean_row, sum_row), count_row) in means.iter_mut().zip(&sums).zip(&counts) { for ((mean, &sum), &count) in mean_row.iter_mut().zip(sum_row).zip(count_row) { *mean = sum / count.max(1.0); total += *mean; } } // Normalised so the tie tolerance below means the same thing at every // exposure. let scale = if total > 0.0 { 36.0 / total } else { 1.0 }; let wb = wb_gains(raw); // Two candidates always fit exactly as well as each other — the tile maps // onto itself under a half-tile shift — so the comparison has to admit a // tie rather than trust the last bit of a float sum. const TIE: f64 = 1e-6; let mut best_spread = f64::INFINITY; let mut best_imbalance = f64::INFINITY; let mut best = (0u32, 0u32); for py in 0..6u32 { for px in 0..6u32 { let mut group_n = [0.0f64; 3]; let mut group_sum = [0.0f64; 3]; let mut group_sq = [0.0f64; 3]; for j in 0..6u32 { for i in 0..6u32 { let k = usize::from(xtrans_colour_at((px, py), i, j)); let m = means[j as usize][i as usize] * scale; group_n[k] += 1.0; group_sum[k] += m; group_sq[k] += m * m; } } let mut spread = 0.0; let mut balanced = [0.0f64; 3]; for k in 0..3 { let n = group_n[k].max(1.0); let mean = group_sum[k] / n; spread += group_sq[k] - n * mean * mean; balanced[k] = f64::from(wb[k]) * mean; } let neutral = (0..3).map(|k| group_n[k] * balanced[k]).sum::() / 36.0; let imbalance = (0..3) .map(|k| group_n[k] * (balanced[k] - neutral).powi(2)) .sum::(); let decisive = spread < best_spread - TIE; let tied_but_more_neutral = spread < best_spread + TIE && imbalance < best_imbalance; if decisive || tied_but_more_neutral { best_spread = best_spread.min(spread); best_imbalance = imbalance; best = (px, py); } } } log::debug!("X-Trans phase detected as {best:?} (spread {best_spread:.4})"); best } /// TRACES: FR-RAW-5 /// Everything the X-Trans shader needs about one image. /// TRACES: FR-RAW-3 /// The hot-pixel repair's pipeline and its three bindings: the readout, the /// uniform block, and the repaired copy it writes. fn hot_pixel_pipeline(ctx: &GpuContext) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) { let shader = ctx .device .create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("hot-pixels"), source: wgpu::ShaderSource::Wgsl(include_str!("shaders/hot_pixels.wgsl").into()), }); let storage = |binding, read_only| wgpu::BindGroupLayoutEntry { binding, visibility: wgpu::ShaderStages::COMPUTE, ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only }, has_dynamic_offset: false, min_binding_size: None, }, count: None, }; let layout = ctx .device .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("hot-pixel-bgl"), entries: &[ storage(0, true), 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, }, storage(2, false), ], }); let pipeline_layout = ctx .device .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("hot-pixel-layout"), bind_group_layouts: &[Some(&layout)], immediate_size: 0, }); let pipeline = ctx .device .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("hot-pixel-pipeline"), layout: Some(&pipeline_layout), module: &shader, entry_point: Some("main"), compilation_options: Default::default(), cache: None, }); (pipeline, layout) } /// The colour of each position of a Bayer cell, row-major, for the pattern /// the decoder reported: 0=R, 1=G, 2=B. `None` for anything that is not a /// 2×2 pattern. fn bayer_cell(pattern: CfaPattern) -> Option<[u32; 4]> { match pattern { CfaPattern::Rggb => Some([0, 1, 1, 2]), CfaPattern::Bggr => Some([2, 1, 1, 0]), CfaPattern::Grbg => Some([1, 0, 2, 1]), CfaPattern::Gbrg => Some([1, 2, 0, 1]), _ => None, } } /// A Bayer cell as the 6×6 sensor-anchored tile the repair indexes. /// /// The decoder's pattern is phased for the *crop* origin, and the tile is /// indexed by sensor coordinate, so each position is shifted by the crop. /// Six is even, so a column's parity modulo 6 is its parity outright and the /// cell repeats cleanly. fn pack_bayer_tile(cell: [u32; 4], crop_x: u32, crop_y: u32) -> [u32; 4] { let mut out = [0u32; 4]; for row in 0..6u32 { for col in 0..6u32 { let i = (((row + crop_y) & 1) * 2 + ((col + crop_x) & 1)) as usize; out[(row >> 1) as usize] |= cell[i] << ((row & 1) * 12 + col * 2); } } out } /// The repair's uniforms for one readout. /// /// `xtrans_tile` is the tile the X-Trans demosaic was given, when it was one; /// anything else must be a Bayer pattern, which `run` has already checked. fn hot_pixel_params( raw: &RawImage, (width, height): (u32, u32), words: u32, row_invocations: u32, xtrans_tile: Option<[u32; 4]>, ) -> HotPixelParams { let (black, inv_range, tile) = match (xtrans_tile, bayer_cell(raw.cfa_pattern)) { (Some(tile), _) => { let (black, inv_range) = xtrans_levels(raw); ([black; 4], [inv_range; 4], tile) } (None, Some(cell)) => ( black_per_cell(raw), inv_range_per_cell(raw), pack_bayer_tile(cell, raw.crop.x, raw.crop.y), ), // Not reached from `run`, which refuses any other pattern before // this. A zero tile judges every photosite against all of its // neighbours, which is right for a sensor with no colour filter. (None, None) => (black_per_cell(raw), inv_range_per_cell(raw), [0; 4]), }; HotPixelParams { crop_x: raw.crop.x, crop_y: raw.crop.y, width, height, stride: raw.width, words, row_invocations, samples: raw.data.len() as u32, black, inv_range, tile, } } fn xtrans_params_for(raw: &RawImage, width: u32, height: u32) -> XTransParams { let (black, inv_range) = xtrans_levels(raw); let wb = wb_gains(raw); XTransParams { width, height, crop_x: raw.crop.x, crop_y: raw.crop.y, stride: raw.width, black, inv_range, _pad0: 0, wb: [wb[0], wb[1], wb[2], 1.0], inv_wb: [1.0 / wb[0], 1.0 / wb[1], 1.0 / wb[2], 1.0], tile: pack_xtrans_tile(detect_xtrans_phase(raw)), } } #[cfg(test)] mod tests { use super::*; use dr_decode::CropRect; fn f16_to_f32(bits: u16) -> f32 { let sign = if bits & 0x8000 != 0 { -1.0 } else { 1.0 }; let e = ((bits >> 10) & 0x1F) as i32; let m = (bits & 0x3FF) as f32; if e == 0 { sign * m * 2f32.powi(-24) } else { sign * (1.0 + m / 1024.0) * 2f32.powi(e - 15) } } /// The repair's tile is indexed by sensor coordinate, the decoder's /// pattern by crop coordinate. A crop at an odd origin must shift one /// into the other, or the repair compares red with green. #[test] fn the_bayer_tile_is_anchored_to_the_sensor_not_the_crop() { let cell = bayer_cell(CfaPattern::Rggb).unwrap(); let colour = |tile: [u32; 4], x: u32, y: u32| { (tile[((y % 6) >> 1) as usize] >> (((y % 6) & 1) * 12 + (x % 6) * 2)) & 3 }; for (cx, cy) in [(0, 0), (1, 0), (0, 1), (1, 1), (7, 4)] { let tile = pack_bayer_tile(cell, cx, cy); // Red is the crop's first photosite, wherever the crop starts. assert_eq!(colour(tile, cx, cy), 0, "crop at ({cx}, {cy})"); assert_eq!(colour(tile, cx + 1, cy + 1), 2, "crop at ({cx}, {cy})"); assert_eq!(colour(tile, cx + 1, cy), 1, "crop at ({cx}, {cy})"); } } #[test] fn unclamped_half_keeps_shadows_signs_and_highlights() { // A 14-bit LSB, normalised: subnormal in f16, and must not be zero. let lsb = 1.0 / 16383.0; let back = f16_to_f32(f32_to_f16_bits_unclamped(lsb)); assert!((back - lsb).abs() / lsb < 0.01, "{back} vs {lsb}"); // A quarter of that, still representable. let tiny = lsb / 4.0; let back = f16_to_f32(f32_to_f16_bits_unclamped(tiny)); assert!((back - tiny).abs() / tiny < 0.05, "{back} vs {tiny}"); // Below zero and above one survive. assert!((f16_to_f32(f32_to_f16_bits_unclamped(-0.01)) + 0.01).abs() < 1e-5); assert!((f16_to_f32(f32_to_f16_bits_unclamped(1.75)) - 1.75).abs() < 1e-3); // Exact values are exact. assert_eq!(f32_to_f16_bits_unclamped(1.0), 0x3C00); assert_eq!(f32_to_f16_bits_unclamped(0.5), 0x3800); assert_eq!(f32_to_f16_bits_unclamped(0.0), 0); // Within one ULP of the clamped one on its domain: that one // truncates the mantissa, this one rounds it. for i in 0..=255 { let v = i as f32 / 255.0; let (a, b) = (f32_to_f16_bits_unclamped(v), f32_to_f16_bits(v)); assert!(a.abs_diff(b) <= 1, "{v}: {a} vs {b}"); } } 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, samples_per_pixel: 1, profile: None, profile_tables: None, make: String::new(), model: String::new(), 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, samples_per_pixel: 1, profile: None, profile_tables: None, make: String::new(), model: String::new(), 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::TexelCopyTextureInfo { texture: img.texture(), mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All, }, 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(..); let (tx, rx) = std::sync::mpsc::channel(); slice.map_async(wgpu::MapMode::Read, move |r| { let _ = tx.send(r); }); ctx.device .poll(wgpu::PollType::wait_indefinitely()) .expect("poll"); 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 a_grey_step_edge_stays_grey() { // A flat patch cannot tell the Malvar kernels from any other set of // weights that sum to zero. An edge can. A grey vertical step, so // every photosite records the same profile, must come back with the // three channels close together on both sides; any spread is false // colour from interpolating across the edge. // // The bound is set by the paper's kernels, which peak at 0.19 here. // With the ±2 terms of the green-site kernels transposed — the bug // this test was written against — the peak is 0.375. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let size = 32u32; let white = 16383u16; let mut raw = flat_cfa(CfaPattern::Rggb, size, [0, 0, 0], 0, white); for y in 0..size { for x in size / 2..size { raw.data[(y * size + x) as usize] = white; } } let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, _) = img.size(); let mut worst = (0.0f32, 0u32, 0u32); for y in 2..size - 2 { for x in 2..size - 2 { let p = px[(y * w + x) as usize]; let spread = (p[0] - p[1]).abs().max((p[2] - p[1]).abs()); if spread > worst.0 { worst = (spread, x, y); } } } assert!( worst.0 < 0.25, "false colour of {} at ({}, {}) on a grey edge — the green-site \ kernels are interpolating across the edge", worst.0, worst.1, worst.2 ); } #[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, samples_per_pixel: 1, profile: None, profile_tables: None, make: String::new(), model: String::new(), 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 an_unidentifiable_cfa_is_refused_rather_than_rendered_wrong() { // Both demosaics need to know which filter each photosite carries. A // file whose pattern dr-decode could not name has no answer to that, // and guessing produces a maze of colour that reads as a corrupt // image rather than as an unsupported one. 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::Unknown; assert!( matches!(d.run(&raw), Err(GpuError::UnsupportedCfa(_))), "an unnamed pattern must report the gap, not render artefacts" ); } // ---- X-Trans ------------------------------------------------------- /// The colours the tile assigns across one period, for comparing two /// phases. /// /// Phases are compared this way rather than as coordinate pairs because /// the tile maps onto itself under a half-tile shift: `(0, 0)` and /// `(3, 3)` are different numbers describing the same sensor. fn labelling(phase: (u32, u32)) -> Vec { (0..6) .flat_map(|y| (0..6).map(move |x| xtrans_colour_at(phase, x, y))) .collect() } /// Build a synthetic X-Trans mosaic of a uniform colour at a known phase. fn flat_xtrans( phase: (u32, u32), 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 = usize::from(xtrans_colour_at(phase, x, y)); data[(y * size + x) as usize] = black + rgb[c]; } } RawImage { width: size, height: size, data, cfa_pattern: CfaPattern::XTrans, black_level: [black; 4], white_level: white, // The gains a camera looking at this colour would have recorded. // Not decoration: the phase detector needs them to tell red from // blue, and every real file carries them. wb_coeffs: [ f32::from(rgb[1]) / f32::from(rgb[0]), 1.0, f32::from(rgb[1]) / f32::from(rgb[2]), 1.0, ], color_matrix: None, samples_per_pixel: 1, profile: None, profile_tables: None, make: String::new(), model: String::new(), crop: CropRect { x: 0, y: 0, width: size, height: size, }, } } #[test] fn the_xtrans_tile_is_blind_to_no_colour_on_any_line() { // The property that makes X-Trans what it is, and the cheapest check // that the 36 transcribed entries are the right 36: every row and // every column carries a red and a blue, which is why the pattern // does not alias along an axis the way Bayer does. One mistyped entry // breaks a line here, and would otherwise mislabel one photosite in // thirty-six across the whole sensor. let mut counts = [0usize; 3]; for row in XTRANS_TILE { for colour in row { counts[usize::from(colour)] += 1; } } assert_eq!(counts, [8, 20, 8], "8 red, 20 green, 8 blue"); for (i, row) in XTRANS_TILE.iter().enumerate() { let column: Vec = XTRANS_TILE.iter().map(|r| r[i]).collect(); for (line, what) in [(row.to_vec(), "row"), (column, "column")] { assert!( line.contains(&0) && line.contains(&2), "{what} {i} carries no red or no blue: {line:?}" ); } } } #[test] fn packing_the_tile_survives_every_phase() { // The shader reads the tile back out of four packed words. If the // packing and the unpacking disagree by so much as one shift, every // photosite is assigned somebody else's filter — and the output is // still a plausible-looking image, just the wrong colour. for py in 0..6u32 { for px in 0..6u32 { let packed = pack_xtrans_tile((px, py)); for y in 0..6u32 { for x in 0..6u32 { let word = packed[(y >> 1) as usize]; let unpacked = (word >> ((y & 1) * 12 + x * 2)) & 3; assert_eq!( unpacked as u8, xtrans_colour_at((px, py), x, y), "phase ({px},{py}) at ({x},{y})" ); } } } } } #[test] fn the_xtrans_phase_is_recovered_from_the_sensor_data() { // dr-decode reports "X-Trans" and not where the tile starts, and the // Fujifilm bodies do not agree on that — rawler's database gives four // different origins. Reading the phase back out of the pixels is the // only thing standing between a Fuji file and every photosite being // assigned the wrong filter. for py in 0..6u32 { for px in 0..6u32 { let raw = flat_xtrans((px, py), 36, [9000, 12000, 4000], 0, 16383); let found = detect_xtrans_phase(&raw); assert_eq!( labelling(found), labelling((px, py)), "phase ({px},{py}) came back as {found:?}" ); } } } #[test] fn white_balance_is_what_tells_the_red_sites_from_the_blue() { // Shifting the tile by half a tile turns it into itself with red and // blue exchanged, so nothing about the geometry can choose between the // two — only the camera's own gains can. This pins that down from both // sides: with the gains the answer is exact, and without them green is // still placed correctly while red and blue become a coin toss. If the // second half ever starts insisting on the right answer too, the // tie-break has been replaced by something that only looks like it // works. let mut raw = flat_xtrans((0, 0), 36, [9000, 12000, 4000], 0, 16383); assert_eq!(labelling(detect_xtrans_phase(&raw)), labelling((0, 0))); raw.wb_coeffs = [1.0, 1.0, 1.0, 1.0]; let blind = detect_xtrans_phase(&raw); for y in 0..6u32 { for x in 0..6u32 { assert_eq!( xtrans_colour_at(blind, x, y) == 1, xtrans_colour_at((0, 0), x, y) == 1, "green at ({x},{y}) must not depend on the white balance" ); } } } #[test] fn a_flat_xtrans_patch_demosaics_to_that_colour() { // The same property the Bayer path is held to, and the one that // catches a wrong tile, a wrong phase, or a colour-difference model // that fails to cancel: a sensor looking at a uniform colour must // reconstruct that colour, at every phase and right to the border. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let (black, white) = (1024u16, 16383u16); let range = f32::from(white - black); let rgb = [8000u16, 12000, 3000]; let expected = [ f32::from(rgb[0]) / range, f32::from(rgb[1]) / range, f32::from(rgb[2]) / range, ]; for phase in [(0, 0), (1, 0), (0, 1), (3, 0), (2, 5), (4, 3)] { let raw = flat_xtrans(phase, 36, rgb, black, white); let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, h) = img.size(); for y in 0..h { for x in 0..w { let p = px[(y * w + x) as usize]; for (ch, &want) in expected.iter().enumerate() { assert!( (p[ch] - want).abs() < 0.01, "phase {phase:?} pixel ({x},{y}) channel {ch}: \ got {}, want {want}", p[ch] ); } } } } } #[test] fn an_xtrans_gradient_carries_no_colour_cast() { // Why the shader fits a plane rather than averaging each channel's // neighbours. The three channels are sampled at different places in // the tile, so a mean compares a red taken slightly to one side of the // pixel with a green taken slightly to the other. On a flat patch that // cancels and the test above passes anyway; on a gradient it is a // colour cast that follows the gradient across the whole frame. A // plane has no such offset and reconstructs a ramp exactly. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let (black, white) = (1024u16, 16383u16); let range = f32::from(white - black); let size = 36u32; // A different slope per channel and per axis, so a cast in any // direction shows. let level = |c: usize, x: u32, y: u32| -> u16 { match c { 0 => 3000 + 20 * x as u16 + 8 * y as u16, 1 => 6000 + 12 * x as u16 + 24 * y as u16, _ => 2000 + 30 * x as u16 + 6 * y as u16, } }; let phase = (2, 1); let mut data = vec![0u16; (size * size) as usize]; for y in 0..size { for x in 0..size { let c = usize::from(xtrans_colour_at(phase, x, y)); data[(y * size + x) as usize] = black + level(c, x, y); } } let mut raw = flat_xtrans(phase, size, [3000, 6000, 2000], black, white); raw.data = data; let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, _) = img.size(); // A two-pixel margin: at the border the window slides inward, and the // clamp to the local sample range can bite where the pixel sits at the // edge of its own window. for y in 2..size - 2 { for x in 2..size - 2 { let p = px[(y * w + x) as usize]; for (c, &got) in p.iter().take(3).enumerate() { let want = f32::from(level(c, x, y)) / range; assert!( (got - want).abs() < 0.005, "pixel ({x},{y}) channel {c}: got {got}, want {want} — \ a ramp must survive the interpolation unbent" ); } } } } #[test] fn an_xtrans_crop_origin_does_not_move_the_tile() { // The 6×6 tile is anchored to the sensor readout, not to the visible // frame — which is why CfaPattern::shifted deliberately leaves X-Trans // alone and says the demosaic handles the offset itself. This is that // handling. Forget to add the crop origin in either the detector or // the shader and an active area starting anywhere but a multiple of // six re-colours the entire image. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let (black, white) = (0u16, 16383u16); let rgb = [9000u16, 13000, 4000]; let expected = [ f32::from(rgb[0]) / f32::from(white), f32::from(rgb[1]) / f32::from(white), f32::from(rgb[2]) / f32::from(white), ]; let mut raw = flat_xtrans((0, 0), 48, rgb, black, white); // Neither offset is a multiple of six, so the visible top-left is a // different filter than the sensor's own origin. raw.crop = CropRect { x: 5, y: 7, width: 34, height: 34, }; let img = d.run(&raw).expect("demosaic"); let px = read_rgba(&ctx, &img); let (w, h) = img.size(); for y in 0..h { for x in 0..w { 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 xtrans_output_is_free_of_nan_and_negatives() { // The same hazard as the Bayer path, reached by a different route: an // f16 NaN propagates silently through every later stage, and a // negative value breaks the ratio-based operations downstream. Here // the risk is the plane solve — a near-singular window divides by a // determinant close to zero — so the input is the noisiest thing a // sensor can produce. let Some(ctx) = ctx() else { return }; let d = Demosaicer::new(&ctx).expect("demosaicer"); let size = 36u32; 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 mut raw = flat_xtrans((0, 0), size, [100, 100, 100], 32, 16383); raw.data = data; 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 an_xtrans_black_level_is_a_single_sensor_wide_value() { // The four-value black level is a 2×2 convention. Indexing it by // position on a 6×6 tile would lift or crush a fifth of the // photosites, so the X-Trans path averages instead — which for the // broadcast case every Fujifilm body produces is exactly the level. let mut raw = raw_for([1024, 0, 0, 0], 16383); raw.cfa_pattern = CfaPattern::XTrans; let (black, inv_range) = xtrans_levels(&raw); assert_eq!(black, 1024.0); assert!((f32::from(16383u16 - 1024) * inv_range - 1.0).abs() < 1e-5); } }