1use rayon::prelude::*;
35use std::f32::consts::PI;
36
37const TWO_PI: f32 = 2.0 * PI;
38
39pub struct Gradient<'a> {
41 g: &'a [f32],
42 w: usize,
43 pub h: usize,
45 pre: Option<(Vec<f32>, Vec<f32>)>,
46}
47
48impl<'a> Gradient<'a> {
49 pub fn lazy(g: &'a [f32], w: usize, h: usize) -> Self {
51 Self { g, w, h, pre: None }
52 }
53
54 pub fn precomputed(g: &'a [f32], w: usize, h: usize) -> Self {
56 let mut mag = vec![0f32; w * h];
57 let mut ang = vec![0f32; w * h];
58 mag.par_chunks_mut(w).zip(ang.par_chunks_mut(w)).enumerate().for_each(|(y, (m, a))| {
59 if y == 0 || y + 1 >= h {
60 return;
61 }
62 for x in 1..w.saturating_sub(1) {
63 let (mm, aa) = grad_at(g, w, x, y);
64 m[x] = mm;
65 a[x] = aa;
66 }
67 });
68 Self { g, w, h, pre: Some((mag, ang)) }
69 }
70
71 pub fn width(&self) -> usize {
73 self.w
74 }
75
76 #[inline]
78 pub fn at(&self, x: usize, y: usize) -> (f32, f32) {
79 match &self.pre {
80 Some((m, a)) => (m[y * self.w + x], a[y * self.w + x]),
81 None => grad_at(self.g, self.w, x, y),
82 }
83 }
84}
85
86#[inline]
87fn grad_at(g: &[f32], w: usize, x: usize, y: usize) -> (f32, f32) {
88 let i = y * w + x;
89 let dx = g[i + 1] - g[i - 1];
90 let dy = g[i + w] - g[i - w];
91 (0.5 * (dx * dx + dy * dy).sqrt(), dy.atan2(dx))
92}
93
94#[derive(Clone, Copy, Debug)]
96pub struct OrientParams {
97 pub max_num_orientations: usize,
99 pub bin_interpolation: bool,
101}
102
103#[inline]
105pub fn quantize_angle(theta: f32) -> u16 {
106 let f = (theta / TWO_PI).rem_euclid(1.0);
107 ((f * 65535.0).floor() as u32).min(65534) as u16
108}
109
110#[inline]
112pub fn dequantize_angle(q: u16) -> f32 {
114 q as f32 / 65535.0 * TWO_PI
115}
116
117pub fn orientations(grad: &Gradient, x: f32, y: f32, sigma: f32, p: &OrientParams) -> Vec<u16> {
119 const NB: usize = 36;
120 let w = grad.width();
121 let h = grad.h;
122 let sw = 1.5 * sigma;
123 let rad = 2.0 * sw;
124 let x0 = ((x - rad).floor() as i64).max(1);
125 let x1 = ((x + rad).floor() as i64).min(w as i64 - 2);
126 let y0 = ((y - rad).floor() as i64).max(1);
127 let y1 = ((y + rad).floor() as i64).min(h as i64 - 2);
128 let r2 = rad * rad + 0.5;
129 let inv = 1.0 / (2.0 * sw * sw);
130 let mut hist = [0f32; NB];
131 let bscale = NB as f32 / TWO_PI;
132 for yi in y0..=y1 {
133 let dy = yi as f32 + 0.5 - y;
134 for xi in x0..=x1 {
135 let dx = xi as f32 + 0.5 - x;
136 let d2 = dx * dx + dy * dy;
137 if d2 >= r2 {
138 continue;
139 }
140 let (m, th) = grad.at(xi as usize, yi as usize);
141 let wgt = m * (-d2 * inv).exp();
142 if p.bin_interpolation {
143 let t = th * bscale - 0.5;
144 let b0 = t.floor();
145 let f = t - b0;
146 let b0 = (b0 as i64).rem_euclid(NB as i64) as usize;
147 hist[b0] += wgt * (1.0 - f);
148 hist[(b0 + 1) % NB] += wgt * f;
149 } else {
150 let mut b = (th * bscale).floor() as i64;
151 if b < 0 {
152 b += NB as i64;
153 }
154 hist[(b as usize).min(NB - 1)] += wgt;
155 }
156 }
157 }
158 for _ in 0..6 {
160 let prev = hist;
161 for i in 0..NB {
162 hist[i] = (prev[(i + NB - 1) % NB] + prev[i] + prev[(i + 1) % NB]) / 3.0;
163 }
164 }
165 let hmax = hist.iter().cloned().fold(0f32, f32::max);
166 let mut peaks: Vec<(f32, usize)> = Vec::new();
167 for i in 0..NB {
168 let (hl, hc, hr) = (hist[(i + NB - 1) % NB], hist[i], hist[(i + 1) % NB]);
169 if hc > 0.8 * hmax && hc > hl && hc > hr {
170 peaks.push((hc, i));
171 }
172 }
173 peaks.sort_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
175 let mut out: Vec<u16> = Vec::with_capacity(p.max_num_orientations);
176 for &(hc, i) in peaks.iter().take(p.max_num_orientations.max(1)) {
177 let (hl, hr) = (hist[(i + NB - 1) % NB], hist[(i + 1) % NB]);
178 let di = 0.5 * (hr - hl) / (2.0 * hc - hr - hl);
179 let th = (i as f32 + di + 0.5) / bscale;
180 let q = quantize_angle(th);
181 if !out.contains(&q) {
182 out.push(q);
183 }
184 }
185 out
186}
187
188#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
190pub enum DescriptorNormalization {
191 #[default]
193 L1Root,
194 L2,
196}
197
198pub fn descriptor(grad: &Gradient, x: f32, y: f32, sigma: f32, theta: f32, norm: DescriptorNormalization) -> [u8; 128] {
200 let w = grad.width();
201 let h = grad.h;
202 let s = 3.0 * sigma;
203 let (sn, c) = theta.sin_cos();
204 let rad = s * 2.5 * std::f32::consts::SQRT_2 + 1.0;
205 let x0 = ((x - rad).floor() as i64).max(1);
206 let x1 = ((x + rad).ceil() as i64).min(w as i64 - 2);
207 let y0 = ((y - rad).floor() as i64).max(1);
208 let y1 = ((y + rad).ceil() as i64).min(h as i64 - 2);
209 let inv_s = 1.0 / s;
210 let bins_per_rad = 4.0 / PI;
211 let mut d = [0f32; 128];
212 for yi in y0..=y1 {
213 let dy = yi as f32 + 0.5 - y;
214 for xi in x0..=x1 {
215 let dx = xi as f32 + 0.5 - x;
216 let rx = (c * dx + sn * dy) * inv_s;
218 let ry = (c * dy - sn * dx) * inv_s;
219 let fx = rx + 1.5;
220 let fy = ry + 1.5;
221 if fx <= -1.0 || fx >= 4.0 || fy <= -1.0 || fy >= 4.0 {
222 continue;
223 }
224 let (m, th) = grad.at(xi as usize, yi as usize);
225 let gw = (-(rx * rx + ry * ry) / 8.0).exp() * m;
226 let t = ((theta - th) * bins_per_rad).rem_euclid(8.0);
227 let tb = t.floor();
228 let tf = t - tb;
229 let b0 = (tb as usize) % 8;
230 let b1 = (b0 + 1) % 8;
231 let ix0 = fx.floor() as i64;
232 let iy0 = fy.floor() as i64;
233 let wx1 = fx - ix0 as f32;
234 let wy1 = fy - iy0 as f32;
235 for (iy, wy) in [(iy0, 1.0 - wy1), (iy0 + 1, wy1)] {
236 if !(0..4).contains(&iy) || wy <= 0.0 {
237 continue;
238 }
239 for (ix, wx) in [(ix0, 1.0 - wx1), (ix0 + 1, wx1)] {
240 if !(0..4).contains(&ix) || wx <= 0.0 {
241 continue;
242 }
243 let base = ((iy * 4 + ix) * 8) as usize;
244 let ww = gw * wx * wy;
245 d[base + b0] += ww * (1.0 - tf);
246 d[base + b1] += ww * tf;
247 }
248 }
249 }
250 }
251 normalize_quantize(&mut d, norm)
252}
253
254pub fn normalize_quantize(d: &mut [f32; 128], norm: DescriptorNormalization) -> [u8; 128] {
256 let l2 = |d: &mut [f32; 128]| {
257 let n = d.iter().map(|v| v * v).sum::<f32>().sqrt();
258 if n > 0.0 {
259 d.iter_mut().for_each(|v| *v /= n);
260 }
261 };
262 l2(d);
263 d.iter_mut().for_each(|v| *v = v.min(0.2));
264 l2(d);
265 match norm {
266 DescriptorNormalization::L1Root => {
267 let s: f32 = d.iter().sum();
268 if s > 0.0 {
269 d.iter_mut().for_each(|v| *v = (*v / s).sqrt());
270 }
271 }
272 DescriptorNormalization::L2 => l2(d),
273 }
274 let mut q = [0u8; 128];
275 for (o, v) in q.iter_mut().zip(d.iter()) {
276 *o = (512.0 * v).round().clamp(0.0, 255.0) as u8;
277 }
278 q
279}