1use 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 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 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 let cam_translation = pose.translation().data().to_vec();
91 let cam_rotation = pose.rotation().data().to_vec();
92
93 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 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 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}