Add ros transform publisher.

extras/ros/docker-compose-3d.yml

@@ -31,3 +31,16 @@ services:
             - QT_X11_NO_MITSHM=1
             - "PYTHONUNBUFFERED=1"
         command: /bin/bash -i -c 'export ROS_DOMAIN_ID=18 && ros2 run rpt3d_wrapper_cpp rpt3d_wrapper'
+
+    skeleton_tfs:
+        image: rapidposetriangulation_ros3d
+        network_mode: "host"
+        ipc: "host"
+        privileged: true
+        volumes:
+            - ../../:/RapidPoseTriangulation/
+            - /tmp/.X11-unix:/tmp/.X11-unix
+            - /dev/shm:/dev/shm
+        environment:
+            - "PYTHONUNBUFFERED=1"
+        command: /bin/bash -i -c 'export ROS_DOMAIN_ID=18 && ros2 run transform_publisher skeleton_tfs'

extras/ros/dockerfile_3d

@@ -26,10 +26,12 @@ COPY ./extras/include/ /RapidPoseTriangulation/extras/include/
 COPY ./rpt/ /RapidPoseTriangulation/rpt/
 COPY ./extras/ros/rpt_msgs/ /RapidPoseTriangulation/extras/ros/rpt_msgs/
 COPY ./extras/ros/rpt3d_wrapper_cpp/ /RapidPoseTriangulation/extras/ros/rpt3d_wrapper_cpp/
+COPY ./extras/ros/transform_publisher/ /RapidPoseTriangulation/extras/ros/transform_publisher/
 
 # Link and build as ros package
 RUN ln -s /RapidPoseTriangulation/extras/ros/rpt_msgs/ /project/dev_ws/src/
 RUN ln -s /RapidPoseTriangulation/extras/ros/rpt3d_wrapper_cpp/ /project/dev_ws/src/
+RUN ln -s /RapidPoseTriangulation/extras/ros/transform_publisher/ /project/dev_ws/src/
 RUN /bin/bash -i -c 'cd /project/dev_ws/; colcon build --symlink-install --cmake-args -DCMAKE_BUILD_TYPE=Release'
 
 # Update ros packages -> autocompletion and check

extras/ros/transform_publisher/CMakeLists.txt

@@ -0,0 +1,45 @@
+cmake_minimum_required(VERSION 3.5)
+project(transform_publisher)
+
+# Default to C99
+if(NOT CMAKE_C_STANDARD)
+  set(CMAKE_C_STANDARD 99)
+endif()
+
+# Default to C++14
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 14)
+endif()
+
+if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+  add_compile_options(-Wall -Wextra -Wpedantic)
+endif()
+
+# find dependencies
+find_package(ament_cmake REQUIRED)
+find_package(rclcpp REQUIRED)
+find_package(rpt_msgs REQUIRED)
+find_package(geometry_msgs REQUIRED)
+find_package(tf2_ros REQUIRED)
+
+add_executable(skeleton_tfs src/skeleton_tfs.cpp)
+ament_target_dependencies(skeleton_tfs rclcpp rpt_msgs geometry_msgs tf2_ros)
+target_include_directories(skeleton_tfs PUBLIC
+  $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+  $<INSTALL_INTERFACE:include>)
+
+install(TARGETS skeleton_tfs
+  DESTINATION lib/${PROJECT_NAME})
+
+if(BUILD_TESTING)
+  find_package(ament_lint_auto REQUIRED)
+  # the following line skips the linter which checks for copyrights
+  # uncomment the line when a copyright and license is not present in all source files
+  #set(ament_cmake_copyright_FOUND TRUE)
+  # the following line skips cpplint (only works in a git repo)
+  # uncomment the line when this package is not in a git repo
+  #set(ament_cmake_cpplint_FOUND TRUE)
+  ament_lint_auto_find_test_dependencies()
+endif()
+
+ament_package()

extras/ros/transform_publisher/package.xml

@@ -0,0 +1,23 @@
+<?xml version="1.0"?>
+<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
+<package format="3">
+  <name>transform_publisher</name>
+  <version>0.0.0</version>
+  <description>TODO: Package description</description>
+  <maintainer email="root@todo.todo">root</maintainer>
+  <license>TODO: License declaration</license>
+
+  <buildtool_depend>ament_cmake</buildtool_depend>
+
+  <depend>rclcpp</depend>
+  <depend>rpt_msgs</depend>
+  <depend>geometry_msgs</depend>
+  <depend>tf2_ros</depend>
+
+  <test_depend>ament_lint_auto</test_depend>
+  <test_depend>ament_lint_common</test_depend>
+
+  <export>
+    <build_type>ament_cmake</build_type>
+  </export>
+</package>

extras/ros/transform_publisher/src/skeleton_tfs.cpp

@@ -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;
+}