Updated image callback to directly calculate poses.

extras/ros/rpt2D_wrapper_cpp/src/rpt2D_wrapper.cpp

@@ -3,9 +3,7 @@
 #include <cmath>
 #include <iostream>
 #include <memory>
-#include <mutex>
 #include <string>
-#include <thread>
 #include <vector>
 
 // ROS2
@@ -49,9 +47,7 @@ public:
     Rpt2DWrapperNode(const std::string &node_name)
         : Node(node_name)
     {
-        this->stop_flag = false;
+        this->is_busy = false;
-        this->last_input_image = cv::Mat();
-        this->last_input_time = 0.0;
 
         // QoS
         rclcpp::QoS qos_profile(1);
@@ -61,7 +57,7 @@ public:
         // Setup subscriber
         image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
             img_input_topic, qos_profile,
-            std::bind(&Rpt2DWrapperNode::callbackImages, this, std::placeholders::_1));
+            std::bind(&Rpt2DWrapperNode::callback_images, this, std::placeholders::_1));
 
         // Setup publisher
         pose_pub_ = this->create_publisher<std_msgs::msg::String>(pose_out_topic, qos_profile);
@@ -72,48 +68,20 @@ public:
             use_wb, min_bbox_score, min_bbox_area, batch_poses);
 
         RCLCPP_INFO(this->get_logger(), "Finished initialization of pose estimator.");
-
-        // Start background prediction thread
-        model_thread = std::thread(&Rpt2DWrapperNode::callbackWrapper, this);
-    }
-
-    ~Rpt2DWrapperNode()
-    {
-        stop_flag = true;
-        if (model_thread.joinable())
-        {
-            model_thread.join();
-        }
     }
 
 private:
     rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub_;
     rclcpp::Publisher<std_msgs::msg::String>::SharedPtr pose_pub_;
+    std::atomic<bool> is_busy;
 
     // Pose model pointer
     std::unique_ptr<utils_2d_pose::PosePredictor> kpt_model;
     const std::vector<std::string> joint_names_2d = utils_pipeline::get_joint_names(whole_body);
 
-    // Threading
+    void callback_images(const sensor_msgs::msg::Image::SharedPtr msg);
-    std::thread model_thread;
-    std::mutex mutex;
-    std::atomic<bool> stop_flag;
 
