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

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.

core/dr-gpu/examples/mosaic_dump.rs

@@ -5,27 +5,37 @@
 //! network trained on one phase runs on another and paints moiré everywhere
 //! (denoise.md §4.4). So it reads this, not LibRaw.
 //!
+//! The photosites are those the demosaic reads: hot and dead ones repaired by
+//! the app's own pass ([`Demosaicer::repair_hot_pixels`], the same shader
+//! `run` dispatches), because the learned stage replaces the demosaic and
+//! takes its input (denoise.md §2). `--unrepaired` skips it.
+//!
 //! Reads `input<TAB>output-prefix` lines on stdin and writes, per line,
 //! `prefix.npy` (the whole readout, masked border included, `u16`, row-major)
 //! and `prefix.json` (what `decode` and `metadata` say about it). The border
-//! is kept because its optically black photosites are a dark frame for free:
-//! read noise and row noise at that ISO.
+//! is kept, and the repair never touches it, because its optically black
+//! photosites are a dark frame for free: read noise and row noise at that ISO.
 //!
 //! ```sh
 //! printf 'IMG_0001.CR2\tout/IMG_0001\n' |
-//!     cargo run --release -p dr-decode --example mosaic_dump
+//!     cargo run --release -p dr-gpu --example mosaic_dump
 //! ```
 
 use std::io::{BufRead, Write};
 
+use dr_gpu::{Demosaicer, GpuContext};
+
 fn main() {
+    let repair = !std::env::args().any(|a| a == "--unrepaired");
+    let ctx = pollster::block_on(GpuContext::new_headless()).expect("a GPU for the hot-pixel pass");
+    let demosaicer = Demosaicer::new(&ctx).expect("demosaicer");
     let mut failed = 0;
     for line in std::io::stdin().lock().lines() {
         let line = line.expect("stdin");
         let Some((input, prefix)) = line.split_once('\t') else {
             continue;
         };
-        match dump(input, prefix) {
+        match dump(input, prefix, repair.then_some(&demosaicer)) {
             Ok(()) => println!("ok\t{input}"),
             Err(e) => {
                 failed += 1;
@@ -37,12 +47,16 @@ fn main() {
     std::process::exit(if failed > 0 { 1 } else { 0 });
 }
 
-fn dump(input: &str, prefix: &str) -> Result<(), String> {
+fn dump(input: &str, prefix: &str, repair: Option<&Demosaicer>) -> Result<(), String> {
     let bytes = std::fs::read(input).map_err(|e| e.to_string())?;
-    let raw = dr_decode::decode(&bytes).map_err(|e| e.to_string())?;
+    let mut raw = dr_decode::decode(&bytes).map_err(|e| e.to_string())?;
     if raw.samples_per_pixel != 1 {
         return Err("linear DNG: no photosites".into());
     }
+    let repaired = match repair {
+        Some(d) => d.repair_hot_pixels(&mut raw).map_err(|e| e.to_string())? as i64,
+        None => -1,
+    };
     let meta = dr_decode::metadata(&bytes).map_err(|e| e.to_string())?;
 
     let mut npy = Vec::with_capacity(raw.data.len() * 2 + 128);
@@ -73,7 +87,8 @@ fn dump(input: &str, prefix: &str) -> Result<(), String> {
             "\"width\": {}, \"height\": {}, ",
             "\"crop\": [{}, {}, {}, {}], \"cfa\": {:?}, ",
             "\"black\": {:?}, \"white\": {}, \"wb\": {:?}, \"cam_to_srgb\": {}, ",
-            "\"iso\": {}, \"shutter\": {}, \"aperture\": {}, \"captured_at\": {}}}\n"
+            "\"iso\": {}, \"shutter\": {}, \"aperture\": {}, \"captured_at\": {}, ",
+            "\"hot_repaired\": {}}}\n"
         ),
         input,
         raw.make,
@@ -94,6 +109,7 @@ fn dump(input: &str, prefix: &str) -> Result<(), String> {
         opt(meta.aperture),
         meta.captured_at
             .map_or("null".to_string(), |v| v.to_string()),
+        repaired,
     );
     std::fs::write(format!("{prefix}.json"), json).map_err(|e| e.to_string())
 }

core/dr-gpu/src/demosaic.rs

@@ -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.

core/dr-gpu/tests/hot_pixels.rs

@@ -140,3 +140,33 @@ fn a_hot_photosite_on_x_trans_is_invisible() {
     let diff = worst(&clean, &hot);
     assert!(diff <= 1, "a hot X-Trans photosite still shows, by {diff}");
 }
+
+/// The repair alone, read back (FR-DEV-3g): the learned demosaic takes the
+/// mosaic this pass leaves, so it must be the same pass and nothing more —
+/// the hot photosite replaced, a real highlight and every other photosite
+/// untouched.
+#[test]
+fn the_repaired_mosaic_reads_back_with_only_the_defect_changed() {
+    let Some(ctx) = ctx() else {
+        eprintln!("no GPU adapter; skipping");
+        return;
+    };
+    let d = Demosaicer::new(&ctx).expect("demosaicer");
+    let mut star = vec![(MIDDLE, MIDDLE, WHITE)];
+    for dy in 0..3 {
+        for dx in 0..3 {
+            star.push((4 + dx, 4 + dy, WHITE));
+        }
+    }
+    let before = frame(CfaPattern::Rggb, 40, &star);
+    let mut raw = before.clone();
+    let changed = d.repair_hot_pixels(&mut raw).expect("repair");
+    assert_eq!(changed, 1, "only the lone hot photosite should change");
+    let at = (MIDDLE * SIZE + MIDDLE) as usize;
+    assert_eq!(
+        raw.data[at], 40,
+        "repaired to its brightest same-colour neighbour"
+    );
+    let others = (0..raw.data.len()).filter(|&i| i != at);
+    assert!(others.into_iter().all(|i| raw.data[i] == before.data[i]));
+}