DarkRoom/core/dr-denoise/examples/denoise_raw.rs

//! Denoise one RAW file end to end, as develop will, and time it.
//!
//! ```sh
//! DARKROOM_ORT_DIR=~/.local/share/darkroom/runtime \
//!     cargo run --release -p dr-denoise --features native --example denoise_raw -- IMG.CR2 out [fast|best]
//! ```
//!
//! Decode, the app's hot-pixel pass, the frame's noise from its best source,
//! 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
//! at an ONNX Runtime build; the engine's cache goes to `DR_ENGINE_CACHE` or
//! a temporary directory. The whole-frame network (`mosaic-hq.onnx` beside
//! the fixed file) runs where the rung takes any size; `DR_PLAN=tiles` keeps
//! the 1408² tiles anyway, to compare the two.

use std::path::PathBuf;
use std::time::{Duration, Instant};

use dr_denoise::onnx::OnnxNet;
use dr_inference_engine::{Config, Role};

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 [fast|best]");
        std::process::exit(2);
    };
    let shipped = match args.next().as_deref() {
        None | Some("best") => dr_denoise::BEST,
        Some("fast") => dr_denoise::FAST,
        Some(other) => {
            eprintln!("no network called {other}: fast or best");
            std::process::exit(2);
        }
    };
    let model = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../../models/denoise")
        .join(shipped.file);
    let whole = model.with_file_name(shipped.whole);
    let tiles_only = std::env::var("DR_PLAN").is_ok_and(|p| p == "tiles");
    let mut models = vec![(Role::Denoiser, model.clone())];
    if whole.is_file() && !tiles_only {
        models.push((Role::WholeDenoiser, whole));
    }
    let cache = std::env::var_os("DR_ENGINE_CACHE")
        .map(PathBuf::from)
        .unwrap_or_else(|| std::env::temp_dir().join("dr-denoise-engines"));
    let started = Instant::now();
    dr_inference_engine::init(Config {
        runtime_dirs: std::env::var_os("DARKROOM_ORT_DIR")
            .map(PathBuf::from)
            .into_iter()
            .collect(),
        cache_dir: cache,
        models,
        embedded: Vec::new(),
        ceiling: None,
        threads: 0,
        decay: Duration::ZERO,
    });
    // Wait for the probe and the engine build, so the timing below is the
    // rung this machine settles on, not the fallback used while it compiles.
    // The probe starts on its own thread; give it a moment to say so.
    std::thread::sleep(Duration::from_secs(1));
    loop {
        let s = dr_inference_engine::status();
        if !s.probing && s.engines.0 >= s.engines.1 {
            println!(
                "engine        {} ({:.1} s to settle)",
                s.line(),
                started.elapsed().as_secs_f64()
            );
            break;
        }
        std::thread::sleep(Duration::from_millis(200));
    }

    let bytes = std::fs::read(&input).expect("read raw");
    let t = Instant::now();
    let mut raw = dr_decode::decode(&bytes).expect("decode");
    let meta = dr_decode::metadata(&bytes).expect("metadata");
    let decode = t.elapsed();

    let t = Instant::now();
    let ctx =
        pollster::block_on(dr_gpu::GpuContext::new_headless()).expect("GPU for the hot-pixel pass");
    let repaired = dr_gpu::Demosaicer::new(&ctx)
        .expect("demosaicer")
        .repair_hot_pixels(&mut raw)
        .expect("repair");
    let repair = t.elapsed();

    let noise = dr_denoise::noise::for_frame(&raw, &bytes, meta.iso)
        .expect("no noise source for this frame");
    println!(
        "frame         {} {} ISO {:?}, {}×{}, {:?}, {repaired} hot photosites repaired",
        raw.make, raw.model, meta.iso, raw.crop.width, raw.crop.height, raw.cfa_pattern
    );
    println!(
        "noise         {} — σ at 10 % grey (G) {:.5}, read {:.5}, row {:.5}, col {:.5}",
        noise.source.label(),
        noise.sigma(1, 0.1),
        noise.o[1].sqrt(),
        noise.row,
        noise.col
    );

    let mut net = if tiles_only {
        OnnxNet::open_tiled(&model, shipped)
    } else {
        OnnxNet::open(&model, shipped)
    }
    .expect("model");
    println!(
        "rung          {} · {}",
        net.rung().map(|r| r.label()).unwrap_or("?"),
        if net.whole_frame() {
            "whole frame"
        } else {
            "1408² tiles"
        }
    );
    let t = Instant::now();
    let rgb = dr_denoise::denoise(&raw, &noise, &mut net, &mut |done, total| {
        eprint!("\rtile          {done}/{total}");
        true
    })
    .expect("denoise")
    .expect("not cancelled");
    let run = t.elapsed();
    eprintln!();
    println!(
        "time          decode {:.2} s · hot pixels {:.2} s · network {:.2} s ({:.1} MP)",
        decode.as_secs_f64(),
        repair.as_secs_f64(),
        run.as_secs_f64(),
        (raw.crop.width * raw.crop.height) as f64 / 1e6
    );

    let (h, w) = (raw.crop.height as usize, raw.crop.width as usize);
    let mut npy = Vec::with_capacity(rgb.len() * 4 + 128);
    let mut header =
        format!("{{'descr': '<f4', 'fortran_order': False, 'shape': ({h}, {w}, 3), }}");
    while (10 + header.len() + 1) % 64 != 0 {
        header.push(' ');
    }
    header.push('\n');
    npy.extend_from_slice(b"\x93NUMPY\x01\x00");
    npy.extend_from_slice(&(header.len() as u16).to_le_bytes());
    npy.extend_from_slice(header.as_bytes());
    for v in &rgb {
        npy.extend_from_slice(&v.to_le_bytes());
    }
    std::fs::write(format!("{out}.npy"), npy).expect("write");
    println!("wrote         {out}.npy");
}