aliked-colmap
Commit: d65ca97
Parent: b0c154a

Implement aliked-colmap CLI: ALIKED+LightGlue feature extraction and matching

Mårten Åsberg committed on 2026-08-03 at 09:34
Adds feature_extractor and exhaustive_matcher subcommands that populate a
COLMAP-compatible SQLite database using ONNX Runtime (OpenVINO EP) for
inference and a from-scratch normalized 8-point RANSAC for geometric
verification, so colmap mapper can consume the output unchanged.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Browse files at this commit
Cargo.lock +1914 -0
Large diff (1914 lines changed) - not displayed
Cargo.toml +11 -0
diff --git a/Cargo.toml b/Cargo.toml
index 726b6f4..224674b 100644
@@ -4,3 +4,14 @@ version = "0.1.0"
edition = "2024"
[dependencies]
anyhow = "1.0.104"
bytemuck = { version = "1.25.2", features = ["derive"] }
clap = { version = "4.6.5", features = ["derive"] }
image = "0.25.10"
nalgebra = "0.35.0"
ndarray = "0.17.2"
ort = { version = "2.0.0-rc.13", features = ["openvino"] }
rand = "0.10.2"
rayon = "1.12.0"
rusqlite = { version = "0.37", features = ["bundled"] }
walkdir = "2.5.0"
src/database.rs +225 -0
diff --git a/src/database.rs b/src/database.rs
new file mode 100644
index 0000000..840c995
@@ -0,0 +1,225 @@
use std::path::Path;
use anyhow::{Context, Result};
use rusqlite::Connection;
/// COLMAP's SIMPLE_RADIAL camera model id.
pub const SIMPLE_RADIAL_MODEL_ID: i64 = 2;
/// COLMAP caps image ids below this value; the pair_id encoding relies on it.
pub const MAX_NUM_IMAGES: i64 = 2_147_483_647;
const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS cameras (
camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
model INTEGER NOT NULL,
width INTEGER NOT NULL,
height INTEGER NOT NULL,
params BLOB,
prior_focal_length INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS images (
image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
name TEXT NOT NULL UNIQUE,
camera_id INTEGER NOT NULL,
prior_qw REAL,
prior_qx REAL,
prior_qy REAL,
prior_qz REAL,
prior_tx REAL,
prior_ty REAL,
prior_tz REAL
);
CREATE TABLE IF NOT EXISTS keypoints (
image_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data BLOB
);
CREATE TABLE IF NOT EXISTS descriptors (
image_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data BLOB
);
CREATE TABLE IF NOT EXISTS matches (
pair_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data BLOB
);
CREATE TABLE IF NOT EXISTS two_view_geometries (
pair_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data BLOB,
config INTEGER NOT NULL,
F BLOB,
E BLOB,
H BLOB,
qvec BLOB,
tvec BLOB
);
"#;
/// Computes COLMAP's canonical pair id for an unordered pair of image ids.
///
/// Must match COLMAP exactly: id1/id2 are ordered smaller-first before encoding.
pub fn image_pair_to_pair_id(id1: i64, id2: i64) -> i64 {
if id1 > id2 {
MAX_NUM_IMAGES * id2 + id1
} else {
MAX_NUM_IMAGES * id1 + id2
}
}
pub fn open(path: &Path) -> Result<Connection> {
let conn = Connection::open(path).with_context(|| format!("opening database at {}", path.display()))?;
conn.execute_batch(SCHEMA).context("creating database schema")?;
Ok(conn)
}
/// Looks up the camera_id for a given (width, height), inserting a new SIMPLE_RADIAL
/// camera record if one doesn't already exist for that resolution.
pub fn get_or_create_camera(conn: &Connection, width: u32, height: u32) -> Result<i64> {
if let Ok(camera_id) = conn.query_row(
"SELECT camera_id FROM cameras WHERE width = ?1 AND height = ?2 LIMIT 1",
(width, height),
|row| row.get::<_, i64>(0),
) {
return Ok(camera_id);
}
let f = (width.max(height) as f64) * 1.2;
let cx = width as f64 / 2.0;
let cy = height as f64 / 2.0;
let k = 0.0_f64;
let params: [f64; 4] = [f, cx, cy, k];
let params_blob: &[u8] = bytemuck::cast_slice(&params);
conn
.execute(
"INSERT INTO cameras (model, width, height, params, prior_focal_length) VALUES (?1, ?2, ?3, ?4, 0)",
(SIMPLE_RADIAL_MODEL_ID, width, height, params_blob),
)
.context("inserting camera")?;
Ok(conn.last_insert_rowid())
}
/// Inserts an image row, returning its assigned image_id.
pub fn insert_image(conn: &Connection, name: &str, camera_id: i64) -> Result<i64> {
conn
.execute(
"INSERT INTO images (name, camera_id) VALUES (?1, ?2)",
(name, camera_id),
)
.with_context(|| format!("inserting image row for {name}"))?;
Ok(conn.last_insert_rowid())
}
/// Returns (width, height) of the camera associated with an image.
pub fn image_dimensions(conn: &Connection, image_id: i64) -> Result<(u32, u32)> {
conn
.query_row(
"SELECT c.width, c.height FROM images i JOIN cameras c ON c.camera_id = i.camera_id WHERE i.image_id = ?1",
[image_id],
|row| Ok((row.get::<_, i64>(0)? as u32, row.get::<_, i64>(1)? as u32)),
)
.with_context(|| format!("reading dimensions for image_id {image_id}"))
}
pub fn all_images(conn: &Connection) -> Result<Vec<(i64, String)>> {
let mut stmt = conn.prepare("SELECT image_id, name FROM images ORDER BY name")?;
let rows = stmt
.query_map((), |row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)))?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
}
/// keypoints.data is float32[N][4]: [x, y, scale, orientation] in original pixel coordinates.
pub fn write_keypoints(conn: &Connection, image_id: i64, keypoints: &[[f32; 4]]) -> Result<()> {
let data: &[u8] = bytemuck::cast_slice(keypoints);
conn
.execute(
"INSERT INTO keypoints (image_id, rows, cols, data) VALUES (?1, ?2, 4, ?3)",
(image_id, keypoints.len() as i64, data),
)
.context("inserting keypoints")?;
Ok(())
}
pub fn read_keypoints(conn: &Connection, image_id: i64) -> Result<Vec<[f32; 4]>> {
let data: Vec<u8> = conn.query_row("SELECT data FROM keypoints WHERE image_id = ?1", [image_id], |row| {
row.get(0)
})?;
let flat: &[f32] = bytemuck::cast_slice(&data);
Ok(flat.chunks_exact(4).map(|c| [c[0], c[1], c[2], c[3]]).collect())
}
/// descriptors.data is float32[N][128], row-major.
pub fn write_descriptors(
conn: &Connection,
image_id: i64,
rows: usize,
cols: usize,
descriptors: &[f32],
) -> Result<()> {
debug_assert_eq!(descriptors.len(), rows * cols);
let data: &[u8] = bytemuck::cast_slice(descriptors);
conn
.execute(
"INSERT INTO descriptors (image_id, rows, cols, data) VALUES (?1, ?2, ?3, ?4)",
(image_id, rows as i64, cols as i64, data),
)
.context("inserting descriptors")?;
Ok(())
}
/// Returns (rows, cols, flat row-major descriptor data).
pub fn read_descriptors(conn: &Connection, image_id: i64) -> Result<(usize, usize, Vec<f32>)> {
let (rows, cols, data): (i64, i64, Vec<u8>) = conn.query_row(
"SELECT rows, cols, data FROM descriptors WHERE image_id = ?1",
[image_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
)?;
let flat: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
Ok((rows as usize, cols as usize, flat))
}
/// matches.data is int32[N][2]: [kpt_idx_in_image_A, kpt_idx_in_image_B].
pub fn write_matches(conn: &Connection, pair_id: i64, matches: &[[i32; 2]]) -> Result<()> {
let data: &[u8] = bytemuck::cast_slice(matches);
conn
.execute(
"INSERT INTO matches (pair_id, rows, cols, data) VALUES (?1, ?2, 2, ?3)",
(pair_id, matches.len() as i64, data),
)
.context("inserting matches")?;
Ok(())
}
/// Writes a two_view_geometries row. `f_matrix` is the row-major 3x3 fundamental matrix.
/// `config` should be 3 (UNCALIBRATED) unless calibrated intrinsics are available.
pub fn write_two_view_geometry(
conn: &Connection,
pair_id: i64,
config: i64,
inlier_matches: &[[i32; 2]],
f_matrix: &[f64; 9],
) -> Result<()> {
let data: &[u8] = bytemuck::cast_slice(inlier_matches);
let f_blob: &[u8] = bytemuck::cast_slice(f_matrix.as_slice());
conn
.execute(
"INSERT INTO two_view_geometries (pair_id, rows, cols, data, config, F, E, H, qvec, tvec)
VALUES (?1, ?2, 2, ?3, ?4, ?5, NULL, NULL, NULL, NULL)",
(pair_id, inlier_matches.len() as i64, data, config, f_blob),
)
.context("inserting two_view_geometries")?;
Ok(())
}
src/exhaustive_matcher.rs +246 -0
diff --git a/src/exhaustive_matcher.rs b/src/exhaustive_matcher.rs
new file mode 100644
index 0000000..0673ff4
@@ -0,0 +1,246 @@
use std::path::PathBuf;
use anyhow::{Context, Result};
use ndarray::ArrayView3;
use ort::session::Session;
use ort::value::TensorRef;
use rayon::prelude::*;
use crate::database::{self, image_pair_to_pair_id};
use crate::models::{self, Device};
use crate::ransac::{self, RansacResult};
pub struct ExhaustiveMatcherConfig {
pub database_path: PathBuf,
pub lightglue_model: PathBuf,
pub min_score: f32,
pub device: Device,
}
const MIN_INLIERS: usize = 15;
/// two_view_geometries.config: UNCALIBRATED (uses F, no known intrinsics).
const UNCALIBRATED_CONFIG: i64 = 3;
struct ImageFeatures {
image_id: i64,
name: String,
/// Original pixel-space [x, y, scale, orientation], as stored in the database.
keypoints: Vec<[f32; 4]>,
/// Flat n*2 keypoints re-normalized to [-1, 1], matching ALIKED's own convention
/// (normalized by max(width, height) of the source image, not per-axis).
kpts_norm: Vec<f32>,
/// Flat row-major n*desc_dim descriptors.
descriptors: Vec<f32>,
desc_dim: usize,
}
fn load_features(conn: &rusqlite::Connection, image_id: i64, name: String) -> Result<ImageFeatures> {
let keypoints = database::read_keypoints(conn, image_id)
.with_context(|| format!("reading keypoints for {name} (run feature_extractor first?)"))?;
let (rows, desc_dim, descriptors) = database::read_descriptors(conn, image_id)?;
if rows != keypoints.len() {
anyhow::bail!(
"{name}: keypoint count ({}) doesn't match descriptor row count ({rows})",
keypoints.len()
);
}
let (w, h) = database::image_dimensions(conn, image_id)?;
let max_dim = w.max(h) as f32;
let kpts_norm: Vec<f32> = keypoints
.iter()
.flat_map(|k| [2.0 * k[0] / max_dim - 1.0, 2.0 * k[1] / max_dim - 1.0])
.collect();
Ok(ImageFeatures {
image_id,
name,
keypoints,
kpts_norm,
descriptors,
desc_dim,
})
}
/// Runs LightGlue on a pair of images' features. Returns (index_in_a, index_in_b) pairs
/// for matches passing `min_score`, with unmatched (-1) entries already excluded.
fn run_lightglue(
session: &mut Session,
a: &ImageFeatures,
b: &ImageFeatures,
min_score: f32,
) -> Result<Vec<(usize, usize)>> {
let n = a.keypoints.len();
let m = b.keypoints.len();
if n == 0 || m == 0 {
return Ok(Vec::new());
}
let kpts0 = ArrayView3::from_shape((1, n, 2), &a.kpts_norm)?;
let kpts1 = ArrayView3::from_shape((1, m, 2), &b.kpts_norm)?;
let desc0 = ArrayView3::from_shape((1, n, a.desc_dim), &a.descriptors)?;
let desc1 = ArrayView3::from_shape((1, m, b.desc_dim), &b.descriptors)?;
let outputs = session.run(ort::inputs! {
"kpts0" => TensorRef::from_array_view(kpts0)?,
"kpts1" => TensorRef::from_array_view(kpts1)?,
"desc0" => TensorRef::from_array_view(desc0)?,
"desc1" => TensorRef::from_array_view(desc1)?,
})?;
let matches_val = outputs
.get("matches0")
.ok_or_else(|| models::missing_output_err("matches0", &outputs))?;
let (_shape, matches_data) = matches_val.try_extract_tensor::<i64>()?;
let scores_val = outputs
.get("mscores0")
.ok_or_else(|| models::missing_output_err("mscores0", &outputs))?;
let (_shape2, scores_data) = scores_val.try_extract_tensor::<f32>()?;
Ok(
matches_data
.iter()
.zip(scores_data.iter())
.enumerate()
.filter_map(|(i, (&m_idx, &score))| (m_idx >= 0 && score >= min_score).then_some((i, m_idx as usize)))
.collect(),
)
}
struct PendingPair {
pair_id: i64,
matched: Vec<(usize, usize)>,
pts0: Vec<(f32, f32)>,
pts1: Vec<(f32, f32)>,
}
pub fn run(cfg: ExhaustiveMatcherConfig) -> Result<()> {
let conn = database::open(&cfg.database_path)?;
let mut session = models::load_session(&cfg.lightglue_model, cfg.device)?;
let images = database::all_images(&conn)?;
if images.len() < 2 {
anyhow::bail!(
"need at least 2 images with extracted features to match; found {}. Run feature_extractor first.",
images.len()
);
}
let mut features = Vec::with_capacity(images.len());
for (image_id, name) in images {
features.push(load_features(&conn, image_id, name)?);
}
let n_images = features.len();
println!(
"Matching {n_images} images exhaustively ({} pairs)...",
n_images * (n_images - 1) / 2
);
let mut pending: Vec<PendingPair> = Vec::new();
for i in 0..n_images {
for j in (i + 1)..n_images {
let matched = run_lightglue(&mut session, &features[i], &features[j], cfg.min_score)?;
if matched.is_empty() {
continue;
}
// COLMAP's pair_id encoding requires the matches columns to be ordered
// (smaller image_id, larger image_id); this normally coincides with insertion
// order but is enforced explicitly here rather than assumed.
let (lo, hi, swapped) = if features[i].image_id <= features[j].image_id {
(&features[i], &features[j], false)
} else {
(&features[j], &features[i], true)
};
let ordered_matched: Vec<(usize, usize)> = if swapped {
matched.iter().map(|&(a, b)| (b, a)).collect()
} else {
matched.clone()
};
let pair_id = image_pair_to_pair_id(lo.image_id, hi.image_id);
let raw_matches: Vec<[i32; 2]> = ordered_matched.iter().map(|&(a, b)| [a as i32, b as i32]).collect();
database::write_matches(&conn, pair_id, &raw_matches)?;
let pts0: Vec<(f32, f32)> = ordered_matched
.iter()
.map(|&(a, _)| (lo.keypoints[a][0], lo.keypoints[a][1]))
.collect();
let pts1: Vec<(f32, f32)> = ordered_matched
.iter()
.map(|&(_, b)| (hi.keypoints[b][0], hi.keypoints[b][1]))
.collect();
println!(
" {} <-> {}: {} raw matches",
features[i].name,
features[j].name,
ordered_matched.len()
);
pending.push(PendingPair {
pair_id,
matched: ordered_matched,
pts0,
pts1,
});
}
}
println!(
"Running RANSAC geometric verification on {} candidate pairs...",
pending.len()
);
let ransac_results: Vec<Option<RansacResult>> = pending
.par_iter()
.map(|p| {
ransac::ransac_fundamental_matrix(
&p.pts0,
&p.pts1,
ransac::DEFAULT_MAX_ITERATIONS,
ransac::DEFAULT_INLIER_THRESHOLD,
ransac::DEFAULT_CONFIDENCE,
)
})
.collect();
let mut verified_pairs = 0usize;
let mut total_inliers = 0usize;
for (pair, result) in pending.iter().zip(ransac_results.iter()) {
let Some(result) = result else { continue };
if result.inliers.len() < MIN_INLIERS {
continue;
}
let inlier_matches: Vec<[i32; 2]> = result
.inliers
.iter()
.map(|&idx| {
let (a, b) = pair.matched[idx];
[a as i32, b as i32]
})
.collect();
database::write_two_view_geometry(
&conn,
pair.pair_id,
UNCALIBRATED_CONFIG,
&inlier_matches,
&result.f_matrix,
)?;
verified_pairs += 1;
total_inliers += inlier_matches.len();
}
let avg_inliers = if verified_pairs > 0 {
total_inliers as f64 / verified_pairs as f64
} else {
0.0
};
println!(
"Done: {} pairs had raw matches, {verified_pairs} passed geometric verification (avg {avg_inliers:.1} inliers/pair)",
pending.len()
);
Ok(())
}
src/feature_extractor.rs +176 -0
diff --git a/src/feature_extractor.rs b/src/feature_extractor.rs
new file mode 100644
index 0000000..5300587
@@ -0,0 +1,176 @@
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use image::imageops::FilterType;
use ndarray::Array4;
use ort::session::Session;
use ort::value::TensorRef;
use walkdir::WalkDir;
use crate::database;
use crate::models::{self, Device};
pub struct FeatureExtractorConfig {
pub database_path: PathBuf,
pub image_path: PathBuf,
pub aliked_model: PathBuf,
pub top_k: usize,
pub image_size: u32,
pub device: Device,
}
struct Preprocessed {
tensor: Array4<f32>,
orig_width: u32,
orig_height: u32,
}
const IMAGE_EXTENSIONS: [&str; 3] = ["jpg", "jpeg", "png"];
fn collect_image_paths(root: &Path) -> Result<Vec<PathBuf>> {
let mut paths: Vec<PathBuf> = WalkDir::new(root)
.into_iter()
.filter_map(|e| e.ok())
.filter(|e| e.file_type().is_file())
.filter(|e| {
e.path()
.extension()
.and_then(|s| s.to_str())
.map(|s| IMAGE_EXTENSIONS.iter().any(|ext| ext.eq_ignore_ascii_case(s)))
.unwrap_or(false)
})
.map(|e| e.path().to_path_buf())
.collect();
// Sorted so image_id assignment order matches filename order, which colmap mapper expects.
paths.sort();
Ok(paths)
}
/// Letterbox-resizes (top-left aligned, matching the ALIKED-LightGlue-ONNX reference
/// implementation) into an `image_size` x `image_size` canvas and produces an NCHW RGB
/// float32 tensor normalized to [0, 1].
fn preprocess(path: &Path, image_size: u32) -> Result<Preprocessed> {
let img = image::open(path)
.with_context(|| format!("opening image {}", path.display()))?
.into_rgb8();
let (w, h) = img.dimensions();
if w == 0 || h == 0 {
anyhow::bail!("image {} has zero dimension", path.display());
}
let resize_scale = (image_size as f32 / h as f32).min(image_size as f32 / w as f32);
let w_new = ((w as f32 * resize_scale) as u32).clamp(1, image_size);
let h_new = ((h as f32 * resize_scale) as u32).clamp(1, image_size);
let resized = image::imageops::resize(&img, w_new, h_new, FilterType::Triangle);
let size = image_size as usize;
let mut tensor = Array4::<f32>::zeros((1, 3, size, size));
for y in 0..h_new {
for x in 0..w_new {
let px = resized.get_pixel(x, y);
for c in 0..3 {
tensor[[0, c, y as usize, x as usize]] = px[c] as f32 / 255.0;
}
}
}
Ok(Preprocessed {
tensor,
orig_width: w,
orig_height: h,
})
}
/// (keypoints in normalized [-1,1] canvas coordinates, flat row-major descriptors [n x 128], per-keypoint scores)
type AlikedOutput = (Vec<[f32; 2]>, Vec<f32>, Vec<f32>);
/// Runs ALIKED on a preprocessed tensor.
fn run_aliked(session: &mut Session, tensor: &Array4<f32>) -> Result<AlikedOutput> {
let input = TensorRef::from_array_view(tensor)?;
let outputs = session.run(ort::inputs![input])?;
let kpts_val = outputs
.get("keypoints")
.ok_or_else(|| models::missing_output_err("keypoints", &outputs))?;
let (_kpts_shape, kpts_data) = kpts_val.try_extract_tensor::<f32>()?;
let keypoints: Vec<[f32; 2]> = kpts_data.chunks_exact(2).map(|c| [c[0], c[1]]).collect();
let desc_val = outputs
.get("descriptors")
.ok_or_else(|| models::missing_output_err("descriptors", &outputs))?;
let (_desc_shape, desc_data) = desc_val.try_extract_tensor::<f32>()?;
let scores_val = outputs
.get("scores")
.ok_or_else(|| models::missing_output_err("scores", &outputs))?;
let (_scores_shape, scores_data) = scores_val.try_extract_tensor::<f32>()?;
Ok((keypoints, desc_data.to_vec(), scores_data.to_vec()))
}
pub fn run(cfg: FeatureExtractorConfig) -> Result<()> {
let conn = database::open(&cfg.database_path)?;
let mut session = models::load_session(&cfg.aliked_model, cfg.device)?;
let paths = collect_image_paths(&cfg.image_path)?;
if paths.is_empty() {
anyhow::bail!("no images found under {}", cfg.image_path.display());
}
println!("Found {} images under {}", paths.len(), cfg.image_path.display());
for path in &paths {
let rel_name = path
.strip_prefix(&cfg.image_path)
.unwrap_or(path)
.to_string_lossy()
.replace('\\', "/");
let pre = preprocess(path, cfg.image_size)?;
let camera_id = database::get_or_create_camera(&conn, pre.orig_width, pre.orig_height)?;
let image_id = database::insert_image(&conn, &rel_name, camera_id)?;
let (keypoints, descriptors_flat, scores) = run_aliked(&mut session, &pre.tensor)?;
let n = keypoints.len();
if descriptors_flat.len() % n.max(1) != 0 {
anyhow::bail!(
"{}: descriptor count ({}) doesn't divide evenly by keypoint count ({})",
rel_name,
descriptors_flat.len(),
n
);
}
let desc_dim = if n == 0 { 128 } else { descriptors_flat.len() / n };
// The ONNX model already caps output count near its export-time top-k; this only
// trims further if more keypoints than requested were returned.
let keep: Vec<usize> = if n > cfg.top_k {
let mut idx: Vec<usize> = (0..n).collect();
idx.sort_unstable_by(|&a, &b| scores[b].partial_cmp(&scores[a]).unwrap());
idx.truncate(cfg.top_k);
idx.sort_unstable();
idx
} else {
(0..n).collect()
};
let max_dim = pre.orig_width.max(pre.orig_height) as f32;
let mut kpt_rows = Vec::with_capacity(keep.len());
let mut desc_rows = Vec::with_capacity(keep.len() * desc_dim);
for &i in &keep {
let [x, y] = keypoints[i];
// ALIKED's exported graph outputs keypoints normalized to [-1, 1] over the padded
// square canvas; this maps directly back to original pixel coordinates.
let x_orig = (x + 1.0) * max_dim / 2.0;
let y_orig = (y + 1.0) * max_dim / 2.0;
kpt_rows.push([x_orig, y_orig, 0.0, 0.0]);
desc_rows.extend_from_slice(&descriptors_flat[i * desc_dim..(i + 1) * desc_dim]);
}
database::write_keypoints(&conn, image_id, &kpt_rows)?;
database::write_descriptors(&conn, image_id, keep.len(), desc_dim, &desc_rows)?;
println!("{rel_name}: {} keypoints", keep.len());
}
Ok(())
}
src/main.rs +84 -2
diff --git a/src/main.rs b/src/main.rs
index 80a1832..50ff664 100644
@@ -1,3 +1,85 @@
fn main() {
println!("Hello, world!");
mod database;
mod exhaustive_matcher;
mod feature_extractor;
mod models;
mod ransac;
use std::path::PathBuf;
use anyhow::Result;
use clap::{Parser, Subcommand};
use models::Device;
/// COLMAP-compatible feature extraction and matching using ALIKED + LightGlue.
#[derive(Parser)]
#[command(name = "aliked-colmap")]
struct Cli {
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
enum Command {
/// Extract ALIKED keypoints/descriptors for all images and populate a COLMAP database.
#[command(name = "feature_extractor")]
FeatureExtractor {
#[arg(long)]
database_path: PathBuf,
#[arg(long)]
image_path: PathBuf,
#[arg(long, default_value = "./models/aliked-n16rot-top2k-640.onnx")]
aliked_model: PathBuf,
#[arg(long, default_value_t = 2048)]
top_k: usize,
#[arg(long, default_value_t = 640)]
image_size: u32,
#[arg(long, value_enum, default_value_t = Device::Gpu)]
device: Device,
},
/// Exhaustively match all image pairs with LightGlue and geometrically verify them.
#[command(name = "exhaustive_matcher")]
ExhaustiveMatcher {
#[arg(long)]
database_path: PathBuf,
#[arg(long, default_value = "./models/lightglue_for_aliked.onnx")]
lightglue_model: PathBuf,
#[arg(long, default_value_t = 0.0)]
min_score: f32,
#[arg(long, value_enum, default_value_t = Device::Gpu)]
device: Device,
},
}
fn main() -> Result<()> {
let cli = Cli::parse();
match cli.command {
Command::FeatureExtractor {
database_path,
image_path,
aliked_model,
top_k,
image_size,
device,
} => feature_extractor::run(feature_extractor::FeatureExtractorConfig {
database_path,
image_path,
aliked_model,
top_k,
image_size,
device,
}),
Command::ExhaustiveMatcher {
database_path,
lightglue_model,
min_score,
device,
} => exhaustive_matcher::run(exhaustive_matcher::ExhaustiveMatcherConfig {
database_path,
lightglue_model,
min_score,
device,
}),
}
}
src/models.rs +55 -0
diff --git a/src/models.rs b/src/models.rs
new file mode 100644
index 0000000..ac17b92
@@ -0,0 +1,55 @@
use std::path::Path;
use anyhow::{Context, Result, anyhow};
use ort::ep::{self, ExecutionProvider};
use ort::session::{Session, SessionOutputs};
#[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
#[value(rename_all = "UPPER")]
pub enum Device {
Cpu,
Gpu,
}
impl Device {
fn openvino_device_type(self) -> &'static str {
match self {
Device::Cpu => "CPU",
Device::Gpu => "GPU",
}
}
}
/// Loads an ONNX model with the OpenVINO execution provider targeting `device`.
/// Registration failures (e.g. no Arc GPU present) are logged and the session
/// silently falls back to the CPU execution provider, per `ort`'s default behavior.
pub fn load_session(model_path: &Path, device: Device) -> Result<Session> {
let openvino = ep::OpenVINO::default().with_device_type(device.openvino_device_type());
match openvino.is_available() {
Ok(true) => {}
Ok(false) => eprintln!(
"OpenVINO execution provider not available; falling back to CPU for {}",
model_path.display()
),
Err(e) => eprintln!("Could not query OpenVINO EP availability: {e}"),
}
Session::builder()?
// `.with_execution_providers` returns an error type that carries the SessionBuilder
// (for `.recover()`), which holds raw FFI pointers and so isn't Send+Sync; anyhow
// requires Send+Sync, so convert explicitly instead of using `?` directly.
.with_execution_providers([openvino.build()])
.map_err(|e| anyhow!("{e}"))?
.commit_from_file(model_path)
.with_context(|| format!("loading ONNX model from {}", model_path.display()))
}
/// Builds a descriptive error when an expected named output is missing from a model's
/// results, listing what output names the model actually produced.
pub fn missing_output_err(name: &str, outputs: &SessionOutputs) -> anyhow::Error {
anyhow!(
"model has no '{name}' output; available outputs: {:?}",
outputs.keys().collect::<Vec<_>>()
)
}
src/ransac.rs +189 -0
diff --git a/src/ransac.rs b/src/ransac.rs
new file mode 100644
index 0000000..b8f6e9b
@@ -0,0 +1,189 @@
use nalgebra::{DMatrix, Matrix3, RowDVector, SVD, Vector3};
pub const DEFAULT_MAX_ITERATIONS: usize = 1000;
pub const DEFAULT_INLIER_THRESHOLD: f64 = 4.0;
pub const DEFAULT_CONFIDENCE: f64 = 0.9999;
pub struct RansacResult {
/// Row-major 3x3 fundamental matrix.
pub f_matrix: [f64; 9],
/// Indices (into the input correspondence slices) of inlier matches.
pub inliers: Vec<usize>,
}
fn normalization_transform(pts: &[(f64, f64)]) -> Matrix3<f64> {
let n = pts.len() as f64;
let (sum_x, sum_y) = pts.iter().fold((0.0, 0.0), |(ax, ay), &(x, y)| (ax + x, ay + y));
let cx = sum_x / n;
let cy = sum_y / n;
let mean_dist = pts
.iter()
.map(|&(x, y)| ((x - cx).powi(2) + (y - cy).powi(2)).sqrt())
.sum::<f64>()
/ n;
let s = if mean_dist > 1e-12 {
std::f64::consts::SQRT_2 / mean_dist
} else {
1.0
};
Matrix3::new(s, 0.0, -s * cx, 0.0, s, -s * cy, 0.0, 0.0, 1.0)
}
fn apply_transform(t: &Matrix3<f64>, p: (f64, f64)) -> (f64, f64) {
let v = t * Vector3::new(p.0, p.1, 1.0);
(v.x / v.z, v.y / v.z)
}
/// Normalized 8-point algorithm. Works with >=8 correspondences via least-squares.
fn fit_fundamental(pts0: &[(f64, f64)], pts1: &[(f64, f64)]) -> Option<Matrix3<f64>> {
let n = pts0.len();
if n < 8 || pts1.len() != n {
return None;
}
let t0 = normalization_transform(pts0);
let t1 = normalization_transform(pts1);
let n0: Vec<(f64, f64)> = pts0.iter().map(|&p| apply_transform(&t0, p)).collect();
let n1: Vec<(f64, f64)> = pts1.iter().map(|&p| apply_transform(&t1, p)).collect();
let mut a = DMatrix::<f64>::zeros(n, 9);
for i in 0..n {
let (x0, y0) = n0[i];
let (x1, y1) = n1[i];
a.set_row(
i,
&RowDVector::from_vec(vec![x1 * x0, x1 * y0, x1, y1 * x0, y1 * y0, y1, x0, y0, 1.0]),
);
}
let svd = SVD::new(a, false, true);
let v_t = svd.v_t?;
let idx = argmin(svd.singular_values.as_slice())?;
let f_row = v_t.row(idx);
let f_normalized = Matrix3::new(
f_row[0], f_row[1], f_row[2], f_row[3], f_row[4], f_row[5], f_row[6], f_row[7], f_row[8],
);
// Enforce rank-2 by zeroing the smallest singular value of F.
let svd_f = SVD::new(f_normalized, true, true);
let u = svd_f.u?;
let v_t2 = svd_f.v_t?;
let mut s = svd_f.singular_values;
let min_idx = argmin(s.as_slice())?;
s[min_idx] = 0.0;
let f_rank2 = u * Matrix3::from_diagonal(&s) * v_t2;
Some(t1.transpose() * f_rank2 * t0)
}
fn argmin(values: &[f64]) -> Option<usize> {
values
.iter()
.enumerate()
.min_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(i, _)| i)
}
fn sampson_distance(f: &Matrix3<f64>, p0: (f64, f64), p1: (f64, f64)) -> f64 {
let x0 = Vector3::new(p0.0, p0.1, 1.0);
let x1 = Vector3::new(p1.0, p1.1, 1.0);
let fx0 = f * x0;
let ftx1 = f.transpose() * x1;
let numerator = x1.dot(&fx0).powi(2);
let denom = fx0.x.powi(2) + fx0.y.powi(2) + ftx1.x.powi(2) + ftx1.y.powi(2);
if denom < 1e-12 {
f64::INFINITY
} else {
numerator / denom
}
}
/// RANSAC estimation of the fundamental matrix relating `pts0` to `pts1` (index i in
/// `pts0` corresponds to index i in `pts1`). Returns `None` if fewer than 8
/// correspondences are given or no valid model is found.
pub fn ransac_fundamental_matrix(
pts0: &[(f32, f32)],
pts1: &[(f32, f32)],
max_iterations: usize,
inlier_threshold: f64,
confidence: f64,
) -> Option<RansacResult> {
let n = pts0.len();
if n < 8 || pts1.len() != n {
return None;
}
let pts0f: Vec<(f64, f64)> = pts0.iter().map(|&(x, y)| (x as f64, y as f64)).collect();
let pts1f: Vec<(f64, f64)> = pts1.iter().map(|&(x, y)| (x as f64, y as f64)).collect();
let mut rng = rand::rng();
let mut best_inliers: Vec<usize> = Vec::new();
let mut best_f: Option<Matrix3<f64>> = None;
let mut iterations_needed = max_iterations;
let mut iter = 0;
while iter < iterations_needed {
iter += 1;
let sample = rand::seq::index::sample(&mut rng, n, 8);
let sample0: Vec<(f64, f64)> = sample.iter().map(|i| pts0f[i]).collect();
let sample1: Vec<(f64, f64)> = sample.iter().map(|i| pts1f[i]).collect();
let Some(f) = fit_fundamental(&sample0, &sample1) else {
continue;
};
let inliers: Vec<usize> = (0..n)
.filter(|&i| sampson_distance(&f, pts0f[i], pts1f[i]) < inlier_threshold)
.collect();
if inliers.len() > best_inliers.len() {
let inlier_ratio = inliers.len() as f64 / n as f64;
best_inliers = inliers;
best_f = Some(f);
if inlier_ratio >= 1.0 {
iterations_needed = iter;
} else if inlier_ratio > 0.0 {
let denom = (1.0 - inlier_ratio.powi(8)).ln();
if denom < 0.0 {
let needed = ((1.0 - confidence).ln() / denom).ceil();
if needed.is_finite() {
iterations_needed = iterations_needed.min((needed as usize).max(1));
}
}
}
}
}
let _ = best_f?;
if best_inliers.len() < 8 {
return None;
}
// Refine F using all inliers found so far, then recompute the inlier set against it.
let in0: Vec<(f64, f64)> = best_inliers.iter().map(|&i| pts0f[i]).collect();
let in1: Vec<(f64, f64)> = best_inliers.iter().map(|&i| pts1f[i]).collect();
let refined_f = fit_fundamental(&in0, &in1)?;
let final_inliers: Vec<usize> = (0..n)
.filter(|&i| sampson_distance(&refined_f, pts0f[i], pts1f[i]) < inlier_threshold)
.collect();
let f_matrix: [f64; 9] = [
refined_f[(0, 0)],
refined_f[(0, 1)],
refined_f[(0, 2)],
refined_f[(1, 0)],
refined_f[(1, 1)],
refined_f[(1, 2)],
refined_f[(2, 0)],
refined_f[(2, 1)],
refined_f[(2, 2)],
];
Some(RansacResult {
f_matrix,
inliers: final_inliers,
})
}