chalkydri/pose/
mod.rs

1use std::{collections::HashMap, sync::Arc, time::Duration};
2
3use field_layout::{AprilTagFieldLayout, Field};
4use sophus_autodiff::{linalg::VecF64, prelude::*};
5use sophus_lie::{HasAverage, Isometry3F64, Rotation3F64, prelude::*};
6use tokio::sync::{Mutex, RwLock, mpsc};
7
8use crate::{Cfg, Nt, error::Error};
9
10pub(crate) mod field_layout;
11
12/// Keeps pose transforms and ...
13#[derive(Clone)]
14pub struct PoseEstimator {
15    ///// Transform registry
16    //reg: Arc<Mutex<Registry>>,
17    //tx: mpsc::Sender<Transform>,
18    layout: Arc<RwLock<Option<AprilTagFieldLayout>>>,
19    tag_mappings: Arc<RwLock<Option<HashMap<usize, Isometry3F64>>>>,
20    poses_rx: Arc<Mutex<mpsc::UnboundedReceiver<Isometry3F64>>>,
21    poses_tx: mpsc::UnboundedSender<Isometry3F64>,
22}
23impl PoseEstimator {
24    pub async fn new() -> Result<Self, Error> {
25        let (poses_tx, poses_rx) = mpsc::unbounded_channel();
26
27        let est = Self {
28            layout: Arc::new(RwLock::new(None)),
29            tag_mappings: Arc::new(RwLock::new(None)),
30            poses_tx,
31            poses_rx: Arc::new(Mutex::new(poses_rx)),
32        };
33
34        est.load_layout().await?;
35
36        Ok(est)
37    }
38
39    /// (Re)load the field layout
40    pub async fn load_layout(&self) -> Result<(), Error> {
41        if let Some(layouts) = &Cfg.read().await.field_layouts {
42            if let Some(layout_name) = &Cfg.read().await.field_layout {
43                if let Some(layout) = layouts.get(layout_name) {
44                    *self.layout.write().await = Some(layout.clone());
45
46                    Ok(())
47                } else {
48                    Err(Error::FieldLayoutDoesNotExist)
49                }
50            } else {
51                Err(Error::FieldLayoutNotSelected)
52            }
53        } else {
54            Err(Error::NoFieldLayouts)
55        }
56    }
57
58    /// Add a transform to the transform registry
59    pub async fn add_transform_from_tag(
60        &self,
61        tag_est_pos: Isometry3F64,
62        tag_id: usize,
63    ) -> Result<(), Error> {
64        if let Some(tag_mappings) = self.tag_mappings.read().await.clone() {
65            if let Some(tag_field_pos) = tag_mappings.get(&tag_id) {
66                let cam_est_rel_pos = tag_est_pos.inverse();
67                let cam_relto_pos = cam_est_rel_pos.translation() / 2.0;
68
69                let cam_relto_x = -cam_relto_pos.x;
70                let cam_relto_y = cam_relto_pos.y;
71                let cam_relto_z = -cam_relto_pos.z;
72
73                let cam_angle = Rotation3F64::try_from_mat(
74                    tag_field_pos.rotation().matrix() - cam_est_rel_pos.rotation().matrix(),
75                )
76                .unwrap();
77
78                let cam_fcs_abs = Isometry3F64::from_translation_and_rotation(
79                    VecF64::<3>::new(
80                        cam_relto_x + tag_field_pos.translation().x,
81                        cam_relto_y + tag_field_pos.translation().y,
82                        cam_relto_z + tag_field_pos.translation().z,
83                    ),
84                    cam_angle,
85                );
86
87                debug!("{cam_fcs_abs}");
88
89                //let cam_fcs_abs_x = cam_fcs_abs.translation().x;
90                //let cam_fcs_abs_y = cam_fcs_abs.translation().y;
91                //let cam_fcs_abs_z = cam_fcs_abs.translation().z;
92
93                //let robot_angle = cam_fcs_abs.rotation() +
94
95                //
96                self.poses_tx.send(cam_fcs_abs).unwrap();
97
98                Ok(())
99            } else {
100                Err(Error::InvalidTag)
101            }
102        } else {
103            Err(Error::FieldLayoutNotSelected)
104        }
105    }
106
107    ///// Interpolate the robot's pose based on transforms
108    //pub async fn get_robot_pose(&self) -> Result<Transform, Error> {
109    //    match self
110    //        .reg
111    //        .lock()
112    //        .await
113    //        .get_transform("robot", "field", Timestamp::now())
114    //    {
115    //        Ok(t) => Ok(t),
116    //        Err(err) => Err(Error::FailedToGetPose(err)),
117    //    }
118    //}
119
120    pub async fn nt_loop(&self) {
121        let est = self.clone();
122        tokio::spawn(async move {
123            let mut t = Nt
124                .publish::<Vec<f64>>(format!(
125                    "/chalkydri/robot_pose/{}/translation",
126                    Cfg.read().await.device_name.clone().unwrap()
127                ))
128                .await
129                .unwrap();
130            let mut r = Nt
131                .publish::<Vec<f64>>(format!(
132                    "/chalkydri/robot_pose/{}/rotation",
133                    Cfg.read().await.device_name.clone().unwrap()
134                ))
135                .await
136                .unwrap();
137            loop {
138                //match est.get_robot_pose().await {
139                //    Ok(pose) => {
140                //        t.set(vec![
141                //            pose.translation.x,
142                //            pose.translation.y,
143                //            pose.translation.z,
144                //        ])
145                //        .await;
146                //        r.set(vec![
147                //            pose.rotation.w,
148                //            pose.rotation.x,
149                //            pose.rotation.y,
150                //            pose.rotation.z,
151                //        ])
152                //        .await;
153                //    }
154                //    Err(err) => {
155                //        error!("failed to get pose");
156                //    }
157                //}
158
159                let mut poses = Vec::new();
160                while let Some(pose) = est.poses_rx.lock().await.recv().await {
161                    poses.push(pose);
162                }
163                let pose = Isometry3F64::average(&poses).unwrap();
164                t.set(vec![
165                    pose.translation().x,
166                    pose.translation().y,
167                    pose.translation().z,
168                ])
169                .await;
170                let rot = pose.rotation().matrix().as_slice().to_vec();
171                r.set(rot).await;
172
173                tokio::time::sleep(Duration::from_millis(20)).await;
174            }
175        });
176    }
177}
chalkydri/pose/mod.rs

chalkydri/pose/
mod.rs
