Run the learned denoise in develop, and export with it
A Bayer photograph keeps its mosaic in the session and is offered the AI Denoise switch. Asked for, the network runs on the decode executor from a hot-pixel-repaired copy — the app's own pass — with the frame's noise from its best source, and its progress in the activity bar; the classical demosaic shows until the result lands, and the finished job says where the noise figures came from. Keep grain is a GrainBlend of the two, made once per value; the render draws it as its source and the adjust pass never knows. demosaiced stays the classical result, so the raw histogram, the white balance picker, masks and segmentation still read the sensor. The develop view reconciles on a 250 ms poll rather than on each way an edit can change (slider, undo, preset, version, a sidecar from another device): two comparisons when nothing changed, and no path that can forget. A failure is not retried until the switch is toggled. An export of a photograph that asks for it waits for a running job or computes it.
This commit is contained in:
@@ -0,0 +1,430 @@
|
||||
//! 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,
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,7 @@
|
||||
//! only methods were ever exported.
|
||||
|
||||
mod curves;
|
||||
mod denoise;
|
||||
mod framing;
|
||||
mod history;
|
||||
mod mask_ops;
|
||||
@@ -29,6 +30,7 @@ mod session;
|
||||
mod tabs;
|
||||
mod white_balance;
|
||||
|
||||
pub use denoise::DenoiseStatus;
|
||||
pub use framing::CropAspect;
|
||||
pub use masks::MASK_COLOURS;
|
||||
pub use segmentation::{Abandon, RefinedInstance, Segmented, SessionId};
|
||||
|
||||
@@ -235,7 +235,10 @@ impl DevelopSession {
|
||||
reach: u32,
|
||||
) -> Result<std::sync::Arc<dr_gpu::DemosaicedImage>, String> {
|
||||
let Some(full) = self.full.clone() else {
|
||||
return Ok(self.demosaiced.clone());
|
||||
// TRACES: FR-DEV-3g
|
||||
// The learned demosaic, with its grain, where the edit asks for
|
||||
// it and it has landed; the classical one otherwise.
|
||||
return Ok(self.developed_source());
|
||||
};
|
||||
let frame = self.demosaiced.size();
|
||||
let ratio = self.graph.render_scale(frame, (w, h)).ratio();
|
||||
@@ -706,6 +709,10 @@ impl DevelopSession {
|
||||
&mut self,
|
||||
space: dr_types::ColourSpace,
|
||||
) -> Result<dr_export::Frame, String> {
|
||||
// TRACES: FR-DEV-3g
|
||||
// A photograph that asks for the learned demosaic is exported with
|
||||
// it, computed now if it is not here (denoise.md §7.1).
|
||||
self.denoise_blocking()?;
|
||||
let (sw, sh) = self.demosaiced.size();
|
||||
let (w, h) = self.graph.output_size(sw, sh);
|
||||
|
||||
|
||||
@@ -118,6 +118,10 @@ pub struct DevelopSession {
|
||||
/// history still stands on that crop. See [`super::framing::CropNotice`].
|
||||
pub(super) crop_notice: Option<super::framing::CropNotice>,
|
||||
pub(super) demosaiced: Arc<DemosaicedImage>,
|
||||
/// TRACES: FR-DEV-3g
|
||||
/// The learned demosaic: the kept mosaic, the job, the result. See
|
||||
/// [`super::denoise`].
|
||||
pub(super) denoise: super::denoise::DenoiseState,
|
||||
/// TRACES: FR-DSP-2 | NFR-RES-2
|
||||
/// The photograph at full resolution, when it is too large to hold in one
|
||||
/// texture. `demosaiced` is then a reduced copy of it, which is all the
|
||||
@@ -376,7 +380,12 @@ impl DevelopSession {
|
||||
let (w, h) = (raw.crop.width.max(1), raw.crop.height.max(1));
|
||||
let edge = PROXY_EDGE.min(DemosaicedImage::max_dimension(ctx));
|
||||
if raw.samples_per_pixel != 3 || w.max(h) <= edge {
|
||||
return Self::open(ctx, &raw, orientation);
|
||||
let mut session = Self::open(ctx, &raw, orientation)?;
|
||||
// TRACES: FR-DEV-3g
|
||||
// The mosaic stays with the session when the learned demosaic
|
||||
// could take it, so asking for it later needs no second decode.
|
||||
session.keep_mosaic(raw);
|
||||
return Ok(session);
|
||||
}
|
||||
let reduce = w.max(h).div_ceil(edge);
|
||||
log::info!("{w}×{h} is larger than one texture; developing from a 1/{reduce} copy");
|
||||
@@ -438,6 +447,7 @@ impl DevelopSession {
|
||||
compared_snapshot: None,
|
||||
crop_notice: None,
|
||||
demosaiced: Arc::new(demosaiced),
|
||||
denoise: Default::default(),
|
||||
full: None,
|
||||
window: None,
|
||||
adjust: AdjustPass::new(ctx),
|
||||
|
||||
Reference in New Issue
Block a user