Skip to main content

iris/aruco/
mod.rs

1use crate::core::types::Point;
2use crate::error::Result;
3use crate::image::Image;
4use burn::tensor::backend::Backend;
5
6/// Predefined `ArUco` marker dictionary types.
7pub enum ArucoDict {
8    Dict4X4_50,
9    Dict6X6_250,
10}
11
12/// Represents a detected `ArUco` marker.
13#[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    /// Detects `ArUco` markers in the image.
30    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        // Return a mock detected marker for structural correctness
35        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    /// Estimates 3D poses (rotation and translation vectors) for detected markers.
47    #[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}