dtourolle
b69fb3e191
The learned-denoise branch merged while this one was open, and two of its test fixtures build a RawImage without the baseline_exposure field this branch added; zero is the no-op value.
432 lines
15 KiB
Rust
432 lines
15 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 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>,
|
||
}
|
||
|
||
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;
|
||
}
|
||
}
|
||
|
||
/// 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) {
|
||
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(model) = crate::library::denoise_model() else {
|
||
self.denoise.failed = Some("the denoise model is not installed".into());
|
||
return;
|
||
};
|
||
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(),
|
||
};
|
||
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 });
|
||
}
|
||
|
||
/// 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(());
|
||
}
|
||
let Some(raw) = self.denoise.mosaic.clone() else {
|
||
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 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)),
|
||
};
|
||
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,
|
||
cancel: Arc<AtomicBool>,
|
||
}
|
||
|
||
impl Work {
|
||
fn run(self, progress: &mut dyn FnMut(usize, usize)) -> Result<Finished, String> {
|
||
let started = std::time::Instant::now();
|
||
// 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).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")?;
|
||
Ok(Finished {
|
||
rgb,
|
||
width: raw.crop.width,
|
||
height: raw.crop.height,
|
||
source: noise.source,
|
||
rung,
|
||
seconds: started.elapsed().as_secs_f64(),
|
||
})
|
||
}
|
||
}
|
||
|
||
#[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);
|
||
|
||
// Off: the classical demosaic, result or no result.
|
||
assert!(same(&s.developed_source(), &classical));
|
||
// On, no grain: the network's result as it is.
|
||
s.graph
|
||
.set_param(learned_denoise::ID, learned_denoise::APPLY, 1.0);
|
||
assert!(same(&s.developed_source(), &result));
|
||
// Grain: a blend, made once per value and reused until it moves.
|
||
s.graph
|
||
.set_param(learned_denoise::ID, learned_denoise::GRAIN, 40.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::GRAIN, 60.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.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"
|
||
);
|
||
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"
|
||
);
|
||
}
|
||
}
|