1use crate::error::{Result, VisionError};
12use crate::reconstruction::sfm::{
13 build_projection_matrix, rodrigues_to_rotation, triangulate_dlt_single, EssentialMatrix,
14 FundamentalMatrix, IntrinsicMatrix, PointCloud, Triangulation,
15};
16use scirs2_core::ndarray::{Array1, Array2};
17use std::collections::{HashMap, HashSet};
18
19pub use crate::reconstruction::sfm::IntrinsicMatrix as CameraIntrinsics;
21
22#[derive(Debug, Clone)]
28pub struct Pose {
29 pub rvec: Array1<f64>,
31 pub tvec: Array1<f64>,
33}
34
35impl Pose {
36 pub fn identity() -> Self {
38 Self {
39 rvec: Array1::zeros(3),
40 tvec: Array1::zeros(3),
41 }
42 }
43
44 pub fn to_projection(&self, k: &IntrinsicMatrix) -> Array2<f64> {
46 let r = rodrigues_to_rotation(&self.rvec);
47 let p_cam = build_projection_matrix(&r, &self.tvec);
48 let km = k.to_matrix();
49 let mut p = Array2::<f64>::zeros((3, 4));
51 for i in 0..3 {
52 for j in 0..4 {
53 for k_idx in 0..3 {
54 p[[i, j]] += km[[i, k_idx]] * p_cam[[k_idx, j]];
55 }
56 }
57 }
58 p
59 }
60}
61
62#[derive(Debug, Clone)]
68pub struct MapPoint {
69 pub id: usize,
71 pub position: Array1<f64>,
73 pub observations: HashMap<usize, [f64; 2]>,
75 pub match_count: usize,
77 pub is_outlier: bool,
79}
80
81impl MapPoint {
82 pub fn new(id: usize, position: Array1<f64>) -> Self {
84 Self {
85 id,
86 position,
87 observations: HashMap::new(),
88 match_count: 0,
89 is_outlier: false,
90 }
91 }
92
93 pub fn add_observation(&mut self, frame_id: usize, pixel: [f64; 2]) {
95 self.observations.insert(frame_id, pixel);
96 }
97
98 pub fn observation_count(&self) -> usize {
100 self.observations.len()
101 }
102}
103
104#[derive(Debug, Clone)]
110pub struct Keyframe {
111 pub id: usize,
113 pub pose: Pose,
115 pub keypoints: Vec<[f64; 2]>,
117 pub map_point_ids: Vec<Option<usize>>,
119 pub bow_descriptor: Vec<f32>,
121}
122
123impl Keyframe {
124 pub fn new(id: usize, pose: Pose, keypoints: Vec<[f64; 2]>) -> Self {
126 let n = keypoints.len();
127 Self {
128 id,
129 pose,
130 keypoints,
131 map_point_ids: vec![None; n],
132 bow_descriptor: Vec::new(),
133 }
134 }
135
136 pub fn get_pixel_for(&self, map_point_id: usize) -> Option<[f64; 2]> {
138 for (i, mp_id) in self.map_point_ids.iter().enumerate() {
139 if *mp_id == Some(map_point_id) {
140 return Some(self.keypoints[i]);
141 }
142 }
143 None
144 }
145}
146
147#[derive(Debug, Clone)]
153pub struct KeyframeSelector {
154 pub min_parallax: f64,
156 pub max_overlap: f64,
158 pub min_tracked: usize,
160 pub min_frame_gap: usize,
162}
163
164impl Default for KeyframeSelector {
165 fn default() -> Self {
166 Self {
167 min_parallax: 20.0,
168 max_overlap: 0.9,
169 min_tracked: 15,
170 min_frame_gap: 5,
171 }
172 }
173}
174
175impl KeyframeSelector {
176 pub fn should_insert(
182 &self,
183 current_pts: &[[f64; 2]],
184 prev_kf_pts: &[[f64; 2]],
185 frames_since_kf: usize,
186 ) -> bool {
187 if frames_since_kf < self.min_frame_gap {
188 return false;
189 }
190 let n = current_pts.len().min(prev_kf_pts.len());
191 if n < self.min_tracked {
192 return true;
194 }
195 let mut parallaxes: Vec<f64> = (0..n)
197 .map(|i| {
198 let dx = current_pts[i][0] - prev_kf_pts[i][0];
199 let dy = current_pts[i][1] - prev_kf_pts[i][1];
200 (dx * dx + dy * dy).sqrt()
201 })
202 .collect();
203 parallaxes.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
204 let median_parallax = parallaxes[n / 2];
205
206 let overlap = n as f64 / prev_kf_pts.len().max(1) as f64;
208
209 median_parallax > self.min_parallax || overlap < (1.0 - self.max_overlap)
210 }
211}
212
213#[derive(Debug, Clone)]
219pub struct Covisibility {
220 edges: HashMap<(usize, usize), usize>,
222 num_keyframes: usize,
224}
225
226impl Covisibility {
227 pub fn new() -> Self {
229 Self {
230 edges: HashMap::new(),
231 num_keyframes: 0,
232 }
233 }
234
235 pub fn update(&mut self, keyframes: &[Keyframe]) {
237 self.num_keyframes = keyframes.len();
238 self.edges.clear();
239
240 #[allow(clippy::needless_range_loop)]
242 for i in 0..keyframes.len() {
243 let ids_i: HashSet<usize> = keyframes[i]
244 .map_point_ids
245 .iter()
246 .filter_map(|&id| id)
247 .collect();
248 for j in (i + 1)..keyframes.len() {
249 let ids_j: HashSet<usize> = keyframes[j]
250 .map_point_ids
251 .iter()
252 .filter_map(|&id| id)
253 .collect();
254 let shared = ids_i.intersection(&ids_j).count();
255 if shared > 0 {
256 self.edges.insert((i, j), shared);
257 }
258 }
259 }
260 }
261
262 pub fn get_neighbours(&self, kf_idx: usize, min_shared: usize) -> Vec<(usize, usize)> {
264 let mut neighbours = Vec::new();
265 for (&(a, b), &count) in &self.edges {
266 if count >= min_shared {
267 if a == kf_idx {
268 neighbours.push((b, count));
269 } else if b == kf_idx {
270 neighbours.push((a, count));
271 }
272 }
273 }
274 neighbours.sort_by_key(|entry| std::cmp::Reverse(entry.1));
275 neighbours
276 }
277
278 pub fn edge_count(&self) -> usize {
280 self.edges.len()
281 }
282}
283
284impl Default for Covisibility {
285 fn default() -> Self {
286 Self::new()
287 }
288}
289
290pub struct BagOfWords {
296 vocabulary: Array2<f32>,
298}
299
300impl BagOfWords {
301 pub fn build_vocabulary(
307 descriptors: &Array2<f32>,
308 num_words: usize,
309 max_iter: usize,
310 ) -> Result<Self> {
311 if descriptors.nrows() < num_words {
312 return Err(VisionError::InvalidParameter(
313 "BagOfWords: fewer descriptors than vocabulary size".to_string(),
314 ));
315 }
316 let dim = descriptors.ncols();
317 let mut centres = Array2::<f32>::zeros((num_words, dim));
319 for i in 0..num_words {
320 for j in 0..dim {
321 centres[[i, j]] = descriptors[[i, j]];
322 }
323 }
324 for _ in 0..max_iter {
326 let mut sums = Array2::<f32>::zeros((num_words, dim));
327 let mut counts = vec![0usize; num_words];
328 for row in 0..descriptors.nrows() {
330 let nearest = Self::nearest_centre(&descriptors.row(row).to_owned(), ¢res);
331 for j in 0..dim {
332 sums[[nearest, j]] += descriptors[[row, j]];
333 }
334 counts[nearest] += 1;
335 }
336 for k in 0..num_words {
338 if counts[k] > 0 {
339 for j in 0..dim {
340 centres[[k, j]] = sums[[k, j]] / counts[k] as f32;
341 }
342 }
343 }
344 }
345 Ok(Self {
346 vocabulary: centres,
347 })
348 }
349
350 fn nearest_centre(descriptor: &Array1<f32>, centres: &Array2<f32>) -> usize {
351 let mut best = 0usize;
352 let mut best_dist = f32::MAX;
353 for k in 0..centres.nrows() {
354 let dist: f32 = (0..descriptor.len())
355 .map(|j| (descriptor[j] - centres[[k, j]]).powi(2))
356 .sum::<f32>();
357 if dist < best_dist {
358 best_dist = dist;
359 best = k;
360 }
361 }
362 best
363 }
364
365 pub fn compute_bow(&self, descriptors: &Array2<f32>) -> Vec<f32> {
367 let nw = self.vocabulary.nrows();
368 let mut hist = vec![0.0f32; nw];
369 for row in 0..descriptors.nrows() {
370 let word = Self::nearest_centre(&descriptors.row(row).to_owned(), &self.vocabulary);
371 hist[word] += 1.0;
372 }
373 let total: f32 = hist.iter().sum();
374 if total > 0.0 {
375 for v in &mut hist {
376 *v /= total;
377 }
378 }
379 hist
380 }
381
382 pub fn similarity(a: &[f32], b: &[f32]) -> f32 {
384 let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
385 let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
386 let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
387 if na < 1e-10 || nb < 1e-10 {
388 0.0
389 } else {
390 dot / (na * nb)
391 }
392 }
393}
394
395pub struct LoopClosure {
397 pub similarity_threshold: f32,
399 pub min_temporal_gap: usize,
401 known_descriptors: Vec<Vec<f32>>,
403}
404
405impl LoopClosure {
406 pub fn new(similarity_threshold: f32, min_temporal_gap: usize) -> Self {
408 Self {
409 similarity_threshold,
410 min_temporal_gap,
411 known_descriptors: Vec::new(),
412 }
413 }
414
415 pub fn add_keyframe_descriptor(&mut self, bow: Vec<f32>) {
417 self.known_descriptors.push(bow);
418 }
419
420 pub fn query(&self, bow: &[f32]) -> Vec<(usize, f32)> {
424 let current_idx = self.known_descriptors.len();
425 let mut candidates = Vec::new();
426 for (i, known) in self.known_descriptors.iter().enumerate() {
427 if current_idx.saturating_sub(i) < self.min_temporal_gap {
428 continue;
429 }
430 let sim = BagOfWords::similarity(bow, known);
431 if sim >= self.similarity_threshold {
432 candidates.push((i, sim));
433 }
434 }
435 candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
436 candidates
437 }
438
439 pub fn num_keyframes(&self) -> usize {
441 self.known_descriptors.len()
442 }
443}
444
445#[derive(Debug, Clone)]
451pub struct PoseConstraint {
452 pub from: usize,
454 pub to: usize,
456 pub relative_rvec: Array1<f64>,
458 pub relative_tvec: Array1<f64>,
460 pub weight: f64,
462 pub is_loop_closure: bool,
464}
465
466pub struct PoseGraph {
471 pub poses: Vec<Pose>,
473 pub constraints: Vec<PoseConstraint>,
475 pub max_iterations: usize,
477 pub learning_rate: f64,
479 pub tolerance: f64,
481}
482
483impl PoseGraph {
484 pub fn new(max_iterations: usize, learning_rate: f64, tolerance: f64) -> Self {
486 Self {
487 poses: Vec::new(),
488 constraints: Vec::new(),
489 max_iterations,
490 learning_rate,
491 tolerance,
492 }
493 }
494
495 pub fn add_pose(&mut self, pose: Pose) -> usize {
497 let idx = self.poses.len();
498 self.poses.push(pose);
499 idx
500 }
501
502 pub fn add_constraint(&mut self, constraint: PoseConstraint) {
504 self.constraints.push(constraint);
505 }
506
507 pub fn optimize(&mut self) -> Result<f64> {
512 if self.poses.is_empty() {
513 return Err(VisionError::InvalidParameter(
514 "PoseGraph: no poses to optimize".to_string(),
515 ));
516 }
517 if self.constraints.is_empty() {
518 return Ok(0.0);
519 }
520
521 let mut total_error = f64::MAX;
522 for _iter in 0..self.max_iterations {
523 let mut gradient_rvec = vec![Array1::<f64>::zeros(3); self.poses.len()];
524 let mut gradient_tvec = vec![Array1::<f64>::zeros(3); self.poses.len()];
525 let mut error = 0.0f64;
526
527 for constraint in &self.constraints {
528 let i = constraint.from;
529 let j = constraint.to;
530 if i >= self.poses.len() || j >= self.poses.len() {
531 continue;
532 }
533 let pred_rel_r = subtract_pose_rot(&self.poses[j].rvec, &self.poses[i].rvec);
535 let pred_rel_t = subtract_vec(&self.poses[j].tvec, &self.poses[i].tvec);
536 let res_r = subtract_vec(&pred_rel_r, &constraint.relative_rvec);
537 let res_t = subtract_vec(&pred_rel_t, &constraint.relative_tvec);
538 let w = constraint.weight;
539 let e: f64 = (res_r.iter().map(|v| v * v).sum::<f64>()
540 + res_t.iter().map(|v| v * v).sum::<f64>())
541 * w;
542 error += e;
543 for k in 0..3 {
545 gradient_rvec[j][k] += 2.0 * w * res_r[k];
546 gradient_tvec[j][k] += 2.0 * w * res_t[k];
547 gradient_rvec[i][k] -= 2.0 * w * res_r[k];
548 gradient_tvec[i][k] -= 2.0 * w * res_t[k];
549 }
550 }
551
552 for idx in 1..self.poses.len() {
554 for k in 0..3 {
555 self.poses[idx].rvec[k] -= self.learning_rate * gradient_rvec[idx][k];
556 self.poses[idx].tvec[k] -= self.learning_rate * gradient_tvec[idx][k];
557 }
558 }
559
560 let delta = (total_error - error).abs();
561 total_error = error;
562 if delta < self.tolerance {
563 break;
564 }
565 }
566 Ok(total_error)
567 }
568}
569
570fn subtract_vec(a: &Array1<f64>, b: &Array1<f64>) -> Array1<f64> {
571 let mut r = Array1::zeros(a.len());
572 for i in 0..a.len() {
573 r[i] = a[i] - b[i];
574 }
575 r
576}
577
578fn subtract_pose_rot(a: &Array1<f64>, b: &Array1<f64>) -> Array1<f64> {
580 subtract_vec(a, b)
581}
582
583pub struct VisualOdometry {
589 pub intrinsics: IntrinsicMatrix,
591 pub ransac_threshold: f64,
593 pub max_ransac_iter: usize,
595 pub min_inliers: usize,
597 current_pose: Pose,
599 frame_count: usize,
601 prev_keypoints: Vec<[f64; 2]>,
603}
604
605impl VisualOdometry {
606 pub fn new(intrinsics: IntrinsicMatrix) -> Self {
608 Self {
609 intrinsics,
610 ransac_threshold: 1.0,
611 max_ransac_iter: 1000,
612 min_inliers: 10,
613 current_pose: Pose::identity(),
614 frame_count: 0,
615 prev_keypoints: Vec::new(),
616 }
617 }
618
619 pub fn process_frame(
627 &mut self,
628 keypoints: &[[f64; 2]],
629 matches: &[(usize, usize)],
630 ) -> Result<Option<Pose>> {
631 self.frame_count += 1;
632
633 if self.frame_count == 1 {
634 self.prev_keypoints = keypoints.to_vec();
636 return Ok(Some(self.current_pose.clone()));
637 }
638
639 if self.prev_keypoints.is_empty() || matches.len() < self.min_inliers {
640 self.prev_keypoints = keypoints.to_vec();
641 return Ok(None);
642 }
643
644 let ki = self.intrinsics.to_inverse();
645 let normalise = |p: &[f64; 2]| -> [f64; 2] {
646 let v = scirs2_core::ndarray::Array1::from(vec![p[0], p[1], 1.0]);
647 let n = mat3_vec3(&ki, &v);
648 let z = n[2].max(1e-14);
649 [n[0] / z, n[1] / z]
650 };
651
652 let pts1_px: Vec<[f64; 2]> = matches
653 .iter()
654 .map(|&(pi, _)| {
655 if pi < self.prev_keypoints.len() {
656 self.prev_keypoints[pi]
657 } else {
658 [0.0, 0.0]
659 }
660 })
661 .collect();
662 let pts2_px: Vec<[f64; 2]> = matches
663 .iter()
664 .map(|&(_, ci)| {
665 if ci < keypoints.len() {
666 keypoints[ci]
667 } else {
668 [0.0, 0.0]
669 }
670 })
671 .collect();
672
673 let pts1_norm: Vec<[f64; 2]> = pts1_px.iter().map(&normalise).collect();
674 let pts2_norm: Vec<[f64; 2]> = pts2_px.iter().map(&normalise).collect();
675
676 let (f, inliers) = match FundamentalMatrix::from_ransac(
678 &pts1_px,
679 &pts2_px,
680 self.ransac_threshold,
681 0.99,
682 self.max_ransac_iter,
683 ) {
684 Ok(result) => result,
685 Err(_) => {
686 self.prev_keypoints = keypoints.to_vec();
687 return Ok(None);
688 }
689 };
690
691 let inlier_count = inliers.iter().filter(|&&b| b).count();
692 if inlier_count < self.min_inliers {
693 self.prev_keypoints = keypoints.to_vec();
694 return Ok(None);
695 }
696
697 let e = EssentialMatrix::from_fundamental(&f, &self.intrinsics, &self.intrinsics);
698
699 let pts1_in: Vec<[f64; 2]> = inliers
700 .iter()
701 .enumerate()
702 .filter(|(_, &b)| b)
703 .map(|(i, _)| pts1_norm[i])
704 .collect();
705 let pts2_in: Vec<[f64; 2]> = inliers
706 .iter()
707 .enumerate()
708 .filter(|(_, &b)| b)
709 .map(|(i, _)| pts2_norm[i])
710 .collect();
711
712 let (rel_r, rel_t) = match e.recover_pose(&pts1_in, &pts2_in) {
713 Ok(rt) => rt,
714 Err(_) => {
715 self.prev_keypoints = keypoints.to_vec();
716 return Ok(None);
717 }
718 };
719
720 let rel_rvec = rotation_to_rodrigues(&rel_r);
722 let new_rvec = compose_rvec(&self.current_pose.rvec, &rel_rvec);
723 let r_curr = rodrigues_to_rotation(&self.current_pose.rvec);
724 let new_tvec_arr = mat3_vec3(&r_curr, &rel_t);
725 let mut new_tvec = Array1::zeros(3);
726 for k in 0..3 {
727 new_tvec[k] = self.current_pose.tvec[k] + new_tvec_arr[k];
728 }
729
730 self.current_pose = Pose {
731 rvec: new_rvec,
732 tvec: new_tvec,
733 };
734 self.prev_keypoints = keypoints.to_vec();
735 Ok(Some(self.current_pose.clone()))
736 }
737
738 pub fn current_pose(&self) -> &Pose {
740 &self.current_pose
741 }
742
743 pub fn reset(&mut self) {
745 self.current_pose = Pose::identity();
746 self.frame_count = 0;
747 self.prev_keypoints.clear();
748 }
749}
750
751fn mat3_vec3(m: &Array2<f64>, v: &Array1<f64>) -> Array1<f64> {
756 let mut out = Array1::zeros(3);
757 for i in 0..3 {
758 for j in 0..3 {
759 out[i] += m[[i, j]] * v[j];
760 }
761 }
762 out
763}
764
765fn rotation_to_rodrigues(r: &Array2<f64>) -> Array1<f64> {
766 let trace = r[[0, 0]] + r[[1, 1]] + r[[2, 2]];
767 let cos_theta = ((trace - 1.0) / 2.0).clamp(-1.0, 1.0);
768 let theta = cos_theta.acos();
769 if theta.abs() < 1e-10 {
770 return Array1::zeros(3);
771 }
772 let scale = theta / (2.0 * theta.sin());
773 Array1::from(vec![
774 (r[[2, 1]] - r[[1, 2]]) * scale,
775 (r[[0, 2]] - r[[2, 0]]) * scale,
776 (r[[1, 0]] - r[[0, 1]]) * scale,
777 ])
778}
779
780fn compose_rvec(r1: &Array1<f64>, r2: &Array1<f64>) -> Array1<f64> {
782 let m1 = rodrigues_to_rotation(r1);
783 let m2 = rodrigues_to_rotation(r2);
784 let mut m = Array2::<f64>::zeros((3, 3));
785 for i in 0..3 {
786 for k in 0..3 {
787 for j in 0..3 {
788 m[[i, j]] += m2[[i, k]] * m1[[k, j]];
789 }
790 }
791 }
792 rotation_to_rodrigues(&m)
793}
794
795#[cfg(test)]
800mod tests {
801 use super::*;
802
803 #[test]
804 fn test_keyframe_selector_parallax() {
805 let sel = KeyframeSelector::default();
806 let pts1: Vec<[f64; 2]> = (0..20).map(|i| [i as f64, i as f64]).collect();
807 let pts2: Vec<[f64; 2]> = (0..20).map(|i| [i as f64 + 25.0, i as f64]).collect();
808 assert!(sel.should_insert(&pts2, &pts1, 10));
809 }
810
811 #[test]
812 fn test_keyframe_selector_no_insert_too_soon() {
813 let sel = KeyframeSelector {
814 min_frame_gap: 10,
815 ..Default::default()
816 };
817 let pts: Vec<[f64; 2]> = (0..20).map(|i| [i as f64, i as f64]).collect();
818 assert!(!sel.should_insert(&pts, &pts, 3));
819 }
820
821 #[test]
822 fn test_covisibility_graph() {
823 let mut kf1 = Keyframe::new(0, Pose::identity(), vec![[0.0, 0.0]; 5]);
824 kf1.map_point_ids = vec![Some(0), Some(1), Some(2), None, None];
825 let mut kf2 = Keyframe::new(1, Pose::identity(), vec![[0.0, 0.0]; 5]);
826 kf2.map_point_ids = vec![Some(0), Some(1), Some(3), None, None];
827 let mut covis = Covisibility::new();
828 covis.update(&[kf1, kf2]);
829 let neighbours = covis.get_neighbours(0, 1);
830 assert_eq!(neighbours.len(), 1);
831 assert_eq!(neighbours[0].0, 1);
832 assert_eq!(neighbours[0].1, 2);
833 }
834
835 #[test]
836 fn test_loop_closure_query() {
837 let mut lc = LoopClosure::new(0.8, 5);
838 for _ in 0..10 {
839 lc.add_keyframe_descriptor(vec![1.0, 0.0, 0.0]);
840 }
841 lc.add_keyframe_descriptor(vec![0.5, 0.5, 0.0]);
842 let bow = vec![1.0, 0.0, 0.0];
843 let candidates = lc.query(&bow);
844 let _ = candidates; }
847
848 #[test]
849 fn test_pose_graph_optimization() {
850 let mut pg = PoseGraph::new(50, 0.01, 1e-6);
851 pg.add_pose(Pose::identity());
852 let mut p2 = Pose::identity();
853 p2.tvec[0] = 1.1; pg.add_pose(p2);
855 pg.add_constraint(PoseConstraint {
856 from: 0,
857 to: 1,
858 relative_rvec: Array1::zeros(3),
859 relative_tvec: Array1::from(vec![1.0, 0.0, 0.0]),
860 weight: 1.0,
861 is_loop_closure: false,
862 });
863 let result = pg.optimize();
864 assert!(result.is_ok());
865 }
866
867 #[test]
868 fn test_visual_odometry_first_frame() {
869 let k = IntrinsicMatrix {
870 fx: 500.0,
871 fy: 500.0,
872 cx: 320.0,
873 cy: 240.0,
874 };
875 let mut vo = VisualOdometry::new(k);
876 let kps: Vec<[f64; 2]> = (0..20).map(|i| [i as f64 * 10.0, 100.0]).collect();
877 let result = vo.process_frame(&kps, &[]);
878 assert!(result.is_ok());
879 assert!(result.expect("process_frame should succeed").is_some());
880 }
881
882 #[test]
883 fn test_map_point() {
884 let pos = Array1::from(vec![1.0, 2.0, 3.0]);
885 let mut mp = MapPoint::new(0, pos);
886 mp.add_observation(0, [100.0, 200.0]);
887 mp.add_observation(1, [110.0, 205.0]);
888 assert_eq!(mp.observation_count(), 2);
889 assert!(mp.observations.contains_key(&0)); }
892}