1use spatialrust_image::{Image, ImageView};
4
5use crate::{
6 detect_fast, gaussian_blur, resize, BorderMode, DescriptorBuffer, FastOptions, FeatureSet2,
7 Interpolation, Keypoint2, VisionError, VisionResult,
8};
9
10const ORB_DESCRIPTOR_BYTES: usize = 32;
11
12#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
14pub enum OrbScoreType {
15 #[default]
17 Harris,
18 Fast,
20}
21
22#[derive(Clone, Copy, Debug, PartialEq)]
24pub struct OrbOptions {
25 pub max_features: usize,
27 pub scale_factor: f32,
29 pub levels: usize,
31 pub edge_threshold: usize,
33 pub fast_threshold: u8,
35 pub patch_size: usize,
37 pub score_type: OrbScoreType,
39}
40
41impl Default for OrbOptions {
42 fn default() -> Self {
43 Self {
44 max_features: 500,
45 scale_factor: 1.2,
46 levels: 8,
47 edge_threshold: 31,
48 fast_threshold: 20,
49 patch_size: 31,
50 score_type: OrbScoreType::Harris,
51 }
52 }
53}
54
55impl OrbOptions {
56 fn validate(self) -> VisionResult<Self> {
57 if !self.scale_factor.is_finite() || self.scale_factor <= 1.0 {
58 return Err(VisionError::InvalidParameter(
59 "ORB scale_factor must be finite and greater than one".into(),
60 ));
61 }
62 if self.levels == 0 {
63 return Err(VisionError::InvalidParameter("ORB levels must be positive".into()));
64 }
65 if self.patch_size < 7 || self.patch_size % 2 == 0 {
66 return Err(VisionError::InvalidParameter(
67 "ORB patch_size must be odd and at least seven".into(),
68 ));
69 }
70 Ok(self)
71 }
72}
73
74#[derive(Clone, Copy)]
75struct Candidate {
76 level: usize,
77 level_x: usize,
78 level_y: usize,
79 scale: f32,
80 response: f32,
81}
82
83pub fn detect_and_describe_orb(
89 input: ImageView<'_, u8, 1>,
90 options: OrbOptions,
91) -> VisionResult<FeatureSet2> {
92 let options = options.validate()?;
93 if input.width() == 0 || input.height() == 0 || options.max_features == 0 {
94 return FeatureSet2::try_new(
95 Vec::new(),
96 DescriptorBuffer::try_binary(0, ORB_DESCRIPTOR_BYTES, Vec::new())?,
97 );
98 }
99
100 let mut pyramid = Vec::with_capacity(options.levels);
101 let mut candidates = Vec::new();
102 for level in 0..options.levels {
103 let scale = options.scale_factor.powi(level as i32);
104 let width = ((input.width() as f32 / scale).round() as usize).max(1);
105 let height = ((input.height() as f32 / scale).round() as usize).max(1);
106 let image = if level == 0 {
107 Image::try_new_with_metadata(
108 input.width(),
109 input.height(),
110 (0..input.height())
111 .flat_map(|y| {
112 (0..input.width())
113 .map(move |x| input.get(x, y).expect("coordinate in bounds")[0])
114 })
115 .collect(),
116 input.metadata(),
117 )?
118 } else {
119 resize(input, width, height, Interpolation::Bilinear)?
120 };
121 let margin = options.edge_threshold.max(options.patch_size / 2 + 1);
122 let fast = detect_fast(
123 image.view(),
124 FastOptions { threshold: options.fast_threshold, nonmax_suppression: true },
125 )?;
126 for point in fast {
127 let x = point.x() as usize;
128 let y = point.y() as usize;
129 if x < margin
130 || y < margin
131 || x.saturating_add(margin) >= image.width()
132 || y.saturating_add(margin) >= image.height()
133 {
134 continue;
135 }
136 let response = match options.score_type {
137 OrbScoreType::Harris => harris_score(image.view(), x, y),
138 OrbScoreType::Fast => point.response(),
139 };
140 candidates.push(Candidate { level, level_x: x, level_y: y, scale, response });
141 }
142 pyramid.push(image);
143 }
144
145 candidates.sort_by(|left, right| {
146 right
147 .response
148 .total_cmp(&left.response)
149 .then_with(|| left.level.cmp(&right.level))
150 .then_with(|| left.level_y.cmp(&right.level_y))
151 .then_with(|| left.level_x.cmp(&right.level_x))
152 });
153 candidates.truncate(options.max_features);
154
155 let pattern = brief_pattern(options.patch_size / 2);
156 let mut blurred = Vec::with_capacity(pyramid.len());
157 for image in &pyramid {
158 blurred.push(gaussian_blur(image.view(), 7, 7, 2.0, 2.0, BorderMode::Reflect101)?);
159 }
160 let mut keypoints = Vec::with_capacity(candidates.len());
161 let mut descriptors = Vec::with_capacity(candidates.len() * ORB_DESCRIPTOR_BYTES);
162 for candidate in candidates {
163 let image = blurred[candidate.level].view();
164 let angle = intensity_centroid_angle(
165 image,
166 candidate.level_x,
167 candidate.level_y,
168 options.patch_size / 2,
169 );
170 keypoints.push(
171 Keypoint2::try_new(
172 candidate.level_x as f32 * candidate.scale,
173 candidate.level_y as f32 * candidate.scale,
174 candidate.response,
175 )?
176 .with_size(options.patch_size as f32 * candidate.scale)?
177 .with_angle_degrees(angle.to_degrees())?
178 .with_octave(candidate.level as i32),
179 );
180 descriptors.extend(describe(image, candidate.level_x, candidate.level_y, angle, &pattern));
181 }
182 FeatureSet2::try_new(
183 keypoints,
184 DescriptorBuffer::try_binary(
185 descriptors.len() / ORB_DESCRIPTOR_BYTES,
186 ORB_DESCRIPTOR_BYTES,
187 descriptors,
188 )?,
189 )
190}
191
192fn harris_score(image: ImageView<'_, u8, 1>, x: usize, y: usize) -> f32 {
193 let mut xx = 0.0_f32;
194 let mut xy = 0.0_f32;
195 let mut yy = 0.0_f32;
196 for dy in -3_isize..=3 {
197 for dx in -3_isize..=3 {
198 let px = (x as isize + dx) as usize;
199 let py = (y as isize + dy) as usize;
200 let gx = f32::from(image.get(px + 1, py).unwrap()[0])
201 - f32::from(image.get(px - 1, py).unwrap()[0]);
202 let gy = f32::from(image.get(px, py + 1).unwrap()[0])
203 - f32::from(image.get(px, py - 1).unwrap()[0]);
204 xx += gx * gx;
205 xy += gx * gy;
206 yy += gy * gy;
207 }
208 }
209 xx.mul_add(yy, -(xy * xy)) - 0.04 * (xx + yy) * (xx + yy)
210}
211
212fn intensity_centroid_angle(image: ImageView<'_, u8, 1>, x: usize, y: usize, radius: usize) -> f32 {
213 let mut m10 = 0_i64;
214 let mut m01 = 0_i64;
215 let radius_squared = (radius * radius) as isize;
216 for dy in -(radius as isize)..=radius as isize {
217 let extent = ((radius_squared - dy * dy) as f64).sqrt().floor() as isize;
218 for dx in -extent..=extent {
219 let intensity = i64::from(
220 image.get((x as isize + dx) as usize, (y as isize + dy) as usize).unwrap()[0],
221 );
222 m10 += dx as i64 * intensity;
223 m01 += dy as i64 * intensity;
224 }
225 }
226 (m01 as f32).atan2(m10 as f32)
227}
228
229fn describe(
230 image: ImageView<'_, u8, 1>,
231 x: usize,
232 y: usize,
233 angle: f32,
234 pattern: &[((i8, i8), (i8, i8))],
235) -> [u8; ORB_DESCRIPTOR_BYTES] {
236 let (sin, cos) = angle.sin_cos();
237 let mut descriptor = [0_u8; ORB_DESCRIPTOR_BYTES];
238 for (bit, &(first, second)) in pattern.iter().enumerate() {
239 let sample = |point: (i8, i8)| {
240 let rx = (f32::from(point.0) * cos - f32::from(point.1) * sin).round() as isize;
241 let ry = (f32::from(point.0) * sin + f32::from(point.1) * cos).round() as isize;
242 image.get((x as isize + rx) as usize, (y as isize + ry) as usize).unwrap()[0]
243 };
244 if sample(first) < sample(second) {
245 descriptor[bit / 8] |= 1 << (bit % 8);
246 }
247 }
248 descriptor
249}
250
251fn brief_pattern(radius: usize) -> Vec<((i8, i8), (i8, i8))> {
252 let radius = radius.min(i8::MAX as usize) as i32;
253 let mut state = 0x6d2b_79f5_u32;
254 let mut next_point = || loop {
255 state ^= state << 13;
256 state ^= state >> 17;
257 state ^= state << 5;
258 let span = (radius * 2 + 1) as u32;
259 let x = (state % span) as i32 - radius;
260 state = state.rotate_left(11).wrapping_mul(0x9e37_79b1);
261 let y = (state % span) as i32 - radius;
262 if x * x + y * y <= radius * radius {
263 return (x as i8, y as i8);
264 }
265 };
266 (0..ORB_DESCRIPTOR_BYTES * 8)
267 .map(|_| {
268 let first = next_point();
269 let mut second = next_point();
270 while second == first {
271 second = next_point();
272 }
273 (first, second)
274 })
275 .collect()
276}
277
278#[cfg(test)]
279mod tests {
280 use super::{detect_and_describe_orb, OrbOptions};
281 use crate::{match_descriptors, MatchOptions};
282 use spatialrust_image::Image;
283
284 fn textured_image() -> Image<u8, 1> {
285 Image::try_new(
286 96,
287 80,
288 (0..80)
289 .flat_map(|y| {
290 (0..96).map(move |x| {
291 (((x * 37 + y * 19) ^ (x * y * 3) ^ ((x / 8 + y / 8) * 127)) & 255) as u8
292 })
293 })
294 .collect(),
295 )
296 .unwrap()
297 }
298
299 #[test]
300 fn orb_is_deterministic_bounded_and_well_formed() {
301 let image = textured_image();
302 let options = OrbOptions { max_features: 80, edge_threshold: 16, ..OrbOptions::default() };
303 let first = detect_and_describe_orb(image.view(), options).unwrap();
304 let second = detect_and_describe_orb(image.view(), options).unwrap();
305 assert_eq!(first, second);
306 assert!(!first.keypoints().is_empty());
307 assert!(first.keypoints().len() <= 80);
308 assert_eq!(first.descriptors().width(), 32);
309 assert!(first.keypoints().iter().all(|point| point.angle_degrees().is_some()));
310 }
311
312 #[test]
313 fn orb_descriptors_self_match_at_zero_distance() {
314 let image = textured_image();
315 let features = detect_and_describe_orb(
316 image.view(),
317 OrbOptions { max_features: 40, edge_threshold: 16, ..OrbOptions::default() },
318 )
319 .unwrap();
320 let matches = match_descriptors(
321 features.descriptors(),
322 features.descriptors(),
323 MatchOptions { cross_check: true, ..MatchOptions::default() },
324 )
325 .unwrap();
326 assert_eq!(matches.len(), features.keypoints().len());
327 assert!(matches.iter().all(|feature_match| feature_match.distance() == 0.0));
328 }
329
330 #[test]
331 fn invalid_or_empty_orb_inputs_are_handled() {
332 let image = textured_image();
333 assert!(detect_and_describe_orb(
334 image.view(),
335 OrbOptions { levels: 0, ..OrbOptions::default() }
336 )
337 .is_err());
338 let empty = detect_and_describe_orb(
339 image.view(),
340 OrbOptions { max_features: 0, ..OrbOptions::default() },
341 )
342 .unwrap();
343 assert!(empty.keypoints().is_empty());
344 assert_eq!(empty.descriptors().width(), 32);
345 }
346}