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::{Matrix3, Matrix6, 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 )
123}
124
125#[must_use]
127#[allow(clippy::missing_panics_doc)]
128#[tracing::instrument(skip_all, name = "pipeline::estimate_tag_pose")]
129pub fn estimate_tag_pose_with_config(
130 intrinsics: &CameraIntrinsics,
131 corners: &[[f64; 2]; 4],
132 tag_size: f64,
133 img: Option<&ImageView>,
134 mode: PoseEstimationMode,
135 config: &crate::config::DetectorConfig,
136) -> (Option<Pose>, Option<[[f64; 6]; 6]>) {
137 let Some(h_poly) = crate::decoder::Homography::square_to_quad(corners) else {
139 return (None, None);
140 };
141 let h_pixel = h_poly.h;
142
143 let k_inv = intrinsics.inv_matrix();
145 let h_norm = k_inv * h_pixel;
146
147 let scaler = 2.0 / tag_size;
153 let mut h_metric = h_norm;
154 let h1 = h_norm.column(0).into_owned();
155 let h2 = h_norm.column(1).into_owned();
156 h_metric.column_mut(0).scale_mut(scaler);
157 h_metric.column_mut(1).scale_mut(scaler);
158 h_metric
159 .column_mut(2)
160 .copy_from(&(h_norm.column(2) - h1 - h2));
161
162 let Some(candidates) = solve_ippe_square(&h_metric) else {
164 return (None, None);
165 };
166
167 let best_pose = find_best_pose(intrinsics, corners, tag_size, &candidates);
169
170 match (mode, img) {
172 (PoseEstimationMode::Accurate, Some(image)) => {
173 let uncertainty = crate::pose_weighted::compute_framework_uncertainty(
175 image,
176 corners,
177 &h_poly,
178 config.tikhonov_alpha_max,
179 config.sigma_n_sq,
180 config.structure_tensor_radius,
181 );
182 let (refined_pose, covariance) = crate::pose_weighted::refine_pose_lm_weighted(
183 intrinsics,
184 corners,
185 tag_size,
186 best_pose,
187 &uncertainty,
188 );
189 (Some(refined_pose), Some(covariance))
190 },
191 _ => {
192 (
194 Some(refine_pose_lm(
195 intrinsics,
196 corners,
197 tag_size,
198 best_pose,
199 config.huber_delta_px,
200 )),
201 None,
202 )
203 },
204 }
205}
206
207fn solve_ippe_square(h: &Matrix3<f64>) -> Option<[Pose; 2]> {
213 let h1 = h.column(0);
216 let h2 = h.column(1);
217
218 let a = h1.dot(&h1);
222 let b = h2.dot(&h2);
223 let c = h1.dot(&h2);
224
225 let trace = a + b;
228 let det = a * b - c * c;
229 let delta = (trace * trace - 4.0 * det).max(0.0).sqrt();
230
231 let s1_sq = (trace + delta) * 0.5;
233 let s2_sq = (trace - delta) * 0.5;
234
235 let s1 = s1_sq.sqrt();
237 let s2 = s2_sq.sqrt();
238
239 if (s1 - s2).abs() < 1e-4 * s1 {
242 let mut r1 = h1.clone_owned();
246 let scale = 1.0 / r1.norm();
247 r1 *= scale;
248
249 let mut r2 = h2 - r1 * (h2.dot(&r1));
251 r2 = r2.normalize();
252
253 let r3 = r1.cross(&r2);
254 let rot = Matrix3::from_columns(&[r1, r2, r3]);
255
256 let gamma = (s1 + s2) * 0.5;
261 if gamma < 1e-8 {
262 return None;
263 } let tz = 1.0 / gamma;
265 let t = h.column(2) * tz;
266
267 let pose = Pose::new(rot, t);
268 return Some([pose, pose]);
269 }
270
271 let v1 = if c.abs() > 1e-8 {
287 let v = nalgebra::Vector2::new(s1_sq - b, c);
288 v.normalize()
289 } else if a >= b {
290 nalgebra::Vector2::new(1.0, 0.0)
291 } else {
292 nalgebra::Vector2::new(0.0, 1.0)
293 };
294
295 let v2 = nalgebra::Vector2::new(-v1.y, v1.x);
297
298 let j_v1 = h1 * v1.x + h2 * v1.y;
302 let j_v2 = h1 * v2.x + h2 * v2.y;
303
304 if s1 < 1e-8 {
306 return None;
307 }
308 let u1 = j_v1 / s1;
309 let u2 = j_v2 / s2.max(1e-8); let r1_a = u1 * v1.x + u2 * v2.x;
319 let r2_a = u1 * v1.y + u2 * v2.y;
320 let r3_a = r1_a.cross(&r2_a);
321 let rot_a = Matrix3::from_columns(&[r1_a, r2_a, r3_a]);
322
323 let gamma = (s1 + s2) * 0.5;
325 let tz = 1.0 / gamma;
326 let t_a = h.column(2) * tz;
327 let pose_a = Pose::new(rot_a, t_a);
328
329 let n_a = rot_a.column(2);
337 let n_b_raw = Vector3::new(-n_a.x, -n_a.y, n_a.z);
339 let n_b = if n_b_raw.norm_squared() > 1e-8 {
340 n_b_raw.normalize()
341 } else {
342 Vector3::z_axis().into_inner()
344 };
345
346 let h1_norm = h1.normalize();
350 let x_b_raw = h1_norm - n_b * h1_norm.dot(&n_b);
351 let x_b = if x_b_raw.norm_squared() > 1e-8 {
352 x_b_raw.normalize()
353 } else {
354 let tangent = if n_b.x.abs() > 0.9 {
356 Vector3::y_axis().into_inner()
357 } else {
358 Vector3::x_axis().into_inner()
359 };
360 tangent.cross(&n_b).normalize()
361 };
362
363 let y_b = n_b.cross(&x_b);
364 let rot_b = Matrix3::from_columns(&[x_b, y_b, n_b]);
365 let pose_b = Pose::new(rot_b, t_a);
366
367 Some([pose_a, pose_b])
368}
369
370#[allow(clippy::too_many_lines)]
387fn refine_pose_lm(
388 intrinsics: &CameraIntrinsics,
389 corners: &[[f64; 2]; 4],
390 tag_size: f64,
391 initial_pose: Pose,
392 huber_delta_px: f64,
393) -> Pose {
394 let huber_delta = huber_delta_px;
395
396 let mut pose = initial_pose;
397 let s = tag_size;
398 let obj_pts = [
400 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), ];
405
406 let mut lambda = 1e-3_f64;
408 let mut nu = 2.0_f64;
409 let mut current_cost = huber_cost(intrinsics, corners, &obj_pts, &pose, huber_delta);
410
411 for _ in 0..20 {
413 let mut jtj = Matrix6::<f64>::zeros();
416 let mut jtr = Vector6::<f64>::zeros();
417
418 for i in 0..4 {
419 let p_world = obj_pts[i];
420 let p_cam = pose.rotation * p_world + pose.translation;
421 let z_inv = 1.0 / p_cam.z;
422 let z_inv2 = z_inv * z_inv;
423
424 let u_est = intrinsics.fx * p_cam.x * z_inv + intrinsics.cx;
425 let v_est = intrinsics.fy * p_cam.y * z_inv + intrinsics.cy;
426
427 let res_u = corners[i][0] - u_est;
428 let res_v = corners[i][1] - v_est;
429
430 let r_norm = (res_u * res_u + res_v * res_v).sqrt();
433 let w = if r_norm <= huber_delta {
434 1.0
435 } else {
436 huber_delta / r_norm
437 };
438
439 let du_dp = Vector3::new(
443 intrinsics.fx * z_inv,
444 0.0,
445 -intrinsics.fx * p_cam.x * z_inv2,
446 );
447 let dv_dp = Vector3::new(
448 0.0,
449 intrinsics.fy * z_inv,
450 -intrinsics.fy * p_cam.y * z_inv2,
451 );
452
453 let mut row_u = Vector6::zeros();
456 row_u[0] = du_dp[0];
457 row_u[1] = du_dp[1];
458 row_u[2] = du_dp[2];
459 row_u[3] = du_dp[1] * p_cam.z - du_dp[2] * p_cam.y;
460 row_u[4] = du_dp[2] * p_cam.x - du_dp[0] * p_cam.z;
461 row_u[5] = du_dp[0] * p_cam.y - du_dp[1] * p_cam.x;
462
463 let mut row_v = Vector6::zeros();
464 row_v[0] = dv_dp[0];
465 row_v[1] = dv_dp[1];
466 row_v[2] = dv_dp[2];
467 row_v[3] = dv_dp[1] * p_cam.z - dv_dp[2] * p_cam.y;
468 row_v[4] = dv_dp[2] * p_cam.x - dv_dp[0] * p_cam.z;
469 row_v[5] = dv_dp[0] * p_cam.y - dv_dp[1] * p_cam.x;
470
471 jtj += w * (row_u * row_u.transpose() + row_v * row_v.transpose());
473 jtr += w * (row_u * res_u + row_v * res_v);
474 }
475
476 if jtr.amax() < 1e-8 {
478 break;
479 }
480
481 let d_diag = Vector6::new(
485 jtj[(0, 0)].max(1e-8),
486 jtj[(1, 1)].max(1e-8),
487 jtj[(2, 2)].max(1e-8),
488 jtj[(3, 3)].max(1e-8),
489 jtj[(4, 4)].max(1e-8),
490 jtj[(5, 5)].max(1e-8),
491 );
492
493 let mut jtj_damped = jtj;
495 for k in 0..6 {
496 jtj_damped[(k, k)] += lambda * d_diag[k];
497 }
498
499 let delta = if let Some(chol) = jtj_damped.cholesky() {
500 chol.solve(&jtr)
501 } else {
502 lambda *= 10.0;
504 nu = 2.0;
505 continue;
506 };
507
508 let predicted_reduction = 0.5 * delta.dot(&(lambda * d_diag.component_mul(&delta) + jtr));
512
513 let twist = Vector3::new(delta[3], delta[4], delta[5]);
515 let trans_update = Vector3::new(delta[0], delta[1], delta[2]);
516 let rot_update = nalgebra::Rotation3::new(twist).matrix().into_owned();
517 let new_pose = Pose::new(
518 rot_update * pose.rotation,
519 rot_update * pose.translation + trans_update,
520 );
521
522 let new_cost = huber_cost(intrinsics, corners, &obj_pts, &new_pose, huber_delta);
523 let actual_reduction = current_cost - new_cost;
524
525 let rho = if predicted_reduction > 1e-12 {
526 actual_reduction / predicted_reduction
527 } else {
528 0.0
529 };
530
531 if rho > 0.0 {
532 pose = new_pose;
534 current_cost = new_cost;
535 lambda *= (1.0 - (2.0 * rho - 1.0).powi(3)).max(1.0 / 3.0);
537 nu = 2.0;
538
539 if delta.norm() < 1e-7 {
541 break;
542 }
543 } else {
544 lambda *= nu;
546 nu *= 2.0;
547 }
548 }
549
550 pose
551}
552
553fn huber_cost(
558 intrinsics: &CameraIntrinsics,
559 corners: &[[f64; 2]; 4],
560 obj_pts: &[Vector3<f64>; 4],
561 pose: &Pose,
562 delta: f64,
563) -> f64 {
564 let mut cost = 0.0;
565 for i in 0..4 {
566 let p = pose.project(&obj_pts[i], intrinsics);
567 let r_u = corners[i][0] - p[0];
568 let r_v = corners[i][1] - p[1];
569 let r_norm = (r_u * r_u + r_v * r_v).sqrt();
571 if r_norm <= delta {
572 cost += 0.5 * r_norm * r_norm;
573 } else {
574 cost += delta * (r_norm - 0.5 * delta);
575 }
576 }
577 cost
578}
579
580fn reprojection_error(
581 intrinsics: &CameraIntrinsics,
582 corners: &[[f64; 2]; 4],
583 obj_pts: &[Vector3<f64>; 4],
584 pose: &Pose,
585) -> f64 {
586 let mut err_sq = 0.0;
587 for i in 0..4 {
588 let p = pose.project(&obj_pts[i], intrinsics);
589 err_sq += (p[0] - corners[i][0]).powi(2) + (p[1] - corners[i][1]).powi(2);
590 }
591 err_sq
592}
593
594fn find_best_pose(
595 intrinsics: &CameraIntrinsics,
596 corners: &[[f64; 2]; 4],
597 tag_size: f64,
598 candidates: &[Pose; 2],
599) -> Pose {
600 let s = tag_size;
602 let obj_pts = [
603 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), ];
608
609 let err0 = reprojection_error(intrinsics, corners, &obj_pts, &candidates[0]);
610 let err1 = reprojection_error(intrinsics, corners, &obj_pts, &candidates[1]);
611
612 if err1 < err0 {
614 candidates[1]
615 } else {
616 candidates[0]
617 }
618}
619
620#[tracing::instrument(skip_all, name = "pipeline::pose_refinement")]
625pub fn refine_poses_soa(
626 batch: &mut DetectionBatch,
627 v: usize,
628 intrinsics: &CameraIntrinsics,
629 tag_size: f64,
630 img: Option<&ImageView>,
631 mode: PoseEstimationMode,
632) {
633 refine_poses_soa_with_config(
634 batch,
635 v,
636 intrinsics,
637 tag_size,
638 img,
639 mode,
640 &crate::config::DetectorConfig::default(),
641 );
642}
643
644#[tracing::instrument(skip_all, name = "pipeline::pose_refinement")]
646pub fn refine_poses_soa_with_config(
647 batch: &mut DetectionBatch,
648 v: usize,
649 intrinsics: &CameraIntrinsics,
650 tag_size: f64,
651 img: Option<&ImageView>,
652 mode: PoseEstimationMode,
653 config: &crate::config::DetectorConfig,
654) {
655 use rayon::prelude::*;
656
657 let poses: Vec<_> = (0..v)
660 .into_par_iter()
661 .map(|i| {
662 let corners = [
663 [
664 f64::from(batch.corners[i][0].x),
665 f64::from(batch.corners[i][0].y),
666 ],
667 [
668 f64::from(batch.corners[i][1].x),
669 f64::from(batch.corners[i][1].y),
670 ],
671 [
672 f64::from(batch.corners[i][2].x),
673 f64::from(batch.corners[i][2].y),
674 ],
675 [
676 f64::from(batch.corners[i][3].x),
677 f64::from(batch.corners[i][3].y),
678 ],
679 ];
680
681 let (pose_opt, _) =
682 estimate_tag_pose_with_config(intrinsics, &corners, tag_size, img, mode, config);
683
684 if let Some(pose) = pose_opt {
685 let q = UnitQuaternion::from_matrix(&pose.rotation);
686 let t = pose.translation;
687
688 let mut data = [0.0f32; 7];
690 data[0] = t.x as f32;
691 data[1] = t.y as f32;
692 data[2] = t.z as f32;
693 data[3] = q.coords.x as f32;
694 data[4] = q.coords.y as f32;
695 data[5] = q.coords.z as f32;
696 data[6] = q.coords.w as f32;
697 Some(data)
698 } else {
699 None
700 }
701 })
702 .collect();
703
704 for (i, pose_data) in poses.into_iter().enumerate() {
705 if let Some(data) = pose_data {
706 batch.poses[i] = Pose6D { data, padding: 0.0 };
707 } else {
708 batch.poses[i] = Pose6D {
709 data: [0.0; 7],
710 padding: 0.0,
711 };
712 }
713 }
714}
715
716#[cfg(test)]
717mod tests {
718 use super::*;
719 use proptest::prelude::*;
720
721 #[test]
722 fn test_pose_projection() {
723 let intrinsics = CameraIntrinsics::new(500.0, 500.0, 320.0, 240.0);
724 let rotation = Matrix3::identity();
725 let translation = Vector3::new(0.0, 0.0, 2.0); let pose = Pose::new(rotation, translation);
727
728 let p_world = Vector3::new(0.1, 0.1, 0.0);
729 let p_img = pose.project(&p_world, &intrinsics);
730
731 assert!((p_img[0] - 345.0).abs() < 1e-6);
734 assert!((p_img[1] - 265.0).abs() < 1e-6);
735 }
736
737 #[test]
738 fn test_perfect_pose_estimation() {
739 let intrinsics = CameraIntrinsics::new(800.0, 800.0, 400.0, 300.0);
740 let gt_rot = Matrix3::identity(); let gt_t = Vector3::new(0.1, -0.2, 1.5);
742 let gt_pose = Pose::new(gt_rot, gt_t);
743
744 let tag_size = 0.16; let s = tag_size;
746 let obj_pts = [
748 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), ];
753
754 let mut img_pts = [[0.0, 0.0]; 4];
755 for i in 0..4 {
756 img_pts[i] = gt_pose.project(&obj_pts[i], &intrinsics);
757 }
758
759 let (est_pose, _) = estimate_tag_pose(
760 &intrinsics,
761 &img_pts,
762 tag_size,
763 None,
764 PoseEstimationMode::Fast,
765 );
766 let est_pose = est_pose.expect("Pose estimation failed");
767
768 assert!((est_pose.translation.x - gt_t.x).abs() < 1e-3);
770 assert!((est_pose.translation.y - gt_t.y).abs() < 1e-3);
771 assert!((est_pose.translation.z - gt_t.z).abs() < 1e-3);
772
773 let diff_rot = est_pose.rotation - gt_rot;
775 assert!(diff_rot.norm() < 1e-3);
776 }
777
778 proptest! {
779 #[test]
780 fn prop_intrinsics_inversion(
781 fx in 100.0..2000.0f64,
782 fy in 100.0..2000.0f64,
783 cx in 0.0..1000.0f64,
784 cy in 0.0..1000.0f64
785 ) {
786 let intrinsics = CameraIntrinsics::new(fx, fy, cx, cy);
787 let k = intrinsics.as_matrix();
788 let k_inv = intrinsics.inv_matrix();
789 let identity = k * k_inv;
790
791 let expected = Matrix3::<f64>::identity();
792 for i in 0..3 {
793 for j in 0..3 {
794 prop_assert!((identity[(i, j)] - expected[(i, j)]).abs() < 1e-9);
795 }
796 }
797 }
798
799 #[test]
800 fn prop_pose_recovery_stability(
801 tx in -0.5..0.5f64,
802 ty in -0.5..0.5f64,
803 tz in 0.5..5.0f64, roll in -0.5..0.5f64,
805 pitch in -0.5..0.5f64,
806 yaw in -0.5..0.5f64,
807 noise in 0.0..0.1f64 ) {
809 let intrinsics = CameraIntrinsics::new(800.0, 800.0, 400.0, 300.0);
810 let translation = Vector3::new(tx, ty, tz);
811
812 let r_obj = nalgebra::Rotation3::from_euler_angles(roll, pitch, yaw);
814 let rotation = r_obj.matrix().into_owned();
815 let gt_pose = Pose::new(rotation, translation);
816
817 let tag_size = 0.16;
818 let s = tag_size;
819 let obj_pts = [
821 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), ];
826
827 let mut img_pts = [[0.0, 0.0]; 4];
828 for i in 0..4 {
829 let p = gt_pose.project(&obj_pts[i], &intrinsics);
830 img_pts[i] = [p[0] + noise, p[1] + noise];
832 }
833
834 if let (Some(est_pose), _) = estimate_tag_pose(&intrinsics, &img_pts, tag_size, None, PoseEstimationMode::Fast) {
835 let t_err = (est_pose.translation - translation).norm();
838 prop_assert!(t_err < 0.1 + noise * 0.1, "Translation error {} too high for noise {}", t_err, noise);
839
840 let r_err = (est_pose.rotation - rotation).norm();
841 prop_assert!(r_err < 0.1 + noise * 0.1, "Rotation error {} too high for noise {}", r_err, noise);
842 }
843 }
844 }
845}