dtourolle/DarkRoom
1
0
Fork 0
Code Issues 39 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Merge branch 'xtrans-demosaic'

Browse Source
This commit is contained in:
dtourolle
2026-08-17 09:25:28 +02:00
parent 8ea92545df 1c0994c807
commit 9f577ed6a2
3 changed files with 940 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-gpu/src/demosaic.rs
+687 -47
View File
@@ -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<DemosaicedImage, GpuError> {
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(&params),
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::<f32>() / 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::<f64>() / 36.0;
let imbalance = (0..3)
.map(|k| group_n[k] * (balanced[k] - neutral).powi(2))
.sum::<f64>();
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<u8> {
(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<u8> = 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);
}
}
core/dr-gpu/src/error.rs
+4 -3
View File
@@ -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),
core/dr-gpu/src/shaders/xtrans.wgsl
+249
View File
@@ -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<f32>,
inv_wb: vec4<f32>,
// 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<u32>,
}
@group(0) @binding(0) var<storage, read> raw: array<u32>;
@group(0) @binding(1) var<uniform> params: XTransParams;
@group(0) @binding(2) var output: texture_storage_2d<rgba16float, write>;
// 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<u32>) {
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<f32>(0.0);
var sx = vec3<f32>(0.0);
var sy = vec3<f32>(0.0);
var sxx = vec3<f32>(0.0);
var sxy = vec3<f32>(0.0);
var syy = vec3<f32>(0.0);
var t0 = vec3<f32>(0.0);
var tx = vec3<f32>(0.0);
var ty = vec3<f32>(0.0);
var lo = vec3<f32>(1.0e30);
var hi = vec3<f32>(-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<f32>(
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<f32>(0.0);
p = select(p, vec3<f32>(p[centre]), empty);
lo = select(lo, vec3<f32>(0.0), empty);
hi = select(hi, vec3<f32>(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<f32>(measured * params.wb[centre]) + p - vec3<f32>(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<f32>(0.0));
textureStore(output, vec2<i32>(x, y), vec4<f32>(rgb, 1.0));
}