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

Repair hot and dead photosites before the demosaic

Browse Source 
A hot photosite went into the demosaic as it was read, and came out as a
coloured cross three pixels wide that nothing later could take back
out. Night and long exposures showed them; the defect-map reader added
for FR-RAW-3 was never wired in, and a CR2 carries no map anyway.

A pass over the mosaic now runs ahead of the demosaic, into a second
buffer. A photosite is hot when it reads more than twice every
same-colour photosite in its 5x5 window plus 2% of the range, and more
than twice each of its eight immediate neighbours of any colour. The
second half keeps stars and glints: real light reaches the sensor
through a lens and an anti-aliasing filter and lights a patch, so the
photosites beside it are lit too, where a hot photosite's are dark. It
is replaced by its brightest same-colour neighbour, which invents
nothing. Dead photosites are the mirror case, judged only where the
neighbourhood is above 5%, so shadow noise clipped at black is left
alone.

The colour of each photosite comes from a 6x6 sensor-anchored tile, so
Bayer and X-Trans share the pass. Export and every other path that
demosaics get it too, and there is no setting: the repair only fires
where a single photosite disagrees with everything around it.

Cost, warm, on a Canon 6D frame (RTX 3050): 91-99 ms to demosaic
before, 94-98 ms after; the extra pass is inside the run-to-run noise.

