diff --git a/Cargo.lock b/Cargo.lock index 7260f40..7915fa9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1498,6 +1498,23 @@ dependencies = [ "zune-jpeg 0.4.21", ] +[[package]] +name = "dr-denoise" +version = "0.20.0" +dependencies = [ + "dr-decode", + "dr-gpu", + "dr-inference-engine", + "env_logger", + "log", + "ndarray", + "ort", + "pollster", + "serde", + "serde_norway", + "thiserror 2.0.20", +] + [[package]] name = "dr-export" version = "0.20.0" diff --git a/Cargo.toml b/Cargo.toml index 0827e21..abe05b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ "core/dr-catalog", "core/dr-thumbs", "core/dr-decode", + "core/dr-denoise", "core/dr-export", "core/dr-face", "core/dr-film", @@ -44,6 +45,7 @@ dr-types = { path = "core/dr-types" } dr-catalog = { path = "core/dr-catalog" } dr-thumbs = { path = "core/dr-thumbs" } dr-decode = { path = "core/dr-decode" } +dr-denoise = { path = "core/dr-denoise" } dr-export = { path = "core/dr-export" } # Stated explicitly for the same reason as `dr-segment` below: no dependant # should drag in an ONNX runtime by accident. Members opt in with diff --git a/core/dr-decode/src/profile.rs b/core/dr-decode/src/profile.rs index f4da951..835812f 100644 --- a/core/dr-decode/src/profile.rs +++ b/core/dr-decode/src/profile.rs @@ -757,7 +757,7 @@ pub fn read_noise_profile(decoder: &dyn rawler::decoders::Decoder) -> Option = (0..n / 2) diff --git a/core/dr-denoise/Cargo.toml b/core/dr-denoise/Cargo.toml new file mode 100644 index 0000000..6211463 --- /dev/null +++ b/core/dr-denoise/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "dr-denoise" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[dependencies] +dr-decode.workspace = true +serde = { workspace = true } +serde_norway.workspace = true +thiserror.workspace = true +log.workspace = true + +# The network runs under the inference engine like every other model +# (docs/dev/inference.md): `ort` is the API, the engine picks the rung. +# Optional so the noise model and the tiling test without a runtime. +ort = { workspace = true, optional = true } +dr-inference-engine = { workspace = true, optional = true } +ndarray = { workspace = true, optional = true } + +[features] +default = ["onnx"] +onnx = ["dep:ort", "dep:dr-inference-engine", "dep:ndarray"] +# A real ONNX Runtime from disk rather than tract alone, as the app links it. +native = ["onnx", "dr-inference-engine/native"] + +[dev-dependencies] +# The example repairs hot photosites with the app's own pass, as develop will. +dr-gpu.workspace = true +pollster.workspace = true +env_logger.workspace = true diff --git a/core/dr-denoise/examples/denoise_raw.rs b/core/dr-denoise/examples/denoise_raw.rs new file mode 100644 index 0000000..22740c4 --- /dev/null +++ b/core/dr-denoise/examples/denoise_raw.rs @@ -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': ' bool { + raw.samples_per_pixel == 1 && tile::rggb_offset(raw.cfa_pattern).is_some() +} + +/// The active area of `raw`, denoised and demosaiced: `crop.height × +/// crop.width` interleaved RGB, linear camera space, normalised black 0 and +/// white 1 per photosite as the classical demosaic normalises. +/// +/// `raw` must be hot-pixel repaired. `None` when `progress` stopped it. +pub fn denoise( + raw: &RawImage, + noise: &NoiseModel, + net: &mut dyn TileNet, + progress: &mut dyn FnMut(usize, usize) -> bool, +) -> Result>, DenoiseError> { + if !eligible(raw) { + return Err(DenoiseError::Unsupported(format!( + "{:?} with {} samples per photosite", + raw.cfa_pattern, raw.samples_per_pixel + ))); + } + let active = noise::active(raw); + let (h, w) = (active.h, active.w); + tile::run_tiled( + net, + h, + w, + raw.cfa_pattern, + &|y, x| active.at(y, x), + &|c, v| noise.sigma(c, v), + progress, + ) +} diff --git a/core/dr-denoise/src/noise.rs b/core/dr-denoise/src/noise.rs new file mode 100644 index 0000000..f10d323 --- /dev/null +++ b/core/dr-denoise/src/noise.rs @@ -0,0 +1,396 @@ +//! TRACES: FR-DEV-3g +//! How noisy each photosite is: the network is told, not left to guess +//! (denoise.md §3.3). +//! +//! The model is `σ² = S·x + O + row² + col²` per photosite, `x` the signal +//! normalised black-to-white the way the demosaic normalises it. Three +//! sources, best first: +//! +//! 1. **A measured table** for the body ([`Source::Table`]) — the Canon EOS 6D +//! today, from the library's own frames. +//! 2. **The DNG's `NoiseProfile`** ([`Source::DngProfile`]) — what Adobe's +//! converter measured for the body at that ISO. +//! 3. **The frame itself** ([`Source::Measured`]) — read, row and column +//! noise from its masked border, which is a dark frame taken in the same +//! instant, and only the shot gain estimated, from the quietest flat +//! patches. Checked against the 6D's table on 130 frames: within ±10 % at +//! ISO 1000 and above, scattered below; the network loses under 0.3 dB for +//! a σ off by 15–20 %, and over-estimating costs half what +//! under-estimating does, so the estimate leans high. +//! +//! Row and column noise come from the masked border whenever the frame has +//! one, whatever the source of the rest. + +use dr_decode::{CfaPattern, RawImage}; +use serde::Deserialize; + +/// Where a frame's noise figures came from, for develop to say. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Source { + Table, + DngProfile, + Measured, +} + +impl Source { + pub fn label(self) -> &'static str { + match self { + Source::Table => "measured for this camera", + Source::DngProfile => "from the DNG's noise profile", + Source::Measured => "estimated from this photograph", + } + } +} + +/// Per-photosite noise in the frame's own normalisation (black 0, white 1). +#[derive(Clone, Debug, PartialEq)] +pub struct NoiseModel { + /// Shot gain per colour, R G B. + pub s: [f32; 3], + /// Read variance per colour, R G B. + pub o: [f32; 3], + /// Standard deviation shared by a whole row, and by a whole column. + pub row: f32, + pub col: f32, + pub source: Source, +} + +impl NoiseModel { + /// σ for a photosite of colour `c` (0 R, 1 G, 2 B) reading `x`. + #[inline] + pub fn sigma(&self, c: usize, x: f32) -> f32 { + (self.s[c] * x.max(0.0) + self.o[c] + self.row * self.row + self.col * self.col).sqrt() + } + + /// The same figures scaled for the Amount the spec describes (§3.3): + /// above 1 tells the network there is more noise than there is. + pub fn scaled(&self, amount: f32) -> NoiseModel { + let a2 = amount * amount; + NoiseModel { + s: self.s.map(|v| v * a2), + o: self.o.map(|v| v * a2), + row: self.row * amount, + col: self.col * amount, + source: self.source, + } + } +} + +/// The frame's noise, from the best source it has. +/// +/// `bytes` is the file (for a DNG's `NoiseProfile`), `iso` its EXIF ISO. +/// `None` only for a frame with no masked border, no profile and no table +/// that is also too dark or too busy to measure. +pub fn for_frame(raw: &RawImage, bytes: &[u8], iso: Option) -> Option { + let dark = dark_border(raw); + let mut model = iso + .and_then(|iso| from_table(raw, iso)) + .or_else(|| dr_decode::noise_profile(bytes).and_then(|p| from_dng_profile(raw, &p))) + .or_else(|| measured(raw, dark.as_ref()))?; + if let Some(d) = dark { + // The border saw this exposure's row and column noise directly. + if model.source != Source::Table { + model.row = d.row; + model.col = d.col; + } + } + Some(model) +} + +#[derive(Deserialize)] +struct Table { + make: String, + model: String, + rows: Vec, +} + +#[derive(Deserialize)] +struct TableRow { + iso: u32, + s_dn: [f32; 4], + o_dn: [f32; 4], + row_dn: f32, + col_dn: f32, +} + +const TABLES: &[&str] = &[include_str!("../tables/canon-eos-6d.yaml")]; + +/// The body's measured table at the nearest ISO it holds, converted from DN +/// to this frame's normalisation. +pub fn from_table(raw: &RawImage, iso: u32) -> Option { + let table = TABLES.iter().find_map(|t| { + let t: Table = serde_norway::from_str(t).ok()?; + (t.make.eq_ignore_ascii_case(&raw.make) && t.model.eq_ignore_ascii_case(&raw.model)) + .then_some(t) + })?; + let row = table.rows.iter().min_by(|a, b| { + let d = |r: &TableRow| ((r.iso as f32).ln() - (iso as f32).ln()).abs(); + d(a).total_cmp(&d(b)) + })?; + let span = span(raw); + // RGGB positions → colours: the greens share. + let s = [row.s_dn[0], 0.5 * (row.s_dn[1] + row.s_dn[2]), row.s_dn[3]].map(|v| v / span); + let o = + [row.o_dn[0], 0.5 * (row.o_dn[1] + row.o_dn[2]), row.o_dn[3]].map(|v| v / (span * span)); + Some(NoiseModel { + s, + o, + row: row.row_dn / span, + col: row.col_dn / span, + source: Source::Table, + }) +} + +/// A DNG's `NoiseProfile`: one pair for every plane, or one per colour plane +/// (R, G, B for a Bayer DNG), already in the file's black-to-white units — +/// which are the units `dr-decode` normalises by. +pub fn from_dng_profile(raw: &RawImage, pairs: &[(f32, f32)]) -> Option { + if raw.cfa_pattern.is_xtrans() || raw.samples_per_pixel != 1 { + return None; + } + let (s, o) = match pairs { + [(s, o)] => ([*s; 3], [*o; 3]), + [r, g, b, ..] => ([r.0, g.0, b.0], [r.1, g.1, b.1]), + _ => return None, + }; + Some(NoiseModel { + s, + o, + row: 0.0, + col: 0.0, + source: Source::DngProfile, + }) +} + +/// Read, row and column noise measured on the masked border, normalised. +#[derive(Clone, Copy, Debug)] +pub struct Dark { + pub read: f32, + pub row: f32, + pub col: f32, +} + +/// The optically black photosites beside and above the active area. +/// +/// Keeps well clear of the active area: on the 6D the dozen columns nearest +/// it see light. Photosites over 8σ are the strip's own hot photosites — the +/// same ones in every frame — and are left out, as the app's hot-pixel pass +/// removes their kin before the network sees them. +pub fn dark_border(raw: &RawImage) -> Option { + let (x0, y0, w, h) = ( + raw.crop.x as usize, + raw.crop.y as usize, + raw.crop.width as usize, + raw.crop.height as usize, + ); + let stride = raw.width as usize; + let span = span(raw); + if x0 < 40 || raw.samples_per_pixel != 1 { + return None; + } + let cols = 4..x0 - 16; + let nc = cols.len() as f32; + // Residual after removing each row's mean and each column's mean. + let mut row_means = Vec::with_capacity(h); + let mut col_sum = vec![0.0f64; cols.len()]; + for y in y0..y0 + h { + let line = &raw.data[y * stride..y * stride + x0]; + let m = cols.clone().map(|x| line[x] as f32).sum::() / nc; + row_means.push(m); + for (k, x) in cols.clone().enumerate() { + col_sum[k] += (line[x] as f32 - m) as f64; + } + } + let col_mean: Vec = col_sum.iter().map(|s| (*s / h as f64) as f32).collect(); + let resid = |y: usize, k: usize, x: usize| { + raw.data[y * stride + x] as f32 - row_means[y - y0] - col_mean[k] + }; + let (mut s1, mut n) = (0.0f64, 0usize); + for y in y0..y0 + h { + for (k, x) in cols.clone().enumerate() { + s1 += (resid(y, k, x) as f64).powi(2); + n += 1; + } + } + let rough = (s1 / n as f64).sqrt() as f32; + let (mut s2, mut n2) = (0.0f64, 0usize); + for y in y0..y0 + h { + for (k, x) in cols.clone().enumerate() { + let r = resid(y, k, x); + if r.abs() < 8.0 * rough { + s2 += (r as f64).powi(2); + n2 += 1; + } + } + } + let read = (s2 / n2.max(1) as f64).sqrt() as f32; + let rm = row_means.iter().sum::() / h as f32; + let row_var = row_means.iter().map(|m| (m - rm).powi(2)).sum::() / h as f32; + let row = (row_var - read * read / nc).max(0.0).sqrt(); + + // Columns: the masked rows above the image span every column. + let col = if y0 >= 24 { + let rows = 4..y0 - 12; + let nr = rows.len() as f32; + let means: Vec = (x0..x0 + w) + .map(|x| { + rows.clone() + .map(|y| raw.data[y * stride + x] as f32) + .sum::() + / nr + }) + .collect(); + let mm = means.iter().sum::() / means.len() as f32; + let var = means.iter().map(|m| (m - mm).powi(2)).sum::() / means.len() as f32; + (var - read * read / nr).max(0.0).sqrt() + } else { + 0.0 + }; + Some(Dark { + read: read / span, + row: row / span, + col: col / span, + }) +} + +/// The quietest-third bias of the patch variance, and the residual bias the +/// estimate showed against the 6D's table (0.91 at the median), in one: the +/// estimate is divided by this. +const QUIET_FACTOR: f32 = 0.85 * 0.91; + +/// The frame's own noise: read noise from the border (or, lacking one, the +/// floor of the quietest patches), shot gain from flat patches of one green +/// plane, the same for every colour, as a sensor's gain is. +pub fn measured(raw: &RawImage, dark: Option<&Dark>) -> Option { + if raw.cfa_pattern.is_xtrans() || raw.samples_per_pixel != 1 { + return None; + } + let m = active(raw); + let (h, w) = (m.h, m.w); + // One green plane at a two-photosite pitch. + let (gy, gx) = green_offset(raw.cfa_pattern)?; + let ph = (h - gy) / 2; + let pw = (w - gx) / 2; + let g = |y: usize, x: usize| m.at(gy + 2 * y, gx + 2 * x); + const B: usize = 8; + let mut patches: Vec<(f32, f32)> = Vec::new(); // (level, variance) + for by in 0..ph / B { + for bx in 0..(pw - 2) / B { + let (mut s, mut s2, mut lv) = (0.0f32, 0.0f32, 0.0f32); + for y in by * B..by * B + B { + for x in bx * B..bx * B + B { + // Second difference: cancels any gradient; var = 6σ². + let d = g(y, x + 2) - 2.0 * g(y, x + 1) + g(y, x); + s += d; + s2 += d * d; + lv += g(y, x + 1); + } + } + let n = (B * B) as f32; + let var = (s2 / n - (s / n).powi(2)) / 6.0; + patches.push((lv / n, var)); + } + } + let floor = dark.map(|d| d.read); + let lo = 4.0 * floor.unwrap_or(0.002); + patches.retain(|(l, _)| *l > lo && *l < 0.7); + if patches.len() < 500 { + return None; + } + patches.sort_by(|a, b| a.0.total_cmp(&b.0)); + let bins = 12; + let per = patches.len() / bins; + let mut ests = Vec::new(); + let mut floors = Vec::new(); + for b in 0..bins { + let mut bin: Vec<(f32, f32)> = patches[b * per..(b + 1) * per].to_vec(); + if bin.len() < 60 { + continue; + } + bin.sort_by(|a, b| a.1.total_cmp(&b.1)); + let quiet = &bin[..bin.len() / 3]; + let read2 = floor.map(|r| r * r); + let mut e: Vec = quiet + .iter() + .map(|(l, v)| (v / QUIET_FACTOR - read2.unwrap_or(0.0)) / l) + .collect(); + e.sort_by(f32::total_cmp); + ests.push(e[e.len() / 2]); + floors.push(quiet[quiet.len() / 2]); + } + ests.sort_by(f32::total_cmp); + let s = *ests.get(ests.len() / 2)?; + if !(s.is_finite() && s > 0.0) { + return None; + } + // No border: the read variance is what the darkest bin leaves unexplained. + let read2 = match floor { + Some(r) => r * r, + None => { + let (l, v) = floors.first().copied()?; + (v / QUIET_FACTOR - s * l).max(1e-9) + } + }; + Some(NoiseModel { + s: [s; 3], + o: [read2; 3], + row: dark.map_or(0.0, |d| d.row), + col: dark.map_or(0.0, |d| d.col), + source: Source::Measured, + }) +} + +/// Black-to-white range of the frame, as the demosaic normalises it. +pub(crate) fn span(raw: &RawImage) -> f32 { + let black = raw.black_level.iter().map(|&b| b as f32).sum::() / 4.0; + (raw.white_level as f32 - black).max(1.0) +} + +/// Where a green photosite sits in the pattern's 2×2 cell, (dy, dx). +fn green_offset(p: CfaPattern) -> Option<(usize, usize)> { + match p { + CfaPattern::Rggb | CfaPattern::Bggr => Some((0, 1)), + CfaPattern::Grbg | CfaPattern::Gbrg => Some((0, 0)), + _ => None, + } +} + +/// The active area, normalised, read lazily. +pub(crate) struct Active<'a> { + raw: &'a RawImage, + black: [f32; 4], + inv: [f32; 4], + pub h: usize, + pub w: usize, +} + +impl Active<'_> { + /// Photosite (y, x) of the active area, black 0, white 1. + #[inline] + pub fn at(&self, y: usize, x: usize) -> f32 { + let c = (y & 1) * 2 + (x & 1); + let v = self.raw.data[(self.raw.crop.y as usize + y) * self.raw.width as usize + + self.raw.crop.x as usize + + x]; + (v as f32 - self.black[c]) * self.inv[c] + } +} + +/// Black levels per position of the crop's 2×2 cell, as the demosaic reads +/// them: one reported level is broadcast. +pub(crate) fn active(raw: &RawImage) -> Active<'_> { + let b = raw.black_level; + let black = if b[1] == 0 && b[2] == 0 && b[3] == 0 { + [b[0] as f32; 4] + } else { + b.map(|v| v as f32) + }; + let inv = black.map(|bl| 1.0 / (raw.white_level as f32 - bl).max(1.0)); + Active { + raw, + black, + inv, + h: raw.crop.height as usize, + w: raw.crop.width as usize, + } +} diff --git a/core/dr-denoise/src/onnx.rs b/core/dr-denoise/src/onnx.rs new file mode 100644 index 0000000..4e903c2 --- /dev/null +++ b/core/dr-denoise/src/onnx.rs @@ -0,0 +1,66 @@ +//! TRACES: FR-DEV-3g +//! The denoise network under the inference engine. +//! +//! The shipped export takes `mosaic` and `sigma`, `1×1×1408×1408`, and +//! returns `rgb`, `1×3×1408×1408` (darkroom-denoise `denoise/export.py`, +//! fixed shape because every model the engine runs is). The engine picks the +//! rung: fp16 on TensorRT and MIGraphX, which measured 0.00 dB from f32; f32 +//! on CUDA and the CPU; never the Hexagon, where int8 lost 6–9 dB. + +use crate::tile::TileNet; +use crate::DenoiseError; +use dr_inference_engine::{Model, Role}; + +/// The edge of the tile the shipped export takes. +pub const TILE: usize = 1408; + +pub struct OnnxNet { + model: Model, + tile: usize, +} + +impl OnnxNet { + pub fn from_path(path: &std::path::Path) -> Result { + let (path, form) = dr_inference_engine::resolve_model(Role::Denoiser, path); + let bytes = std::fs::read(&path)?; + Ok(OnnxNet { + model: dr_inference_engine::open(Role::Denoiser, form, &bytes)?, + tile: TILE, + }) + } + + /// Where it runs, for a status line. + pub fn rung(&self) -> Result { + Ok(self.model.acquire()?.rung()) + } +} + +impl TileNet for OnnxNet { + fn tile(&self) -> usize { + self.tile + } + + fn run(&mut self, mosaic: &[f32], sigma: &[f32]) -> Result, DenoiseError> { + let n = self.tile; + let shape = ndarray::IxDyn(&[1, 1, n, n]); + let m = ort::value::Tensor::from_array( + ndarray::Array::from_shape_vec(shape.clone(), mosaic.to_vec()) + .map_err(|e| DenoiseError::Model(e.to_string()))?, + )?; + let s = ort::value::Tensor::from_array( + ndarray::Array::from_shape_vec(shape, sigma.to_vec()) + .map_err(|e| DenoiseError::Model(e.to_string()))?, + )?; + let acquired = self.model.acquire()?; + let mut session = acquired.lock(); + let outputs = session.run(ort::inputs!["mosaic" => m, "sigma" => s])?; + let (shape, data) = outputs[0].try_extract_tensor::()?; + let dims: Vec = shape.iter().copied().collect(); + if dims != [1, 3, n as i64, n as i64] { + return Err(DenoiseError::Model(format!( + "output is {dims:?}, expected [1, 3, {n}, {n}]" + ))); + } + Ok(data.to_vec()) + } +} diff --git a/core/dr-denoise/src/tile.rs b/core/dr-denoise/src/tile.rs new file mode 100644 index 0000000..cc14ad0 --- /dev/null +++ b/core/dr-denoise/src/tile.rs @@ -0,0 +1,295 @@ +//! TRACES: FR-DEV-3g +//! A whole frame through a fixed-shape network, exactly (denoise.md §3.4). +//! +//! The network sees `TILE_IN`² photosites and its output is exact in the +//! central `TILE_IN − 2·HALO`: the halo is wider than its receptive field +//! (185 photosites, counted from the layers), so a tile's centre equals the +//! whole frame's at the same place. The frame is extended by reflection +//! about its edge photosites, which keeps every photosite's CFA colour, so +//! edge tiles see real context too. +//! +//! **Phase.** The network was trained on RGGB. A frame whose pattern starts +//! on another colour is read from one photosite up and/or left — the +//! reflection supplies that row or column — so its top-left is red, and the +//! output is read back from the same offset. Nothing is cropped. + +use dr_decode::CfaPattern; + +/// Photosites of context beyond a tile's kept centre, on every side. +pub const HALO: usize = 192; + +/// A fixed-shape network: `mosaic` and `sigma`, `n×n` RGGB, in; `3×n×n` +/// planar linear camera RGB out. +pub trait TileNet { + /// The edge `n` of the square tile the network takes. + fn tile(&self) -> usize; + fn run(&mut self, mosaic: &[f32], sigma: &[f32]) -> Result, crate::DenoiseError>; +} + +/// Index into `0..n` by reflection about the end photosites, any distance +/// out: …2 1 [0 1 2 … n−1] n−2 n−3…, period `2(n−1)`. Parity is kept, which +/// is what keeps a CFA colour. +#[inline] +pub fn reflect(i: isize, n: usize) -> usize { + if n == 1 { + return 0; + } + let p = 2 * (n as isize - 1); + let m = i.rem_euclid(p); + (if m < n as isize { m } else { p - m }) as usize +} + +/// How far up and left to start reading so the first photosite is red. +pub fn rggb_offset(p: CfaPattern) -> Option<(usize, usize)> { + match p { + CfaPattern::Rggb => Some((0, 0)), + CfaPattern::Grbg => Some((0, 1)), + CfaPattern::Gbrg => Some((1, 0)), + CfaPattern::Bggr => Some((1, 1)), + _ => None, + } +} + +/// Run `net` over an `h×w` mosaic given by `at(y, x)`, with σ from +/// `sigma(colour, value)`, and return `h×w` interleaved RGB. +/// +/// `progress(done, total)` is called after each tile and stops the run by +/// returning `false`, in which case the result is `Ok(None)`. +#[allow(clippy::too_many_arguments)] +pub fn run_tiled( + net: &mut dyn TileNet, + h: usize, + w: usize, + pattern: CfaPattern, + at: &dyn Fn(usize, usize) -> f32, + sigma: &dyn Fn(usize, f32) -> f32, + progress: &mut dyn FnMut(usize, usize) -> bool, +) -> Result>, crate::DenoiseError> { + let (dy, dx) = rggb_offset(pattern).ok_or_else(|| { + crate::DenoiseError::Unsupported(format!("{pattern:?} is not a Bayer pattern")) + })?; + let n = net.tile(); + if n <= 2 * HALO || !(n - 2 * HALO).is_multiple_of(2) { + return Err(crate::DenoiseError::Model(format!( + "tile {n} leaves no even centre past a {HALO} halo" + ))); + } + let core = n - 2 * HALO; + // In unified coordinates the frame spans u ∈ [dy, dy + h), v ∈ [dx, dx + w). + let (uh, uw) = (h + dy, w + dx); + let (ty, tx) = (uh.div_ceil(core), uw.div_ceil(core)); + let total = ty * tx; + let mut out = vec![0.0f32; h * w * 3]; + let mut mos = vec![0.0f32; n * n]; + let mut sig = vec![0.0f32; n * n]; + // RGGB colour of unified position (u, v). + let colour = |u: usize, v: usize| [[0, 1], [1, 2]][u & 1][v & 1]; + for (k, (i, j)) in (0..ty) + .flat_map(|i| (0..tx).map(move |j| (i, j))) + .enumerate() + { + let (u0, v0) = (i * core, j * core); + for r in 0..n { + // Unified row u = u0 + r − HALO; frame row y = u − dy, reflected. + let u = u0 as isize + r as isize - HALO as isize; + let y = reflect(u - dy as isize, h); + for c in 0..n { + let v = v0 as isize + c as isize - HALO as isize; + let x = reflect(v - dx as isize, w); + let val = at(y, x); + mos[r * n + c] = val; + sig[r * n + c] = sigma(colour(r, c), val); + } + } + let rgb = net.run(&mos, &sig)?; + if rgb.len() != 3 * n * n { + return Err(crate::DenoiseError::Model(format!( + "network returned {} values for a {n}² tile", + rgb.len() + ))); + } + for r in HALO..HALO + core { + let u = u0 + r - HALO; + if u < dy || u >= uh { + continue; + } + let y = u - dy; + for c in HALO..HALO + core { + let v = v0 + c - HALO; + if v < dx || v >= uw { + continue; + } + let x = v - dx; + let o = (y * w + x) * 3; + for ch in 0..3 { + out[o + ch] = rgb[ch * n * n + r * n + c]; + } + } + } + if !progress(k + 1, total) { + return Ok(None); + } + } + Ok(Some(out)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reflection_keeps_parity_any_distance_out() { + let n = 7; + for i in -40isize..40 { + let r = reflect(i, n); + assert!(r < n); + assert_eq!( + r % 2, + i.rem_euclid(2) as usize, + "index {i} reflected to {r}" + ); + } + assert_eq!(reflect(-1, n), 1); + assert_eq!(reflect(7, n), 5); + } + + /// A stand-in network with a known, finite reach: each output photosite + /// is its 2×2 quad's (R, mean G, B), averaged over the quads within + /// `reach` quads. Purely a function of the tile, like the real one. + struct BoxNet { + n: usize, + reach: usize, + } + + impl TileNet for BoxNet { + fn tile(&self) -> usize { + self.n + } + fn run(&mut self, m: &[f32], _s: &[f32]) -> Result, crate::DenoiseError> { + let n = self.n; + let q = n / 2; + let quad = |qy: usize, qx: usize| { + let (y, x) = (2 * qy, 2 * qx); + [ + m[y * n + x], + 0.5 * (m[y * n + x + 1] + m[(y + 1) * n + x]), + m[(y + 1) * n + x + 1], + ] + }; + let mut out = vec![0.0; 3 * n * n]; + for qy in 0..q { + for qx in 0..q { + let mut acc = [0.0f32; 3]; + let mut cnt = 0.0; + for a in qy.saturating_sub(self.reach)..(qy + self.reach + 1).min(q) { + for b in qx.saturating_sub(self.reach)..(qx + self.reach + 1).min(q) { + let v = quad(a, b); + for c in 0..3 { + acc[c] += v[c]; + } + cnt += 1.0; + } + } + for (dy, dx) in [(0, 0), (0, 1), (1, 0), (1, 1)] { + for c in 0..3 { + out[c * n * n + (2 * qy + dy) * n + 2 * qx + dx] = acc[c] / cnt; + } + } + } + } + Ok(out) + } + } + + /// The mosaic of a smooth colour field in `pattern`, read at (y, x). + fn field(pattern: CfaPattern) -> impl Fn(usize, usize) -> f32 { + move |y, x| { + let rgb = [0.2 + 0.0004 * x as f32, 0.5, 0.1 + 0.0003 * y as f32]; + rgb[pattern.colour_at(x as u32, y as u32) as usize] + } + } + + #[test] + fn every_bayer_phase_comes_back_as_its_own_colours() { + // A frame of each pattern, its colours known: the network must see + // red where the frame's red photosites are, whatever the phase. + for p in [ + CfaPattern::Rggb, + CfaPattern::Grbg, + CfaPattern::Gbrg, + CfaPattern::Bggr, + ] { + let (h, w) = (300, 410); + let at = field(p); + let mut net = BoxNet { + n: 2 * HALO + 64, + reach: 0, + }; + let out = run_tiled(&mut net, h, w, p, &at, &|_, _| 0.01, &mut |_, _| true) + .unwrap() + .unwrap(); + for (y, x) in [(10, 10), (150, 201), (299, 409), (0, 0), (77, 333)] { + let o = &out[(y * w + x) * 3..(y * w + x) * 3 + 3]; + let want = [0.2 + 0.0004 * x as f32, 0.5, 0.1 + 0.0003 * y as f32]; + for c in 0..3 { + // Within the quad the binned value is at most a photosite away. + assert!( + (o[c] - want[c]).abs() < 0.0012, + "{p:?} at ({y},{x}) channel {c}: {} vs {}", + o[c], + want[c] + ); + } + } + } + } + + #[test] + fn tiles_reproduce_one_pass_over_the_reflected_frame() { + // A network whose reach is inside the halo gives the same answer + // tiled small as in one tile covering everything. + let (h, w) = (230, 170); + for p in [CfaPattern::Rggb, CfaPattern::Bggr] { + let at = |y: usize, x: usize| ((y * 7919 + x * 104729) % 1000) as f32 / 1000.0; + let mut small = BoxNet { + n: 2 * HALO + 32, + reach: 20, + }; + let mut big = BoxNet { + n: 2 * HALO + 256, + reach: 20, + }; + let a = run_tiled(&mut small, h, w, p, &at, &|_, _| 0.0, &mut |_, _| true) + .unwrap() + .unwrap(); + let b = run_tiled(&mut big, h, w, p, &at, &|_, _| 0.0, &mut |_, _| true) + .unwrap() + .unwrap(); + let worst = a + .iter() + .zip(&b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max); + assert!(worst < 1e-5, "{p:?}: tiled and whole differ by {worst}"); + } + } + + #[test] + fn a_cancelled_run_returns_nothing() { + let mut net = BoxNet { + n: 2 * HALO + 32, + reach: 0, + }; + let r = run_tiled( + &mut net, + 100, + 100, + CfaPattern::Rggb, + &|_, _| 0.5, + &|_, _| 0.0, + &mut |done, _| done < 2, + ) + .unwrap(); + assert!(r.is_none()); + } +} diff --git a/core/dr-denoise/tables/canon-eos-6d.yaml b/core/dr-denoise/tables/canon-eos-6d.yaml new file mode 100644 index 0000000..6d3c684 --- /dev/null +++ b/core/dr-denoise/tables/canon-eos-6d.yaml @@ -0,0 +1,36 @@ +# Canon EOS 6D noise, measured from the library's own frames (denoise.md §5). +# Shot gain S and read variance O per RGGB position from Adobe's NoiseProfile in +# converted DNGs, in DN at the ISO's own white level; read noise checked against +# the masked border (within 2-3 %); row and column noise from the masked border. +# ISO 50 and 100 are extrapolated (S proportional to ISO). Generated by +# darkroom-denoise tools/profile.py; regenerate there, never edit by hand. +make: Canon +model: EOS 6D +black: 2048 +rows: + - {iso: 50, white: 15000, s_dn: [0.0854021, 0.085467, 0.085467, 0.0839724], o_dn: [38.2741, 38.6675, 38.6675, 38.9649], row_dn: 0.3423, col_dn: 0.505} + - {iso: 100, white: 15000, s_dn: [0.170804, 0.170934, 0.170934, 0.167945], o_dn: [38.339, 38.7332, 38.7332, 39.031], row_dn: 0.3423, col_dn: 0.505} + - {iso: 125, white: 15035, s_dn: [0.228108, 0.230361, 0.230361, 0.228345], o_dn: [36.9455, 37.8995, 37.8995, 38.2358], row_dn: 0.3423, col_dn: 0.505} + - {iso: 160, white: 12373, s_dn: [0.289653, 0.294915, 0.294915, 0.286887], o_dn: [15.3717, 16.1346, 16.1346, 16.0413], row_dn: 0.212, col_dn: 0.07151} + - {iso: 200, white: 15035, s_dn: [0.370969, 0.369922, 0.369922, 0.361443], o_dn: [24.2761, 24.0847, 24.0847, 24.2414], row_dn: 0.2692, col_dn: 0} + - {iso: 250, white: 15035, s_dn: [0.461889, 0.457975, 0.457975, 0.449318], o_dn: [38.0975, 37.5041, 37.5041, 37.7424], row_dn: 0.3345, col_dn: 0.4786} + - {iso: 320, white: 12323, s_dn: [0.590765, 0.59755, 0.59755, 0.576843], o_dn: [18.5426, 18.9163, 18.9163, 19.1202], row_dn: 0.3158, col_dn: 0.5174} + - {iso: 400, white: 15035, s_dn: [0.753591, 0.740586, 0.740586, 0.729874], o_dn: [29.3028, 29.8496, 29.8496, 29.7961], row_dn: 0.4378, col_dn: 0.2691} + - {iso: 500, white: 15035, s_dn: [0.937458, 0.920836, 0.920836, 0.899293], o_dn: [45.2196, 46.3954, 46.3954, 46.1012], row_dn: 0.5473, col_dn: 0.4328} + - {iso: 640, white: 12323, s_dn: [1.12726, 1.13527, 1.13527, 1.10159], o_dn: [24.9951, 25.1029, 25.1029, 25.7183], row_dn: 0.316, col_dn: 0.4544} + - {iso: 800, white: 15035, s_dn: [1.44048, 1.42299, 1.42299, 1.40795], o_dn: [38.7891, 39.302, 39.302, 40.079], row_dn: 0.3877, col_dn: 0.2132} + - {iso: 1000, white: 15000, s_dn: [1.77595, 1.75662, 1.75662, 1.74584], o_dn: [63.9499, 64.4203, 64.4203, 65.0739], row_dn: 0.4593, col_dn: 0.3307} + - {iso: 1250, white: 12346, s_dn: [2.18211, 2.18313, 2.18313, 2.11979], o_dn: [41.6124, 42.9483, 42.9483, 43.2075], row_dn: 0.3979, col_dn: 0.4496} + - {iso: 1600, white: 15035, s_dn: [2.75544, 2.74633, 2.74633, 2.69951], o_dn: [66.3905, 66.4104, 66.4104, 67.253], row_dn: 0.4944, col_dn: 0.4593} + - {iso: 2000, white: 15035, s_dn: [3.42754, 3.40445, 3.40445, 3.36808], o_dn: [104.349, 103.648, 103.648, 106.404], row_dn: 0.6094, col_dn: 0.3602} + - {iso: 2500, white: 12330, s_dn: [4.17112, 4.17551, 4.17551, 4.17175], o_dn: [94.4289, 91.8508, 91.8508, 96.3598], row_dn: 0.5671, col_dn: 0} + - {iso: 3200, white: 15035, s_dn: [5.30088, 5.25742, 5.25742, 5.21782], o_dn: [147.421, 147.302, 147.302, 147.01], row_dn: 0.748, col_dn: 0.8611} + - {iso: 4000, white: 15035, s_dn: [6.62037, 6.59922, 6.59922, 6.60871], o_dn: [224.765, 232.408, 232.408, 231.419], row_dn: 0.9335, col_dn: 1.125} + - {iso: 5000, white: 12323, s_dn: [8.49542, 8.48265, 8.48265, 8.41176], o_dn: [232.672, 233.922, 233.922, 256.059], row_dn: 1.085, col_dn: 1.852} + - {iso: 6400, white: 15035, s_dn: [10.6956, 10.7417, 10.7417, 10.6503], o_dn: [360.311, 368.198, 368.198, 362.848], row_dn: 1.326, col_dn: 2.277} + - {iso: 8000, white: 15035, s_dn: [13.1307, 13.3864, 13.3864, 13.147], o_dn: [615.02, 566.666, 566.666, 611.738], row_dn: 1.768, col_dn: 3.141} + - {iso: 10000, white: 12365, s_dn: [16.5338, 16.7603, 16.7603, 16.4739], o_dn: [914.064, 904.583, 904.583, 938.024], row_dn: 2.214, col_dn: 3.605} + - {iso: 12800, white: 15000, s_dn: [18.4717, 20.9315, 20.9315, 19.3821], o_dn: [1431.85, 1432.33, 1432.33, 1477.82], row_dn: 2.568, col_dn: 4.661} + - {iso: 16000, white: 15000, s_dn: [20.527, 26.0841, 26.0841, 21.8866], o_dn: [2203.77, 2357.78, 2357.78, 2193.1], row_dn: 3.521, col_dn: 5.805} + - {iso: 20000, white: 13000, s_dn: [25.1517, 32.5307, 32.5307, 26.2303], o_dn: [3490.34, 3647.93, 3647.93, 3423.33], row_dn: 4.336, col_dn: 7.143} + - {iso: 25600, white: 15000, s_dn: [22.8743, 40.4641, 40.4641, 23.5537], o_dn: [5184.57, 5690.43, 5690.43, 5286.07], row_dn: 5.682, col_dn: 9.193} diff --git a/models/LICENCE.md b/models/LICENCE.md index e6aaa64..df68400 100644 --- a/models/LICENCE.md +++ b/models/LICENCE.md @@ -122,3 +122,16 @@ position of any model here. The training set is Places2, a research dataset, but the weights are released under the repository's licence without a data-derived restriction (contrast the gaze models §7 of the requirements declined, and the InsightFace grant of D13). + +## `denoise/` — the mosaic denoiser, the project's own + +| File | Source | Trained on | Used by | +|---|---|---|---| +| `denoise/mosaic-1408.onnx` | trained from scratch in the `darkroom-denoise` repository (2026-10-03, run `m2`, 60 000 steps) | 427 of the maintainer's own base-ISO Canon EOS 6D raws, with the 6D's measured noise added | the learned demosaic and denoise (FR-DEV-3g) | + +A U-Net of plain 3×3 convolutions, ReLU, strided and transposed +convolutions and additive skips — no third-party architecture code or +weights — at a fixed `1×1×1408×1408` for `mosaic` and `sigma`, exported by +`python -m denoise.export` in `darkroom-denoise`. Trained only on +photographs the maintainer owns, so the weights carry no grant but the +project's own: GPL-3.0-or-later, like the code (denoise.md §10). diff --git a/models/denoise/mosaic-1408.onnx b/models/denoise/mosaic-1408.onnx new file mode 100644 index 0000000..eb8a22a --- /dev/null +++ b/models/denoise/mosaic-1408.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:02c9b0efdc5fbbf69d947a8e345ef340f51c2d899c2d3d87f9c4efd3c2c77277 +size 12624295