wip

Author: MichaelGrupp

crates/store/re_data_loader/src/loader_urdf.rs

@@ -6,13 +6,13 @@ use std::{
 use ahash::{HashMap, HashMapExt as _, HashSet, HashSetExt as _};
 use anyhow::{Context as _, bail};
 use itertools::Itertools as _;
-use urdf_rs::{Geometry, Joint, Link, Material, Robot, Vec3, Vec4};
+use urdf_rs::{Geometry, Joint, Link, LinkName, Material, Robot, Vec3, Vec4};
 
 use re_chunk::{ChunkBuilder, ChunkId, EntityPath, RowId, TimePoint};
 use re_log_types::{EntityPathPart, StoreId};
 use re_types::{
     AsComponents, Component as _, ComponentDescriptor, SerializedComponentBatch,
-    archetypes::{Asset3D, Transform3D},
+    archetypes::{Asset3D, CoordinateFrame, Transform3D, TransformAxes3D},
     datatypes::Vec3D,
     external::glam,
 };
@@ -282,6 +282,15 @@ fn log_robot(
         .map(|prefix| prefix / EntityPath::from_single_string(urdf_tree.name.clone()))
         .unwrap_or_else(|| EntityPath::from_single_string(urdf_tree.name.clone()));
 
+    // Log the robot's root coordinate frame_id.
+    send_archetype(
+        tx,
+        store_id,
+        entity_path.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame_id(urdf_tree.root.name.clone()),
+    )?;
+
     walk_tree(
         &urdf_tree,
         tx,
@@ -345,8 +354,8 @@ fn log_joint(
         name: _,
         joint_type,
         origin,
-        parent: _,
-        child: _,
+        parent,
+        child,
         axis,
         limit,
         calibration,
@@ -355,7 +364,24 @@ fn log_joint(
         safety_controller,
     } = joint;
 
-    send_transform(tx, store_id, joint_path.clone(), origin, timepoint)?;
+    // A joint's own coordinate frame is that of its parent link.
+    send_archetype(
+        tx,
+        store_id,
+        joint_path.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame_id(parent.link.clone()),
+    )?;
+    // Send the joint origin, i.e. the default transform from parent link to child link.
+    send_transform(
+        tx,
+        store_id,
+        joint_path.clone(),
+        origin,
+        timepoint,
+        Some(parent),
+        Some(child),
+    )?;
 
     log_debug_format(
         tx,
@@ -404,13 +430,23 @@ fn log_joint(
     Ok(())
 }
 
-fn transform_from_pose(origin: &urdf_rs::Pose) -> Transform3D {
+fn transform_from_pose(
+    origin: &urdf_rs::Pose,
+    parent: Option<&LinkName>,
+    child: Option<&LinkName>,
+) -> Transform3D {
     let urdf_rs::Pose { xyz, rpy } = origin;
     let translation = [xyz[0] as f32, xyz[1] as f32, xyz[2] as f32];
     let quaternion = quat_xyzw_from_roll_pitch_yaw(rpy[0] as f32, rpy[1] as f32, rpy[2] as f32);
-    Transform3D::update_fields()
+    let transform = Transform3D::update_fields()
         .with_translation(translation)
-        .with_quaternion(quaternion)
+        .with_quaternion(quaternion);
+    if let (Some(parent), Some(child)) = (parent, child) {
+        return transform
+            .with_parent_frame(parent.link.clone())
+            .with_child_frame(child.link.clone());
+    }
+    transform
 }
 
 fn send_transform(
@@ -419,19 +455,29 @@ fn send_transform(
     entity_path: EntityPath,
     origin: &urdf_rs::Pose,
     timepoint: &TimePoint,
+    parent: Option<&LinkName>,
+    child: Option<&LinkName>,
 ) -> anyhow::Result<()> {
     let urdf_rs::Pose { xyz, rpy } = origin;
     let is_identity = xyz.0 == [0.0, 0.0, 0.0] && rpy.0 == [0.0, 0.0, 0.0];
 
     if is_identity {
         Ok(()) // avoid noise
     } else {
+        // TODO: remove axis log this after debugging
+        send_archetype(
+            tx,
+            store_id,
+            entity_path.clone(),
+            timepoint,
+            &TransformAxes3D::update_fields().with_axis_length(0.1),
+        )?;
         send_archetype(
             tx,
             store_id,
             entity_path,
             timepoint,
-            &transform_from_pose(origin),
+            &transform_from_pose(origin, parent, child),
         )
     }
 }
@@ -485,6 +531,19 @@ fn log_link(
         timepoint,
     )?;
 
+    // Log coordinate frame ID of the link.
+    send_archetype(
+        tx,
+        store_id,
+        link_entity.clone(),
+        timepoint,
+        &CoordinateFrame::update_fields().with_frame_id(link.name.clone()),
+    )?;
+
+    let link_parent = urdf_tree
+        .get_parent_of_link(&link.name)
+        .map(|joint| &joint.child);
+
     for (i, visual) in visual.iter().enumerate() {
         let urdf_rs::Visual {
             name,
@@ -493,7 +552,7 @@ fn log_link(
             material,
         } = visual;
         let name = name.clone().unwrap_or_else(|| format!("visual_{i}"));
-        let vis_entity = link_entity / EntityPathPart::new(name);
+        let vis_entity = link_entity / EntityPathPart::new(name.clone());
 
         // Prefer inline defined material properties if present, otherwise fall back to global material.
         let material = material.as_ref().and_then(|mat| {
@@ -504,7 +563,29 @@ fn log_link(
             }
         });
 
-        send_transform(tx, store_id, vis_entity.clone(), origin, timepoint)?;
+        let link_child = urdf_rs::LinkName { link: name };
+        // TODO
+        send_transform(
+            tx,
+            store_id,
+            vis_entity.clone(),
+            origin,
+            timepoint,
+            link_parent,
+            Some(&link_child),
+        )?;
+
+        if let Some(parent) = link_parent {
+            let coordinate_frame =
+                CoordinateFrame::update_fields().with_frame_id(parent.link.clone());
+            send_archetype(
+                tx,
+                store_id,
+                vis_entity.clone(),
+                timepoint,
+                &coordinate_frame,
+            )?;
+        }
 
         log_geometry(
             urdf_tree, tx, store_id, vis_entity, geometry, material, timepoint,
@@ -518,9 +599,30 @@ fn log_link(
             geometry,
         } = collision;
         let name = name.clone().unwrap_or_else(|| format!("collision_{i}"));
-        let collision_entity = link_entity / EntityPathPart::new(name);
+        let collision_entity = link_entity / EntityPathPart::new(name.clone());
 
-        send_transform(tx, store_id, collision_entity.clone(), origin, timepoint)?;
+        let link_child = urdf_rs::LinkName { link: name };
+        send_transform(
+            tx,
+            store_id,
+            collision_entity.clone(),
+            origin,
+            timepoint,
+            link_parent,
+            Some(&link_child),
+        )?;
+
+        if let Some(parent) = link_parent {
+            let coordinate_frame =
+                CoordinateFrame::update_fields().with_frame_id(parent.link.clone());
+            send_archetype(
+                tx,
+                store_id,
+                collision_entity.clone(),
+                timepoint,
+                &coordinate_frame,
+            )?;
+        }
 
         log_geometry(
             urdf_tree,