dtourolle/DarkRoom
1
0
Fork 0
Code Issues 39 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Offer three denoise networks and a method to choose between them

Browse Source 
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.
This commit is contained in:
dtourolle
2026-10-04 08:02:25 -04:00
parent 14f08a565f
commit 06422a07db
26 changed files with 540 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-denoise/examples/denoise_raw.rs
+16 -6
View File
@@ -2,11 +2,11 @@
//!
//! ```sh
//! DARKROOM_ORT_DIR=~/.local/share/darkroom/runtime \
//! cargo run --release -p dr-denoise --features native --example denoise_raw -- IMG.CR2 out
//! cargo run --release -p dr-denoise --features native --example denoise_raw -- IMG.CR2 out [fast|medium|best]
//! ```
//!
//! Decode, the app's hot-pixel pass, the frame's noise from its best source,
//! then the shipped network under the inference engine on whatever rung this
//! then one of the shipped networks (`best` unless named) under the inference engine on whatever rung this
//! machine probes to. Writes `out.npy` — the active area, `h×w×3` f32 linear
//! camera RGB — for comparison with the training repo's own path
//! (`tools/compare_rust.py` in darkroom-denoise). `DARKROOM_ORT_DIR` points
@@ -23,11 +23,21 @@ fn main() {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("warn")).init();
let mut args = std::env::args().skip(1);
let (Some(input), Some(out)) = (args.next(), args.next()) else {
eprintln!("usage: denoise_raw RAW OUT_PREFIX");
eprintln!("usage: denoise_raw RAW OUT_PREFIX [fast|medium|best]");
std::process::exit(2);
};
let model =
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../models/denoise/mosaic-1408.onnx");
let shipped = match args.next().as_deref() {
None | Some("best") => dr_denoise::BEST,
Some("medium") => dr_denoise::MEDIUM,
Some("fast") => dr_denoise::FAST,
Some(other) => {
eprintln!("no network called {other}: fast, medium or best");
std::process::exit(2);
}
};
let model = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("../../models/denoise")
.join(shipped.file);
let cache = std::env::var_os("DR_ENGINE_CACHE")
.map(PathBuf::from)
.unwrap_or_else(|| std::env::temp_dir().join("dr-denoise-engines"));
@@ -91,7 +101,7 @@ fn main() {
noise.col
);
let mut net = OnnxNet::from_path(&model).expect("model");
let mut net = OnnxNet::from_path(&model, shipped.halo).expect("model");
println!(
"rung {}",
net.rung().map(|r| r.label()).unwrap_or("?")
core/dr-denoise/src/lib.rs
+27
View File
@@ -27,6 +27,33 @@ 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}")]
core/dr-denoise/src/onnx.rs
+9 -1
View File
@@ -21,15 +21,19 @@ pub const TILE: usize = 1408;
pub struct OnnxNet {
model: Model,
tile: usize,
halo: usize,
}
impl OnnxNet {
pub fn from_path(path: &std::path::Path) -> Result<Self, DenoiseError> {
/// The network at `path`, which needs `halo` photosites of context
/// ([`crate::Shipped::halo`]).
pub fn from_path(path: &std::path::Path, halo: usize) -> 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,
halo,
})
}
@@ -44,6 +48,10 @@ impl TileNet for OnnxNet {
self.tile
}
fn halo(&self) -> usize {
self.halo
}
fn run(
&mut self,
mosaic: Vec<f32>,
core/dr-denoise/src/tile.rs
+18 -11
View File
@@ -15,7 +15,9 @@
use dr_decode::CfaPattern;
/// Photosites of context beyond a tile's kept centre, on every side.
/// Photosites of context beyond a tile's kept centre, on every side, for a
/// single network; a mixture reaches further and says so through
/// [`TileNet::halo`].
pub const HALO: usize = 192;
/// A fixed-shape network: `mosaic` and `sigma`, `n×n` RGGB, in; `3×n×n`
@@ -28,6 +30,11 @@ pub const HALO: usize = 192;
pub trait TileNet {
/// The edge `n` of the square tile the network takes.
fn tile(&self) -> usize;
/// Photosites of context it needs past a tile's kept centre: at least
/// its receptive field. [`HALO`] unless the network says otherwise.
fn halo(&self) -> usize {
HALO
}
fn run(
&mut self,
mosaic: Vec<f32>,
@@ -78,13 +85,13 @@ pub fn run_tiled(
let (dy, dx) = rggb_offset(pattern).ok_or_else(|| {
crate::DenoiseError::Unsupported(format!("{pattern:?} is not a Bayer pattern"))
})?;
let n = net.tile();
if n <= 2 * HALO || !(n - 2 * HALO).is_multiple_of(2) {
let (n, halo) = (net.tile(), net.halo());
if n <= 2 * halo || !(n - 2 * halo).is_multiple_of(2) {
return Err(crate::DenoiseError::Model(format!(
"tile {n} leaves no even centre past a {HALO} halo"
"tile {n} leaves no even centre past a {halo} halo"
)));
}
let core = n - 2 * HALO;
let core = n - 2 * halo;
// In unified coordinates the frame spans u ∈ [dy, dy + h), v ∈ [dx, dx + w).
let (uh, uw) = (h + dy, w + dx);
let (ty, tx) = (uh.div_ceil(core), uw.div_ceil(core));
@@ -108,11 +115,11 @@ pub fn run_tiled(
scope.spawn(move || {
for (i, (mrow, srow)) in m.chunks_mut(n).zip(s.chunks_mut(n)).enumerate() {
let r = chunk * rows_per + i;
// Unified row u = u0 + r − HALO; frame row y = u − dy, reflected.
let u = u0 as isize + r as isize - HALO as isize;
// Unified row u = u0 + r − halo; frame row y = u − dy, reflected.
let u = u0 as isize + r as isize - halo as isize;
let y = reflect(u - dy as isize, h);
for c in 0..n {
let v = v0 as isize + c as isize - HALO as isize;
let v = v0 as isize + c as isize - halo as isize;
let x = reflect(v - dx as isize, w);
let val = at(y, x);
mrow[c] = val;
@@ -169,10 +176,10 @@ pub fn run_tiled(
scope.spawn(move || {
for (i, row) in block.chunks_mut(w * 3).enumerate() {
let y = y_lo + chunk * per + i;
// Tile row of frame row y: u = y + dy = u0 + r − HALO.
let r = y + dy + HALO - u0;
// Tile row of frame row y: u = y + dy = u0 + r − halo.
let r = y + dy + halo - u0;
for x in x_lo..x_hi {
let c = x + dx + HALO - v0;
let c = x + dx + halo - v0;
for ch in 0..3 {
row[x * 3 + ch] = rgb[ch * n * n + r * n + c];
}