use crate::decode::Gray;
use crate::timed;
pub const DESC_LEN: usize = 128;
const D: usize = 4; const N: usize = 8; const ORI_BINS: usize = 36;
const ORI_SIG_FCTR: f32 = 1.5;
const ORI_RADIUS: f32 = 3.0 * ORI_SIG_FCTR;
const ORI_PEAK_RATIO: f32 = 0.8;
const DESCR_SCL_FCTR: f32 = 3.0;
const DESCR_MAG_THR: f32 = 0.2;
const INT_DESCR_FCTR: f32 = 512.0;
const IMG_BORDER: i32 = 5;
const MAX_INTERP_STEPS: usize = 5;
#[derive(Clone, Copy, Debug)]
pub struct Params {
pub n_layers: usize,
pub sigma: f32,
pub contrast: f32,
pub edge: f32,
pub max_features: usize,
pub upsample_below: usize,
pub candidate_pool: usize,
}
impl Default for Params {
fn default() -> Self {
Params {
n_layers: 3,
sigma: 1.6,
contrast: 2.0 / 255.0,
edge: 10.0,
max_features: 800,
upsample_below: 512,
candidate_pool: 3,
}
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct Keypoint {
pub x: f32,
pub y: f32,
pub sigma: f32,
pub angle: f32,
pub response: f32,
}
#[derive(Clone, Debug, Default)]
pub struct Features {
pub w: u32,
pub h: u32,
pub kps: Vec<Keypoint>,
pub desc: Vec<u8>,
}
impl Features {
#[inline]
pub fn d(&self, i: usize) -> &[u8] {
&self.desc[i * DESC_LEN..(i + 1) * DESC_LEN]
}
pub fn len(&self) -> usize {
self.kps.len()
}
}
const EXP_RANGE: f32 = 40.0;
const EXP_N: usize = 8192;
struct ExpTable([f32; EXP_N]);
static EXP_TABLE: std::sync::LazyLock<ExpTable> = std::sync::LazyLock::new(|| {
let mut t = [0f32; EXP_N];
for (i, v) in t.iter_mut().enumerate() {
*v = (-(i as f32 + 0.5) * EXP_RANGE / EXP_N as f32).exp();
}
ExpTable(t)
});
impl ExpTable {
#[inline(always)]
fn at(&self, t: f32) -> f32 {
if t >= EXP_RANGE {
return 0.0;
}
self.at_index(unsafe { Self::index(t) })
}
#[inline(always)]
unsafe fn index(t: f32) -> u32 {
let f = t * (EXP_N as f32 / EXP_RANGE);
unsafe { f.to_int_unchecked::<u32>() }
}
#[inline(always)]
fn at_index(&self, i: u32) -> f32 {
debug_assert!((i as usize) < EXP_N);
unsafe { *self.0.get_unchecked(i as usize) }
}
}
#[inline]
pub fn fast_atan2_deg(y: f32, x: f32) -> f32 {
const P1: f32 = 0.999_787_8 * (180.0 / std::f32::consts::PI);
const P3: f32 = -0.325_808_4 * (180.0 / std::f32::consts::PI);
const P5: f32 = 0.155_578_65 * (180.0 / std::f32::consts::PI);
const P7: f32 = -0.044_326_555 * (180.0 / std::f32::consts::PI);
let ax = x.abs();
let ay = y.abs();
let steep = ax < ay;
let num = if steep { ax } else { ay };
let den = if steep { ay } else { ax };
let c = num / (den + f32::EPSILON);
let c2 = c * c;
let a = (((P7 * c2 + P5) * c2 + P3) * c2 + P1) * c;
let a = if steep { 90.0 - a } else { a };
let a = if x < 0.0 { 180.0 - a } else { a };
if y < 0.0 { 360.0 - a } else { a }
}
struct Grad {
w: usize,
px: Vec<[f32; 2]>,
}
impl Grad {
fn of(l: &Layer) -> Grad {
let (w, h) = (l.w, l.h);
let n = w * h;
let mut px: Vec<[f32; 2]> = Vec::with_capacity(n);
{
let spare = &mut px.spare_capacity_mut()[..n];
for y in 0..h {
let urow = &mut spare[y * w..(y + 1) * w];
if y == 0 || y + 1 >= h || w < 3 {
for u in urow.iter_mut() {
u.write([0.0, 0.0]);
}
continue;
}
let up = &l.px[(y - 1) * w..y * w];
let row = &l.px[y * w..(y + 1) * w];
let dn = &l.px[(y + 1) * w..(y + 2) * w];
urow[0].write([0.0, 0.0]);
urow[w - 1].write([0.0, 0.0]);
for x in 1..w - 1 {
let dx = row[x + 1] - row[x - 1];
let dy = up[x] - dn[x];
urow[x].write([(dx * dx + dy * dy).sqrt(), fast_atan2_deg(dy, dx)]);
}
}
}
unsafe { px.set_len(n) };
Grad { w, px }
}
}
struct Layer {
w: usize,
h: usize,
px: Vec<f32>,
}
impl Layer {
#[inline]
fn at(&self, x: i32, y: i32) -> f32 {
self.px[y as usize * self.w + x as usize]
}
}
fn gaussian_kernel(sigma: f32) -> Vec<f32> {
let radius = (sigma * 3.0).ceil().max(1.0) as usize;
let mut k = vec![0.0f32; 2 * radius + 1];
let mut sum = 0.0;
for i in 0..k.len() {
let x = i as f32 - radius as f32;
let v = (-x * x / (2.0 * sigma * sigma)).exp();
k[i] = v;
sum += v;
}
for v in k.iter_mut() {
*v /= sum;
}
k
}
struct BlurScratch {
ring: Vec<f32>,
padded: Vec<f32>,
row: Vec<f32>,
}
thread_local! {
static BLUR_SCRATCH: std::cell::RefCell<BlurScratch> =
const { std::cell::RefCell::new(BlurScratch { ring: Vec::new(), padded: Vec::new(), row: Vec::new() }) };
}
pub fn release_scratch() {
let _ = BLUR_SCRATCH.try_with(|s| {
let mut s = s.borrow_mut();
s.ring = Vec::new();
s.padded = Vec::new();
s.row = Vec::new();
});
}
fn blur(src: &Layer, sigma: f32) -> Layer {
blur_plane(src.w, src.h, &src.px, sigma)
}
fn blur_plane(w: usize, h: usize, px: &[f32], sigma: f32) -> Layer {
BLUR_SCRATCH.with(|s| blur_into(w, h, px, sigma, &mut s.borrow_mut(), false, true).0.unwrap())
}
fn blur_dog(src: &Layer, sigma: f32) -> (Layer, Layer) {
let (g, d) = BLUR_SCRATCH.with(|s| blur_into(src.w, src.h, &src.px, sigma, &mut s.borrow_mut(), true, true));
(g.unwrap(), d.unwrap())
}
fn blur_top(src: &Layer, sigma: f32) -> Layer {
BLUR_SCRATCH.with(|s| blur_into(src.w, src.h, &src.px, sigma, &mut s.borrow_mut(), true, false)).1.unwrap()
}
#[inline]
fn blur_row(row: &[f32], padded: &mut [f32], out: &mut [f32], kc: f32, ks: &[f32], r: usize, w: usize) {
for i in 0..r {
padded[i] = row[reflect101(i as i32 - r as i32, w)];
padded[w + r + i] = row[reflect101((w + i) as i32, w)];
}
padded[r..r + w].copy_from_slice(row);
let mut x = 0;
while x + 8 <= w {
let mut acc = [0f32; 8];
let c = &padded[x + r..x + r + 8];
for i in 0..8 {
acc[i] = c[i] * kc;
}
for (t, &kv) in ks.iter().enumerate() {
let l = &padded[x + r - t - 1..x + r - t - 1 + 8];
let rr = &padded[x + r + t + 1..x + r + t + 1 + 8];
for i in 0..8 {
acc[i] += (l[i] + rr[i]) * kv;
}
}
out[x..x + 8].copy_from_slice(&acc);
x += 8;
}
while x < w {
let mut acc = padded[x + r] * kc;
for (t, &kv) in ks.iter().enumerate() {
acc += (padded[x + r - t - 1] + padded[x + r + t + 1]) * kv;
}
out[x] = acc;
x += 1;
}
}
fn blur_into(w: usize, h: usize, src: &[f32], sigma: f32, s: &mut BlurScratch, want_dog: bool, want_gauss: bool) -> (Option<Layer>, Option<Layer>) {
let k = gaussian_kernel(sigma);
let r = k.len() / 2;
let kc = k[r];
let ks: &[f32] = &k[r + 1..];
let ring_rows = (2 * r + 1).min(h);
if s.ring.len() < ring_rows * w {
s.ring.resize(ring_rows * w, 0.0);
}
if s.padded.len() < w + 2 * r {
s.padded.resize(w + 2 * r, 0.0);
}
if !want_gauss && s.row.len() < w {
s.row.resize(w, 0.0);
}
let ring = &mut s.ring[..ring_rows * w];
let padded = &mut s.padded[..w + 2 * r];
let scratch_row = &mut s.row[..if want_gauss { 0 } else { w }];
let mut dst: Vec<f32> = Vec::with_capacity(if want_gauss { w * h } else { 0 });
let mut dog: Vec<f32> = Vec::with_capacity(if want_dog { w * h } else { 0 });
let mut filtered = 0usize; {
let spare = dst.spare_capacity_mut();
for y in 0..h {
let want = (y + r).min(h - 1);
while filtered <= want {
let slot = filtered % ring_rows;
blur_row(
&src[filtered * w..(filtered + 1) * w],
padded,
&mut ring[slot * w..(slot + 1) * w],
kc,
ks,
r,
w,
);
filtered += 1;
}
let base = y % ring_rows;
let c = &ring[base * w..(base + 1) * w];
let acc: &mut [f32] = if want_gauss {
let urow = &mut spare[y * w..(y + 1) * w];
for x in 0..w {
urow[x].write(c[x] * kc);
}
unsafe { &mut *(urow as *mut [std::mem::MaybeUninit<f32>] as *mut [f32]) }
} else {
for x in 0..w {
scratch_row[x] = c[x] * kc;
}
&mut scratch_row[..]
};
let interior = y >= r && y + r < h;
for (t, &kv) in ks.iter().enumerate() {
let (ya, yb) = if interior {
let mut ya = base + ring_rows - t - 1;
if ya >= ring_rows {
ya -= ring_rows;
}
let mut yb = base + t + 1;
if yb >= ring_rows {
yb -= ring_rows;
}
(ya, yb)
} else {
(
reflect101(y as i32 - t as i32 - 1, h) % ring_rows,
reflect101(y as i32 + t as i32 + 1, h) % ring_rows,
)
};
debug_assert_eq!(ya, reflect101(y as i32 - t as i32 - 1, h) % ring_rows);
debug_assert_eq!(yb, reflect101(y as i32 + t as i32 + 1, h) % ring_rows);
let a = &ring[ya * w..(ya + 1) * w];
let b = &ring[yb * w..(yb + 1) * w];
for x in 0..w {
acc[x] += (a[x] + b[x]) * kv;
}
}
if want_dog {
let below = &src[y * w..(y + 1) * w];
dog.extend(acc.iter().zip(below).map(|(a, b)| a - b));
}
}
}
let g = want_gauss.then(|| {
unsafe { dst.set_len(w * h) };
Layer { w, h, px: dst }
});
let d = want_dog.then(|| Layer { w, h, px: dog });
(g, d)
}
#[inline]
fn reflect101(i: i32, n: usize) -> usize {
let n = n as i32;
if n == 1 {
return 0;
}
let mut i = i;
while i < 0 || i >= n {
if i < 0 {
i = -i;
}
if i >= n {
i = 2 * (n - 1) - i;
}
}
i as usize
}
fn upsample(g: &Gray, f: usize) -> Layer {
let (w, h) = (g.w * f, g.h * f);
let ff = f as f32;
let mut cols: Vec<(usize, usize, f32)> = Vec::with_capacity(w);
for x in 0..w {
let sx = (x as f32 + 0.5) / ff - 0.5;
let x0 = sx.floor().max(0.0) as usize;
let x1 = (x0 + 1).min(g.w - 1);
cols.push((x0, x1, (sx - x0 as f32).clamp(0.0, 1.0)));
}
let mut px: Vec<f32> = Vec::with_capacity(w * h);
for y in 0..h {
let sy = (y as f32 + 0.5) / ff - 0.5;
let y0 = sy.floor().max(0.0) as usize;
let y1 = (y0 + 1).min(g.h - 1);
let fy = (sy - y0 as f32).clamp(0.0, 1.0);
let r0 = &g.px[y0 * g.w..(y0 + 1) * g.w];
let r1 = &g.px[y1 * g.w..(y1 + 1) * g.w];
px.extend(cols.iter().map(|&(x0, x1, fx)| {
let a = r0[x0] * (1.0 - fx) + r0[x1] * fx;
let b = r1[x0] * (1.0 - fx) + r1[x1] * fx;
a * (1.0 - fy) + b * fy
}));
}
Layer { w, h, px }
}
fn halve(src: &Layer) -> Layer {
let (w, h) = ((src.w / 2).max(1), (src.h / 2).max(1));
let mut px: Vec<f32> = Vec::with_capacity(w * h);
for y in 0..h {
let row = &src.px[(y * 2) * src.w..];
px.extend((0..w).map(|x| row[x * 2]));
}
Layer { w, h, px }
}
pub fn extract(g: &Gray, p: &Params) -> Features {
let mut feats = Features { w: g.w as u32, h: g.h as u32, ..Default::default() };
if g.w < 8 || g.h < 8 {
return feats;
}
let mut factor = 1usize;
while g.w.max(g.h) * factor * 2 <= p.upsample_below.max(2) {
factor *= 2;
}
let (init_sigma, coord_scale) = if factor > 1 { (0.5 * factor as f32, 1.0 / factor as f32) } else { (0.5f32, 1.0f32) };
let sig_diff = (p.sigma * p.sigma - init_sigma * init_sigma).max(0.01).sqrt();
let base = timed!(5, {
if factor > 1 {
blur(&upsample(g, factor), sig_diff)
} else {
blur_plane(g.w, g.h, &g.px, sig_diff)
}
});
let min_side = base.w.min(base.h) as f32;
let n_octaves = ((min_side.ln() / 2f32.ln()).round() as i32 - 2).max(1) as usize;
let s = p.n_layers;
let k = 2f32.powf(1.0 / s as f32);
let mut sig = vec![p.sigma; s + 3];
for i in 1..s + 3 {
let prev = p.sigma * k.powi(i as i32 - 1);
let total = prev * k;
sig[i] = (total * total - prev * prev).sqrt();
}
let thr_pre = 0.5 * p.contrast / s as f32;
let mut cands: Vec<Cand> = Vec::new();
let mut grads: Vec<Vec<Option<Grad>>> = Vec::with_capacity(n_octaves);
let mut heights: Vec<usize> = Vec::with_capacity(n_octaves);
let mut octave_base = base;
for o in 0..n_octaves {
let height = octave_base.h;
let mut gauss: Vec<Option<Layer>> = Vec::with_capacity(s + 3);
gauss.push(Some(std::mem::replace(&mut octave_base, Layer { w: 0, h: 0, px: vec![] })));
let mut dog: Vec<Layer> = Vec::with_capacity(s + 2);
for i in 1..s + 3 {
if i == s + 2 {
dog.push(timed!(6, blur_top(gauss[i - 1].as_ref().unwrap(), sig[i])));
gauss.push(None);
gauss[i - 1] = None;
break;
}
let (l, d) = timed!(6, blur_dog(gauss[i - 1].as_ref().unwrap(), sig[i]));
gauss.push(Some(l));
dog.push(d);
if !(1..=s).contains(&(i - 1)) {
gauss[i - 1] = None;
}
}
gauss[s + 2] = None;
timed!(7, find_extrema(&dog, o, p, thr_pre, coord_scale, &mut cands));
drop(dog);
heights.push(height);
if o + 1 < n_octaves {
octave_base = timed!(32, halve(gauss[s].as_ref().unwrap()));
}
grads.push(timed!(8,
(0..s + 3)
.map(|i| gauss[i].take().filter(|_| (1..=s).contains(&i)).map(|l| Grad::of(&l)))
.collect::<Vec<_>>()
));
}
cands.sort_by(|a, b| {
(a.octave, key3(&a.kp))
.cmp(&(b.octave, key3(&b.kp)))
.then(b.kp.response.partial_cmp(&a.kp.response).unwrap())
});
cands.dedup_by(|a, b| a.octave == b.octave && key3(&a.kp) == key3(&b.kp));
cands.sort_by(|a, b| {
b.kp.response
.partial_cmp(&a.kp.response)
.unwrap()
.then((a.octave, key3(&a.kp)).cmp(&(b.octave, key3(&b.kp))))
});
cands.truncate(p.max_features * p.candidate_pool + 8);
const KEEP_MARGIN: usize = 8;
let stop_at = p.max_features + KEEP_MARGIN;
timed!(9, {
for c in cands.iter() {
if feats.kps.len() >= stop_at && c.kp.response < feats.kps[stop_at - 1].response {
break;
}
let oct_scale = (1u32 << c.octave) as f32 * coord_scale;
let grad = grads[c.octave][c.layer].as_ref().unwrap();
let h = heights[c.octave];
let scl_octv = c.kp.sigma / oct_scale;
let px = c.kp.x / oct_scale;
let py = c.kp.y / oct_scale;
let mut hist = [0f32; ORI_BINS];
let radius = (ORI_RADIUS * scl_octv).round() as i32;
let omax = timed!(30, orientation_hist(grad, h, px, py, radius, ORI_SIG_FCTR * scl_octv, &mut hist));
let mag_thr = omax * ORI_PEAK_RATIO;
for j in 0..ORI_BINS {
let l = if j > 0 { j - 1 } else { ORI_BINS - 1 };
let r2 = if j < ORI_BINS - 1 { j + 1 } else { 0 };
if hist[j] > hist[l] && hist[j] > hist[r2] && hist[j] >= mag_thr {
let mut bin = j as f32 + 0.5 * (hist[l] - hist[r2]) / (hist[l] - 2.0 * hist[j] + hist[r2]);
if bin < 0.0 {
bin += ORI_BINS as f32;
} else if bin >= ORI_BINS as f32 {
bin -= ORI_BINS as f32;
}
let mut angle = 360.0 - (360.0 / ORI_BINS as f32) * bin;
if (angle - 360.0).abs() < 1e-5 {
angle = 0.0;
}
let mut kp = c.kp;
kp.angle = angle;
let mut d = [0u8; DESC_LEN];
timed!(31, descriptor(grad, h, px, py, angle, scl_octv, &mut d));
feats.kps.push(kp);
feats.desc.extend_from_slice(&d);
}
}
}
});
timed!(10, retain_best(&mut feats, p.max_features));
feats
}
#[inline(always)]
fn fmax(a: f32, b: f32) -> f32 {
if a > b { a } else { b }
}
#[inline(always)]
fn fmin(a: f32, b: f32) -> f32 {
if a < b { a } else { b }
}
#[inline]
fn rows3(l: &Layer, y: usize, w: usize) -> (&[f32], &[f32], &[f32]) {
(&l.px[(y - 1) * w..y * w], &l.px[y * w..(y + 1) * w], &l.px[(y + 1) * w..(y + 2) * w])
}
#[inline]
fn is_extreme(v: f32, x: usize, rows: [&[f32]; 6], max: bool) -> bool {
if max {
for r in rows {
if v < r[x - 1] || v < r[x] || v < r[x + 1] {
return false;
}
}
} else {
for r in rows {
if v > r[x - 1] || v > r[x] || v > r[x + 1] {
return false;
}
}
}
true
}
struct Cand {
kp: Keypoint,
octave: usize,
layer: usize,
}
#[inline]
fn key3(k: &Keypoint) -> (i32, i32, i32) {
((k.x * 4.0) as i32, (k.y * 4.0) as i32, (k.sigma * 16.0) as i32)
}
fn find_extrema(dog: &[Layer], octave: usize, p: &Params, thr_pre: f32, coord_scale: f32, out: &mut Vec<Cand>) {
let s = p.n_layers;
let (w, h) = (dog[0].w as i32, dog[0].h as i32);
if w <= 2 * IMG_BORDER || h <= 2 * IMG_BORDER {
return;
}
let oct_scale = (1u32 << octave) as f32 * coord_scale;
let wu = w as usize;
let (lo, hi) = (IMG_BORDER as usize, (w - IMG_BORDER) as usize);
let span = hi - lo;
let mut alive = vec![false; span];
for layer in 1..=s {
let cur = &dog[layer];
let prv = &dog[layer - 1];
let nxt = &dog[layer + 1];
for y in IMG_BORDER..h - IMG_BORDER {
let yu = y as usize;
let (c0, c1, c2) = rows3(cur, yu, wu);
let (p0, p1, p2) = rows3(prv, yu, wu);
let (n0, n1, n2) = rows3(nxt, yu, wu);
let (vc, cl, cr) = (&c1[lo..hi], &c1[lo - 1..hi - 1], &c1[lo + 1..hi + 1]);
let (ul, um, ur) = (&c0[lo - 1..hi - 1], &c0[lo..hi], &c0[lo + 1..hi + 1]);
let (dl, dm, dr) = (&c2[lo - 1..hi - 1], &c2[lo..hi], &c2[lo + 1..hi + 1]);
for i in 0..span {
let v = vc[i];
let biggest = fmax(
fmax(fmax(cl[i], cr[i]), fmax(ul[i], um[i])),
fmax(fmax(ur[i], dl[i]), fmax(dm[i], dr[i])),
);
let smallest = fmin(
fmin(fmin(cl[i], cr[i]), fmin(ul[i], um[i])),
fmin(fmin(ur[i], dl[i]), fmin(dm[i], dr[i])),
);
alive[i] = ((v > thr_pre) & (v >= biggest)) | ((v < -thr_pre) & (v <= smallest));
}
let mut i = 0usize;
while i < span {
if i + 8 <= span {
let eight = unsafe { std::ptr::read_unaligned(alive.as_ptr().add(i) as *const u64) };
if eight == 0 {
i += 8;
continue;
}
}
let here = i;
i += 1;
if !alive[here] {
continue;
}
let i = here;
let xu = lo + i;
let v = c1[xu];
let positive = v > 0.0;
if !is_extreme(v, xu, [p0, p1, p2, n0, n1, n2], positive) {
continue;
}
if let Some((kp, lay)) = adjust(dog, octave, layer, xu as i32, y, p, oct_scale) {
out.push(Cand { kp, octave, layer: lay });
}
}
}
}
}
fn adjust(
dog: &[Layer],
_octave: usize,
layer0: usize,
x0: i32,
y0: i32,
p: &Params,
oct_scale: f32,
) -> Option<(Keypoint, usize)> {
let s = p.n_layers;
let (mut layer, mut x, mut y) = (layer0 as i32, x0, y0);
let (w, h) = (dog[0].w as i32, dog[0].h as i32);
let mut xi = 0.0f32;
let mut xr = 0.0f32;
let mut xc = 0.0f32;
let mut converged = false;
for _ in 0..MAX_INTERP_STEPS {
let cur = &dog[layer as usize];
let prv = &dog[layer as usize - 1];
let nxt = &dog[layer as usize + 1];
let dx = (cur.at(x + 1, y) - cur.at(x - 1, y)) * 0.5;
let dy = (cur.at(x, y + 1) - cur.at(x, y - 1)) * 0.5;
let ds = (nxt.at(x, y) - prv.at(x, y)) * 0.5;
let v2 = cur.at(x, y) * 2.0;
let dxx = cur.at(x + 1, y) + cur.at(x - 1, y) - v2;
let dyy = cur.at(x, y + 1) + cur.at(x, y - 1) - v2;
let dss = nxt.at(x, y) + prv.at(x, y) - v2;
let dxy = (cur.at(x + 1, y + 1) - cur.at(x - 1, y + 1) - cur.at(x + 1, y - 1) + cur.at(x - 1, y - 1)) * 0.25;
let dxs = (nxt.at(x + 1, y) - nxt.at(x - 1, y) - prv.at(x + 1, y) + prv.at(x - 1, y)) * 0.25;
let dys = (nxt.at(x, y + 1) - nxt.at(x, y - 1) - prv.at(x, y + 1) + prv.at(x, y - 1)) * 0.25;
let hm = [[dxx, dxy, dxs], [dxy, dyy, dys], [dxs, dys, dss]];
let g = [dx, dy, ds];
let sol = solve3(hm, g)?;
xc = -sol[0];
xr = -sol[1];
xi = -sol[2];
if xc.abs() < 0.5 && xr.abs() < 0.5 && xi.abs() < 0.5 {
converged = true;
break;
}
if xc.abs() > 1e6 || xr.abs() > 1e6 || xi.abs() > 1e6 {
return None;
}
x += xc.round() as i32;
y += xr.round() as i32;
layer += xi.round() as i32;
if layer < 1 || layer > s as i32 || x < IMG_BORDER || x >= w - IMG_BORDER || y < IMG_BORDER || y >= h - IMG_BORDER {
return None;
}
}
if !converged {
return None;
}
let cur = &dog[layer as usize];
let prv = &dog[layer as usize - 1];
let nxt = &dog[layer as usize + 1];
let dx = (cur.at(x + 1, y) - cur.at(x - 1, y)) * 0.5;
let dy = (cur.at(x, y + 1) - cur.at(x, y - 1)) * 0.5;
let ds = (nxt.at(x, y) - prv.at(x, y)) * 0.5;
let contr = cur.at(x, y) + 0.5 * (dx * xc + dy * xr + ds * xi);
if contr.abs() * (s as f32) < p.contrast {
return None;
}
let v2 = cur.at(x, y) * 2.0;
let dxx = cur.at(x + 1, y) + cur.at(x - 1, y) - v2;
let dyy = cur.at(x, y + 1) + cur.at(x, y - 1) - v2;
let dxy = (cur.at(x + 1, y + 1) - cur.at(x - 1, y + 1) - cur.at(x + 1, y - 1) + cur.at(x - 1, y - 1)) * 0.25;
let tr = dxx + dyy;
let det = dxx * dyy - dxy * dxy;
if det <= 0.0 || tr * tr * p.edge >= (p.edge + 1.0) * (p.edge + 1.0) * det {
return None;
}
let kp = Keypoint {
x: (x as f32 + xc) * oct_scale,
y: (y as f32 + xr) * oct_scale,
sigma: p.sigma * 2f32.powf((layer as f32 + xi) / s as f32) * oct_scale,
angle: 0.0,
response: contr.abs(),
};
Some((kp, layer as usize))
}
fn solve3(a: [[f32; 3]; 3], b: [f32; 3]) -> Option<[f32; 3]> {
let det = a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1])
- a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0])
+ a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0]);
if det.abs() < 1e-12 {
return None;
}
let inv = 1.0 / det;
let mut x = [0f32; 3];
for i in 0..3 {
let mut m = a;
for r in 0..3 {
m[r][i] = b[r];
}
let d = m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]);
x[i] = d * inv;
}
Some(x)
}
const ORI_SWEEP: usize = 32;
fn orientation_hist(g: &Grad, h: usize, px: f32, py: f32, radius: i32, sigma: f32, hist: &mut [f32; ORI_BINS]) -> f32 {
let expf_scale = -1.0 / (2.0 * sigma * sigma);
let mut temphist = [0f32; ORI_BINS];
let (w, h) = (g.w as i32, h as i32);
let cx = px.round() as i32;
let cy = py.round() as i32;
let neg_scale = -expf_scale;
let tbl = &*EXP_TABLE;
for i in -radius..=radius {
let y = cy + i;
if y <= 0 || y >= h - 1 {
continue;
}
let j0 = (-radius).max(1 - cx);
let j1 = radius.min(w - 2 - cx);
if j1 < j0 {
continue;
}
let row = y as usize * g.w;
let span = &g.px[row + (cx + j0) as usize..row + (cx + j1) as usize + 1];
let mut jj = j0;
for block in span.chunks(ORI_SWEEP) {
let mut sw_t = [0f32; ORI_SWEEP];
let mut sw_bin = [0u32; ORI_SWEEP];
for (u, &[_, ori]) in block.iter().enumerate() {
let j = jj + u as i32;
sw_t[u] = (i * i + j * j) as f32 * neg_scale;
let b = unsafe { (ori * ORI_BINS as f32 / 360.0).round().to_int_unchecked::<i32>() };
sw_bin[u] = if b >= ORI_BINS as i32 { b - ORI_BINS as i32 } else { b } as u32;
}
jj += block.len() as i32;
for (u, &[mag, _]) in block.iter().enumerate() {
let bin = sw_bin[u] as usize;
debug_assert!(bin < ORI_BINS);
unsafe { *temphist.get_unchecked_mut(bin) += tbl.at(sw_t[u]) * mag };
}
}
}
let n = ORI_BINS;
let mut maxval = 0.0f32;
for i in 0..n {
let v = (temphist[(i + n - 2) % n] + temphist[(i + 2) % n]) * (1.0 / 16.0)
+ (temphist[(i + n - 1) % n] + temphist[(i + 1) % n]) * (4.0 / 16.0)
+ temphist[i] * (6.0 / 16.0);
hist[i] = v;
maxval = maxval.max(v);
}
maxval
}
#[inline]
fn j_span(a: f32, l: f32, u: f32) -> (f32, f32) {
let (p, q) = (l / a, u / a);
if p <= q { (p, q) } else { (q, p) }
}
fn descriptor(g: &Grad, h: usize, px: f32, py: f32, kp_angle: f32, scl: f32, dst: &mut [u8; DESC_LEN]) {
let mut ori = 360.0 - kp_angle;
if (ori - 360.0).abs() < 1e-5 {
ori = 0.0;
}
let (rows, cols) = (h as i32, g.w as i32);
let pt_x = px.round() as i32;
let pt_y = py.round() as i32;
let mut cos_t = ori.to_radians().cos();
let mut sin_t = ori.to_radians().sin();
let bins_per_rad = N as f32 / 360.0;
let neg_exp_scale = 1.0 / (D as f32 * D as f32 * 0.5);
let hist_width = DESCR_SCL_FCTR * scl;
let mut radius = (hist_width * 2f32.sqrt() * (D as f32 + 1.0) * 0.5).round() as i32;
let diag = ((cols * cols + rows * rows) as f32).sqrt() as i32;
radius = radius.min(diag);
cos_t /= hist_width;
sin_t /= hist_width;
let hlen = (D + 2) * (D + 2) * (N + 2);
let mut hist = [0f32; (D + 2) * (D + 2) * (N + 2)];
debug_assert_eq!(hlen, hist.len());
let tbl = &*EXP_TABLE;
const SWEEP: usize = 64;
const RAMP: [f32; SWEEP] = [
0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0,
24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0,
32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0,
40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0,
48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0,
56.0, 57.0, 58.0, 59.0, 60.0, 61.0, 62.0, 63.0,
];
let mut sw_w = [0u32; SWEEP];
let mut sw_mag = [0f32; SWEEP];
let mut sw_rb = [0f32; SWEEP];
let mut sw_cb = [0f32; SWEEP];
let mut sw_ob = [0f32; SWEEP];
let mut sw_idx = [0i32; SWEEP];
let mut sw_in = [0u32; SWEEP];
let mut sw_hit = [0u32; SWEEP];
for i in -radius..=radius {
let r = pt_y + i;
if r <= 0 || r >= rows - 1 {
continue;
}
let fi = i as f32;
let (p1, q1) = j_span(cos_t, -2.5 + fi * sin_t, 2.5 + fi * sin_t);
let (p2, q2) = j_span(sin_t, -2.5 - fi * cos_t, 2.5 - fi * cos_t);
let lo = p1.max(p2).max(-radius as f32);
let hi = q1.min(q2).min(radius as f32);
if !(hi >= lo) {
continue;
}
let j0 = (lo.floor() as i32 - 1).max(-radius).max(1 - pt_x);
let j1 = (hi.ceil() as i32 + 1).min(radius).min(cols - 2 - pt_x);
if j1 < j0 {
continue;
}
let row = r as usize * g.w;
let span = &g.px[row + (pt_x + j0) as usize..row + (pt_x + j1) as usize + 1];
let mut jf = j0 as f32;
for block in span.chunks(SWEEP) {
for (u, &[_, o]) in block.iter().enumerate() {
let jj = jf + RAMP[u];
let c_rot = jj * cos_t - fi * sin_t;
let r_rot = jj * sin_t + fi * cos_t;
let rbin = r_rot + (D / 2) as f32 - 0.5;
let cbin = c_rot + (D / 2) as f32 - 0.5;
let t = (c_rot * c_rot + r_rot * r_rot) * neg_exp_scale;
debug_assert!((0.0..EXP_RANGE).contains(&t));
sw_w[u] = unsafe { ExpTable::index(t) };
let obin = (o - ori) * bins_per_rad;
let r0 = rbin.floor();
let c0 = cbin.floor();
let o0 = obin.floor();
sw_rb[u] = rbin - r0;
sw_cb[u] = cbin - c0;
sw_ob[u] = obin - o0;
debug_assert!(rbin.abs() < 7.0 && cbin.abs() < 7.0 && obin.abs() < N as f32);
let (r0i, c0i, o0i) = unsafe {
(
r0.to_int_unchecked::<i32>(),
c0.to_int_unchecked::<i32>(),
o0.to_int_unchecked::<i32>(),
)
};
let o0i = o0i & (N as i32 - 1);
sw_idx[u] = ((r0i + 1) * (D as i32 + 2) + c0i + 1) * (N as i32 + 2) + o0i;
sw_in[u] = (rbin > -1.0 && rbin < D as f32 && cbin > -1.0 && cbin < D as f32) as u32;
}
jf += block.len() as f32;
let mut n_hit = 0usize;
for u in 0..block.len() {
sw_mag[u] = block[u][0] * tbl.at_index(sw_w[u]);
debug_assert!(n_hit <= u && u < SWEEP);
unsafe { *sw_hit.get_unchecked_mut(n_hit) = u as u32 };
n_hit += (sw_in[u] != 0) as usize;
}
for &u in &sw_hit[..n_hit] {
let u = u as usize;
debug_assert!(u < SWEEP);
let (mag, rb, cb, ob, sidx) = unsafe {
(
*sw_mag.get_unchecked(u),
*sw_rb.get_unchecked(u),
*sw_cb.get_unchecked(u),
*sw_ob.get_unchecked(u),
*sw_idx.get_unchecked(u),
)
};
let v_r1 = mag * rb;
let v_r0 = mag - v_r1;
let v_rc11 = v_r1 * cb;
let v_rc10 = v_r1 - v_rc11;
let v_rc01 = v_r0 * cb;
let v_rc00 = v_r0 - v_rc01;
let v_rco111 = v_rc11 * ob;
let v_rco110 = v_rc11 - v_rco111;
let v_rco101 = v_rc10 * ob;
let v_rco100 = v_rc10 - v_rco101;
let v_rco011 = v_rc01 * ob;
let v_rco010 = v_rc01 - v_rco011;
let v_rco001 = v_rc00 * ob;
let v_rco000 = v_rc00 - v_rco001;
let idx = sidx as usize;
let stride_c = N + 2;
let stride_r = (D + 2) * (N + 2);
debug_assert!(idx + stride_r + stride_c + 1 < hist.len());
unsafe {
let h = hist.as_mut_ptr().add(idx);
*h += v_rco000;
*h.add(1) += v_rco001;
*h.add(stride_c) += v_rco010;
*h.add(stride_c + 1) += v_rco011;
*h.add(stride_r) += v_rco100;
*h.add(stride_r + 1) += v_rco101;
*h.add(stride_r + stride_c) += v_rco110;
*h.add(stride_r + stride_c + 1) += v_rco111;
}
}
}
}
let mut out = [0f32; DESC_LEN];
for i in 0..D {
for j in 0..D {
let idx = ((i + 1) * (D + 2) + (j + 1)) * (N + 2);
hist[idx] += hist[idx + N];
hist[idx + 1] += hist[idx + N + 1];
for k2 in 0..N {
out[(i * D + j) * N + k2] = hist[idx + k2];
}
}
}
let nrm2: f32 = out.iter().map(|v| v * v).sum();
let thr = nrm2.sqrt() * DESCR_MAG_THR;
let mut nrm2b = 0.0f32;
for v in out.iter_mut() {
if *v > thr {
*v = thr;
}
nrm2b += *v * *v;
}
let scale = INT_DESCR_FCTR / nrm2b.sqrt().max(1e-12);
for (i, v) in out.iter().enumerate() {
dst[i] = (v * scale).round().clamp(0.0, 255.0) as u8;
}
}
fn retain_best(f: &mut Features, n: usize) {
let mut idx: Vec<usize> = (0..f.kps.len()).collect();
idx.sort_by(|&a, &b| {
let (ka, kb) = (&f.kps[a], &f.kps[b]);
kb.response
.partial_cmp(&ka.response)
.unwrap()
.then(ka.x.partial_cmp(&kb.x).unwrap())
.then(ka.y.partial_cmp(&kb.y).unwrap())
.then(ka.sigma.partial_cmp(&kb.sigma).unwrap())
.then(ka.angle.partial_cmp(&kb.angle).unwrap())
});
let mut kps = Vec::with_capacity(n.min(idx.len()));
let mut desc = Vec::with_capacity(n.min(idx.len()) * DESC_LEN);
let mut last: Option<(i32, i32, i32, i32)> = None;
for &i in &idx {
let k = f.kps[i];
let key = ((k.x * 4.0) as i32, (k.y * 4.0) as i32, (k.sigma * 16.0) as i32, (k.angle * 2.0) as i32);
if last == Some(key) {
continue;
}
last = Some(key);
kps.push(k);
desc.extend_from_slice(f.d(i));
if kps.len() >= n {
break;
}
}
f.kps = kps;
f.desc = desc;
}
pub fn mirror_perm() -> [u8; DESC_LEN] {
let mut p = [0u8; DESC_LEN];
for r in 0..D {
for c in 0..D {
for o in 0..N {
let src = ((D - 1 - r) * D + c) * N + ((N - o) % N);
p[(r * D + c) * N + o] = src as u8;
}
}
}
p
}
pub fn invert_perm() -> [u8; DESC_LEN] {
let mut p = [0u8; DESC_LEN];
for r in 0..D {
for c in 0..D {
for o in 0..N {
let src = ((D - 1 - r) * D + (D - 1 - c)) * N + o;
p[(r * D + c) * N + o] = src as u8;
}
}
}
p
}
pub fn permute(desc: &[u8], perm: &[u8; DESC_LEN], out: &mut [u8]) {
for i in 0..DESC_LEN {
out[i] = desc[perm[i] as usize];
}
}
#[cfg(test)]
mod bench {
use super::*;
fn synthetic(w: usize, h: usize) -> Layer {
let mut px = vec![0f32; w * h];
let mut s = 0x1234_5678u32;
for v in px.iter_mut() {
s = s.wrapping_mul(1664525).wrapping_add(1013904223);
*v = ((s >> 16) & 0xff) as f32 / 255.0;
}
for y in 0..h {
for x in 0..w {
px[y * w + x] = px[y * w + x] * 0.3
+ (((x / 17 + y / 13) % 5) as f32) * 0.15
+ ((x as f32 * 0.05).sin() * (y as f32 * 0.03).cos()) * 0.2;
}
}
Layer { w, h, px }
}
fn ms(f: impl Fn()) -> f64 {
let mut best = f64::MAX;
for _ in 0..9 {
let t = std::time::Instant::now();
f();
best = best.min(t.elapsed().as_secs_f64() * 1000.0);
}
best
}
#[test]
#[ignore]
fn extract_threads() {
let img = |w: usize, h: usize, seed: u32| {
let mut px = vec![0f32; w * h];
let mut s = seed;
for (i, v) in px.iter_mut().enumerate() {
s = s.wrapping_mul(1664525).wrapping_add(1013904223);
let (x, y) = ((i % w) as f32, (i / w) as f32);
*v = ((s >> 16) & 0xff) as f32 / 255.0 * 0.15
+ (((i % w) / (11 + seed as usize % 7) + (i / w) / 13) % 5) as f32 * 0.12
+ ((x * 0.07 + seed as f32).sin() * (y * 0.045).cos()) * 0.25
+ 0.3;
}
crate::decode::Gray { w, h, px }
};
let p = Params { max_features: 600, ..Params::default() };
for (w, h) in [(224usize, 224usize), (384, 288)] {
let imgs: Vec<crate::decode::Gray> = (0..8).map(|k| img(w, h, 7 + k as u32 * 101)).collect();
let mut sum = 0u64;
for g in imgs.iter() {
let f = extract(g, &p);
for k in f.kps.iter() {
for v in [k.x, k.y, k.sigma, k.angle, k.response] {
sum = sum.wrapping_mul(0x100000001b3).wrapping_add(v.to_bits() as u64);
}
}
for &b in f.desc.iter() {
sum = sum.wrapping_mul(0x100000001b3).wrapping_add(b as u64);
}
}
let reps = 1;
let t1 = ms(|| {
for _ in 0..reps {
for g in imgs.iter() {
std::hint::black_box(extract(g, &p));
}
}
}) / (reps * imgs.len()) as f64;
let threads = rayon::current_num_threads();
let t8 = ms(|| {
use rayon::prelude::*;
(0..threads).into_par_iter().for_each(|_| {
for _ in 0..reps {
for g in imgs.iter() {
std::hint::black_box(extract(g, &p));
}
}
});
}) / (reps * imgs.len()) as f64;
println!("extract {w}x{h}: 1 thread {t1:.3} ms, {threads} threads {t8:.3} ms each (x{:.2}); checksum {sum:016x}", t8 / t1);
}
}
#[test]
#[ignore]
fn blur_threads() {
let base = synthetic(448, 448);
let sig = [1.226f32, 1.545, 1.946, 2.452, 3.089];
let mut sum = 0u64;
{
let mut g = Layer { w: base.w, h: base.h, px: base.px.clone() };
for (k, &s) in sig.iter().enumerate() {
if k + 1 == sig.len() {
for v in blur_top(&g, s).px.iter() {
sum = sum.wrapping_mul(0x100000001b3).wrapping_add(v.to_bits() as u64);
}
break;
}
let (l, d) = blur_dog(&g, s);
for v in l.px.iter().chain(d.px.iter()) {
sum = sum.wrapping_mul(0x100000001b3).wrapping_add(v.to_bits() as u64);
}
g = l;
}
}
let once = || {
let mut g = Layer { w: base.w, h: base.h, px: base.px.clone() };
for (k, &s) in sig.iter().enumerate() {
if k + 1 == sig.len() {
std::hint::black_box(blur_top(&g, s));
break;
}
let (l, d) = blur_dog(&g, s);
std::hint::black_box(&d);
g = l;
}
};
let reps = 40;
let t1 = ms(|| {
for _ in 0..reps {
once();
}
}) / reps as f64;
let threads = rayon::current_num_threads();
let t8 = ms(|| {
use rayon::prelude::*;
(0..threads).into_par_iter().for_each(|_| {
for _ in 0..reps {
once();
}
});
}) / reps as f64;
println!("blur octave 448x448: 1 thread {t1:.3} ms, {threads} threads {t8:.3} ms each (x{:.2}); checksum {sum:016x}", t8 / t1);
}
#[test]
#[ignore]
fn kernel_timings() {
let base = synthetic(640, 480);
for sigma in [1.226f32, 1.545, 1.946, 2.452, 3.089] {
let t = ms(|| {
std::hint::black_box(blur_dog(&base, sigma));
});
println!("blur_dog sigma {sigma:.3} (r={}): {t:8.3} ms", gaussian_kernel(sigma).len() / 2);
}
let g = crate::decode::Gray { w: 640, h: 480, px: base.px.clone() };
let p = Params::default();
let t = ms(|| {
std::hint::black_box(extract(&g, &p));
});
println!("extract 640x480: {t:8.3} ms ({} features)", extract(&g, &p).len());
let gauss = blur(&base, 1.0);
let grad = Grad::of(&gauss);
let t = ms(|| {
let mut d = [0u8; DESC_LEN];
let mut acc = 0f32;
for i in 0..2000 {
let x = 60.0 + ((i * 37) % 500) as f32;
let y = 60.0 + ((i * 53) % 340) as f32;
descriptor(&grad, 480, x, y, (i % 360) as f32, 1.6 + (i % 5) as f32 * 0.3, &mut d);
acc += d[0] as f32;
}
std::hint::black_box(acc);
});
println!("2000 descriptors: {t:8.3} ms");
let t = ms(|| {
let mut hist = [0f32; ORI_BINS];
let mut acc = 0f32;
for i in 0..2000 {
let x = 60.0 + ((i * 37) % 500) as f32;
let y = 60.0 + ((i * 53) % 340) as f32;
let scl = 1.6 + (i % 5) as f32 * 0.3;
acc += orientation_hist(&grad, 480, x, y, (ORI_RADIUS * scl).round() as i32, ORI_SIG_FCTR * scl, &mut hist);
}
std::hint::black_box(acc);
});
println!("2000 orientation hists: {t:8.3} ms");
let t = ms(|| {
std::hint::black_box(Grad::of(&gauss));
});
println!("Grad::of 640x480: {t:8.3} ms");
}
}