dtourolle
06422a07db
AI Denoise's Apply switch becomes Method: Bilinear, Fast, Medium, Best, default Best, so an untouched raw writes nothing and develops through the mixture. `apply` is still read and never written: 0 is Bilinear, 1 keeps a network already chosen. - Best is the mixture of a flat and an edge expert with a learned gate; Medium and Fast are students distilled from it. 2.48 s, 0.79 s and 0.57 s for a 20 MP frame on TensorRT fp16. - Each network carries its own tile border (256 for the mixture, 192 for the students) through `dr_denoise::Shipped` and `TileNet::halo`. - The file is hashed once at open and each network keys its own cached result; Bilinear keeps the result in memory for the way back. - Each has an .a16w16 sibling for the Hexagon: 0.00 dB on the 6D gate, at most 0.11 dB with the noise scaled x0.5 to x4. - APK BUNDLED 19 -> 23; the PKGBUILD installs all three.
119 lines
4.1 KiB
Rust
119 lines
4.1 KiB
Rust
//! 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 repair;
|
||
pub mod tile;
|
||
|
||
use dr_decode::RawImage;
|
||
|
||
pub use noise::{NoiseModel, Source};
|
||
pub use tile::{TileNet, HALO};
|
||
|
||
/// TRACES: FR-DEV-3g
|
||
/// A network the app ships in `models/denoise/`: its file, and the context
|
||
/// it needs past a tile's kept centre (docs/dev/denoise.md §13).
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub struct Shipped {
|
||
pub file: &'static str,
|
||
pub halo: usize,
|
||
}
|
||
|
||
/// The smallest student: 0.9 M parameters, 11 GMAC a megapixel.
|
||
pub const FAST: Shipped = Shipped {
|
||
file: "mosaic-fast-1408.onnx",
|
||
halo: HALO,
|
||
};
|
||
/// A student of the mixture with the first release's shape: 3.2 M
|
||
/// parameters, 48 GMAC a megapixel.
|
||
pub const MEDIUM: Shipped = Shipped {
|
||
file: "mosaic-medium-1408.onnx",
|
||
halo: HALO,
|
||
};
|
||
/// The mixture: a flat expert, an edge expert and the gate that blends them.
|
||
/// It reaches further than either, so it keeps a smaller centre of each tile.
|
||
pub const BEST: Shipped = Shipped {
|
||
file: "mosaic-best-1408.onnx",
|
||
halo: 256,
|
||
};
|
||
|
||
#[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);
|
||
// The active area laid out once, then the noise-aware repair the model
|
||
// was trained behind (see `repair`).
|
||
let mut mosaic: Vec<f32> = (0..h * w).map(|i| active.at(i / w, i % w)).collect();
|
||
let pattern = raw.cfa_pattern;
|
||
let repaired = repair::repair(&mut mosaic, h, w, repair::REPAIR_K, &|y, x, v| {
|
||
noise.sigma(pattern.colour_at(x as u32, y as u32) as usize, v)
|
||
});
|
||
log::info!(
|
||
"learned denoise: {repaired} photosites beyond {}σ of every neighbour repaired",
|
||
repair::REPAIR_K
|
||
);
|
||
tile::run_tiled(
|
||
net,
|
||
h,
|
||
w,
|
||
raw.cfa_pattern,
|
||
&|y, x| mosaic[y * w + x],
|
||
&|c, v| noise.sigma(c, v),
|
||
progress,
|
||
)
|
||
}
|