-    cv::Mat last_input_image;
+    std::vector<std::vector<std::array<float, 3>>> call_model(const cv::Mat &image);
-    double last_input_time;
-
-    void callbackImages(const sensor_msgs::msg::Image::SharedPtr msg);
-    void callbackModel();
-
-    void callbackWrapper()
-    {
-        using namespace std::chrono_literals;
-        while (!stop_flag)
-        {
-            callbackModel();
-            std::this_thread::sleep_for(std::chrono::microseconds(100));
-        }
-    }
 
     void publish(const json &data)
     {
@@ -125,8 +93,16 @@ private:
 
 // =================================================================================================
 
-void Rpt2DWrapperNode::callbackImages(const sensor_msgs::msg::Image::SharedPtr msg)
+void Rpt2DWrapperNode::callback_images(const sensor_msgs::msg::Image::SharedPtr msg)
 {
+    if (this->is_busy)
+    {
+        RCLCPP_WARN(this->get_logger(), "Skipping frame, still processing...");
+        return;
+    }
+    this->is_busy = true;
+    auto ts_image = std::chrono::high_resolution_clock::now();
+
     // Load or convert image to Bayer format
     cv::Mat bayer_image;
     try
@@ -158,43 +134,51 @@ void Rpt2DWrapperNode::callbackImages(const sensor_msgs::msg::Image::SharedPtr m
         return;
     }
 
-    // Get timestamp
+    // Call model
-    double time_stamp = msg->header.stamp.sec + msg->header.stamp.nanosec / 1.0e9;
+    const auto &valid_poses = this->call_model(bayer_image);
 
-    // Store in member variables with lock
+    // Calculate timings
-    {
+    double time_stamp = msg->header.stamp.sec + msg->header.stamp.nanosec / 1.0e9;
-        std::lock_guard<std::mutex> lk(mutex);
+    auto ts_image_sec = std::chrono::duration<double>(ts_image.time_since_epoch()).count();
-        this->last_input_image = std::move(bayer_image);
+    auto ts_pose = std::chrono::high_resolution_clock::now();
-        this->last_input_time = time_stamp;
+    double ts_pose_sec = std::chrono::duration<double>(ts_pose.time_since_epoch()).count();
-    }
+    double z_trigger_image = ts_image_sec - time_stamp;
+    double z_trigger_pose = ts_pose_sec - time_stamp;
+    double z_image_pose = ts_pose_sec - ts_image_sec;
+
+    // Build JSON
+    json poses_msg;
+    poses_msg["bodies2D"] = valid_poses; // shape: persons x keypoints x 3
+    poses_msg["joints"] = joint_names_2d;
+    poses_msg["num_persons"] = valid_poses.size();
+    poses_msg["timestamps"] = {
+        {"trigger", time_stamp},
+        {"image", ts_image_sec},
+        {"pose", ts_pose_sec},
+        {"z-trigger-image", z_trigger_image},
+        {"z-image-pose", z_image_pose},
+        {"z-trigger-pose", z_trigger_pose}};
+
+    // Publish
+    publish(poses_msg);
+
+    // Print info
+    double elapsed_time = std::chrono::duration<double>(
+                              std::chrono::high_resolution_clock::now() - ts_image)
+                              .count();
+    std::cout << "Detected persons: " << valid_poses.size()
+              << " - Prediction time: " << elapsed_time << "s" << std::endl;
+
+    this->is_busy = false;
 }
 
 // =================================================================================================
 
-void Rpt2DWrapperNode::callbackModel()
+std::vector<std::vector<std::array<float, 3>>> Rpt2DWrapperNode::call_model(const cv::Mat &image)
 {
-    auto ptime = std::chrono::high_resolution_clock::now();
-
-    // Check if we have an image
-    cv::Mat local_image;
-    double local_timestamp = 0.0;
-    {
-        std::lock_guard<std::mutex> lk(mutex);
-        if (last_input_time == 0.0)
-        {
-            return;
-        }
-        local_image = std::move(last_input_image);
-        local_timestamp = last_input_time;
-
-        last_input_image = cv::Mat();
-        last_input_time = 0.0;
-    }
-
     // Create image vector
-    cv::Mat rgb_image = utils_pipeline::bayer2rgb(local_image);
+    cv::Mat rgb_image = utils_pipeline::bayer2rgb(image);
-    std::vector<cv::Mat> images_2d;
+    std::vector<cv::Mat> images_2d = {rgb_image};
-    images_2d.push_back(rgb_image);
 
     // Predict 2D poses
     auto poses_2d_all = kpt_model->predict(images_2d);
@@ -228,29 +212,7 @@ void Rpt2DWrapperNode::callbackModel()
         }
     }
 
-    // Build JSON
+    return valid_poses;
-    auto ts_pose = std::chrono::high_resolution_clock::now();
-    double ts_pose_sec = std::chrono::duration<double>(ts_pose.time_since_epoch()).count();
-    double z_images_pose = ts_pose_sec - local_timestamp;
-
-    json poses_msg;
-    poses_msg["bodies2D"] = valid_poses; // shape: persons x keypoints x 3
-    poses_msg["joints"] = joint_names_2d;
-    poses_msg["num_persons"] = valid_poses.size();
-    poses_msg["timestamps"] = {
-        {"image", local_timestamp},
-        {"pose", ts_pose_sec},
-        {"z-images-pose", z_images_pose}};
-
-    // Publish
-    publish(poses_msg);
-
-    // Print info
-    double elapsed_time = std::chrono::duration<double>(
-                              std::chrono::high_resolution_clock::now() - ptime)
-                              .count();
-    std::cout << "Detected persons: " << valid_poses.size()
-              << " - Prediction time: " << elapsed_time << "s" << std::endl;
 }
 
 // =================================================================================================