diff --git a/core/dr-gpu/src/demosaic.rs b/core/dr-gpu/src/demosaic.rs index 6e84077..3b79fb6 100644 --- a/core/dr-gpu/src/demosaic.rs +++ b/core/dr-gpu/src/demosaic.rs @@ -1,4 +1,4 @@ -//! Raw upload, black/white normalisation, and Bayer demosaic. +//! 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 @@ -31,6 +31,32 @@ struct DemosaicParams { 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], +} + /// A demosaiced image living on the GPU. /// /// RGBA16Float, scene-referred, camera colour space. This is the input every @@ -222,10 +248,17 @@ fn f32_to_f16_bits(v: f32) -> u16 { ((exp as u16) << 10) | mantissa as u16 } -/// Runs the demosaic pass. Holds the pipeline so repeated images reuse it. +/// 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, } @@ -237,7 +270,17 @@ impl Demosaicer { 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 { @@ -296,32 +339,31 @@ impl Demosaicer { 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, + }); + Ok(Self { ctx: ctx.clone(), pipeline, + xtrans_pipeline, bind_group_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.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(); @@ -332,6 +374,38 @@ impl Demosaicer { ))); } + // 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; + let (pipeline, params_bytes) = if raw.cfa_pattern.is_xtrans() { + xtrans_params = xtrans_params_for(raw, width, height); + (&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); @@ -344,24 +418,12 @@ impl Demosaicer { 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), + contents: params_bytes, usage: wgpu::BufferUsages::UNIFORM, }); @@ -418,7 +480,7 @@ impl Demosaicer { label: Some("demosaic-pass"), timestamp_writes: None, }); - pass.set_pipeline(&self.pipeline); + pass.set_pipeline(pipeline); pass.set_bind_group(0, &bind_group, &[]); pass.dispatch_workgroups(width.div_ceil(8), height.div_ceil(8), 1); } @@ -478,25 +540,259 @@ fn black_per_cell(raw: &RawImage) -> [f32; 4] { } /// 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_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 - }; +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. +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::*; @@ -886,15 +1182,359 @@ mod tests { } #[test] - fn xtrans_is_refused_rather_than_rendered_wrong() { + 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::XTrans; + raw.cfa_pattern = CfaPattern::Unknown; assert!( matches!(d.run(&raw), Err(GpuError::UnsupportedCfa(_))), - "an X-Trans file must report the gap, not render artefacts" + "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, + 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); + } } diff --git a/core/dr-gpu/src/error.rs b/core/dr-gpu/src/error.rs index b3cc8ed..5fdd8b7 100644 --- a/core/dr-gpu/src/error.rs +++ b/core/dr-gpu/src/error.rs @@ -22,9 +22,10 @@ pub enum GpuError { #[error("readback failed: {0}")] Readback(String), - /// The X-Trans demosaic is not implemented (FR-RAW-5). Reported rather - /// than approximated with the Bayer path, which would produce a maze of - /// colour artefacts and look like a corrupt file. + /// The CFA layout is one no demosaic here handles — in practice a file + /// `dr-decode` could not identify the pattern of. Reported rather than + /// approximated with the Bayer path, which would produce a maze of colour + /// artefacts and look like a corrupt file. #[error("unsupported CFA pattern: {0}")] UnsupportedCfa(String), diff --git a/core/dr-gpu/src/shaders/xtrans.wgsl b/core/dr-gpu/src/shaders/xtrans.wgsl new file mode 100644 index 0000000..b00acba --- /dev/null +++ b/core/dr-gpu/src/shaders/xtrans.wgsl @@ -0,0 +1,249 @@ +// Black/white normalisation and X-Trans demosaic, in one pass. +// +// The Fujifilm counterpart to demosaic.wgsl (FR-RAW-5). The input and the +// output are the same — packed u16 photosites in, linear camera-space +// RGBA16Float out — but the colour filter array is a 6x6 tile rather than a +// 2x2 one, and none of the Bayer kernels survive that. In a Bayer cell every +// pixel has the missing channels at a fixed offset; in X-Trans the offsets +// differ at all 36 positions, so a fixed kernel per site would need 36 of +// them and would still say nothing about which neighbours to trust. +// +// The method here is *local plane fitting on the colour difference*: +// +// 1. Take a 5x5 window and sort its photosites by the channel each one +// measures. Every window holds at least four red, four blue and thirteen +// green samples, whatever the phase — checked exhaustively, not assumed. +// 2. Fit a weighted least-squares plane through each channel's samples and +// evaluate all three planes at the pixel. A plane rather than a mean +// because the three channels are sampled at *different* places: a mean +// would compare a red average taken slightly left of the pixel with a +// green average taken slightly right of it, and the difference of those +// two offsets is a colour cast that follows every gradient in the frame. +// A plane has no such bias — it reconstructs any linear gradient exactly. +// 3. Keep the pixel's own measured value, and carry the other two channels +// across as the *difference* between the fitted planes. +// +// Step 3 is the standard constant-colour-difference model, and the reason it +// is applied in white-balanced space is that the model is exact only where +// the channel difference is locally constant. On a neutral subject that is +// true after white balance and false before it, so the gains go on before the +// fit and come off after — which costs one multiply and halves the error at a +// luminance edge (measured on a synthetic step: 0.34 -> 0.20 max error). +// The pixels written out are still as-shot, unbalanced camera space; nothing +// downstream sees the difference. +// +// **What this is not.** It is not Markesteijn. It has no directional +// hypotheses and no homogeneity map, so it does not resolve detail finer than +// the CFA period, and a hard edge arrives about two pixels wide. It does not +// produce the "worms" that FR-RAW-5 exists to avoid — the output is bounded +// by the local sample range, so it cannot ring — but the Markesteijn-class +// quality that requirement asks for is still owed. + +struct XTransParams { + // Dimensions of the *cropped* output, in pixels. + width: u32, + height: u32, + // Origin of the crop within the sensor readout, in photosites. Added to + // every read, and to every pattern lookup: the 6x6 tile is anchored to the + // sensor, not to the visible frame. + crop_x: u32, + crop_y: u32, + // Row stride of the input, in samples. + stride: u32, + // One black level and one reciprocal range for the whole sensor. The + // four-value form the Bayer path uses is a 2x2 convention with no meaning + // on a 6x6 tile. + black: f32, + inv_range: f32, + _pad0: u32, + // As-shot white balance gains, green-normalised, and their reciprocals. + wb: vec4, + inv_wb: vec4, + // The 6x6 tile, already rotated to this sensor's phase on the CPU, packed + // two bits per photosite: word k holds row 2k in its low 12 bits and row + // 2k+1 in the next 12. The fourth word is padding. + tile: vec4, +} + +@group(0) @binding(0) var raw: array; +@group(0) @binding(1) var params: XTransParams; +@group(0) @binding(2) var output: texture_storage_2d; + +// Half-width of the fitting window. Two is the smallest radius for which +// every phase of the tile still offers enough red and blue samples to pin a +// plane down; three would be smoother and blurrier. +const RADIUS: i32 = 2; + +// Colour of the photosite at absolute sensor coordinates: 0=R, 1=G, 2=B. +fn colour_at(sx: u32, sy: u32) -> u32 { + let row = sy % 6u; + let col = sx % 6u; + let word = params.tile[row >> 1u]; + return (word >> ((row & 1u) * 12u + col * 2u)) & 3u; +} + +// Read one photosite, normalised to [0, 1] against the black level. +// +// Coordinates are relative to the crop origin and must already be in range; +// unlike the Bayer pass there is no reflection here, because the window is +// slid inside the image instead (see `main`) and so never asks for a +// photosite that does not exist. +fn sample(cx: i32, cy: i32) -> f32 { + let sx = u32(cx) + params.crop_x; + let sy = u32(cy) + params.crop_y; + let index = sy * params.stride + sx; + + let word = raw[index >> 1u]; + let raw_value = select(word & 0xFFFFu, word >> 16u, (index & 1u) == 1u); + + // Sensor noise puts real signal below the black point, so subtracting it + // can go negative; clamped rather than allowed to wrap. + return max((f32(raw_value) - params.black) * params.inv_range, 0.0); +} + +// Solve the 3x3 weighted least-squares normal equations for a plane +// `c0 + c1*dx + c2*dy` and evaluate it at `(ex, ey)`. +// +// The accumulators are the usual moments: `n` is the summed weight, `sx`..`syy` +// the first and second moments of the sample positions, `t0`..`ty` the same +// moments weighted by value. Cramer's rule rather than a factorisation — the +// matrix is 3x3 and symmetric, and this keeps the whole solve in registers. +fn plane_at( + n: f32, sx: f32, sy: f32, sxx: f32, sxy: f32, syy: f32, + t0: f32, tx: f32, ty: f32, ex: f32, ey: f32, +) -> f32 { + if (n <= 0.0) { + return 0.0; + } + let det = n * (sxx * syy - sxy * sxy) + - sx * (sx * syy - sxy * sy) + + sy * (sx * sxy - sxx * sy); + + // Degenerate only if a channel's samples in this window are collinear, + // which the 5x5 geometry rules out for every phase — but an image a few + // photosites across is clipped down to fewer samples than that, and a + // division by a near-zero determinant there would put NaN in the texture. + // Falling back to the plain weighted mean loses the gradient term and + // nothing else. + if (abs(det) < 1e-6 * n * n * n) { + return t0 / n; + } + + let inv = 1.0 / det; + let c0 = (t0 * (sxx * syy - sxy * sxy) + - sx * (tx * syy - sxy * ty) + + sy * (tx * sxy - sxx * ty)) * inv; + let c1 = (n * (tx * syy - sxy * ty) + - t0 * (sx * syy - sxy * sy) + + sy * (sx * ty - tx * sy)) * inv; + let c2 = (n * (sxx * ty - tx * sxy) + - sx * (sx * ty - tx * sy) + + t0 * (sx * sxy - sxx * sy)) * inv; + return c0 + c1 * ex + c2 * ey; +} + +@compute @workgroup_size(8, 8, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + if (gid.x >= params.width || gid.y >= params.height) { + return; + } + + let x = i32(gid.x); + let y = i32(gid.y); + let w = i32(params.width); + let h = i32(params.height); + + // Near an edge the window slides inward rather than reflecting. Reflection + // is right for the Bayer pass, whose kernels only need the mirrored sample + // to be the same colour; here the fit needs real geometry, and a mirrored + // photosite sitting at a position it does not occupy tilts the plane. A + // slid window is entirely real data, so the plane stays exact right up to + // the border — the pixel is still inside the window, just not centred in + // it, which is what `ex`/`ey` below account for. + let wx = clamp(x, RADIUS, max(RADIUS, w - 1 - RADIUS)); + let wy = clamp(y, RADIUS, max(RADIUS, h - 1 - RADIUS)); + + // Per-channel moments, indexed 0=R, 1=G, 2=B. + var n = vec3(0.0); + var sx = vec3(0.0); + var sy = vec3(0.0); + var sxx = vec3(0.0); + var sxy = vec3(0.0); + var syy = vec3(0.0); + var t0 = vec3(0.0); + var tx = vec3(0.0); + var ty = vec3(0.0); + var lo = vec3(1.0e30); + var hi = vec3(-1.0e30); + + for (var dy = -RADIUS; dy <= RADIUS; dy = dy + 1) { + for (var dx = -RADIUS; dx <= RADIUS; dx = dx + 1) { + // The clamp only bites on an image narrower than the window, where + // a duplicated photosite is better than a missing channel. + let px = clamp(wx + dx, 0, w - 1); + let py = clamp(wy + dy, 0, h - 1); + + let v = sample(px, py); + let k = colour_at(u32(px) + params.crop_x, u32(py) + params.crop_y); + let u = v * params.wb[k]; + + let fx = f32(px - wx); + let fy = f32(py - wy); + // Nearer photosites describe this pixel better. 1/(1+r^2) rather + // than a Gaussian because it needs no width to tune and leaves the + // normal equations well conditioned at every phase. + let g = 1.0 / (1.0 + fx * fx + fy * fy); + + n[k] = n[k] + g; + sx[k] = sx[k] + g * fx; + sy[k] = sy[k] + g * fy; + sxx[k] = sxx[k] + g * fx * fx; + sxy[k] = sxy[k] + g * fx * fy; + syy[k] = syy[k] + g * fy * fy; + t0[k] = t0[k] + g * u; + tx[k] = tx[k] + g * u * fx; + ty[k] = ty[k] + g * u * fy; + + lo[k] = min(lo[k], v); + hi[k] = max(hi[k], v); + } + } + + let ex = f32(x - wx); + let ey = f32(y - wy); + var p = vec3( + plane_at(n.r, sx.r, sy.r, sxx.r, sxy.r, syy.r, t0.r, tx.r, ty.r, ex, ey), + plane_at(n.g, sx.g, sy.g, sxx.g, sxy.g, syy.g, t0.g, tx.g, ty.g, ex, ey), + plane_at(n.b, sx.b, sy.b, sxx.b, sxy.b, syy.b, t0.b, tx.b, ty.b, ex, ey), + ); + + // The pixel's own channel always has a sample — itself — so its plane is + // always real, and it is the reference the other two are carried across + // from. + let centre = colour_at(u32(x) + params.crop_x, u32(y) + params.crop_y); + let measured = sample(x, y); + + // A channel with no sample at all cannot happen in a 5x5 window; it can on + // an image a few photosites across, where the window collapses. Such a + // channel is given the reference plane, which renders the pixel grey + // rather than arbitrary. + let empty = n <= vec3(0.0); + p = select(p, vec3(p[centre]), empty); + lo = select(lo, vec3(0.0), empty); + hi = select(hi, vec3(1.0e30), empty); + + // Constant colour difference, in white-balanced space, undone on the way + // out. The measured channel comes back bit-for-bit: its own plane cancels. + var rgb = (vec3(measured * params.wb[centre]) + p - vec3(p[centre])) + * params.inv_wb.rgb; + + // Bound each channel by what was actually measured nearby. The plane + // difference overshoots wherever the colour itself changes across the + // window — a red edge against green — and an overshoot here is a coloured + // halo. The pixel is always inside its own window, so a true value can + // never be clipped away by this on smooth content. + rgb = clamp(rgb, lo, hi); + rgb = max(rgb, vec3(0.0)); + + textureStore(output, vec2(x, y), vec4(rgb, 1.0)); +}