use anyhow::Result;
use nalgebra as na;
use rerun as rr;
use std::{collections::HashMap, path::PathBuf};
use urdf_rs as urdf;
pub struct RerunUrdfLoader {
pub robot: urdf::Robot,
pub mesh_path: PathBuf,
pub root_prefix: String,
link_entity: HashMap<String, String>,
joint_entity: HashMap<String, String>,
}
impl RerunUrdfLoader {
pub fn from_urdf(
urdf_path: impl Into<PathBuf>,
mesh_path: Option<impl Into<PathBuf>>,
root_prefix: impl Into<String>,
) -> Result<Self> {
let urdf_path = urdf_path.into();
let mesh_path = mesh_path
.map(Into::into)
.unwrap_or(PathBuf::clone(&urdf_path));
let robot = urdf::read_file(&urdf_path)?;
let root_prefix = root_prefix.into();
let mut link_entity = HashMap::new();
for link in &robot.links {
let ent = format!("{}/{}", root_prefix, link.name);
link_entity.insert(link.name.clone(), ent.clone());
}
let mut joint_entity = HashMap::new();
for joint in &robot.joints {
let ent = format!("{}/{}", root_prefix, joint.name);
joint_entity.insert(joint.name.clone(), ent.clone());
}
Ok(Self {
robot,
mesh_path,
root_prefix,
link_entity,
joint_entity,
})
}
pub fn root_prefix(&self) -> &str {
&self.root_prefix
}
fn load_joint(
&self,
rec: &rr::RecordingStream,
joint: &urdf::Joint,
entity_path: &str,
parent_origin: &mut na::Isometry3<f64>,
) -> Result<urdf::LinkName> {
*parent_origin *= origin_to_isometry3(&joint.origin);
let tran = isometry3_to_transform3d(parent_origin);
rec.log(entity_path, &tran)?;
Ok(joint.child.clone())
}
fn load_link(
&self,
rec: &rr::RecordingStream,
link: &urdf::Link,
entity_path: &str,
parent_origin: &na::Isometry3<f64>,
) -> Result<String> {
rec.log(entity_path, &isometry3_to_transform3d(parent_origin))?;
for (i, visual) in link.visual.iter().enumerate() {
let vroot = format!("{entity_path}/visual_{i}");
self.load_visual(rec, visual, &vroot)?;
}
Ok(link.name.clone())
}
fn load_visual(
&self,
rec: &rr::RecordingStream,
visual: &urdf::Visual,
entity_path: &str,
) -> Result<()> {
let mut tran = origin_to_transform3d(&visual.origin);
match &visual.geometry {
urdf::Geometry::Mesh { filename, scale } => {
let asset = rr::Asset3D::from_file_path(self.mesh_path.join(filename))?;
rec.log_static(entity_path, &asset)?;
if let Some(scale) = scale {
tran = tran.with_scale([scale[0], scale[1], scale[2]]);
}
rec.log_static(entity_path, &tran)?;
}
urdf::Geometry::Box { size } => {
let hs = [size[0] / 2.0, size[1] / 2.0, size[2] / 2.0];
rec.log_static(entity_path, &rr::Boxes3D::from_half_sizes([hs]))?;
rec.log_static(entity_path, &tran)?;
}
urdf::Geometry::Cylinder { radius, length } => {
rec.log_static(
entity_path,
&rr::Cylinders3D::from_lengths_and_radii([*length as f32], [*radius as f32]),
)?;
rec.log_static(entity_path, &tran)?;
}
urdf::Geometry::Sphere { radius } => {
rec.log_static(
entity_path,
&rr::Points3D::new([visual.origin.xyz.0]).with_radii([*radius as f32]),
)?;
rec.log_static(entity_path, &tran)?;
}
_ => {
rec.log_static(entity_path, &rr::TextLog::new("Unsupported geometry"))?;
}
}
Ok(())
}
pub fn register_statics(
&self,
rec: &rr::RecordingStream,
base: na::Isometry3<f64>,
) -> Result<()> {
rec.log_static("world", &rr::ViewCoordinates::RIGHT_HAND_Z_UP())?;
let mut link_map: HashMap<String, &urdf::Link> = self
.robot
.links
.iter()
.map(|link| (link.name.clone(), link))
.collect();
let mut joint_map: HashMap<String, &urdf::Joint> = self
.robot
.joints
.iter()
.map(|joint| (joint.parent.link.clone(), joint))
.collect();
let mut parent = self.robot.links[0].name.clone();
let mut pose = base;
self.load_link(
rec,
link_map.remove(&parent).unwrap(),
&self.link_entity[&parent],
&pose,
)?;
while let Some(joint) = joint_map.remove(&parent) {
let child = self.load_joint(rec, joint, &self.joint_entity[&joint.name], &mut pose)?;
parent = self.load_link(
rec,
link_map.remove(&child.link).unwrap(),
&self.link_entity[&child.link],
&pose,
)?;
}
Ok(())
}
pub fn log_frame(
&self,
rec: &rr::RecordingStream,
frame_idx: i64,
link_world_poses: &HashMap<String, na::Isometry3<f64>>,
timeline: &str, ) {
rec.set_time_sequence(timeline, frame_idx);
for (link_name, ent_path) in &self.link_entity {
if let Some(pose) = link_world_poses.get(link_name) {
let _ = rec.log(ent_path.clone(), &isometry3_to_transform3d(pose));
}
}
}
pub fn link_names(&self) -> impl Iterator<Item = &str> {
self.link_entity.keys().map(|s| s.as_str())
}
pub fn link_entity_path(&self, link: &str) -> Option<&str> {
self.link_entity.get(link).map(|s| s.as_str())
}
pub fn all_visual_entity_paths(&self) -> Result<Vec<String>> {
let mut out = Vec::new();
for link in &self.robot.links {
if let Some(link_ent) = self.link_entity.get(&link.name) {
for (i, _visual) in link.visual.iter().enumerate() {
out.push(format!("{link_ent}/visual_{i}"));
}
}
}
Ok(out)
}
pub fn ordered_links_urdf(&self) -> Result<Vec<String>> {
use std::collections::HashSet;
let child_set: HashSet<&str> = self
.robot
.joints
.iter()
.map(|j| j.child.link.as_str())
.collect();
let root = self
.robot
.links
.iter()
.find(|l| !child_set.contains(l.name.as_str()))
.ok_or_else(|| anyhow::anyhow!("Cannot find root link"))?
.name
.clone();
let mut ordered = Vec::with_capacity(self.robot.links.len());
ordered.push(root);
for j in &self.robot.joints {
ordered.push(j.child.link.clone());
}
Ok(ordered)
}
}
pub fn origin_to_transform3d(pose: &urdf::Pose) -> rr::Transform3D {
let urdf::Pose {
xyz: urdf::Vec3(xyz),
rpy: urdf::Vec3(rpy),
} = pose;
let rot = na::UnitQuaternion::from_euler_angles(rpy[0], rpy[1], rpy[2]);
let rot = rr::Quaternion::from_wxyz([rot.w as f32, rot.i as f32, rot.j as f32, rot.k as f32]);
rr::Transform3D::from_translation_rotation(xyz, rot)
}
pub fn origin_to_isometry3(pose: &urdf::Pose) -> na::Isometry3<f64> {
let urdf::Pose {
xyz: urdf::Vec3(xyz),
rpy: urdf::Vec3(rpy),
} = pose;
let rot = na::UnitQuaternion::from_euler_angles(rpy[0], rpy[1], rpy[2]);
na::Isometry3::from_parts(na::Translation3::new(xyz[0], xyz[1], xyz[2]), rot)
}
pub fn isometry3_to_transform3d(pose: &na::Isometry3<f64>) -> rr::Transform3D {
let tran = pose.translation.vector;
let rot = pose.rotation;
rr::Transform3D::from_translation_rotation(
[tran.x, tran.y, tran.z],
rr::Quaternion::from_wxyz([rot.w as f32, rot.i as f32, rot.j as f32, rot.k as f32]),
)
}