Tests render a frame with and without the defect and compare the
finished pixels. Without the repair a hot photosite showed by 230 and a
dead one by 168; with it neither shows, and a 3x3 highlight at white
survives.
This commit is contained in:
dtourolle
2026-09-27 06:20:23 -04:00
parent 6b99f67f47
commit c50d96e949
4 changed files with 588 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-gpu/src/demosaic.rs
+255 -1
View File
@@ -57,6 +57,32 @@ struct XTransParams {
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
@@ -437,6 +463,8 @@ pub struct Demosaicer {
pipeline: wgpu::ComputePipeline,
xtrans_pipeline: wgpu::ComputePipeline,
bind_group_layout: wgpu::BindGroupLayout,
hot_pixel_pipeline: wgpu::ComputePipeline,
hot_pixel_layout: wgpu::BindGroupLayout,
}
impl Demosaicer {
@@ -527,11 +555,15 @@ impl Demosaicer {
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,
})
}
@@ -558,8 +590,12 @@ impl Demosaicer {
// 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 {
@@ -598,6 +634,64 @@ impl Demosaicer {
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
@@ -636,7 +730,7 @@ impl Demosaicer {
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: raw_buf.as_entire_binding(),
resource: repaired.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
@@ -655,6 +749,18 @@ impl Demosaicer {
.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"),
@@ -960,6 +1066,136 @@ fn detect_xtrans_phase(raw: &RawImage) -> (u32, u32) {
/// 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);
@@ -994,6 +1230,24 @@ mod tests {
}
}
/// 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.
core/dr-gpu/src/shaders/hot_pixels.wgsl
+185
View File
@@ -0,0 +1,185 @@
// Hot and dead photosite repair, on the raw mosaic, before demosaic.
//
// A hot photosite reads far above anything the light put there — a leaky
// well, lit by its own dark current on a long or high-ISO exposure. Left in,
// the demosaic spreads it into its neighbours' interpolated channels and it
// becomes a coloured cross, three pixels wide, that no later stage can take
// back out: by then it is five pixels of plausible colour rather than one
// photosite of nonsense. So it is repaired here, where it is still one value.
//
// **What counts as hot.** A photosite far above *every* photosite of its own
// colour in its 5x5 window, and also far above every one of its eight
// immediate neighbours whatever their colour. The second half is what keeps a
// star or a glint: real light arrives through a lens and an anti-aliasing
// filter, so even the sharpest point lands on a patch of photosites, and the
// ones beside it are lit too. A hot photosite's neighbours are as dark as the
// rest of the frame. Dead photosites are the mirror image and are handled the
// same way.
//
// **What it becomes.** The brightest (for a hot photosite) or darkest (for a
// dead one) same-colour neighbour — the value nearest to what it read that
// the neighbourhood can vouch for. An average would soften the one case this
// gets wrong, a real highlight that happened to pass both tests; a clamp to
// the neighbourhood's range cannot invent anything.
//
// Written for either colour filter array: the colour of a photosite comes from
// a 6x6 tile anchored to the sensor, which holds the X-Trans pattern as it is
// and a Bayer 2x2 cell repeated nine times.
struct HotPixelParams {
// The cropped area, in sensor photosites. Only photosites inside it are
// judged, and only photosites inside it are asked as neighbours — the
// masked border sits at black and would make everything look hot.
crop_x: u32,
crop_y: u32,
width: u32,
height: u32,
// Row stride of the readout, in samples, and the number of u32 words.
stride: u32,
words: u32,
// How many invocations one row of the dispatch grid holds, so a frame
// wider than a dispatch dimension can be addressed as two.
row_invocations: u32,
// Samples in the readout. One less than twice `words` when the count is
// odd, and the padding half of the last word is never judged.
samples: u32,
// Per-position black levels and reciprocal ranges, indexed by the
// photosite's parity within the *crop*: (y&1)*2 + (x&1) counted from its
// origin, as the demosaic counts them.
black: vec4<f32>,
inv_range: vec4<f32>,
// The 6x6 colour tile, two bits per photosite, indexed by sensor
// coordinate modulo 6: 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: HotPixelParams;
@group(0) @binding(2) var<storage, read_write> repaired: array<u32>;
// How far above its brightest neighbour a photosite must read to be hot, as a
// ratio and a margin in normalised units. Twice the neighbourhood and two
// percent of the range above it: far enough that shot noise in a lit area
// never qualifies, near enough that a hot photosite in a night sky — reading
// a third of the range over a sky at one percent — always does.
const HOT_RATIO: f32 = 2.0;
const HOT_MARGIN: f32 = 0.02;
// A dead photosite reads under half its darkest neighbour, and only counts
// where that neighbour is at least this bright: in the shadows, a photosite
// at zero is noise that clipped at the black point, not a defect.
const DEAD_RATIO: f32 = 0.5;
const DEAD_FLOOR: f32 = 0.05;
fn value_at(index: u32) -> u32 {
let word = raw[index >> 1u];
return select(word & 0xFFFFu, word >> 16u, (index & 1u) == 1u);
}
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;
}
// A raw value against its own black level and range. Compared rather than
// stored, so it is left unclamped at the top: a hot photosite above white is
// still more above white than its neighbours are.
fn level(sx: u32, sy: u32, v: u32) -> f32 {
let cell = ((sy - params.crop_y) & 1u) * 2u + ((sx - params.crop_x) & 1u);
return max(f32(v) - params.black[cell], 0.0) * params.inv_range[cell];
}
// The value to store for the photosite at `index`.
fn repair(index: u32) -> u32 {
let v = value_at(index);
let sx = index % params.stride;
let sy = index / params.stride;
if (sx < params.crop_x || sy < params.crop_y
|| sx >= params.crop_x + params.width || sy >= params.crop_y + params.height) {
return v;
}
let centre = level(sx, sy, v);
let colour = colour_at(sx, sy);
var same_hi = -1.0;
var same_lo = 1.0e9;
var same_hi_raw = v;
var same_lo_raw = v;
var same_count = 0u;
var adjacent_hi = 0.0;
var adjacent_lo = 1.0e9;
for (var dy = -2; dy <= 2; dy++) {
for (var dx = -2; dx <= 2; dx++) {
if (dx == 0 && dy == 0) {
continue;
}
let nx = i32(sx) + dx;
let ny = i32(sy) + dy;
if (nx < i32(params.crop_x) || ny < i32(params.crop_y)
|| nx >= i32(params.crop_x + params.width)
|| ny >= i32(params.crop_y + params.height)) {
continue;
}
let nsx = u32(nx);
let nsy = u32(ny);
let nv = value_at(nsy * params.stride + nsx);
let n = level(nsx, nsy, nv);
if (abs(dx) <= 1 && abs(dy) <= 1) {
adjacent_hi = max(adjacent_hi, n);
adjacent_lo = min(adjacent_lo, n);
}
if (colour_at(nsx, nsy) == colour) {
same_count += 1u;
if (n > same_hi) {
same_hi = n;
same_hi_raw = nv;
}
if (n < same_lo) {
same_lo = n;
same_lo_raw = nv;
}
}
}
}
// A corner of the crop can leave a photosite with a single same-colour
// neighbour, and one witness is not a neighbourhood.
if (same_count < 2u) {
return v;
}
let hot_line_same = same_hi * HOT_RATIO + HOT_MARGIN;
let hot_line_adjacent = adjacent_hi * HOT_RATIO + HOT_MARGIN;
if (centre > hot_line_same && centre > hot_line_adjacent) {
return same_hi_raw;
}
if (same_lo >= DEAD_FLOOR && centre < same_lo * DEAD_RATIO
&& centre < adjacent_lo * DEAD_RATIO) {
return same_lo_raw;
}
return v;
}
// One invocation per u32 word: two photosites, packed as the demosaic reads
// them. A word may straddle two rows when the stride is odd, which `repair`
// does not mind — it addresses by sample index.
@compute @workgroup_size(64, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let word = gid.y * params.row_invocations + gid.x;
if (word >= params.words) {
return;
}
let lo = repair(word * 2u);
// The padding half of an odd-length readout is copied, not judged: it is
// not a photosite, and the demosaic never addresses it.
var hi = raw[word] >> 16u;
if (word * 2u + 1u < params.samples) {
hi = repair(word * 2u + 1u);
}
repaired[word] = (lo & 0xFFFFu) | (hi << 16u);
}