DarkRoom/ui/dr-ui/examples/merge.rs

//! Merge a set of frames into a panorama DNG, headless.
//!
//! ```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 --fill frame1.CR2 …
//! ```
//!
//! The job the interface runs (`dr_ui::merge`), driven from a loop instead
//! of a timer: the same code, the same events, a file at the end. What it
//! proves that the alignment example cannot is the second half — the tap,
//! the chunked warp, the resolve to sixteen bits and the DNG — and that the
//! result opens in the application as one RAW.

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 out_dir = None;
    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] [--fill] [--wait-engines] <frame> <frame>...");
        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<dr_ui::merge::MergeInput> = args
        .iter()
        .map(|a| dr_ui::merge::MergeInput::read(std::path::Path::new(a)).expect("read"))
        .collect();
    let dir = out_dir.unwrap_or_else(|| {
        PathBuf::from(&args[0])
            .parent()
            .map(std::path::Path::to_path_buf)
            .unwrap_or_else(|| PathBuf::from("."))
    });
    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();
    let t = Instant::now();
    let worker = std::thread::spawn(move || dr_ui::merge::run(gpu, request, tx, decision, cancel));

    let mut last_stage = "";
    for event in rx {
        use dr_ui::merge::MergeEvent::*;
        match event {
            Progress { stage, done, total } => {
                if stage != last_stage {
                    println!("{stage}…");
                    last_stage = stage;
                }
                if done == total || done % 4 == 0 {
                    println!("  {done}/{total}");
                }
            }
            Aligned(r) => {
                println!(
                    "aligned: {:.1} mm equivalent, rms {:.2} px, {} links, {:?} {}×{}",
                    r.focal_mm, r.rms_px, r.links, r.projection, r.width, r.height
                );
                for (k, f) in r.frames.iter().enumerate() {
                    match f {
                        Ok((yaw, pitch)) => {
                            println!("  frame {k}: yaw {yaw:7.2}° pitch {pitch:6.2}°")
                        }
                        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,
                    fill: fill.then(dr_ui::merge::FillSettings::default),
                });
            }
            Done {
                path,
                width,
                height,
                ..
            } => {
                println!(
                    "wrote {} ({width}×{height}) in {:?}",
                    path.display(),
                    t.elapsed()
                );
            }
            Failed(e) => {
                eprintln!("failed: {e}");
                std::process::exit(1);
            }
            Cancelled => println!("cancelled"),
        }
    }
    let _ = worker.join();
}