Skip to main content

spatialrust_vision/
orb.rs

1//! Oriented FAST and rotated BRIEF binary features.
2
3use 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/// Response used to rank ORB keypoints.
13#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
14pub enum OrbScoreType {
15    /// Rank FAST candidates by a local Harris response.
16    #[default]
17    Harris,
18    /// Retain the FAST segment-test score.
19    Fast,
20}
21
22/// Multi-scale ORB detector and descriptor configuration.
23#[derive(Clone, Copy, Debug, PartialEq)]
24pub struct OrbOptions {
25    /// Maximum number of descriptor rows across every pyramid level.
26    pub max_features: usize,
27    /// Scale multiplier between adjacent levels; must be greater than one.
28    pub scale_factor: f32,
29    /// Number of pyramid levels.
30    pub levels: usize,
31    /// Minimum distance from a level-image edge.
32    pub edge_threshold: usize,
33    /// FAST intensity threshold.
34    pub fast_threshold: u8,
35    /// Odd intensity-centroid and descriptor patch diameter.
36    pub patch_size: usize,
37    /// Candidate ranking response.
38    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
83/// Detects oriented multi-scale keypoints and computes 256-bit rotated BRIEF descriptors.
84///
85/// The fixed-seed BRIEF sampling pattern is stable across platforms and SpatialRust
86/// releases. It intentionally does not promise bit identity with OpenCV's private
87/// learned sampling table.
88pub 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}