From 8baa46ff4965e28d42250bb2270d615c0fc558d7 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sat, 19 Sep 2026 20:34:13 +0200 Subject: [PATCH] Let the merge example fill, wait for engines and dump the filler's input; add a fill example that re-runs it stage by stage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A fill that went wrong took a seven-minute merge to look at again. Now DR_FILL_DUMP=dir makes the merge write what the filler was given, and the fill example runs fill_border on that, or a crop of it, on the engine and writes coarse, each band and the feathered result as PPMs — seconds per attempt on TensorRT. Both examples take DARKROOM_ORT_DIR as the app does, and --wait-engines lets a compiling rung finish before timing. --- ui/dr-ui/examples/fill.rs | 183 +++++++++++++++++++++++++++++++++++++ ui/dr-ui/examples/merge.rs | 73 +++++++++++++-- 2 files changed, 249 insertions(+), 7 deletions(-) create mode 100644 ui/dr-ui/examples/fill.rs diff --git a/ui/dr-ui/examples/fill.rs b/ui/dr-ui/examples/fill.rs new file mode 100644 index 0000000..a5055ef --- /dev/null +++ b/ui/dr-ui/examples/fill.rs @@ -0,0 +1,183 @@ +//! Re-run the panorama border fill on what a merge dumped, stage by stage. +//! +//! DR_FILL_DUMP=/var/tmp/fill cargo run --release -p dr-ui --example merge -- --fill a.CR2 b.CR2 +//! cargo run --release -p dr-ui --example fill -- /var/tmp/fill [--crop X,Y,W,H] [--feather N] [--mirror N] [--band N] [--coarse N] +//! +//! Reads `input.ppm` and `known.pgm` from the directory (`dr_ui::merge` +//! writes them when `DR_FILL_DUMP` is set), runs `fill_border` on the +//! engine as the merge would, and writes `coarse.ppm`, `band-N.ppm` and +//! `feathered.ppm` beside them — so a fill that went wrong can be looked +//! at one stage at a time, in seconds, without the merge around it. +//! `--crop` works on a window of the dump; the coverage edge must be in it. + +use std::path::PathBuf; +use std::time::Instant; + +fn main() { + env_logger::init(); + let mut args: Vec = std::env::args().skip(1).collect(); + let mut crop: Option<[usize; 4]> = None; + let mut params = dr_ui::merge::FillSettings::default().params; + let mut wait_engines = false; + let mut i = 0; + while i < args.len() { + match args[i].as_str() { + "--crop" => { + let v: Vec = args[i + 1] + .split(',') + .map(|s| s.parse().expect("--crop X,Y,W,H")) + .collect(); + crop = Some([v[0], v[1], v[2], v[3]]); + args.drain(i..i + 2); + } + "--feather" => { + params.feather = args[i + 1].parse().expect("--feather N"); + args.drain(i..i + 2); + } + "--mirror" => { + params.mirror_depth = args[i + 1].parse().expect("--mirror N"); + args.drain(i..i + 2); + } + "--band" => { + params.band = args[i + 1].parse().expect("--band N"); + args.drain(i..i + 2); + } + "--coarse" => { + params.coarse = args[i + 1].parse().expect("--coarse N"); + args.drain(i..i + 2); + } + "--wait-engines" => { + wait_engines = true; + args.remove(i); + } + _ => i += 1, + } + } + let Some(dir) = args.first().map(PathBuf::from) else { + eprintln!("usage: fill [--crop X,Y,W,H] [--feather N] [--mirror N] [--band N] [--coarse N] [--wait-engines] DIR"); + std::process::exit(2); + }; + + let (w, h, mut rgb) = read_ppm(&dir.join("input.ppm")); + let (kw, kh, known) = read_pgm(&dir.join("known.pgm")); + assert_eq!((w, h), (kw, kh), "input and mask differ in size"); + let mut known: Vec = known.into_iter().map(|v| v > 127).collect(); + let (w, h) = if let Some([x0, y0, cw, ch]) = crop { + let mut r = Vec::with_capacity(cw * ch * 3); + let mut k = Vec::with_capacity(cw * ch); + for y in y0..y0 + ch { + r.extend_from_slice(&rgb[(y * w + x0) * 3..(y * w + x0 + cw) * 3]); + k.extend_from_slice(&known[y * w + x0..y * w + x0 + cw]); + } + rgb = r; + known = k; + (cw, ch) + } else { + (w, h) + }; + eprintln!( + "{w}×{h}, {:.1}% to fill", + 100.0 * known.iter().filter(|k| !**k).count() as f64 / known.len() as f64 + ); + + let mut runtime_dirs = Vec::new(); + if let Some(d) = std::env::var_os("DARKROOM_ORT_DIR") { + runtime_dirs.push(PathBuf::from(d)); + } + runtime_dirs.push(PathBuf::from("/usr/lib/darkroom")); + runtime_dirs.push(PathBuf::from("/usr/lib")); + dr_ui::inference::init(runtime_dirs); + if wait_engines { + loop { + let s = dr_inference_engine::status(); + eprintln!( + "inference: {} ({}), engines {}/{}", + s.line(), + s.reason, + s.engines.0, + s.engines.1 + ); + if !s.probing && s.engines.0 >= s.engines.1 { + break; + } + std::thread::sleep(std::time::Duration::from_secs(5)); + } + } + let path = dr_ui::inpaint_model().expect("no migan-512.onnx in the model directory"); + let mut model = dr_pano::migan::MiGan::from_path(&path).expect("open the filler"); + eprintln!( + "filler on {}", + model.rung().map(|r| r.label()).unwrap_or("?") + ); + + struct Dump { + dir: PathBuf, + started: Instant, + } + impl dr_pano::Observer for Dump { + fn progress(&mut self, done: usize, total: usize) { + if done.is_multiple_of(8) || done == total { + eprintln!( + " {done}/{total} {:.1}s", + self.started.elapsed().as_secs_f64() + ); + } + } + fn stage(&mut self, name: &str, rgb: &[f32], width: usize, height: usize) { + eprintln!( + "{name}: {width}×{height} at {:.1}s", + self.started.elapsed().as_secs_f64() + ); + write_ppm(&self.dir.join(format!("{name}.ppm")), rgb, width, height); + } + } + let mut dump = Dump { + dir: dir.clone(), + started: Instant::now(), + }; + let tiles = + dr_pano::fill_border(&mut rgb, w, h, &known, &mut model, params, &mut dump).expect("fill"); + eprintln!( + "{tiles} tiles in {:.1}s", + dump.started.elapsed().as_secs_f64() + ); +} + +fn read_header(bytes: &[u8], magic: &str) -> (usize, usize, usize) { + // P6/P5, width, height, maxval, one whitespace, then the data. + let text = std::str::from_utf8(&bytes[..64.min(bytes.len())]).unwrap_or(""); + let mut it = text.split_ascii_whitespace(); + assert_eq!(it.next(), Some(magic), "not a {magic} file"); + let w: usize = it.next().unwrap().parse().unwrap(); + let h: usize = it.next().unwrap().parse().unwrap(); + let maxval: &str = it.next().unwrap(); + assert_eq!(maxval, "255"); + // The data starts one byte after the maxval token. + let header_len = text.find(maxval).unwrap() + maxval.len() + 1; + (w, h, header_len) +} + +fn read_ppm(path: &std::path::Path) -> (usize, usize, Vec) { + let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("{}: {e}", path.display())); + let (w, h, start) = read_header(&bytes, "P6"); + let rgb = bytes[start..start + w * h * 3] + .iter() + .map(|&v| v as f32 / 255.0) + .collect(); + (w, h, rgb) +} + +fn read_pgm(path: &std::path::Path) -> (usize, usize, Vec) { + let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("{}: {e}", path.display())); + let (w, h, start) = read_header(&bytes, "P5"); + (w, h, bytes[start..start + w * h].to_vec()) +} + +fn write_ppm(path: &std::path::Path, rgb: &[f32], w: usize, h: usize) { + let mut out = format!("P6\n{w} {h}\n255\n").into_bytes(); + out.extend( + rgb.iter() + .map(|v| (v.clamp(0.0, 1.0) * 255.0).round() as u8), + ); + std::fs::write(path, out).expect("write stage"); +} diff --git a/ui/dr-ui/examples/merge.rs b/ui/dr-ui/examples/merge.rs index 7bed683..668c56c 100644 --- a/ui/dr-ui/examples/merge.rs +++ b/ui/dr-ui/examples/merge.rs @@ -2,7 +2,7 @@ //! //! ```sh //! cargo run -p dr-ui --example merge --release -- fixtures/pano/2025-08-05/*.CR2 -//! cargo run -p dr-ui --example merge --release -- --out /var/tmp frame1.CR2 … +//! cargo run -p dr-ui --example merge --release -- --out /var/tmp --fill frame1.CR2 … //! ``` //! //! The job the interface runs (`dr_ui::merge`), driven from a loop instead @@ -18,18 +18,65 @@ fn main() { env_logger::init(); let mut args: Vec = std::env::args().skip(1).collect(); let mut out_dir = None; - if args.first().map(String::as_str) == Some("--out") { - args.remove(0); - out_dir = Some(PathBuf::from(args.remove(0))); + let mut fill = false; + let mut wait_engines = false; + loop { + match args.first().map(String::as_str) { + Some("--out") => { + args.remove(0); + out_dir = Some(PathBuf::from(args.remove(0))); + } + Some("--fill") => { + args.remove(0); + fill = true; + } + Some("--wait-engines") => { + args.remove(0); + wait_engines = true; + } + _ => break, + } } if args.len() < 2 { - eprintln!("usage: merge [--out DIR] ..."); + eprintln!("usage: merge [--out DIR] [--fill] [--wait-engines] ..."); std::process::exit(2); } let Ok(gpu) = pollster::block_on(dr_gpu::GpuContext::new_headless()) else { eprintln!("no GPU"); std::process::exit(1); }; + // The inference engine as the desktop app starts it, so the border fill + // runs where the app would run it — TensorRT, ONNX Runtime's CPU pool — + // rather than on tract, which is a hundred times slower for this model. + // `DARKROOM_ORT_DIR` first, as the desktop app takes it: a runtime with + // the TensorRT provider where the distribution's has none. + let mut runtime_dirs = Vec::new(); + if let Some(dir) = std::env::var_os("DARKROOM_ORT_DIR") { + runtime_dirs.push(std::path::PathBuf::from(dir)); + } + runtime_dirs.push(std::path::PathBuf::from("/usr/lib/darkroom")); + runtime_dirs.push(std::path::PathBuf::from("/usr/lib")); + dr_ui::inference::init(runtime_dirs); + // A compiling rung (TensorRT) serves models from the fallback until + // their engines are built, in the background, on first run; the app + // lives long enough for that, this example does not. `--wait-engines` + // waits for the cache to fill so the fill's timing is the rung's. + if wait_engines { + loop { + let s = dr_inference_engine::status(); + eprintln!( + "inference: {} ({}), engines {}/{}", + s.line(), + s.reason, + s.engines.0, + s.engines.1 + ); + if !s.probing && s.engines.0 >= s.engines.1 { + break; + } + std::thread::sleep(std::time::Duration::from_secs(5)); + } + } let frames: Vec = args .iter() @@ -41,8 +88,10 @@ fn main() { .map(std::path::Path::to_path_buf) .unwrap_or_else(|| PathBuf::from(".")) }); - let request = + let mut request = dr_ui::merge::MergeRequest::new(frames, dr_ui::merge::MergeDestination::Local(dir)); + // The filler from the shared models directory, as the app finds it. + request.inpaint_model = dr_ui::merge::inpaint_model_path(); let (tx, rx) = std::sync::mpsc::channel(); let (decide, decision) = std::sync::mpsc::channel(); let cancel = dr_ui::merge::Cancel::default(); @@ -75,8 +124,18 @@ fn main() { Err(why) => println!(" frame {k}: NOT ALIGNED — {why}"), } } + println!( + " border fill: {}", + match &r.filler { + Ok(rung) => format!("available, runs on {rung}"), + Err(why) => why.clone(), + } + ); // Headless: the confirmation the dialog would ask for. - let _ = decide.send(dr_ui::merge::Decision::Merge { projection: None }); + let _ = decide.send(dr_ui::merge::Decision::Merge { + projection: None, + fill: fill.then(dr_ui::merge::FillSettings::default), + }); } Done { path,