1use crate::error::{AlignError, Result};
34use crate::umeyama::{robust_umeyama, RobustUmeyamaOptions, RobustUmeyamaResult};
35use cumulus3d_core::{Point3DId, Reconstruction, Vec3};
36use std::collections::HashMap;
37use std::ops::Range;
38use std::sync::Arc;
39
40pub type NameFilter = Arc<dyn Fn(&str) -> bool + Send + Sync>;
42
43#[derive(Clone, Default)]
45pub struct SharedPointOptions {
46 pub position_range: Option<Range<u32>>,
48 pub name_filter: Option<NameFilter>,
50}
51
52impl std::fmt::Debug for SharedPointOptions {
53 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54 f.debug_struct("SharedPointOptions")
55 .field("position_range", &self.position_range)
56 .field("name_filter", &self.name_filter.is_some())
57 .finish()
58 }
59}
60
61impl SharedPointOptions {
62 fn accepts(&self, name: &str) -> bool {
63 if let Some(r) = &self.position_range {
64 match position_index_from_name(name) {
65 Some(p) if r.contains(&p) => {}
66 _ => return false,
67 }
68 }
69 self.name_filter.as_ref().is_none_or(|f| f(name))
70 }
71}
72
73pub fn position_index_from_name(name: &str) -> Option<u32> {
75 let file = name.rsplit(['/', '\\']).next().unwrap_or(name);
76 let stem = file.rsplit_once('.').map_or(file, |(s, _)| s);
77 let end = stem.rfind(|c: char| c.is_ascii_digit())? + 1;
78 let start = stem[..end].rfind(|c: char| !c.is_ascii_digit()).map_or(0, |i| i + 1);
79 stem[start..end].parse().ok()
80}
81
82pub fn shared_point_correspondences(a: &Reconstruction, b: &Reconstruction, opts: &SharedPointOptions) -> Vec<(Point3DId, Point3DId)> {
84 let b_by_name: HashMap<&str, _> = b.images().map(|im| (im.name.as_str(), im)).collect();
85 let mut pairs = Vec::new();
86 for ia in a.images() {
87 if !opts.accepts(&ia.name) {
88 continue;
89 }
90 let Some(ib) = b_by_name.get(ia.name.as_str()) else { continue };
91 for (pa, pb) in ia.points2d().iter().zip(ib.points2d()) {
94 if pa.has_point3d() && pb.has_point3d() {
95 pairs.push((pa.point3d_id, pb.point3d_id));
96 }
97 }
98 }
99 pairs.sort_unstable();
100 pairs.dedup();
101 pairs
102}
103
104pub fn align_reconstructions(
107 src: &Reconstruction,
108 dst: &Reconstruction,
109 opts: &SharedPointOptions,
110 robust: &RobustUmeyamaOptions,
111) -> Result<(RobustUmeyamaResult, Vec<(Point3DId, Point3DId)>)> {
112 let pairs = shared_point_correspondences(src, dst, opts);
113 if pairs.len() < 3 {
114 return Err(AlignError::TooFewCorrespondences { found: pairs.len(), required: 3 });
115 }
116 let (xs, ys): (Vec<Vec3>, Vec<Vec3>) =
117 pairs.iter().map(|(i, j)| (src.point3d(*i).expect("track id").xyz, dst.point3d(*j).expect("track id").xyz)).unzip();
118 let r = robust_umeyama(&xs, &ys, robust).ok_or(AlignError::Degenerate)?;
119 Ok((r, pairs))
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125 use crate::umeyama::tests::{random_sim3, sim_close};
126 use cumulus3d_core::{Camera, CameraModelKind, Image, Rigid3, TrackEntry, Vec2};
127
128 fn base(n_img: u32, n_pts: u32) -> Reconstruction {
129 let mut rec = Reconstruction::new();
130 rec.add_camera_own_rig(camera(1)).unwrap();
131 for i in 1..=n_img {
132 let pts: Vec<Vec2> = (0..n_pts).map(|k| Vec2::new(k as f64, 0.0)).collect();
133 rec.add_image_own_frame(Image::new(i, format!("camF/camF_{:04}.jpg", i - 1), 1, pts), Some(Rigid3::identity())).unwrap();
134 rec.register_image(i).unwrap();
135 }
136 rec
137 }
138
139 fn camera(id: u32) -> Camera {
140 let mut c = Camera::from_focal(CameraModelKind::Pinhole, 500.0, 640, 480);
141 c.camera_id = id;
142 c
143 }
144
145 #[test]
146 fn name_position_index() {
147 assert_eq!(position_index_from_name("camF/camF_0012.jpg"), Some(12));
148 assert_eq!(position_index_from_name("a7/x_3_0040.png"), Some(40));
149 assert_eq!(position_index_from_name("nodigits.jpg"), None);
150 }
151
152 #[test]
153 fn correspondences_and_alignment() {
154 let truth = random_sim3(&mut cumulus3d_core::ransac::make_rng(Some(4)));
157 let (mut a, mut b) = (base(3, 30), base(3, 30));
158 for k in 0..30u32 {
159 let x = Vec3::new(k as f64, (k * k % 7) as f64, (k % 5) as f64 * 2.0);
160 let ta = vec![TrackEntry::new(1, k), TrackEntry::new(2, k)];
161 a.add_point3d_with_id(k as u64 + 1, cumulus3d_core::Point3D { xyz: x, color: [0; 3], error: -1.0, track: ta }).unwrap();
162 if k < 25 {
164 let tb = vec![TrackEntry::new(2, k), TrackEntry::new(3, k)];
165 let mut y = truth.transform_point(&x);
166 if k == 0 {
167 y += Vec3::new(1000.0, 0.0, 0.0); }
169 b.add_point3d_with_id(k as u64 + 100, cumulus3d_core::Point3D { xyz: y, color: [0; 3], error: -1.0, track: tb }).unwrap();
170 }
171 }
172 let pairs = shared_point_correspondences(&a, &b, &Default::default());
173 assert_eq!(pairs.len(), 25);
174 assert!(pairs.iter().all(|(i, j)| *j == *i + 99));
175 let none = shared_point_correspondences(&a, &b, &SharedPointOptions { position_range: Some(2..3), ..Default::default() });
177 assert!(none.is_empty());
178 let one = shared_point_correspondences(&a, &b, &SharedPointOptions { position_range: Some(1..2), ..Default::default() });
179 assert_eq!(one.len(), 25);
180 let opts = RobustUmeyamaOptions { min_threshold: 1e-6, ..Default::default() };
181 let (r, _) = align_reconstructions(&a, &b, &Default::default(), &opts).unwrap();
182 sim_close(&r.sim3, &truth, 1e-8);
183 assert_eq!(r.num_inliers, 24);
184 assert!(!r.inlier_mask[0]);
185 }
186
187 #[test]
188 fn dedup_pairs() {
189 let (mut a, mut b) = (base(2, 5), base(2, 5));
191 for k in 0..5u32 {
192 let t = vec![TrackEntry::new(1, k), TrackEntry::new(2, k)];
193 let p = cumulus3d_core::Point3D { xyz: Vec3::new(k as f64, 1.0, 2.0), color: [0; 3], error: -1.0, track: t };
194 a.add_point3d_with_id(k as u64 + 1, p.clone()).unwrap();
195 b.add_point3d_with_id(k as u64 + 1, p).unwrap();
196 }
197 assert_eq!(shared_point_correspondences(&a, &b, &Default::default()).len(), 5);
198 }
199}