From 6a589be5a8be8d3d2e474a0861679c9a84d317d0 Mon Sep 17 00:00:00 2001
From: Lucas Walter <wsacul@gmail.com>
Date: Sat, 21 Sep 2024 13:39:14 -0700
Subject: [PATCH] tf_capture outputs a toml of current transforms after looking
 up transforms from an input toml

---
 tf_roslibrust/Cargo.toml         |  4 ++
 tf_roslibrust/src/bin/capture.rs | 68 ++++++++++++++++++++++++++++++++
 tf_roslibrust/src/tf_util.rs     | 40 ++++++++++++++++++-
 3 files changed, 110 insertions(+), 2 deletions(-)
 create mode 100644 tf_roslibrust/src/bin/capture.rs

diff --git a/tf_roslibrust/Cargo.toml b/tf_roslibrust/Cargo.toml
index d7e6ac45..3083fb1d 100644
--- a/tf_roslibrust/Cargo.toml
+++ b/tf_roslibrust/Cargo.toml
@@ -44,6 +44,10 @@ branch = "subscribe_any"
 # path = "../../roslibrust/roslibrust_codegen_macro"
 # version="0.10"
 
+[[bin]]
+name = "tf_capture"
+path = "src/bin/capture.rs"
+
 [[bin]]
 name = "tf_tree"
 path = "src/bin/tree.rs"
diff --git a/tf_roslibrust/src/bin/capture.rs b/tf_roslibrust/src/bin/capture.rs
new file mode 100644
index 00000000..f4214fd6
--- /dev/null
+++ b/tf_roslibrust/src/bin/capture.rs
@@ -0,0 +1,68 @@
+use tf_roslibrust::tf_util;
+use tf_roslibrust::transforms::tf2_msgs;
+
+/// Load a toml file of a list of parent/child frames and then find all those transforms in
+/// a live system and output a new toml listing of their current values
+/// tf_capture examples/transforms.toml > current_transforms.toml
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    use roslibrust::ros1::NodeHandle;
+
+    // need to have leading slash on node name and topic to function properly
+    // so figure out namespace then prefix it to name and topics
+    let mut ns = String::from("");
+    let args = std::env::args();
+    let mut args2 = Vec::new();
+    {
+        // get namespace
+        for arg in args {
+            if arg.starts_with("__ns:=") {
+                ns = arg.replace("__ns:=", "");
+            } else {
+                args2.push(arg);
+            }
+        }
+    }
+
+    let full_node_name = &format!("/{ns}/tf_capture").replace("//", "/");
+    // println!("{}", format!("full ns and node name: {full_node_name}"));
+
+    let config_file = &args2[1];
+    let tfm = tf_util::get_transforms_from_toml(config_file)?;
+
+    let nh = NodeHandle::new(&std::env::var("ROS_MASTER_URI")?, full_node_name)
+        .await
+        .unwrap();
+
+    print!(
+        "# {:.3} getting transforms...",
+        tf_util::stamp_to_f64(tf_util::stamp_now())
+    );
+    let listener = tf_roslibrust::TfListener::new(&nh).await;
+
+    // let some transforms arrive
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    let mut new_tfm = tf2_msgs::TFMessage::default();
+    for old_tfs in &tfm.transforms {
+        // TODO(lucasw) make a lookup_most_recent_transform and then don't need the None
+        let res =
+            listener.lookup_transform(&old_tfs.header.frame_id, &old_tfs.child_frame_id, None);
+        match res {
+            Ok(new_tfs) => {
+                new_tfm.transforms.push(new_tfs);
+                // TODO(lucasw) compare old values to current values, warn if large
+            }
+            Err(err) => {
+                eprintln!("# {err:?} -> going to use old version for this transform");
+                new_tfm.transforms.push(old_tfs.clone());
+            }
+        }
+    }
+
+    let toml = tf_util::transforms_to_toml(tfm)?;
+    print!("\n\n{toml}");
+
+    Ok(())
+}
diff --git a/tf_roslibrust/src/tf_util.rs b/tf_roslibrust/src/tf_util.rs
index 5846ca39..f7c3f7a8 100644
--- a/tf_roslibrust/src/tf_util.rs
+++ b/tf_roslibrust/src/tf_util.rs
@@ -1,6 +1,6 @@
 use chrono::TimeDelta;
 use roslibrust_codegen::Time;
-use serde_derive::Deserialize;
+use serde_derive::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::time::SystemTime;
 
@@ -39,7 +39,7 @@ pub fn stamp_to_f64(stamp: Time) -> f64 {
 }
 
 /// use for loading from a toml
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Serialize, Debug)]
 struct TransformRaw {
     frame: String,
     child_frame: String,
@@ -51,6 +51,42 @@ struct TransformRaw {
     yaw: Option<f64>,
 }
 
+impl TransformRaw {
+    fn from_transform_stamped(tfs: geometry_msgs::TransformStamped) -> Self {
+        let rot = tfs.transform.rotation;
+        let quat = nalgebra::UnitQuaternion::new_normalize(nalgebra::geometry::Quaternion::new(
+            rot.w, rot.x, rot.y, rot.z,
+        ));
+        let (roll, pitch, yaw) = quat.euler_angles();
+
+        let tr = tfs.transform.translation;
+        let (x, y, z) = (tr.x, tr.y, tr.z);
+
+        Self {
+            frame: tfs.header.frame_id,
+            child_frame: tfs.child_frame_id,
+            x: Some(x),
+            y: Some(y),
+            z: Some(z),
+            roll: Some(roll),
+            pitch: Some(pitch),
+            yaw: Some(yaw),
+        }
+    }
+}
+
+pub fn transforms_to_toml(tfm: tf2_msgs::TFMessage) -> Result<String, anyhow::Error> {
+    let mut tf_vec = Vec::new();
+    for tfs in tfm.transforms {
+        tf_vec.push(TransformRaw::from_transform_stamped(tfs));
+    }
+
+    let mut tf_data = HashMap::new();
+    tf_data.insert("tf", tf_vec);
+
+    Ok(toml::to_string(&tf_data).unwrap())
+}
+
 pub fn get_transforms_from_toml(filename: &str) -> Result<tf2_msgs::TFMessage, anyhow::Error> {
     let tf_data = {
         let contents = match std::fs::read_to_string(filename) {