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,79 @@
|
||||
//! 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,
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user