Add dr-denoise: the learned demosaic and denoise, without the UI

The noise model takes the best source the frame has: the body's measured
table (the Canon EOS 6D's, from the library), the DNG's NoiseProfile, or
the frame itself — read, row and column noise from its masked border, and
only the shot gain estimated, from the quietest flat patches. Checked on
130 6D frames, the estimate is within 10 % from ISO 1000 up; the network
loses under 0.3 dB for a sigma off by 15-20 %, so every Bayer body is
eligible.

Tiles of 1408 keep their central 1024 behind a 192-photosite halo, past the
185-photosite receptive field, and the frame is extended by reflection,
which keeps every photosite's colour; a pattern that starts on another
colour is read from one photosite up or left so the network sees RGGB, and
nothing is cropped. The tests run every Bayer phase, tiled against whole,
with a stand-in network of known reach.

The model ships as models/denoise/mosaic-1408.onnx (LFS), trained in
darkroom-denoise on the maintainer's own photographs, GPL like the code.
denoise_raw runs a file end to end: on a 6D frame at ISO 8000 the result
matches the training repository's own path to 2.5e-4 at worst, and takes
3.1 s on TensorRT fp16 (75 dB from f32) or 14.4 s on the CPU.
This commit is contained in:
2026-10-03 11:15:50 -04:00
parent 20b7bd7663
commit d8304d7c82
12 changed files with 1072 additions and 1 deletions
+132
View File
@@ -0,0 +1,132 @@
//! Denoise one RAW file end to end, as develop will, and time it.
//!
//! ```sh
//! DARKROOM_ORT_DIR=~/.local/share/darkroom/runtime \
//! cargo run --release -p dr-denoise --features native --example denoise_raw -- IMG.CR2 out
//! ```
//!
//! Decode, the app's hot-pixel pass, the frame's noise from its best source,
//! then the shipped network under the inference engine on whatever rung this
//! machine probes to. Writes `out.npy` — the active area, `h×w×3` f32 linear
//! camera RGB — for comparison with the training repo's own path
//! (`tools/compare_rust.py` in darkroom-denoise). `DARKROOM_ORT_DIR` points
//! at an ONNX Runtime build; the engine's cache goes to `DR_ENGINE_CACHE` or
//! a temporary directory.
use std::path::PathBuf;
use std::time::{Duration, Instant};
use dr_denoise::onnx::OnnxNet;
use dr_inference_engine::{Config, Role};
fn main() {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("warn")).init();
let mut args = std::env::args().skip(1);
let (Some(input), Some(out)) = (args.next(), args.next()) else {
eprintln!("usage: denoise_raw RAW OUT_PREFIX");
std::process::exit(2);
};
let model =
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../models/denoise/mosaic-1408.onnx");
let cache = std::env::var_os("DR_ENGINE_CACHE")
.map(PathBuf::from)
.unwrap_or_else(|| std::env::temp_dir().join("dr-denoise-engines"));
let started = Instant::now();
dr_inference_engine::init(Config {
runtime_dirs: std::env::var_os("DARKROOM_ORT_DIR")
.map(PathBuf::from)
.into_iter()
.collect(),
cache_dir: cache,
models: vec![(Role::Denoiser, model.clone())],
embedded: Vec::new(),
ceiling: None,
threads: 0,
decay: Duration::ZERO,
});
// Wait for the probe and the engine build, so the timing below is the
// rung this machine settles on, not the fallback used while it compiles.
// The probe starts on its own thread; give it a moment to say so.
std::thread::sleep(Duration::from_secs(1));
loop {
let s = dr_inference_engine::status();
if !s.probing && s.engines.0 >= s.engines.1 {
println!(
"engine {} ({:.1} s to settle)",
s.line(),
started.elapsed().as_secs_f64()
);
break;
}
std::thread::sleep(Duration::from_millis(200));
}
let bytes = std::fs::read(&input).expect("read raw");
let t = Instant::now();
let mut raw = dr_decode::decode(&bytes).expect("decode");
let meta = dr_decode::metadata(&bytes).expect("metadata");
let decode = t.elapsed();
let t = Instant::now();
let ctx =
pollster::block_on(dr_gpu::GpuContext::new_headless()).expect("GPU for the hot-pixel pass");
let repaired = dr_gpu::Demosaicer::new(&ctx)
.expect("demosaicer")
.repair_hot_pixels(&mut raw)
.expect("repair");
let repair = t.elapsed();
let noise = dr_denoise::noise::for_frame(&raw, &bytes, meta.iso)
.expect("no noise source for this frame");
println!(
"frame {} {} ISO {:?}, {}×{}, {:?}, {repaired} hot photosites repaired",
raw.make, raw.model, meta.iso, raw.crop.width, raw.crop.height, raw.cfa_pattern
);
println!(
"noise {} — σ at 10 % grey (G) {:.5}, read {:.5}, row {:.5}, col {:.5}",
noise.source.label(),
noise.sigma(1, 0.1),
noise.o[1].sqrt(),
noise.row,
noise.col
);
let mut net = OnnxNet::from_path(&model).expect("model");
println!(
"rung {}",
net.rung().map(|r| r.label()).unwrap_or("?")
);
let t = Instant::now();
let rgb = dr_denoise::denoise(&raw, &noise, &mut net, &mut |done, total| {
eprint!("\rtile {done}/{total}");
true
})
.expect("denoise")
.expect("not cancelled");
let run = t.elapsed();
eprintln!();
println!(
"time decode {:.2} s · hot pixels {:.2} s · network {:.2} s ({:.1} MP)",
decode.as_secs_f64(),
repair.as_secs_f64(),
run.as_secs_f64(),
(raw.crop.width * raw.crop.height) as f64 / 1e6
);
let (h, w) = (raw.crop.height as usize, raw.crop.width as usize);
let mut npy = Vec::with_capacity(rgb.len() * 4 + 128);
let mut header =
format!("{{'descr': '<f4', 'fortran_order': False, 'shape': ({h}, {w}, 3), }}");
while (10 + header.len() + 1) % 64 != 0 {
header.push(' ');
}
header.push('\n');
npy.extend_from_slice(b"\x93NUMPY\x01\x00");
npy.extend_from_slice(&(header.len() as u16).to_le_bytes());
npy.extend_from_slice(header.as_bytes());
for v in &rgb {
npy.extend_from_slice(&v.to_le_bytes());
}
std::fs::write(format!("{out}.npy"), npy).expect("write");
println!("wrote {out}.npy");
}