dtourolle
d8304d7c82
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.
67 lines
2.3 KiB
Rust
67 lines
2.3 KiB
Rust
//! 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<Self, DenoiseError> {
|
||
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<dr_inference_engine::Rung, DenoiseError> {
|
||
Ok(self.model.acquire()?.rung())
|
||
}
|
||
}
|
||
|
||
impl TileNet for OnnxNet {
|
||
fn tile(&self) -> usize {
|
||
self.tile
|
||
}
|
||
|
||
fn run(&mut self, mosaic: &[f32], sigma: &[f32]) -> Result<Vec<f32>, 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::<f32>()?;
|
||
let dims: Vec<i64> = 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())
|
||
}
|
||
}
|