Files
DarkRoom/core/dr-denoise/src/lib.rs
T
dtourolle d8304d7c82 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.
2026-10-03 11:15:50 -04:00

80 lines
2.7 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! TRACES: FR-DEV-3g
//! Learned demosaic and denoise on the raw mosaic (docs/dev/denoise.md).
//!
//! A network trained on the library's own base-ISO raws with the 6D's
//! measured noise added takes the repaired, normalised mosaic and a σ for
//! every photosite, and returns linear camera RGB at full resolution — the
//! texture the classical demosaic would have produced, with the noise gone.
//! It replaces the demosaic box; nothing downstream changes (§2).
//!
//! - [`noise`] says how noisy each photosite is, from the best source the
//! frame has.
//! - [`tile`] runs a fixed-shape network over a whole frame, exactly.
//! - [`onnx`] is that network under the inference engine.
//!
//! The input must already have been through the app's hot-pixel pass
//! (`dr_gpu::Demosaicer::repair_hot_pixels`): the noise model was fitted
//! with what that pass removes left out.
pub mod noise;
#[cfg(feature = "onnx")]
pub mod onnx;
pub mod tile;
use dr_decode::RawImage;
pub use noise::{NoiseModel, Source};
pub use tile::{TileNet, HALO};
#[derive(Debug, thiserror::Error)]
pub enum DenoiseError {
#[error("the network cannot take this photograph: {0}")]
Unsupported(String),
#[error("the denoise model misbehaved: {0}")]
Model(String),
#[error("could not read the denoise model: {0}")]
ModelRead(#[from] std::io::Error),
#[cfg(feature = "onnx")]
#[error(transparent)]
Engine(#[from] dr_inference_engine::Error),
#[cfg(feature = "onnx")]
#[error(transparent)]
Ort(#[from] ort::Error),
}
/// Whether the learned stage can take this frame at all: a Bayer mosaic.
/// X-Trans needs its own model (§9); a linear DNG has no photosites.
pub fn eligible(raw: &RawImage) -> 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<Option<Vec<f32>>, 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,
)
}