loftr 0.1.1

Native Rust/tch implementation of LoFTR feature matching
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

use std::{env, fs, path::Path};

use image::{DynamicImage, imageops::FilterType};
use loftr::{LoftrConfig, LoftrModel};
use tch::{Device, Kind, Tensor};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = env::args().collect();
    if args.len() != 5 {
        eprintln!(
            "usage: cargo run -p loftr --example dump_stages -- <weights> <left> <right> <output>"
        );
        std::process::exit(2);
    }

    let mut model = LoftrModel::new(Device::Cpu, LoftrConfig::outdoor())?;
    model.load_weights(&args[1])?;

    let left = load_grayscale(Path::new(&args[2]))?;
    let right = load_grayscale(Path::new(&args[3]))?;
    let debug = model.forward_debug(&left, &right)?;
    fs::write(&args[4], serde_json::to_vec_pretty(&debug)?)?;
    Ok(())
}

fn load_grayscale(path: &Path) -> Result<Tensor, Box<dyn std::error::Error>> {
    let image = image::open(path)?;
    let image = resize_for_loftr(&image);
    let image = image.to_luma32f();
    let height = i64::from(image.height());
    let width = i64::from(image.width());
    let data = image.into_raw();
    let tensor = Tensor::from_slice(&data)
        .view([1, height, width])
        .unsqueeze(0)
        .to_kind(Kind::Float);
    Ok(tensor)
}

fn resize_for_loftr(image: &DynamicImage) -> DynamicImage {
    image.resize_exact(960, 540, FilterType::Triangle)
}