Add ros transform publisher.

2025-09-01 16:57:03 +02:00
parent a428a16c74
commit 86d5a749b5
5 changed files with 306 additions and 0 deletions
@@ -0,0 +1,223 @@
+#include <array>
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+#include <rclcpp/rclcpp.hpp>
+#include <tf2_ros/transform_broadcaster.h>
+#include <geometry_msgs/msg/transform_stamped.hpp>
+
+#include "rpt_msgs/msg/poses.hpp"
+
+// =================================================================================================
+
+const std::string input_topic = "/poses/humans3d";
+
+// =================================================================================================
+// =================================================================================================
+
+class SkeletonTFPublisher : public rclcpp::Node
+{
+public:
+  SkeletonTFPublisher() : Node("skeleton_tf_publisher")
+  {
+    broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>(this);
+    subscriber_ = this->create_subscription<rpt_msgs::msg::Poses>(
+        input_topic, 1,
+        std::bind(&SkeletonTFPublisher::listener_callback, this, std::placeholders::_1));
+
+    pc_connections = {
+        {"hip_middle", "shoulder_middle"},
+        {"shoulder_middle", "head"},
+        {"shoulder_middle", "nose"},
+        {"nose", "eye_left"},
+        {"nose", "eye_right"},
+        {"nose", "ear_left"},
+        {"nose", "ear_right"},
+        {"shoulder_left", "shoulder_right"},
+        {"hip_middle", "shoulder_left"},
+        {"hip_middle", "shoulder_right"},
+        {"shoulder_left", "elbow_left"},
+        {"elbow_left", "wrist_left"},
+        {"hip_middle", "hip_left"},
+        {"hip_left", "knee_left"},
+        {"knee_left", "ankle_left"},
+        {"shoulder_right", "elbow_right"},
+        {"elbow_right", "wrist_right"},
+        {"hip_middle", "hip_right"},
+        {"hip_right", "knee_right"},
+        {"knee_right", "ankle_right"},
+    };
+    for (auto &pair : pc_connections)
+    {
+      cp_map[pair[1]] = pair[0];
+    }
+  }
+
+private:
+  std::vector<std::array<std::string, 2>> pc_connections;
+  std::map<std::string, std::string> cp_map;
+  std::shared_ptr<tf2_ros::TransformBroadcaster> broadcaster_;
+  rclcpp::Subscription<rpt_msgs::msg::Poses>::SharedPtr subscriber_;
+
+  void listener_callback(const rpt_msgs::msg::Poses::SharedPtr msg);
+  std::vector<geometry_msgs::msg::TransformStamped> generate_msg(
+      const std::vector<std::vector<std::array<float, 4>>> &poses,
+      const std::vector<std::string> &joint_names);
+};
+
+// =================================================================================================
+
+void SkeletonTFPublisher::listener_callback(const rpt_msgs::msg::Poses::SharedPtr msg)
+{
+  std::vector<std::vector<std::array<float, 4>>> poses;
+  const std::vector<std::string> &joint_names = msg->joint_names;
+
+  // Unflatten poses
+  size_t idx = 0;
+  std::vector<int32_t> &bshape = msg->bodies_shape;
+  for (int32_t i = 0; i < bshape[0]; ++i)
+  {
+    std::vector<std::array<float, 4>> body;
+    for (int32_t j = 0; j < bshape[1]; ++j)
+    {
+      std::array<float, 4> joint;
+      for (int32_t k = 0; k < bshape[2]; ++k)
+      {
+        joint[k] = msg->bodies_flat[idx];
+        ++idx;
+      }
+      body.push_back(std::move(joint));
+    }
+    poses.push_back(std::move(body));
+  }
+
+  if (poses.empty())
+  {
+    return;
+  }
+
+  auto tf_msgs = generate_msg(poses, joint_names);
+  if (!tf_msgs.empty())
+  {
+    broadcaster_->sendTransform(tf_msgs);
+  }
+}
+
+// =================================================================================================
+
+std::vector<geometry_msgs::msg::TransformStamped> SkeletonTFPublisher::generate_msg(
+    const std::vector<std::vector<std::array<float, 4>>> &poses,
+    const std::vector<std::string> &joint_names)
+{
+  std::vector<geometry_msgs::msg::TransformStamped> tf_msgs;
+  tf_msgs.reserve(poses.size() * pc_connections.size());
+  rclcpp::Time now = this->get_clock()->now();
+
+  for (size_t i = 0; i < poses.size(); ++i)
+  {
+    const auto &body = poses[i];
+
+    // Find index of "hip_middle" in joint_names
+    int hip_mid_idx = -1;
+    for (size_t k = 0; k < joint_names.size(); ++k)
+    {
+      if (joint_names[k] == "hip_middle")
+      {
+        hip_mid_idx = static_cast<int>(k);
+        break;
+      }
+    }
+
+    // Set "hip_middle" as child of "world"
+    const auto &joint = body[hip_mid_idx];
+    geometry_msgs::msg::TransformStamped tf_stamped;
+    tf_stamped.header.stamp = now;
+    tf_stamped.header.frame_id = "world";
+    tf_stamped.child_frame_id = "s" + std::to_string(i) + "-hip_middle";
+    tf_stamped.transform.translation.x = joint[0];
+    tf_stamped.transform.translation.y = joint[1];
+    tf_stamped.transform.translation.z = joint[2];
+    tf_stamped.transform.rotation.x = 0.0;
+    tf_stamped.transform.rotation.y = 0.0;
+    tf_stamped.transform.rotation.z = 0.0;
+    tf_stamped.transform.rotation.w = 1.0;
+    tf_msgs.push_back(std::move(tf_stamped));
+
+    // Add other joints
+    for (size_t j = 0; j < body.size(); ++j)
+    {
+      // Skip "hip_middle"
+      if (static_cast<int>(j) == hip_mid_idx)
+      {
+        continue;
+      }
+
+      const auto &joint_child = body[j];
+      const std::string &child_name = joint_names[j];
+      float conf = joint_child[3];
+      if (conf <= 0.0f)
+      {
+        continue;
+      }
+
+      // Look up its parent
+      auto it = cp_map.find(child_name);
+      if (it == cp_map.end())
+      {
+        continue;
+      }
+      const std::string &parent_name = it->second;
+
+      // Ensure the parent frame was actually published
+      int parent_idx = -1;
+      for (size_t k = 0; k < joint_names.size(); ++k)
+      {
+        if (joint_names[k] == parent_name)
+        {
+          parent_idx = static_cast<int>(k);
+          break;
+        }
+      }
+      if (parent_idx < 0)
+      {
+        // Should never happen if cp_map is consistent
+        continue;
+      }
+      if (parent_name != "hip_middle" && body[parent_idx][3] <= 0.0f)
+      {
+        // Parent not visible, skip this child
+        continue;
+      }
+      const auto &joint_parent = body[parent_idx];
+
+      // Publish child frame
+      geometry_msgs::msg::TransformStamped tf_stamped;
+      tf_stamped.header.stamp = now;
+      tf_stamped.header.frame_id = "s" + std::to_string(i) + "-" + parent_name;
+      tf_stamped.child_frame_id = "s" + std::to_string(i) + "-" + child_name;
+      tf_stamped.transform.translation.x = joint_child[0] - joint_parent[0];
+      tf_stamped.transform.translation.y = joint_child[1] - joint_parent[1];
+      tf_stamped.transform.translation.z = joint_child[2] - joint_parent[2];
+      tf_stamped.transform.rotation.x = 0.0;
+      tf_stamped.transform.rotation.y = 0.0;
+      tf_stamped.transform.rotation.z = 0.0;
+      tf_stamped.transform.rotation.w = 1.0;
+      tf_msgs.push_back(std::move(tf_stamped));
+    }
+  }
+
+  return tf_msgs;
+}
+
+// =================================================================================================
+// =================================================================================================
+
+int main(int argc, char *argv[])
+{
+  rclcpp::init(argc, argv);
+  rclcpp::spin(std::make_shared<SkeletonTFPublisher>());
+  rclcpp::shutdown();
+  return 0;
+}