1#![allow(clippy::many_single_char_names, clippy::similar_names)]
9use crate::batch::{DetectionBatch, Pose6D};
10use crate::config::PoseEstimationMode;
11use crate::image::ImageView;
12use nalgebra::{Matrix2, Matrix3, Matrix6, Rotation3, UnitQuaternion, Vector3, Vector6};
13
14#[derive(Debug, Clone, Copy, PartialEq)]
16#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
17pub struct CameraIntrinsics {
18 pub fx: f64,
20 pub fy: f64,
22 pub cx: f64,
24 pub cy: f64,
26}
27
28impl CameraIntrinsics {
29 #[must_use]
31 pub fn new(fx: f64, fy: f64, cx: f64, cy: f64) -> Self {
32 Self { fx, fy, cx, cy }
33 }
34
35 #[must_use]
37 pub fn as_matrix(&self) -> Matrix3<f64> {
38 Matrix3::new(self.fx, 0.0, self.cx, 0.0, self.fy, self.cy, 0.0, 0.0, 1.0)
39 }
40
41 #[must_use]
43 pub fn inv_matrix(&self) -> Matrix3<f64> {
44 Matrix3::new(
45 1.0 / self.fx,
46 0.0,
47 -self.cx / self.fx,
48 0.0,
49 1.0 / self.fy,
50 -self.cy / self.fy,
51 0.0,
52 0.0,
53 1.0,
54 )
55 }
56}
57
58#[derive(Debug, Clone, Copy)]
61#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
62pub struct Pose {
63 pub rotation: Matrix3<f64>,
65 pub translation: Vector3<f64>,
67}
68
69impl Pose {
70 #[must_use]
72 pub fn new(rotation: Matrix3<f64>, translation: Vector3<f64>) -> Self {
73 Self {
74 rotation,
75 translation,
76 }
77 }
78
79 #[must_use]
81 pub fn project(&self, point: &Vector3<f64>, intrinsics: &CameraIntrinsics) -> [f64; 2] {
82 let p_cam = self.rotation * point + self.translation;
83 let x = (p_cam.x / p_cam.z) * intrinsics.fx + intrinsics.cx;
84 let y = (p_cam.y / p_cam.z) * intrinsics.fy + intrinsics.cy;
85 [x, y]
86 }
87}
88
89#[must_use]
106#[allow(clippy::missing_panics_doc)]
107#[tracing::instrument(skip_all, name = "pipeline::estimate_tag_pose")]
108pub fn estimate_tag_pose(
109 intrinsics: &CameraIntrinsics,
110 corners: &[[f64; 2]; 4],
111 tag_size: f64,
112 img: Option<&ImageView>,
113 mode: PoseEstimationMode,
114) -> (Option<Pose>, Option<[[f64; 6]; 6]>) {
115 estimate_tag_pose_with_config(
116 intrinsics,
117 corners,
118 tag_size,
119 img,
120 mode,
121 &crate::config::DetectorConfig::default(),
122 None,
123 )
124}
125
126#[must_use]
128#[allow(clippy::missing_panics_doc)]
129#[tracing::instrument(skip_all, name = "pipeline::estimate_tag_pose")]
130pub fn estimate_tag_pose_with_config(
131 intrinsics: &CameraIntrinsics,
132 corners: &[[f64; 2]; 4],
133 tag_size: f64,
134 img: Option<&ImageView>,
135 mode: PoseEstimationMode,
136 config: &crate::config::DetectorConfig,
137 external_covariances: Option<&[Matrix2<f64>; 4]>,
138) -> (Option<Pose>, Option<[[f64; 6]; 6]>) {
139 let Some(h_poly) = crate::decoder::Homography::square_to_quad(corners) else {
141 return (None, None);
142 };
143 let h_pixel = h_poly.h;
144
145 let k_inv = intrinsics.inv_matrix();
147 let h_norm = k_inv * h_pixel;
148
149 let scaler = 2.0 / tag_size;
155 let mut h_metric = h_norm;
156 let h1 = h_norm.column(0).into_owned();
157 let h2 = h_norm.column(1).into_owned();
158 h_metric.column_mut(0).scale_mut(scaler);
159 h_metric.column_mut(1).scale_mut(scaler);
160 h_metric
161 .column_mut(2)
162 .copy_from(&(h_norm.column(2) - h1 - h2));
163
164 let Some(candidates) = solve_ippe_square(&h_metric) else {
166 return (None, None);
167 };
168
169 let best_pose = find_best_pose(intrinsics, corners, tag_size, &candidates);
171
172 match (mode, img, external_covariances) {
174 (PoseEstimationMode::Accurate, _, Some(ext_covs)) => {
175 let (refined_pose, covariance) = crate::pose_weighted::refine_pose_lm_weighted(
177 intrinsics, corners, tag_size, best_pose, ext_covs,
178 );
179 (Some(refined_pose), Some(covariance))
180 },
181 (PoseEstimationMode::Accurate, Some(image), None) => {
182 let uncertainty = crate::pose_weighted::compute_framework_uncertainty(
184 image,
185 corners,
186 &h_poly,
187 config.tikhonov_alpha_max,
188 config.sigma_n_sq,
189 config.structure_tensor_radius,
190 );
191 let (refined_pose, covariance) = crate::pose_weighted::refine_pose_lm_weighted(
192 intrinsics,
193 corners,
194 tag_size,
195 best_pose,
196 &uncertainty,
197 );
198 (Some(refined_pose), Some(covariance))
199 },
200 _ => {
201 (
203 Some(refine_pose_lm(
204 intrinsics,
205 corners,
206 tag_size,
207 best_pose,
208 config.huber_delta_px,
209 )),
210 None,
211 )
212 },
213 }
214}
215
216fn solve_ippe_square(h: &Matrix3<f64>) -> Option<[Pose; 2]> {
222 let h1 = h.column(0);
225 let h2 = h.column(1);
226
227 let a = h1.dot(&h1);
231 let b = h2.dot(&h2);
232 let c = h1.dot(&h2);
233
234 let trace = a + b;
237 let det = a * b - c * c;
238 let delta = (trace * trace - 4.0 * det).max(0.0).sqrt();
239
240 let s1_sq = (trace + delta) * 0.5;
242 let s2_sq = (trace - delta) * 0.5;
243
244 let s1 = s1_sq.sqrt();
246 let s2 = s2_sq.sqrt();
247
248 if (s1 - s2).abs() < 1e-4 * s1 {
251 let mut r1 = h1.clone_owned();
255 let scale = 1.0 / r1.norm();
256 r1 *= scale;
257
258 let mut r2 = h2 - r1 * (h2.dot(&r1));
260 r2 = r2.normalize();
261
262 let r3 = r1.cross(&r2);
263 let rot = Matrix3::from_columns(&[r1, r2, r3]);
264
265 let gamma = (s1 + s2) * 0.5;
270 if gamma < 1e-8 {
271 return None;
272 } let tz = 1.0 / gamma;
274 let t = h.column(2) * tz;
275
276 let pose = Pose::new(rot, t);
277 return Some([pose, pose]);
278 }
279
280 let v1 = if c.abs() > 1e-8 {
296 let v = nalgebra::Vector2::new(s1_sq - b, c);
297 v.normalize()
298 } else if a >= b {
299 nalgebra::Vector2::new(1.0, 0.0)
300 } else {
301 nalgebra::Vector2::new(0.0, 1.0)
302 };
303
304 let v2 = nalgebra::Vector2::new(-v1.y, v1.x);
306
307 let j_v1 = h1 * v1.x + h2 * v1.y;
311 let j_v2 = h1 * v2.x + h2 * v2.y;
312
313 if s1 < 1e-8 {
315 return None;
316 }
317 let u1 = j_v1 / s1;
318 let u2 = j_v2 / s2.max(1e-8); let r1_a = u1 * v1.x + u2 * v2.x;
328 let r2_a = u1 * v1.y + u2 * v2.y;
329 let r3_a = r1_a.cross(&r2_a);
330 let rot_a = Matrix3::from_columns(&[r1_a, r2_a, r3_a]);
331
332 let gamma = (s1 + s2) * 0.5;
334 let tz = 1.0 / gamma;
335 let t_a = h.column(2) * tz;
336 let pose_a = Pose::new(rot_a, t_a);
337
338 let n_a = rot_a.column(2);
346 let n_b_raw = Vector3::new(-n_a.x, -n_a.y, n_a.z);
348 let n_b = if n_b_raw.norm_squared() > 1e-8 {
349 n_b_raw.normalize()
350 } else {
351 Vector3::z_axis().into_inner()
353 };
354
355 let h1_norm = h1.normalize();
359 let x_b_raw = h1_norm - n_b * h1_norm.dot(&n_b);
360 let x_b = if x_b_raw.norm_squared() > 1e-8 {
361 x_b_raw.normalize()
362 } else {
363 let tangent = if n_b.x.abs() > 0.9 {
365 Vector3::y_axis().into_inner()
366 } else {
367 Vector3::x_axis().into_inner()
368 };
369 tangent.cross(&n_b).normalize()
370 };
371
372 let y_b = n_b.cross(&x_b);
373 let rot_b = Matrix3::from_columns(&[x_b, y_b, n_b]);
374 let pose_b = Pose::new(rot_b, t_a);
375
376 Some([pose_a, pose_b])
377}
378
379#[allow(clippy::too_many_lines)]
396fn refine_pose_lm(
397 intrinsics: &CameraIntrinsics,
398 corners: &[[f64; 2]; 4],
399 tag_size: f64,
400 initial_pose: Pose,
401 huber_delta_px: f64,
402) -> Pose {
403 let huber_delta = huber_delta_px;
404
405 let mut pose = initial_pose;
406 let s = tag_size;
407 let obj_pts = [
409 Vector3::new(0.0, 0.0, 0.0), Vector3::new(s, 0.0, 0.0), Vector3::new(s, s, 0.0), Vector3::new(0.0, s, 0.0), ];
414
415 let mut lambda = 1e-3_f64;
417 let mut nu = 2.0_f64;
418 let mut current_cost = huber_cost(intrinsics, corners, &obj_pts, &pose, huber_delta);
419
420 for _ in 0..20 {
422 let mut jtj = Matrix6::<f64>::zeros();
425 let mut jtr = Vector6::<f64>::zeros();
426
427 for i in 0..4 {
428 let p_world = obj_pts[i];
429 let p_cam = pose.rotation * p_world + pose.translation;
430 let z_inv = 1.0 / p_cam.z;
431 let z_inv2 = z_inv * z_inv;
432
433 let u_est = intrinsics.fx * p_cam.x * z_inv + intrinsics.cx;
434 let v_est = intrinsics.fy * p_cam.y * z_inv + intrinsics.cy;
435
436 let res_u = corners[i][0] - u_est;
437 let res_v = corners[i][1] - v_est;
438
439 let r_norm = (res_u * res_u + res_v * res_v).sqrt();
442 let w = if r_norm <= huber_delta {
443 1.0
444 } else {
445 huber_delta / r_norm
446 };
447
448 let du_dp = Vector3::new(
452 intrinsics.fx * z_inv,
453 0.0,
454 -intrinsics.fx * p_cam.x * z_inv2,
455 );
456 let dv_dp = Vector3::new(
457 0.0,
458 intrinsics.fy * z_inv,
459 -intrinsics.fy * p_cam.y * z_inv2,
460 );
461
462 let mut row_u = Vector6::zeros();
465 row_u[0] = du_dp[0];
466 row_u[1] = du_dp[1];
467 row_u[2] = du_dp[2];
468 row_u[3] = du_dp[1] * p_cam.z - du_dp[2] * p_cam.y;
469 row_u[4] = du_dp[2] * p_cam.x - du_dp[0] * p_cam.z;
470 row_u[5] = du_dp[0] * p_cam.y - du_dp[1] * p_cam.x;
471
472 let mut row_v = Vector6::zeros();
473 row_v[0] = dv_dp[0];
474 row_v[1] = dv_dp[1];
475 row_v[2] = dv_dp[2];
476 row_v[3] = dv_dp[1] * p_cam.z - dv_dp[2] * p_cam.y;
477 row_v[4] = dv_dp[2] * p_cam.x - dv_dp[0] * p_cam.z;
478 row_v[5] = dv_dp[0] * p_cam.y - dv_dp[1] * p_cam.x;
479
480 jtj += w * (row_u * row_u.transpose() + row_v * row_v.transpose());
482 jtr += w * (row_u * res_u + row_v * res_v);
483 }
484
485 if jtr.amax() < 1e-8 {
487 break;
488 }
489
490 let d_diag = Vector6::new(
494 jtj[(0, 0)].max(1e-8),
495 jtj[(1, 1)].max(1e-8),
496 jtj[(2, 2)].max(1e-8),
497 jtj[(3, 3)].max(1e-8),
498 jtj[(4, 4)].max(1e-8),
499 jtj[(5, 5)].max(1e-8),
500 );
501
502 let mut jtj_damped = jtj;
504 for k in 0..6 {
505 jtj_damped[(k, k)] += lambda * d_diag[k];
506 }
507
508 let delta = if let Some(chol) = jtj_damped.cholesky() {
509 chol.solve(&jtr)
510 } else {
511 lambda *= 10.0;
513 nu = 2.0;
514 continue;
515 };
516
517 let predicted_reduction = 0.5 * delta.dot(&(lambda * d_diag.component_mul(&delta) + jtr));
521
522 let twist = Vector3::new(delta[3], delta[4], delta[5]);
524 let trans_update = Vector3::new(delta[0], delta[1], delta[2]);
525 let rot_update = nalgebra::Rotation3::new(twist).matrix().into_owned();
526 let new_pose = Pose::new(
527 rot_update * pose.rotation,
528 rot_update * pose.translation + trans_update,
529 );
530
531 let new_cost = huber_cost(intrinsics, corners, &obj_pts, &new_pose, huber_delta);
532 let actual_reduction = current_cost - new_cost;
533
534 let rho = if predicted_reduction > 1e-12 {
535 actual_reduction / predicted_reduction
536 } else {
537 0.0
538 };
539
540 if rho > 0.0 {
541 pose = new_pose;
543 current_cost = new_cost;
544 lambda *= (1.0 - (2.0 * rho - 1.0).powi(3)).max(1.0 / 3.0);
546 nu = 2.0;
547
548 if delta.norm() < 1e-7 {
550 break;
551 }
552 } else {
553 lambda *= nu;
555 nu *= 2.0;
556 }
557 }
558
559 pose
560}
561
562fn huber_cost(
567 intrinsics: &CameraIntrinsics,
568 corners: &[[f64; 2]; 4],
569 obj_pts: &[Vector3<f64>; 4],
570 pose: &Pose,
571 delta: f64,
572) -> f64 {
573 let mut cost = 0.0;
574 for i in 0..4 {
575 let p = pose.project(&obj_pts[i], intrinsics);
576 let r_u = corners[i][0] - p[0];
577 let r_v = corners[i][1] - p[1];
578 let r_norm = (r_u * r_u + r_v * r_v).sqrt();
580 if r_norm <= delta {
581 cost += 0.5 * r_norm * r_norm;
582 } else {
583 cost += delta * (r_norm - 0.5 * delta);
584 }
585 }
586 cost
587}
588
589fn reprojection_error(
590 intrinsics: &CameraIntrinsics,
591 corners: &[[f64; 2]; 4],
592 obj_pts: &[Vector3<f64>; 4],
593 pose: &Pose,
594) -> f64 {
595 let mut err_sq = 0.0;
596 for i in 0..4 {
597 let p = pose.project(&obj_pts[i], intrinsics);
598 err_sq += (p[0] - corners[i][0]).powi(2) + (p[1] - corners[i][1]).powi(2);
599 }
600 err_sq
601}
602
603fn find_best_pose(
604 intrinsics: &CameraIntrinsics,
605 corners: &[[f64; 2]; 4],
606 tag_size: f64,
607 candidates: &[Pose; 2],
608) -> Pose {
609 let s = tag_size;
611 let obj_pts = [
612 Vector3::new(0.0, 0.0, 0.0), Vector3::new(s, 0.0, 0.0), Vector3::new(s, s, 0.0), Vector3::new(0.0, s, 0.0), ];
617
618 let err0 = reprojection_error(intrinsics, corners, &obj_pts, &candidates[0]);
619 let err1 = reprojection_error(intrinsics, corners, &obj_pts, &candidates[1]);
620
621 if err1 < err0 {
623 candidates[1]
624 } else {
625 candidates[0]
626 }
627}
628
629#[tracing::instrument(skip_all, name = "pipeline::pose_refinement")]
634pub fn refine_poses_soa(
635 batch: &mut DetectionBatch,
636 v: usize,
637 intrinsics: &CameraIntrinsics,
638 tag_size: f64,
639 img: Option<&ImageView>,
640 mode: PoseEstimationMode,
641) {
642 refine_poses_soa_with_config(
643 batch,
644 v,
645 intrinsics,
646 tag_size,
647 img,
648 mode,
649 &crate::config::DetectorConfig::default(),
650 );
651}
652
653#[tracing::instrument(skip_all, name = "pipeline::pose_refinement")]
655pub fn refine_poses_soa_with_config(
656 batch: &mut DetectionBatch,
657 v: usize,
658 intrinsics: &CameraIntrinsics,
659 tag_size: f64,
660 img: Option<&ImageView>,
661 mode: PoseEstimationMode,
662 config: &crate::config::DetectorConfig,
663) {
664 use rayon::prelude::*;
665
666 let poses: Vec<_> = (0..v)
669 .into_par_iter()
670 .map(|i| {
671 let corners = [
672 [
673 f64::from(batch.corners[i][0].x),
674 f64::from(batch.corners[i][0].y),
675 ],
676 [
677 f64::from(batch.corners[i][1].x),
678 f64::from(batch.corners[i][1].y),
679 ],
680 [
681 f64::from(batch.corners[i][2].x),
682 f64::from(batch.corners[i][2].y),
683 ],
684 [
685 f64::from(batch.corners[i][3].x),
686 f64::from(batch.corners[i][3].y),
687 ],
688 ];
689
690 let mut ext_covs = [Matrix2::zeros(); 4];
693 let mut has_ext_covs = false;
694 #[allow(clippy::needless_range_loop)]
695 for j in 0..4 {
696 ext_covs[j] = Matrix2::new(
697 f64::from(batch.corner_covariances[i][j * 4]),
698 f64::from(batch.corner_covariances[i][j * 4 + 1]),
699 f64::from(batch.corner_covariances[i][j * 4 + 2]),
700 f64::from(batch.corner_covariances[i][j * 4 + 3]),
701 );
702 if ext_covs[j].norm_squared() > 1e-12 {
703 has_ext_covs = true;
704 }
705 }
706
707 let (pose_opt, _) = estimate_tag_pose_with_config(
708 intrinsics,
709 &corners,
710 tag_size,
711 img,
712 mode,
713 config,
714 if has_ext_covs { Some(&ext_covs) } else { None },
715 );
716
717 if let Some(pose) = pose_opt {
718 let rot = Rotation3::from_matrix_unchecked(pose.rotation);
728 let q = UnitQuaternion::from_rotation_matrix(&rot);
729 let t = pose.translation;
730
731 let mut data = [0.0f32; 7];
733 data[0] = t.x as f32;
734 data[1] = t.y as f32;
735 data[2] = t.z as f32;
736 data[3] = q.coords.x as f32;
737 data[4] = q.coords.y as f32;
738 data[5] = q.coords.z as f32;
739 data[6] = q.coords.w as f32;
740 Some(data)
741 } else {
742 None
743 }
744 })
745 .collect();
746
747 for (i, pose_data) in poses.into_iter().enumerate() {
748 if let Some(data) = pose_data {
749 batch.poses[i] = Pose6D { data, padding: 0.0 };
750 } else {
751 batch.poses[i] = Pose6D {
752 data: [0.0; 7],
753 padding: 0.0,
754 };
755 }
756 }
757}
758
759#[cfg(test)]
760mod tests {
761 use super::*;
762 use proptest::prelude::*;
763
764 #[test]
765 fn test_pose_projection() {
766 let intrinsics = CameraIntrinsics::new(500.0, 500.0, 320.0, 240.0);
767 let rotation = Matrix3::identity();
768 let translation = Vector3::new(0.0, 0.0, 2.0); let pose = Pose::new(rotation, translation);
770
771 let p_world = Vector3::new(0.1, 0.1, 0.0);
772 let p_img = pose.project(&p_world, &intrinsics);
773
774 assert!((p_img[0] - 345.0).abs() < 1e-6);
777 assert!((p_img[1] - 265.0).abs() < 1e-6);
778 }
779
780 #[test]
781 fn test_perfect_pose_estimation() {
782 let intrinsics = CameraIntrinsics::new(800.0, 800.0, 400.0, 300.0);
783 let gt_rot = Matrix3::identity(); let gt_t = Vector3::new(0.1, -0.2, 1.5);
785 let gt_pose = Pose::new(gt_rot, gt_t);
786
787 let tag_size = 0.16; let s = tag_size;
789 let obj_pts = [
791 Vector3::new(0.0, 0.0, 0.0), Vector3::new(s, 0.0, 0.0), Vector3::new(s, s, 0.0), Vector3::new(0.0, s, 0.0), ];
796
797 let mut img_pts = [[0.0, 0.0]; 4];
798 for i in 0..4 {
799 img_pts[i] = gt_pose.project(&obj_pts[i], &intrinsics);
800 }
801
802 let (est_pose, _) = estimate_tag_pose(
803 &intrinsics,
804 &img_pts,
805 tag_size,
806 None,
807 PoseEstimationMode::Fast,
808 );
809 let est_pose = est_pose.expect("Pose estimation failed");
810
811 assert!((est_pose.translation.x - gt_t.x).abs() < 1e-3);
813 assert!((est_pose.translation.y - gt_t.y).abs() < 1e-3);
814 assert!((est_pose.translation.z - gt_t.z).abs() < 1e-3);
815
816 let diff_rot = est_pose.rotation - gt_rot;
818 assert!(diff_rot.norm() < 1e-3);
819 }
820
821 proptest! {
822 #[test]
823 fn prop_intrinsics_inversion(
824 fx in 100.0..2000.0f64,
825 fy in 100.0..2000.0f64,
826 cx in 0.0..1000.0f64,
827 cy in 0.0..1000.0f64
828 ) {
829 let intrinsics = CameraIntrinsics::new(fx, fy, cx, cy);
830 let k = intrinsics.as_matrix();
831 let k_inv = intrinsics.inv_matrix();
832 let identity = k * k_inv;
833
834 let expected = Matrix3::<f64>::identity();
835 for i in 0..3 {
836 for j in 0..3 {
837 prop_assert!((identity[(i, j)] - expected[(i, j)]).abs() < 1e-9);
838 }
839 }
840 }
841
842 #[test]
843 fn prop_pose_recovery_stability(
844 tx in -0.5..0.5f64,
845 ty in -0.5..0.5f64,
846 tz in 0.5..5.0f64, roll in -0.5..0.5f64,
848 pitch in -0.5..0.5f64,
849 yaw in -0.5..0.5f64,
850 noise in 0.0..0.1f64 ) {
852 let intrinsics = CameraIntrinsics::new(800.0, 800.0, 400.0, 300.0);
853 let translation = Vector3::new(tx, ty, tz);
854
855 let r_obj = nalgebra::Rotation3::from_euler_angles(roll, pitch, yaw);
857 let rotation = r_obj.matrix().into_owned();
858 let gt_pose = Pose::new(rotation, translation);
859
860 let tag_size = 0.16;
861 let s = tag_size;
862 let obj_pts = [
864 Vector3::new(0.0, 0.0, 0.0), Vector3::new(s, 0.0, 0.0), Vector3::new(s, s, 0.0), Vector3::new(0.0, s, 0.0), ];
869
870 let mut img_pts = [[0.0, 0.0]; 4];
871 for i in 0..4 {
872 let p = gt_pose.project(&obj_pts[i], &intrinsics);
873 img_pts[i] = [p[0] + noise, p[1] + noise];
875 }
876
877 if let (Some(est_pose), _) = estimate_tag_pose(&intrinsics, &img_pts, tag_size, None, PoseEstimationMode::Fast) {
878 let t_err = (est_pose.translation - translation).norm();
881 prop_assert!(t_err < 0.1 + noise * 0.1, "Translation error {} too high for noise {}", t_err, noise);
882
883 let r_err = (est_pose.rotation - rotation).norm();
884 prop_assert!(r_err < 0.1 + noise * 0.1, "Rotation error {} too high for noise {}", r_err, noise);
885 }
886 }
887 }
888}