DarkRoom/core/dr-pano/src/matching.rs

//! Descriptor matching between two images.
//!
//! Mutual nearest neighbour on cosine similarity, with a floor on the
//! similarity — the reference XFeat's own matcher (`match_mkpts`,
//! `min_cossim = 0.82`). For a panorama that is enough: one lens, one
//! scene, near-pure rotation and 20–40 % overlap make the matching problem
//! easy, and what is hard — sky, repeated structure, exposure drift — is
//! handled by the detector's descriptors and by RANSAC downstream, not by a
//! cleverer matcher. A learned matcher (LightGlue) is the step after this
//! one fails on a real set, and it has not (panorama.md §6).
//!
//! Brute force. `4096 × 4096 × 64` multiply-adds is a billion per pair,
//! and a twelve-frame set has sixty-six pairs: a minute single-threaded
//! and scalar (measured 2026-09-19: 51 s), a few seconds vectorised across
//! the cores. Not worth an index, but worth doing properly.

use crate::features::{Features, DESCRIPTOR_LEN};

const _: () = assert!(DESCRIPTOR_LEN % 8 == 0);

/// A correspondence: keypoint `a` in the first image matches keypoint `b`
/// in the second, with the cosine similarity of their descriptors.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Match {
    pub a: usize,
    pub b: usize,
    pub similarity: f32,
}

/// Match two sets of features.
///
/// A pair is kept when each is the other's nearest neighbour and their
/// similarity is at least `min_similarity`.
pub fn match_features(a: &Features, b: &Features, min_similarity: f32) -> Vec<Match> {
    if a.is_empty() || b.is_empty() {
        return Vec::new();
    }
    let (na, nb) = (a.len(), b.len());

    // The whole similarity matrix, once. Both nearest-neighbour directions
    // read it, which halves the multiply-adds against computing each
    // direction on its own; 4096 × 4096 × f32 is 64 MB, transient.
    let mut sim = vec![0.0f32; na * nb];
    let threads = std::thread::available_parallelism()
        .map(usize::from)
        .unwrap_or(1)
        .clamp(1, 16);
    let rows_per = na.div_ceil(threads);
    std::thread::scope(|scope| {
        for (t, chunk) in sim.chunks_mut(rows_per * nb).enumerate() {
            scope.spawn(move || {
                let first = t * rows_per;
                for (r, row) in chunk.chunks_mut(nb).enumerate() {
                    let da = a.descriptor(first + r);
                    for (j, cell) in row.iter_mut().enumerate() {
                        *cell = dot(da, b.descriptor(j));
                    }
                }
            });
        }
    });

    // Best in `b` for each `a`, and best in `a` for each `b`.
    let best_ab: Vec<(usize, f32)> = sim
        .chunks_exact(nb)
        .map(|row| {
            row.iter()
                .enumerate()
                .fold((0usize, f32::MIN), |acc, (j, &s)| if s > acc.1 { (j, s) } else { acc })
        })
        .collect();
    let mut best_ba = vec![(0usize, f32::MIN); nb];
    for (i, row) in sim.chunks_exact(nb).enumerate() {
        for (j, &s) in row.iter().enumerate() {
            if s > best_ba[j].1 {
                best_ba[j] = (i, s);
            }
        }
    }

    best_ab
        .iter()
        .enumerate()
        .filter_map(|(ia, &(ib, s))| {
            (best_ba[ib].0 == ia && s >= min_similarity).then_some(Match {
                a: ia,
                b: ib,
                similarity: s,
            })
        })
        .collect()
}

#[inline]
fn dot(a: &[f32], b: &[f32]) -> f32 {
    // Eight independent accumulators over exact 8-lane chunks: the shape
    // the compiler turns into one vector multiply-add per chunk, and no
    // bounds checks inside the loop. `DESCRIPTOR_LEN` is a multiple of 8.
    let (a, b) = (&a[..DESCRIPTOR_LEN], &b[..DESCRIPTOR_LEN]);
    let mut acc = [0.0f32; 8];
    for (ca, cb) in a.chunks_exact(8).zip(b.chunks_exact(8)) {
        for k in 0..8 {
            acc[k] += ca[k] * cb[k];
        }
    }
    acc.iter().sum()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::features::Keypoint;

    /// Features whose descriptors are unit vectors along the given axes.
    fn along(axes: &[usize]) -> Features {
        let mut descriptors = vec![0.0; axes.len() * DESCRIPTOR_LEN];
        for (i, &ax) in axes.iter().enumerate() {
            descriptors[i * DESCRIPTOR_LEN + ax] = 1.0;
        }
        Features {
            keypoints: axes
                .iter()
                .map(|_| Keypoint {
                    x: 0.0,
                    y: 0.0,
                    score: 1.0,
                })
                .collect(),
            descriptors,
            width: 1,
            height: 1,
        }
    }

    #[test]
    fn identical_descriptors_match_mutually() {
        let a = along(&[0, 1, 2]);
        let b = along(&[2, 0, 1]);
        let m = match_features(&a, &b, 0.8);
        let mut pairs: Vec<(usize, usize)> = m.iter().map(|m| (m.a, m.b)).collect();
        pairs.sort();
        assert_eq!(pairs, vec![(0, 1), (1, 2), (2, 0)]);
        assert!(m.iter().all(|m| (m.similarity - 1.0).abs() < 1e-6));
    }

    #[test]
    fn a_descriptor_with_no_counterpart_is_unmatched() {
        let a = along(&[0, 1, 5]);
        let b = along(&[0, 1]);
        let m = match_features(&a, &b, 0.8);
        assert_eq!(m.len(), 2);
        assert!(m.iter().all(|m| m.a != 2));
    }

    #[test]
    fn mutuality_breaks_a_one_sided_match() {
        // b0 is the nearest to both a0 and a1, but a0 is its nearest — a1
        // must not be matched to it.
        let mut a = along(&[0, 0]);
        a.descriptors[DESCRIPTOR_LEN] = 0.9;
        a.descriptors[DESCRIPTOR_LEN + 1] = (1.0f32 - 0.81).sqrt();
        let b = along(&[0]);
        let m = match_features(&a, &b, 0.0);
        assert_eq!(m.len(), 1);
        assert_eq!((m[0].a, m[0].b), (0, 0));
    }

    #[test]
    fn empty_input_is_empty_output() {
        assert!(match_features(&along(&[]), &along(&[1]), 0.5).is_empty());
    }
}