Configurable input topics.

extras/ros/rpt2d_wrapper_cpp/src/rpt2d_wrapper.cpp

@@ -24,9 +24,8 @@ using json = nlohmann::json;
 
 // =================================================================================================
 
-static const std::string cam_id = "camera01";
-static const std::string img_input_topic = "/" + cam_id + "/pylon_ros2_camera_node/image_raw";
-static const std::string pose_out_topic = "/poses/" + cam_id;
+static const std::string img_input_topic = "/{}/pylon_ros2_camera_node/image_raw";
+static const std::string pose_out_topic = "/poses/{}";
 
 static const float min_bbox_score = 0.4;
 static const float min_bbox_area = 0.1 * 0.1;
@@ -44,10 +43,22 @@ static const std::map<std::string, bool> whole_body = {
 class Rpt2DWrapperNode : public rclcpp::Node
 {
 public:
-    Rpt2DWrapperNode(const std::string &node_name)
-        : Node(node_name)
+    Rpt2DWrapperNode(const std::string &cam_id, const std::string &rotate)
+        : Node("rpt2d_wrapper_" + cam_id)
     {
         this->is_busy = false;
+        std::string img_topic = std::string(img_input_topic)
+                                    .replace(img_input_topic.find("{}"), 2, cam_id);
+        std::string pose_topic = std::string(pose_out_topic)
+                                     .replace(pose_out_topic.find("{}"), 2, cam_id);
+
+        // Check rotation
+        if (rotate != "0" && rotate != "90" && rotate != "180" && rotate != "-90")
+        {
+            RCLCPP_ERROR(this->get_logger(), "Invalid rotation value: %s", rotate.c_str());
+            return;
+        }
+        this->rotate_deg = std::stoi(rotate);
 
         // QoS
         rclcpp::QoS qos_profile(1);
@@ -56,11 +67,11 @@ public:
 
         // Setup subscriber
         image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
-            img_input_topic, qos_profile,
+            img_topic, qos_profile,
             std::bind(&Rpt2DWrapperNode::callback_images, this, std::placeholders::_1));
 
         // Setup publisher
-        pose_pub_ = this->create_publisher<rpt_msgs::msg::Poses>(pose_out_topic, qos_profile);
+        pose_pub_ = this->create_publisher<rpt_msgs::msg::Poses>(pose_topic, qos_profile);
 
         // Load model
         bool use_wb = utils_pipeline::use_whole_body(whole_body);
@@ -74,6 +85,7 @@ private:
     rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub_;
     rclcpp::Publisher<rpt_msgs::msg::Poses>::SharedPtr pose_pub_;
     std::atomic<bool> is_busy;
+    int rotate_deg;
 
     // Pose model pointer
     std::unique_ptr<utils_2d_pose::PosePredictor> kpt_model;
@@ -153,11 +165,11 @@ void Rpt2DWrapperNode::callback_images(const sensor_msgs::msg::Image::SharedPtr
     std::vector<int32_t> bshape = {(int)valid_poses.size(), (int)joint_names_2d.size(), 3};
     pose_msg.bodies_shape = bshape;
     pose_msg.bodies_flat.reserve(bshape[0] * bshape[1] * bshape[2]);
-    for (size_t i = 0; i < bshape[0]; i++)
+    for (int32_t i = 0; i < bshape[0]; i++)
     {
-        for (size_t j = 0; j < bshape[1]; j++)
+        for (int32_t j = 0; j < bshape[1]; j++)
         {
-            for (size_t k = 0; k < bshape[2]; k++)
+            for (int32_t k = 0; k < bshape[2]; k++)
             {
                 pose_msg.bodies_flat.push_back(valid_poses[i][j][k]);
             }
@@ -185,6 +197,24 @@ std::vector<std::vector<std::array<float, 3>>> Rpt2DWrapperNode::call_model(cons
     cv::Mat rgb_image = utils_pipeline::bayer2rgb(image);
     std::vector<cv::Mat> images_2d = {rgb_image};
 
+    // Optionally rotate image
+    if (this->rotate_deg == 0)
+    {
+        // pass
+    }
+    else if (this->rotate_deg == 90)
+    {
+        cv::rotate(images_2d[0], images_2d[0], cv::ROTATE_90_CLOCKWISE);
+    }
+    else if (this->rotate_deg == 180)
+    {
+        cv::rotate(images_2d[0], images_2d[0], cv::ROTATE_180);
+    }
+    else if (this->rotate_deg == -90)
+    {
+        cv::rotate(images_2d[0], images_2d[0], cv::ROTATE_90_COUNTERCLOCKWISE);
+    }
+
     // Predict 2D poses
     auto poses_2d_all = kpt_model->predict(images_2d);
     auto poses_2d_upd = utils_pipeline::update_keypoints(
@@ -217,6 +247,54 @@ std::vector<std::vector<std::array<float, 3>>> Rpt2DWrapperNode::call_model(cons
         }
     }
 
+    // Rotate poses back
+    float img_width = image.cols;
+    float img_height = image.rows;
+    if (this->rotate_deg == 0)
+    {
+        // pass
+    }
+    else if (this->rotate_deg == 90)
+    {
+        for (auto &person : valid_poses)
+        {
+            for (auto &kp : person)
+            {
+                float x = kp[0];
+                float y = kp[1];
+                kp[0] = y;
+                kp[1] = img_height - 1 - x;
+            }
+        }
+    }
+    else if (this->rotate_deg == 180)
+    {
+        for (auto &person : valid_poses)
+        {
+            for (auto &kp : person)
+            {
+                float x = kp[0];
+                float y = kp[1];
+                kp[0] = img_width - 1 - x;
+                kp[1] = img_height - 1 - y;
+            }
+        }
+    }
+    else if (this->rotate_deg == -90)
+    {
+        for (auto &person : valid_poses)
+        {
+            for (auto &kp : person)
+            {
+                float x = kp[0];
+                float y = kp[1];
+                kp[0] = img_width - 1 - y;
+                kp[1] = x;
+
+            }
+        }
+    }
+
     return valid_poses;
 }
 
@@ -227,7 +305,10 @@ int main(int argc, char **argv)
 {
     rclcpp::init(argc, argv);
 
-    auto node = std::make_shared<Rpt2DWrapperNode>("rpt2d_wrapper");
+    const std::string cam_id = std::getenv("CAMID");
+    const std::string rotate = std::getenv("ROTATE") ? std::getenv("ROTATE") : "0";
+
+    auto node = std::make_shared<Rpt2DWrapperNode>(cam_id, rotate);
     rclcpp::spin(node);
 
     rclcpp::shutdown();