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
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.
This commit is contained in:
@@ -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<String> = 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<usize> = 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<bool> = 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<f32>) {
|
||||||
|
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<u8>) {
|
||||||
|
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");
|
||||||
|
}
|
||||||
@@ -2,7 +2,7 @@
|
|||||||
//!
|
//!
|
||||||
//! ```sh
|
//! ```sh
|
||||||
//! cargo run -p dr-ui --example merge --release -- fixtures/pano/2025-08-05/*.CR2
|
//! 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
|
//! The job the interface runs (`dr_ui::merge`), driven from a loop instead
|
||||||
@@ -18,18 +18,65 @@ fn main() {
|
|||||||
env_logger::init();
|
env_logger::init();
|
||||||
let mut args: Vec<String> = std::env::args().skip(1).collect();
|
let mut args: Vec<String> = std::env::args().skip(1).collect();
|
||||||
let mut out_dir = None;
|
let mut out_dir = None;
|
||||||
if args.first().map(String::as_str) == Some("--out") {
|
let mut fill = false;
|
||||||
args.remove(0);
|
let mut wait_engines = false;
|
||||||
out_dir = Some(PathBuf::from(args.remove(0)));
|
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 {
|
if args.len() < 2 {
|
||||||
eprintln!("usage: merge [--out DIR] <frame> <frame>...");
|
eprintln!("usage: merge [--out DIR] [--fill] [--wait-engines] <frame> <frame>...");
|
||||||
std::process::exit(2);
|
std::process::exit(2);
|
||||||
}
|
}
|
||||||
let Ok(gpu) = pollster::block_on(dr_gpu::GpuContext::new_headless()) else {
|
let Ok(gpu) = pollster::block_on(dr_gpu::GpuContext::new_headless()) else {
|
||||||
eprintln!("no GPU");
|
eprintln!("no GPU");
|
||||||
std::process::exit(1);
|
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<dr_ui::merge::MergeInput> = args
|
let frames: Vec<dr_ui::merge::MergeInput> = args
|
||||||
.iter()
|
.iter()
|
||||||
@@ -41,8 +88,10 @@ fn main() {
|
|||||||
.map(std::path::Path::to_path_buf)
|
.map(std::path::Path::to_path_buf)
|
||||||
.unwrap_or_else(|| PathBuf::from("."))
|
.unwrap_or_else(|| PathBuf::from("."))
|
||||||
});
|
});
|
||||||
let request =
|
let mut request =
|
||||||
dr_ui::merge::MergeRequest::new(frames, dr_ui::merge::MergeDestination::Local(dir));
|
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 (tx, rx) = std::sync::mpsc::channel();
|
||||||
let (decide, decision) = std::sync::mpsc::channel();
|
let (decide, decision) = std::sync::mpsc::channel();
|
||||||
let cancel = dr_ui::merge::Cancel::default();
|
let cancel = dr_ui::merge::Cancel::default();
|
||||||
@@ -75,8 +124,18 @@ fn main() {
|
|||||||
Err(why) => println!(" frame {k}: NOT ALIGNED — {why}"),
|
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.
|
// 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 {
|
Done {
|
||||||
path,
|
path,
|
||||||
|
|||||||
Reference in New Issue
Block a user