long_screen/kernels/
features.rs

1//! Corner detection and binary descriptors on analysis-resolution grey images.
2//!
3//! The scoring is the minimum eigenvalue of the 3×3 structure tensor computed from exact integer
4//! rolling row sums; ties and ordering follow the historical implementation so keyframe words, matches
5//! and downstream motion hypotheses are reproducible across releases.
6
7use crate::kernels::geometry::{js_round, Rect, Rng};
8use std::sync::OnceLock;
9
10pub const DESCRIPTOR_WORDS: usize = 8;
11pub const MAX_FEATURES_DEFAULT: usize = 480;
12
13#[derive(Clone, Debug, PartialEq)]
14pub struct Feature {
15    pub x: i32,
16    pub y: i32,
17    pub score: f64,
18    pub descriptor: [u32; DESCRIPTOR_WORDS],
19}
20
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub struct Match {
23    pub a: u32,
24    pub b: u32,
25    pub distance: u16,
26    pub unique: bool,
27}
28
29/// Deterministic BRIEF sampling pairs: 256 × (ax, ay, bx, by), each coordinate in −9..=9.
30fn pairs() -> &'static [[i32; 4]; 256] {
31    static PAIRS: OnceLock<[[i32; 4]; 256]> = OnceLock::new();
32    PAIRS.get_or_init(|| {
33        let mut rng = Rng::new(0x0abc_7361);
34        let mut table = [[0i32; 4]; 256];
35        for pair in table.iter_mut() {
36            for coord in pair.iter_mut() {
37                let sum = rng.next_f64() + rng.next_f64() + rng.next_f64() - 1.5;
38                *coord = js_round(sum * 6.0);
39            }
40        }
41        table
42    })
43}
44
45/// 3×3 binomial blur; the one-pixel border is copied unchanged.
46pub fn smooth(gray: &[u8], width: usize, height: usize, out: &mut [u8]) {
47    out[..width * height].copy_from_slice(&gray[..width * height]);
48    if width < 3 || height < 3 {
49        return;
50    }
51    for y in 1..height - 1 {
52        for x in 1..width - 1 {
53            let i = y * width + x;
54            let a = &gray;
55            let sum = a[i - width - 1] as u32
56                + 2 * a[i - width] as u32
57                + a[i - width + 1] as u32
58                + 2 * a[i - 1] as u32
59                + 4 * a[i] as u32
60                + 2 * a[i + 1] as u32
61                + a[i + width - 1] as u32
62                + 2 * a[i + width] as u32
63                + a[i + width + 1] as u32;
64            out[i] = (sum >> 4) as u8;
65        }
66    }
67}
68
69struct Candidate {
70    x: i32,
71    y: i32,
72    score: f64,
73}
74
75/// Horizontal three-gradient sums for one interior row, indexed by the centre pixel.
76fn tensor_row(d: &[u8], width: usize, y: usize, out: &mut [[i32; 3]]) {
77    let gradient = |x: usize| {
78        let i = y * width + x;
79        let gx = d[i + 1] as i32 - d[i - 1] as i32;
80        let gy = d[i + width] as i32 - d[i - width] as i32;
81        [gx * gx, gy * gy, gx * gy]
82    };
83    let mut left = gradient(10);
84    let mut centre = gradient(11);
85    for (x, sum) in out.iter_mut().enumerate().take(width - 11).skip(11) {
86        let right = gradient(x + 1);
87        *sum = [
88            left[0] + centre[0] + right[0],
89            left[1] + centre[1] + right[1],
90            left[2] + centre[2] + right[2],
91        ];
92        left = centre;
93        centre = right;
94    }
95}
96
97/// Spatially balanced minimum-eigenvalue corners with 256-bit BRIEF descriptors, strongest first.
98pub fn extract_features(
99    gray: &[u8],
100    width: usize,
101    height: usize,
102    max_features: usize,
103    roi: Option<Rect>,
104) -> Vec<Feature> {
105    let mut d = vec![0u8; width * height];
106    smooth(gray, width, height, &mut d);
107    let (w, h) = (width, height);
108    if w < 23 || h < 23 {
109        return Vec::new();
110    }
111    // Only 3×3 sums are queried. Three rows of horizontal tensors plus one cell-height score
112    // band preserve the historical cell-major order without three full-frame i64 integral planes.
113    // |gradient| <= 255, so even a 3×3 sum fits i32 (9 * 255² = 585225).
114    const CELL: usize = 28;
115    let mut tensors = vec![[0i32; 3]; 3 * w];
116    let mut scores = vec![0.0; CELL * w];
117    tensor_row(&d, w, 10, &mut tensors[w..2 * w]);
118    tensor_row(&d, w, 11, &mut tensors[2 * w..3 * w]);
119    let mut candidates: Vec<Candidate> = Vec::new();
120    let mut local: Vec<(i32, i32, f64)> = Vec::with_capacity(CELL * CELL);
121    let mut by = 11;
122    while by < h - 11 {
123        let end = (h - 11).min(by + CELL);
124        for y in by..end {
125            let slot = (y + 1) % 3 * w;
126            tensor_row(&d, w, y + 1, &mut tensors[slot..slot + w]);
127            for x in 11..w - 11 {
128                let xx = (tensors[x][0] + tensors[w + x][0] + tensors[2 * w + x][0]) as f64;
129                let yy = (tensors[x][1] + tensors[w + x][1] + tensors[2 * w + x][1]) as f64;
130                let xy = (tensors[x][2] + tensors[w + x][2] + tensors[2 * w + x][2]) as f64;
131                scores[(y - by) * w + x] =
132                    (xx + yy - ((xx - yy) * (xx - yy) + 4.0 * xy * xy).sqrt()) / 2.0;
133            }
134        }
135        let mut bx = 11;
136        while bx < w - 11 {
137            local.clear();
138            for y in by..(h - 11).min(by + CELL) {
139                for x in bx..(w - 11).min(bx + CELL) {
140                    if let Some(r) = roi {
141                        if !r.contains(x as f64, y as f64) {
142                            continue;
143                        }
144                    }
145                    let score = scores[(y - by) * w + x];
146                    if score > 100.0 {
147                        local.push((x as i32, y as i32, score));
148                    }
149                }
150            }
151            // Three stable argmax passes with a 6px suppression radius, first maximum wins ties.
152            for _ in 0..3 {
153                let mut best: Option<usize> = None;
154                let mut score = 100.0;
155                for (i, c) in local.iter().enumerate() {
156                    if c.2 > score {
157                        best = Some(i);
158                        score = c.2;
159                    }
160                }
161                let Some(best) = best else { break };
162                let (x, y) = (local[best].0, local[best].1);
163                candidates.push(Candidate { x, y, score });
164                for c in local.iter_mut() {
165                    if (c.0 - x).pow(2) + (c.1 - y).pow(2) <= 36 {
166                        c.2 = 0.0;
167                    }
168                }
169            }
170            bx += CELL;
171        }
172        by += CELL;
173    }
174    candidates.sort_by(|a, b| {
175        b.score
176            .partial_cmp(&a.score)
177            .unwrap_or(std::cmp::Ordering::Equal)
178    });
179    candidates.truncate(max_features);
180    let table = pairs();
181    candidates
182        .into_iter()
183        .map(|c| {
184            let mut descriptor = [0u32; DESCRIPTOR_WORDS];
185            for (bit, [ax, ay, bx, by]) in table.iter().enumerate() {
186                let p = ((c.y + ay) as usize) * w + (c.x + ax) as usize;
187                let q = ((c.y + by) as usize) * w + (c.x + bx) as usize;
188                if d[p] < d[q] {
189                    descriptor[bit >> 5] |= 1 << (bit & 31);
190                }
191            }
192            Feature {
193                x: c.x,
194                y: c.y,
195                score: c.score,
196                descriptor,
197            }
198        })
199        .collect()
200}
201
202#[inline]
203pub fn hamming(a: &[u32; DESCRIPTOR_WORDS], b: &[u32; DESCRIPTOR_WORDS]) -> u32 {
204    a.iter().zip(b).map(|(x, y)| (x ^ y).count_ones()).sum()
205}
206
207/// Mutual-nearest distinct matches plus explicitly flagged ambiguous alternatives, in `a` order.
208pub fn match_features(a: &[Feature], b: &[Feature], include_ambiguous: bool) -> Vec<Match> {
209    if a.is_empty() || b.is_empty() {
210        return Vec::new();
211    }
212    let mut best_a = vec![-1i32; a.len()];
213    let mut best_dist = vec![257u32; a.len()];
214    let mut second = vec![257u32; a.len()];
215    let mut alt = vec![-1i32; a.len()];
216    let mut best_b = vec![-1i32; b.len()];
217    let mut dist_b = vec![257u32; b.len()];
218    for (i, fa) in a.iter().enumerate() {
219        for (j, fb) in b.iter().enumerate() {
220            let d = hamming(&fa.descriptor, &fb.descriptor);
221            if d < best_dist[i] {
222                second[i] = best_dist[i];
223                alt[i] = best_a[i];
224                best_dist[i] = d;
225                best_a[i] = j as i32;
226            } else if d < second[i] {
227                second[i] = d;
228                alt[i] = j as i32;
229            }
230            if d < dist_b[j] {
231                dist_b[j] = d;
232                best_b[j] = i as i32;
233            }
234        }
235    }
236    let mut out = Vec::new();
237    for i in 0..a.len() {
238        let j = best_a[i];
239        if j < 0 || best_dist[i] > 72 {
240            continue;
241        }
242        let unique = best_b[j as usize] == i as i32
243            && (best_dist[i] as f64) < second[i] as f64 * 0.82
244            && second[i] - best_dist[i] >= 4;
245        if unique {
246            out.push(Match {
247                a: i as u32,
248                b: j as u32,
249                distance: best_dist[i] as u16,
250                unique: true,
251            });
252        } else if include_ambiguous && best_dist[i] < 45 {
253            out.push(Match {
254                a: i as u32,
255                b: j as u32,
256                distance: best_dist[i] as u16,
257                unique: false,
258            });
259            let k = alt[i];
260            if k >= 0 && second[i] < best_dist[i] + 8 {
261                out.push(Match {
262                    a: i as u32,
263                    b: k as u32,
264                    distance: second[i] as u16,
265                    unique: false,
266                });
267            }
268        }
269    }
270    out
271}
272
273/// Four independent 12-bit vocabulary bands per descriptor; deduplicated in first-seen order.
274pub fn feature_words(descriptors: &[[u32; DESCRIPTOR_WORDS]]) -> Vec<u32> {
275    let mut words = Vec::new();
276    let mut seen = [0u64; 256]; // 4 bands × 4096 words
277    for d in descriptors {
278        for k in 0..4 {
279            let word = ((k as u32) << 12) | ((d[k * 2] ^ (d[k * 2 + 1] >> 8)) & 4095);
280            let (slot, bit) = ((word >> 6) as usize, 1u64 << (word & 63));
281            if seen[slot] & bit == 0 {
282                seen[slot] |= bit;
283                words.push(word);
284            }
285        }
286    }
287    words
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    #[test]
295    fn pair_table_is_bounded_and_stable() {
296        let table = pairs();
297        assert!(table.iter().flatten().all(|c| (-9..=9).contains(c)));
298        assert_eq!(table[0], pairs()[0]);
299    }
300
301    #[test]
302    fn textured_image_yields_ordered_features() {
303        let (w, h) = (200usize, 160usize);
304        let mut rng = Rng::new(97);
305        let gray: Vec<u8> = (0..w * h).map(|_| (rng.next_f64() * 256.0) as u8).collect();
306        let features = extract_features(&gray, w, h, 480, None);
307        assert!(features.len() > 50);
308        assert!(features.windows(2).all(|p| p[0].score >= p[1].score));
309        let matches = match_features(&features, &features, true);
310        assert!(matches.iter().filter(|m| m.unique).count() >= features.len() / 2);
311    }
312}
long_screen/kernels/features.rs

long_screen/kernels/
features.rs
