chalkydri/subsys/
capriltags.rs

1//!
2//! Subsystem for the official AprilTags C library
3//!
4
5// TODO: implement this
6// There's actually already a decent Rust binding we can use.
7// There's an example here: <https://github.com/jerry73204/apriltag-rust/blob/master/apriltag/examples/detector.rs>
8//
9// <https://www.chiefdelphi.com/t/frc-blog-technology-updates-past-present-future-and-beyond-apriltags-and-new-radio/440931>
10// According to this post on CD, we're doing the 36h11 tag family now.
11
12use std::collections::HashMap;
13use std::fs::File;
14use std::path::Path;
15
16use apriltag::{Detector, Family, Image, TagParams};
17use apriltag_image::image::{DynamicImage, RgbImage};
18use apriltag_image::prelude::*;
19use camera_intrinsic_model::{GenericModel, OpenCVModel5};
20use rapier3d::math::{Matrix, Rotation, Translation};
21use rapier3d::na::Quaternion;
22#[cfg(feature = "rerun")]
23use re_sdk::external::re_types_core;
24#[cfg(feature = "rerun")]
25use re_types::{
26    archetypes::{Boxes2D, Points2D},
27    components::{PinholeProjection, PoseRotationQuat, Position2D, ViewCoordinates},
28};
29
30use crate::calibration::CalibratedModel;
31use crate::Subsystem;
32
33const TAG_SIZE: f64 = 165.1;
34
35pub struct CApriltagsDetector {
36    det: apriltag::Detector,
37    layout: HashMap<u64, (Translation<f64>, Rotation<f64>)>,
38    model: CalibratedModel,
39}
40impl<'fr> Subsystem<'fr> for CApriltagsDetector {
41    type Output = (Vec<f64>, Vec<f64>);
42    type Error = Box<dyn std::error::Error + Send>;
43
44    async fn init() -> Result<Self, Self::Error> {
45        let model = CalibratedModel::new();
46
47        let mut path = Path::new("/boot/layout.json");
48        if !path.exists() {
49            path = Path::new("./layout.json");
50        }
51
52        let layout = AprilTagFieldLayout::load(path);
53        let det = Detector::builder()
54            .add_family_bits(Family::tag_36h11(), 3)
55            .build()
56            .unwrap();
57
58        Ok(Self { det, layout, model })
59    }
60    fn process(&mut self, buf: crate::subsystem::Buffer) -> Result<Self::Output, Self::Error> {
61        let img_rgb = DynamicImage::ImageRgb8(RgbImage::from_vec(1280, 720, buf.to_vec()).unwrap());
62        let img_gray = img_rgb.grayscale();
63        let buf = img_gray.as_luma8().unwrap();
64        let img = Image::from_image_buffer(buf);
65        let dets = self.det.detect(&img);
66
67        let poses: Vec<_> = dets
68            .iter()
69            .filter_map(|det| {
70                // Extract camera calibration values from the [CalibratedModel]
71                let OpenCVModel5 { fx, fy, cx, cy, .. } =
72                    if let GenericModel::OpenCVModel5(model) = self.model.inner_model() {
73                        model
74                    } else {
75                        panic!("camera model type not supported yet");
76                    };
77
78                // Estimate tag pose with the camera calibration values
79                let pose = det
80                    .estimate_tag_pose(&TagParams {
81                        fx,
82                        fy,
83                        cx,
84                        cy,
85                        tagsize: TAG_SIZE,
86                    })
87                    .unwrap();
88
89                // Extract the camera's translation and rotation matrices from the [Pose]
90                let cam_translation = pose.translation().data().to_vec();
91                let cam_rotation = pose.rotation().data().to_vec();
92
93                // Convert the camera's translation and rotation matrices into proper Rust datatypes
94                let cam_translation =
95                    Translation::new(cam_translation[0], cam_translation[1], cam_translation[2]);
96                let cam_rotation = Rotation::from_matrix(&Matrix::from_vec(cam_rotation));
97
98                // Try to get the tag's pose from the field layout
99                if let Some((tag_translation, tag_rotation)) = self.layout.get(&(det.id() as u64)) {
100                    let translation = tag_translation * cam_translation;
101                    let rotation = tag_rotation * cam_rotation;
102                    return Some((translation, rotation, det.decision_margin() as f64));
103                }
104
105                None
106            })
107            .collect();
108
109        let mut avg_translation = Translation::new(0.0f64, 0.0, 0.0);
110        let mut avg_rotation = Quaternion::new(0.0f64, 0.0, 0.0, 0.0);
111
112        for pose in poses.iter() {
113            avg_translation.x += pose.0.x;
114            avg_translation.y += pose.0.y;
115            avg_translation.z += pose.0.z;
116
117            avg_rotation.w += pose.1.w;
118            avg_rotation.i += pose.1.i;
119            avg_rotation.j += pose.1.j;
120            avg_rotation.k += pose.1.k;
121        }
122
123        avg_translation.x /= poses.len() as f64;
124        avg_translation.x /= poses.len() as f64;
125        avg_translation.x /= poses.len() as f64;
126
127        avg_rotation.w /= poses.len() as f64;
128        avg_rotation.i /= poses.len() as f64;
129        avg_rotation.j /= poses.len() as f64;
130        avg_rotation.k /= poses.len() as f64;
131
132        Ok((
133            avg_translation.vector.data.as_slice().to_vec(),
134            avg_rotation.vector().data.into_slice().to_vec(),
135        ))
136    }
137}
138
139#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
140#[serde(rename_all = "camelCase")]
141pub struct AprilTagFieldLayout {
142    pub tags: Vec<LayoutTag>,
143    pub field: Field,
144}
145impl AprilTagFieldLayout {
146    pub fn load(path: impl AsRef<Path>) -> HashMap<u64, (Translation<f64>, Rotation<f64>)> {
147        let f = File::open(path).unwrap();
148        let layout: Self = serde_json::from_reader(f).unwrap();
149
150        let mut tags = HashMap::new();
151        for LayoutTag {
152            id,
153            pose:
154                LayoutPose {
155                    translation,
156                    rotation: LayoutRotation { quaternion },
157                },
158        } in layout.tags.clone()
159        {
160            // Turn the field layout values into Rust datatypes
161            let tag_translation = Translation::new(translation.x, translation.y, translation.z);
162            let tag_rotation = Rotation::from_quaternion(Quaternion::new(
163                quaternion.w,
164                quaternion.x,
165                quaternion.y,
166                quaternion.z,
167            ));
168
169            tags.insert(id as u64, (tag_translation, tag_rotation));
170        }
171
172        tags
173    }
174}
175
176#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
177#[serde(rename_all = "camelCase")]
178pub struct LayoutTag {
179    #[serde(rename = "ID")]
180    pub id: i64,
181    pub pose: LayoutPose,
182}
183
184#[derive(Default, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
185#[serde(rename_all = "camelCase")]
186pub struct LayoutPose {
187    pub translation: LayoutTranslation,
188    pub rotation: LayoutRotation,
189}
190
191#[derive(Default, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
192#[serde(rename_all = "camelCase")]
193pub struct LayoutTranslation {
194    pub x: f64,
195    pub y: f64,
196    pub z: f64,
197}
198
199#[derive(Default, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
200#[serde(rename_all = "camelCase")]
201pub struct LayoutRotation {
202    pub quaternion: LayoutQuaternion,
203}
204
205#[derive(Default, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
206#[serde(rename_all = "camelCase")]
207pub struct LayoutQuaternion {
208    #[serde(rename = "W")]
209    pub w: f64,
210    #[serde(rename = "X")]
211    pub x: f64,
212    #[serde(rename = "Y")]
213    pub y: f64,
214    #[serde(rename = "Z")]
215    pub z: f64,
216}
217
218#[derive(Default, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
219#[serde(rename_all = "camelCase")]
220pub struct Field {
221    pub length: f64,
222    pub width: f64,
223}