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

Run the app's hot-pixel pass alone, and dump through it

Browse Source 
The learned demosaic replaces the classical one and takes its input, the
mosaic hot_pixels.wgsl leaves (denoise.md §2), so its training data and its
input in the app must come through that pass and not a lookalike. The pass
was recorded inline in Demosaicer::run; it is now built by hot_pass and
recorded by record_hot_pass, which run still uses unchanged, and
Demosaicer::repair_hot_pixels runs it on its own and reads the mosaic back.

mosaic_dump moves to dr-gpu to call it, records how many photosites changed,
and keeps --unrepaired for a raw readout.
This commit is contained in:
dtourolle
2026-10-03 10:25:28 -04:00
parent 1f266a4478
commit d4248bc0dd
4 changed files with 230 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-gpu/src/demosaic.rs
+169 -62
View File
@@ -761,62 +761,14 @@ impl Demosaicer {
// 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(
// dead photosites repaired.
let hot = self.hot_pass(
raw,
(width, height),
words,
groups_x * HOT_PIXEL_GROUP,
&raw_buf,
packed.len() as u32,
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
@@ -855,7 +807,7 @@ impl Demosaicer {
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: repaired.as_entire_binding(),
resource: hot.repaired.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
@@ -877,15 +829,7 @@ impl Demosaicer {
// 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);
}
self.record_hot_pass(&mut enc, &hot);
{
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("demosaic-pass"),
@@ -921,6 +865,169 @@ impl Demosaicer {
}
}
/// The hot-pixel pass's resources for one frame, ready to record.
struct HotPass {
repaired: wgpu::Buffer,
bind_group: wgpu::BindGroup,
groups: (u32, u32),
}
impl Demosaicer {
/// Buffers and bindings for the hot and dead photosite repair of `raw`,
/// whose packed samples are in `raw_buf`.
fn hot_pass(
&self,
raw: &RawImage,
(width, height): (u32, u32),
raw_buf: &wgpu::Buffer,
words: u32,
xtrans_tile: Option<[u32; 4]>,
) -> HotPass {
// 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 | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
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 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(),
},
],
});
HotPass {
repaired,
bind_group,
groups: (groups_x, groups_y),
}
}
fn record_hot_pass(&self, enc: &mut wgpu::CommandEncoder, hot: &HotPass) {
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(hot.groups.0, hot.groups.1, 1);
}
/// TRACES: FR-RAW-3 | FR-DEV-3g
/// Repair `raw`'s hot and dead photosites in place, exactly as [`Self::run`]
/// does before it demosaics, and return how many changed.
///
/// For the learned demosaic (denoise.md §2), which reads the same repaired
/// mosaic the classical one does: its training data and its input in the
/// app must have been through this one pass, not a lookalike.
pub fn repair_hot_pixels(&self, raw: &mut RawImage) -> Result<usize, GpuError> {
if raw.samples_per_pixel != 1 {
return Ok(0);
}
let (width, height) = (raw.crop.width.max(1), raw.crop.height.max(1));
let xtrans_tile = raw
.cfa_pattern
.is_xtrans()
.then(|| xtrans_params_for(raw, width, height).tile);
let packed = pack_samples(&raw.data);
let raw_buf = self
.ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("raw-samples"),
contents: bytemuck::cast_slice(&packed),
usage: wgpu::BufferUsages::STORAGE,
});
let hot = self.hot_pass(
raw,
(width, height),
&raw_buf,
packed.len() as u32,
xtrans_tile,
);
let readback = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("raw-repaired-readback"),
size: hot.repaired.size(),
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let mut enc = self
.ctx
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("hot-pixel-encoder"),
});
self.record_hot_pass(&mut enc, &hot);
enc.copy_buffer_to_buffer(&hot.repaired, 0, &readback, 0, hot.repaired.size());
self.ctx.queue.submit(Some(enc.finish()));
let slice = readback.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
self.ctx
.device
.poll(wgpu::PollType::wait_indefinitely())
.map_err(|e| GpuError::Readback(e.to_string()))?;
rx.recv()
.map_err(|e| GpuError::Readback(e.to_string()))?
.map_err(|e| GpuError::Readback(e.to_string()))?;
let words: Vec<u32> = bytemuck::cast_slice(&slice.get_mapped_range()).to_vec();
readback.unmap();
let mut changed = 0;
for (i, v) in raw.data.iter_mut().enumerate() {
let w = words[i / 2];
let new = if i % 2 == 0 { w & 0xFFFF } else { w >> 16 } as u16;
changed += usize::from(new != *v);
*v = new;
}
Ok(changed)
}
}
const IDENTITY_3X3: [f32; 9] = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
/// Pack u16 samples two per u32, little-endian within the word.