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
ui/dr-ui/src/develop/denoise.rs
+107 -43
View File
@@ -18,6 +18,7 @@ use std::sync::Arc;
use dr_decode::RawImage;
use dr_gpu::{DemosaicedImage, GrainBlend};
use dr_pipeline::learned_denoise::Method;
use super::session::DevelopSession;
@@ -41,9 +42,12 @@ pub(crate) struct DenoiseState {
failed: Option<String>,
/// Where the noise figures came from, for the panel.
source: Option<dr_denoise::Source>,
/// The result's key in the on-disk cache: the file's bytes and the model
/// (`denoise_cache::key`). `None` where there is no model to key on.
cache_key: Option<String>,
/// The file's bytes, hashed for the on-disk cache, which keys each
/// network's result on them and the model (`denoise_cache::FileHash`).
file_hash: Option<super::denoise_cache::FileHash>,
/// The method the result, the job and the failure above are for. A
/// different one asked for discards them.
method: Option<Method>,
}
struct Job {
@@ -93,10 +97,10 @@ impl DevelopSession {
self.denoise.profile = dr_decode::noise_profile(bytes);
self.denoise.iso = meta.iso;
// TRACES: FR-DEV-3g
// The bytes are only here now, so the cache key is made now: a
// reopened or exported photograph finds its result on disk.
self.denoise.cache_key = crate::library::denoise_model()
.map(|model| super::denoise_cache::key(bytes, &model));
// The bytes are only here now, so they are hashed now: a
// reopened or exported photograph finds its result on disk,
// under whichever method it asks for.
self.denoise.file_hash = Some(super::denoise_cache::FileHash::of(bytes));
}
}
@@ -105,6 +109,7 @@ impl DevelopSession {
/// Cheap when nothing changed; the develop view calls it after every
/// change to the edit, whatever made it — a slider, undo, a version.
pub fn reconcile_denoise(&mut self) {
self.forget_other_method();
let wanted = self.graph.denoise_applied() && self.denoise.mosaic.is_some();
if !wanted {
if let Some(job) = self.denoise.job.take() {
@@ -120,21 +125,11 @@ impl DevelopSession {
{
return;
}
let Some(model) = crate::library::denoise_model() else {
self.denoise.failed = Some("the denoise model is not installed".into());
let Some(work) = self.work(Arc::new(AtomicBool::new(false))) else {
return;
};
let cancel = work.cancel.clone();
let (tx, rx) = mpsc::channel();
let cancel = Arc::new(AtomicBool::new(false));
let work = Work {
ctx: self.ctx.clone(),
raw: self.denoise.mosaic.clone().expect("checked above"),
profile: self.denoise.profile.clone(),
iso: self.denoise.iso,
model,
cancel: cancel.clone(),
cache_key: self.denoise.cache_key.clone(),
};
crate::executors::spawn(crate::executors::Executor::Decode, "denoise", move || {
let result = work.run(&mut |done, total| {
let _ = tx.send(Msg::Progress(done, total));
@@ -144,6 +139,51 @@ impl DevelopSession {
self.denoise.job = Some(Job { rx, cancel });
}
/// Drop what was computed, or is being computed, for a network the edit
/// no longer asks for — a choice in the panel, an undo, a version.
/// Each network's result stays in the on-disk cache, so going back to
/// one is a read, not a run. The classical demosaic asks for no network
/// and drops nothing: the result is kept for the way back.
fn forget_other_method(&mut self) {
let asked = self.graph.denoise_method();
if !asked.learned() || self.denoise.method == Some(asked) {
return;
}
if self.denoise.method.is_none() {
// The first network asked for: nothing computed is another's.
self.denoise.method = Some(asked);
return;
}
if let Some(job) = self.denoise.job.take() {
job.cancel.store(true, Ordering::Relaxed);
}
self.denoise.result = None;
self.denoise.blended = None;
self.denoise.failed = None;
self.denoise.source = None;
self.denoise.method = Some(asked);
}
/// The job for the method the edit asks for, or `None` — with the reason
/// kept as the failure — where its network is not installed.
fn work(&mut self, cancel: Arc<AtomicBool>) -> Option<Work> {
let raw = self.denoise.mosaic.clone()?;
let Some((model, net)) = crate::library::denoise_model(self.graph.denoise_method()) else {
self.denoise.failed = Some("the denoise model is not installed".into());
return None;
};
Some(Work {
ctx: self.ctx.clone(),
raw,
profile: self.denoise.profile.clone(),
iso: self.denoise.iso,
cache_key: self.denoise.file_hash.as_ref().map(|h| h.key(&model)),
model,
halo: net.halo,
cancel,
})
}
/// Collect what the job sent since the last poll.
pub fn poll_denoise(&mut self) -> DenoiseStatus {
let Some(job) = &self.denoise.job else {
@@ -182,9 +222,10 @@ impl DevelopSession {
if !self.graph.denoise_applied() || self.denoise.result.is_some() {
return Ok(());
}
let Some(raw) = self.denoise.mosaic.clone() else {
self.forget_other_method();
if self.denoise.mosaic.is_none() {
return Ok(());
};
}
// Already under way: wait for it rather than start again.
if let Some(job) = self.denoise.job.take() {
for msg in job.rx.iter() {
@@ -193,15 +234,8 @@ impl DevelopSession {
}
}
}
let model = crate::library::denoise_model().ok_or("the denoise model is not installed")?;
let work = Work {
ctx: self.ctx.clone(),
raw,
profile: self.denoise.profile.clone(),
iso: self.denoise.iso,
model,
cancel: Arc::new(AtomicBool::new(false)),
cache_key: self.denoise.cache_key.clone(),
let Some(work) = self.work(Arc::new(AtomicBool::new(false))) else {
return Err(self.denoise.failed.clone().unwrap_or_default());
};
let finished = work.run(&mut |_, _| {})?;
self.land(finished)
@@ -282,6 +316,8 @@ struct Work {
profile: Option<Vec<(f32, f32)>>,
iso: Option<u32>,
model: std::path::PathBuf,
/// The context `model` needs past a tile's kept centre.
halo: usize,
cancel: Arc<AtomicBool>,
cache_key: Option<String>,
}
@@ -315,8 +351,8 @@ impl Work {
.map_err(|e| e.to_string())?;
let noise = dr_denoise::noise::for_frame_with(&raw, self.profile.as_deref(), self.iso)
.ok_or("this photograph gives no way to measure its noise")?;
let mut net =
dr_denoise::onnx::OnnxNet::from_path(&self.model).map_err(|e| e.to_string())?;
let mut net = dr_denoise::onnx::OnnxNet::from_path(&self.model, self.halo)
.map_err(|e| e.to_string())?;
let rung = net
.rung()
.map(|r| r.label().to_string())
@@ -426,12 +462,18 @@ mod tests {
// On by default, at full strength: the network's result as it is.
assert!(same(&s.developed_source(), &result));
// Off: the classical demosaic, result or no result.
s.graph
.set_param(learned_denoise::ID, learned_denoise::APPLY, 0.0);
// Bilinear: the classical demosaic, result or no result.
s.graph.set_param(
learned_denoise::ID,
learned_denoise::METHOD,
Method::Bilinear.index(),
);
assert!(same(&s.developed_source(), &classical));
s.graph
.set_param(learned_denoise::ID, learned_denoise::APPLY, 1.0);
s.graph.set_param(
learned_denoise::ID,
learned_denoise::METHOD,
Method::DEFAULT.index(),
);
// Less strength: a blend, made once per value and reused until it
// moves.
s.graph
@@ -455,19 +497,41 @@ mod tests {
let mut s =
DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
s.denoise.failed = Some("no model".into());
s.graph
.set_param(learned_denoise::ID, learned_denoise::APPLY, 1.0);
s.reconcile_denoise();
assert!(
s.denoise.job.is_none(),
"a failure is not retried while the switch stays on"
"a failure is not retried while the method stays"
);
s.graph.set_param(
learned_denoise::ID,
learned_denoise::METHOD,
Method::Bilinear.index(),
);
s.graph
.set_param(learned_denoise::ID, learned_denoise::APPLY, 0.0);
s.reconcile_denoise();
assert!(
s.denoise.failed.is_none(),
"toggling off is how a failure is retried"
"choosing bilinear is how a failure is retried"
);
}
#[test]
fn another_network_discards_the_result_and_bilinear_keeps_it() {
let Some(ctx) = headless() else { return };
let mut s =
DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
s.denoise.method = Some(Method::DEFAULT);
landed(&mut s);
let to = |s: &mut DevelopSession, m: Method| {
s.graph
.set_param(learned_denoise::ID, learned_denoise::METHOD, m.index());
s.forget_other_method();
};
to(&mut s, Method::Bilinear);
assert!(s.denoise.result.is_some(), "kept for the way back");
to(&mut s, Method::DEFAULT);
assert!(s.denoise.result.is_some());
to(&mut s, Method::Fast);
assert!(s.denoise.result.is_none(), "another network's picture");
assert_eq!(s.denoise.method, Some(Method::Fast));
}
}