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.
538 lines
19 KiB
Rust
538 lines
19 KiB
Rust
//! TRACES: FR-DEV-3g
|
||
//! The learned denoise in a develop session (docs/dev/denoise.md §7).
|
||
//!
|
||
//! The classical demosaic shows at once; when the photograph asks for the
|
||
//! learned one, it is computed off the UI thread from the mosaic the
|
||
//! session kept, and swapped in when it lands. Grain is a blend of the two
|
||
//! results, made once per slider value by [`dr_gpu::GrainBlend`] and handed
|
||
//! to the render as its source — the adjust pass never knows.
|
||
//!
|
||
//! `demosaiced` stays the classical result for the session's life: the raw
|
||
//! histogram, the white balance picker, masks and segmentation measure the
|
||
//! sensor data, and only [`DevelopSession::developed_source`] — what the
|
||
//! render draws — changes.
|
||
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
use std::sync::mpsc;
|
||
use std::sync::Arc;
|
||
|
||
use dr_decode::RawImage;
|
||
use dr_gpu::{DemosaicedImage, GrainBlend};
|
||
use dr_pipeline::learned_denoise::Method;
|
||
|
||
use super::session::DevelopSession;
|
||
|
||
/// What the session holds for the learned denoise.
|
||
#[derive(Default)]
|
||
pub(crate) struct DenoiseState {
|
||
/// The mosaic as decoded, kept only for a photograph that can take the
|
||
/// learned stage. The job repairs a copy.
|
||
mosaic: Option<Arc<RawImage>>,
|
||
/// The file's `NoiseProfile` and ISO, read from the header at open.
|
||
profile: Option<Vec<(f32, f32)>>,
|
||
iso: Option<u32>,
|
||
/// The network's result, once it has landed.
|
||
result: Option<Arc<DemosaicedImage>>,
|
||
/// The last grain blend made, and the grain it was made at.
|
||
blended: Option<(f32, Arc<DemosaicedImage>)>,
|
||
blend: Option<GrainBlend>,
|
||
job: Option<Job>,
|
||
/// Why the last attempt failed; not retried until the switch is
|
||
/// toggled, so a photograph that cannot be denoised does not loop.
|
||
failed: Option<String>,
|
||
/// Where the noise figures came from, for the panel.
|
||
source: Option<dr_denoise::Source>,
|
||
/// 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 {
|
||
rx: mpsc::Receiver<Msg>,
|
||
cancel: Arc<AtomicBool>,
|
||
}
|
||
|
||
enum Msg {
|
||
Progress(usize, usize),
|
||
Done(Result<Finished, String>),
|
||
}
|
||
|
||
struct Finished {
|
||
rgb: Vec<f32>,
|
||
width: u32,
|
||
height: u32,
|
||
source: dr_denoise::Source,
|
||
rung: String,
|
||
seconds: f64,
|
||
}
|
||
|
||
/// What a poll found, for the develop view's status line and redraw.
|
||
#[derive(Debug, PartialEq)]
|
||
pub enum DenoiseStatus {
|
||
/// Nothing running and nothing new.
|
||
Idle,
|
||
/// Tiles done of tiles.
|
||
Running(usize, usize),
|
||
/// The result landed this poll: redraw.
|
||
Landed,
|
||
Failed(String),
|
||
}
|
||
|
||
impl DevelopSession {
|
||
/// Keep the mosaic for the learned denoise, if it can take this frame.
|
||
pub(super) fn keep_mosaic(&mut self, raw: RawImage) {
|
||
if dr_denoise::eligible(&raw) {
|
||
self.denoise.mosaic = Some(Arc::new(raw));
|
||
}
|
||
self.graph
|
||
.set_denoise_available(self.denoise.mosaic.is_some());
|
||
}
|
||
|
||
/// The header's part: the DNG's measured noise and the ISO.
|
||
pub fn prepare_denoise(&mut self, bytes: &[u8], meta: &dr_decode::Metadata) {
|
||
if self.denoise.mosaic.is_some() {
|
||
self.denoise.profile = dr_decode::noise_profile(bytes);
|
||
self.denoise.iso = meta.iso;
|
||
// TRACES: FR-DEV-3g
|
||
// 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));
|
||
}
|
||
}
|
||
|
||
/// Bring what is computed in line with what the edit asks for: start the
|
||
/// network when it is wanted and has not run, stop it when it is not.
|
||
/// 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() {
|
||
job.cancel.store(true, Ordering::Relaxed);
|
||
}
|
||
// Toggling off is how a failure is retried.
|
||
self.denoise.failed = None;
|
||
return;
|
||
}
|
||
if self.denoise.result.is_some()
|
||
|| self.denoise.job.is_some()
|
||
|| self.denoise.failed.is_some()
|
||
{
|
||
return;
|
||
}
|
||
let Some(work) = self.work(Arc::new(AtomicBool::new(false))) else {
|
||
return;
|
||
};
|
||
let cancel = work.cancel.clone();
|
||
let (tx, rx) = mpsc::channel();
|
||
crate::executors::spawn(crate::executors::Executor::Decode, "denoise", move || {
|
||
let result = work.run(&mut |done, total| {
|
||
let _ = tx.send(Msg::Progress(done, total));
|
||
});
|
||
let _ = tx.send(Msg::Done(result));
|
||
});
|
||
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 {
|
||
return DenoiseStatus::Idle;
|
||
};
|
||
let mut last = None;
|
||
let mut done = None;
|
||
while let Ok(msg) = job.rx.try_recv() {
|
||
match msg {
|
||
Msg::Progress(d, t) => last = Some((d, t)),
|
||
Msg::Done(r) => done = Some(r),
|
||
}
|
||
}
|
||
match done {
|
||
Some(Ok(f)) => {
|
||
self.denoise.job = None;
|
||
match self.land(f) {
|
||
Ok(()) => DenoiseStatus::Landed,
|
||
Err(e) => self.fail(e),
|
||
}
|
||
}
|
||
Some(Err(e)) => {
|
||
self.denoise.job = None;
|
||
self.fail(e)
|
||
}
|
||
None => last.map_or(DenoiseStatus::Running(0, 1), |(d, t)| {
|
||
DenoiseStatus::Running(d, t)
|
||
}),
|
||
}
|
||
}
|
||
|
||
/// Compute the result now, on this thread, if the edit wants it and it
|
||
/// is not here: for an export, which must not write the classical
|
||
/// picture of a photograph that asks for the learned one (§7.1).
|
||
pub fn denoise_blocking(&mut self) -> Result<(), String> {
|
||
if !self.graph.denoise_applied() || self.denoise.result.is_some() {
|
||
return Ok(());
|
||
}
|
||
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() {
|
||
if let Msg::Done(result) = msg {
|
||
return result.and_then(|f| self.land(f));
|
||
}
|
||
}
|
||
}
|
||
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)
|
||
}
|
||
|
||
/// What the render draws: the learned result with the asked-for grain
|
||
/// where there is one, else the classical demosaic.
|
||
pub(super) fn developed_source(&mut self) -> Arc<DemosaicedImage> {
|
||
if !self.graph.denoise_applied() {
|
||
return self.demosaiced.clone();
|
||
}
|
||
let Some(result) = self.denoise.result.clone() else {
|
||
return self.demosaiced.clone();
|
||
};
|
||
let grain = self.graph.denoise_grain();
|
||
if grain <= 0.0 {
|
||
return result;
|
||
}
|
||
if let Some((g, image)) = &self.denoise.blended {
|
||
if (*g - grain).abs() < 1e-4 {
|
||
return image.clone();
|
||
}
|
||
}
|
||
let blend = self
|
||
.denoise
|
||
.blend
|
||
.get_or_insert_with(|| GrainBlend::new(&self.ctx));
|
||
match blend.blend(&result, &self.demosaiced, grain) {
|
||
Ok(image) => {
|
||
self.denoise.blended = Some((grain, image.clone()));
|
||
image
|
||
}
|
||
Err(e) => {
|
||
log::warn!("grain blend failed, showing the denoised result without grain: {e}");
|
||
result
|
||
}
|
||
}
|
||
}
|
||
|
||
/// For the panel: where the noise figures came from, once computed.
|
||
pub fn denoise_source(&self) -> Option<dr_denoise::Source> {
|
||
self.denoise.source
|
||
}
|
||
|
||
pub fn denoise_failure(&self) -> Option<&str> {
|
||
self.denoise.failed.as_deref()
|
||
}
|
||
|
||
fn land(&mut self, f: Finished) -> Result<(), String> {
|
||
let image =
|
||
DemosaicedImage::from_rgb_f32(&self.ctx, &self.demosaiced, f.width, f.height, &f.rgb)
|
||
.map_err(|e| e.to_string())?;
|
||
log::info!(
|
||
"learned denoise: {}×{} on {} in {:.1} s, noise {}",
|
||
f.width,
|
||
f.height,
|
||
f.rung,
|
||
f.seconds,
|
||
f.source.label()
|
||
);
|
||
self.denoise.result = Some(Arc::new(image));
|
||
self.denoise.blended = None;
|
||
self.denoise.source = Some(f.source);
|
||
Ok(())
|
||
}
|
||
|
||
fn fail(&mut self, e: String) -> DenoiseStatus {
|
||
log::warn!("learned denoise failed: {e}");
|
||
self.denoise.failed = Some(e.clone());
|
||
DenoiseStatus::Failed(e)
|
||
}
|
||
}
|
||
|
||
/// Everything the job needs, owned, so it can leave the UI thread.
|
||
struct Work {
|
||
ctx: dr_gpu::GpuContext,
|
||
raw: Arc<RawImage>,
|
||
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>,
|
||
}
|
||
|
||
impl Work {
|
||
fn run(self, progress: &mut dyn FnMut(usize, usize)) -> Result<Finished, String> {
|
||
let started = std::time::Instant::now();
|
||
// TRACES: FR-DEV-3g
|
||
// A result computed before — this photograph opened earlier, or
|
||
// developed and now exported — is read back rather than recomputed.
|
||
let cache_dir = super::denoise_cache::dir();
|
||
if let Some(hit) = self
|
||
.cache_key
|
||
.as_deref()
|
||
.and_then(|key| super::denoise_cache::load(&cache_dir, key))
|
||
{
|
||
return Ok(Finished {
|
||
rgb: hit.rgb,
|
||
width: hit.width,
|
||
height: hit.height,
|
||
source: hit.source,
|
||
rung: "the cache".into(),
|
||
seconds: started.elapsed().as_secs_f64(),
|
||
});
|
||
}
|
||
// The app's own hot-pixel pass, on a copy: the classical source was
|
||
// repaired by the same pass inside `Demosaicer::run`.
|
||
let mut raw = (*self.raw).clone();
|
||
dr_gpu::Demosaicer::new(&self.ctx)
|
||
.and_then(|d| d.repair_hot_pixels(&mut raw))
|
||
.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, self.halo)
|
||
.map_err(|e| e.to_string())?;
|
||
let rung = net
|
||
.rung()
|
||
.map(|r| r.label().to_string())
|
||
.unwrap_or_default();
|
||
let cancel = self.cancel;
|
||
let rgb = dr_denoise::denoise(&raw, &noise, &mut net, &mut |done, total| {
|
||
progress(done, total);
|
||
!cancel.load(Ordering::Relaxed)
|
||
})
|
||
.map_err(|e| e.to_string())?
|
||
.ok_or("stopped")?;
|
||
let finished = Finished {
|
||
rgb,
|
||
width: raw.crop.width,
|
||
height: raw.crop.height,
|
||
source: noise.source,
|
||
rung,
|
||
seconds: started.elapsed().as_secs_f64(),
|
||
};
|
||
if let Some(key) = self.cache_key.as_deref() {
|
||
let entry = super::denoise_cache::Entry {
|
||
rgb: finished.rgb.clone(),
|
||
width: finished.width,
|
||
height: finished.height,
|
||
source: finished.source,
|
||
};
|
||
super::denoise_cache::store(&cache_dir, key, &entry, super::denoise_cache::BUDGET);
|
||
}
|
||
Ok(finished)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::develop::test_support::headless;
|
||
use dr_pipeline::learned_denoise;
|
||
|
||
fn bayer(w: u32, h: u32) -> RawImage {
|
||
RawImage {
|
||
width: w,
|
||
height: h,
|
||
data: vec![800; (w * h) as usize],
|
||
cfa_pattern: dr_decode::CfaPattern::Rggb,
|
||
black_level: [0; 4],
|
||
white_level: 4095,
|
||
wb_coeffs: [2.0, 1.0, 1.5, 1.0],
|
||
color_matrix: Some([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]),
|
||
samples_per_pixel: 1,
|
||
profile: None,
|
||
profile_tables: None,
|
||
baseline_exposure: 0.0,
|
||
make: String::new(),
|
||
model: String::new(),
|
||
crop: dr_decode::CropRect {
|
||
x: 0,
|
||
y: 0,
|
||
width: w,
|
||
height: h,
|
||
},
|
||
}
|
||
}
|
||
|
||
fn landed(session: &mut DevelopSession) {
|
||
let (w, h) = (64, 64);
|
||
session
|
||
.land(Finished {
|
||
rgb: vec![0.2; (w * h * 3) as usize],
|
||
width: w,
|
||
height: h,
|
||
source: dr_denoise::Source::Measured,
|
||
rung: "test".into(),
|
||
seconds: 0.0,
|
||
})
|
||
.expect("upload");
|
||
}
|
||
|
||
#[test]
|
||
fn a_bayer_raw_offers_the_switch_and_an_rgb_image_does_not() {
|
||
let Some(ctx) = headless() else { return };
|
||
let raw =
|
||
DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
|
||
assert!(raw
|
||
.graph
|
||
.capabilities()
|
||
.iter()
|
||
.any(|c| c.id == learned_denoise::ID));
|
||
let rgba: Vec<u8> = (0..64 * 64).flat_map(|_| [128u8, 128, 128, 255]).collect();
|
||
let jpeg =
|
||
DevelopSession::open_rgb(&ctx, &rgba, 64, 64, dr_types::Orientation::NORMAL).unwrap();
|
||
assert!(!jpeg
|
||
.graph
|
||
.capabilities()
|
||
.iter()
|
||
.any(|c| c.id == learned_denoise::ID));
|
||
}
|
||
|
||
#[test]
|
||
fn the_render_draws_what_the_edit_asks_for() {
|
||
let Some(ctx) = headless() else { return };
|
||
let mut s =
|
||
DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
|
||
let classical = s.demosaiced.clone();
|
||
landed(&mut s);
|
||
let result = s.denoise.result.clone().unwrap();
|
||
let same = |a: &Arc<DemosaicedImage>, b: &Arc<DemosaicedImage>| Arc::ptr_eq(a, b);
|
||
|
||
// On by default, at full strength: the network's result as it is.
|
||
assert!(same(&s.developed_source(), &result));
|
||
// 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::METHOD,
|
||
Method::DEFAULT.index(),
|
||
);
|
||
// Less strength: a blend, made once per value and reused until it
|
||
// moves.
|
||
s.graph
|
||
.set_param(learned_denoise::ID, learned_denoise::STRENGTH, 60.0);
|
||
let blended = s.developed_source();
|
||
assert!(!same(&blended, &result) && !same(&blended, &classical));
|
||
assert!(
|
||
same(&s.developed_source(), &blended),
|
||
"the same grain must not blend again"
|
||
);
|
||
s.graph
|
||
.set_param(learned_denoise::ID, learned_denoise::STRENGTH, 40.0);
|
||
assert!(!same(&s.developed_source(), &blended));
|
||
// The sensor's own reading stays the classical one throughout.
|
||
assert!(same(&s.demosaiced, &classical));
|
||
}
|
||
|
||
#[test]
|
||
fn switching_off_stops_and_forgets_a_failure() {
|
||
let Some(ctx) = headless() else { return };
|
||
let mut s =
|
||
DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
|
||
s.denoise.failed = Some("no model".into());
|
||
s.reconcile_denoise();
|
||
assert!(
|
||
s.denoise.job.is_none(),
|
||
"a failure is not retried while the method stays"
|
||
);
|
||
s.graph.set_param(
|
||
learned_denoise::ID,
|
||
learned_denoise::METHOD,
|
||
Method::Bilinear.index(),
|
||
);
|
||
s.reconcile_denoise();
|
||
assert!(
|
||
s.denoise.failed.is_none(),
|
||
"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));
|
||
}
|
||
}
|