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

//! 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");
}