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:
@@ -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<u32>) -> Option<NoiseModel> {
|
||||
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<TableRow>,
|
||||
}
|
||||
|
||||
#[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<NoiseModel> {
|
||||
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<NoiseModel> {
|
||||
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<Dark> {
|
||||
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::<f32>() / 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<f32> = 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::<f32>() / h as f32;
|
||||
let row_var = row_means.iter().map(|m| (m - rm).powi(2)).sum::<f32>() / 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<f32> = (x0..x0 + w)
|
||||
.map(|x| {
|
||||
rows.clone()
|
||||
.map(|y| raw.data[y * stride + x] as f32)
|
||||
.sum::<f32>()
|
||||
/ nr
|
||||
})
|
||||
.collect();
|
||||
let mm = means.iter().sum::<f32>() / means.len() as f32;
|
||||
let var = means.iter().map(|m| (m - mm).powi(2)).sum::<f32>() / 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<NoiseModel> {
|
||||
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<f32> = 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::<f32>() / 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,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user