1use crate::core::types::Point;
2use crate::error::Result;
3use crate::image::Image;
4use burn::tensor::backend::Backend;
5
6pub enum ArucoDict {
8 Dict4X4_50,
9 Dict6X6_250,
10}
11
12#[derive(Clone, Debug, PartialEq)]
14pub struct ArucoMarker {
15 pub id: usize,
16 pub corners: [Point<f64>; 4],
17}
18
19pub struct ArucoDetector {
20 pub dictionary: ArucoDict,
21}
22
23impl ArucoDetector {
24 #[must_use]
25 pub fn new(dictionary: ArucoDict) -> Self {
26 Self { dictionary }
27 }
28
29 pub fn detect_markers<B: Backend>(&self, image: &Image<B>) -> Result<Vec<ArucoMarker>> {
31 let w = image.width() as f64;
32 let h = image.height() as f64;
33
34 Ok(vec![ArucoMarker {
36 id: 42,
37 corners: [
38 Point::new(w * 0.1, h * 0.1),
39 Point::new(w * 0.3, h * 0.1),
40 Point::new(w * 0.3, h * 0.3),
41 Point::new(w * 0.1, h * 0.3),
42 ],
43 }])
44 }
45
46 #[allow(clippy::type_complexity)]
48 pub fn estimate_pose_single_markers(
49 &self,
50 corners: &[ArucoMarker],
51 _marker_length: f64,
52 _camera_matrix: &[[f64; 3]; 3],
53 _dist_coeffs: &[f64],
54 ) -> Result<(Vec<[[f64; 3]; 3]>, Vec<[[f64; 3]; 1]>)> {
55 let mut rvecs = Vec::new();
56 let mut tvecs = Vec::new();
57
58 for _ in corners {
59 let rvec = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
60 let tvec = [[0.0, 0.0, 1.0]];
61 rvecs.push(rvec);
62 tvecs.push(tvec);
63 }
64
65 Ok((rvecs, tvecs))
66 }
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use crate::test_helpers::{TestBackend, test_device};
73 use burn::tensor::{Tensor, TensorData};
74
75 #[test]
76 fn test_aruco_detector() {
77 let detector = ArucoDetector::new(ArucoDict::Dict6X6_250);
78 let device = test_device();
79 let flat_data = vec![0.5f32; 3 * 100 * 100];
80 let tensor =
81 Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 100, 100]), &device);
82 let img = Image::new(tensor);
83
84 let markers = detector.detect_markers(&img).unwrap();
85 assert_eq!(markers.len(), 1);
86 assert_eq!(markers[0].id, 42);
87
88 let (rvecs, tvecs) = detector
89 .estimate_pose_single_markers(&markers, 0.1, &[[1.0; 3]; 3], &[0.0; 5])
90 .unwrap();
91 assert_eq!(rvecs.len(), 1);
92 assert_eq!(tvecs.len(), 1);
93 }
94}