269 lines
8.0 KiB
C++
269 lines
8.0 KiB
C++
#include <atomic>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
// ROS2
|
|
#include <rclcpp/rclcpp.hpp>
|
|
#include <sensor_msgs/msg/image.hpp>
|
|
#include <std_msgs/msg/string.hpp>
|
|
|
|
// OpenCV / cv_bridge
|
|
#include <cv_bridge/cv_bridge.h>
|
|
#include <opencv2/opencv.hpp>
|
|
|
|
// JSON library
|
|
#include "/RapidPoseTriangulation/extras/include/nlohmann/json.hpp"
|
|
using json = nlohmann::json;
|
|
|
|
#include "/RapidPoseTriangulation/scripts/utils_2d_pose.hpp"
|
|
#include "/RapidPoseTriangulation/scripts/utils_pipeline.hpp"
|
|
|
|
// =================================================================================================
|
|
|
|
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 float min_bbox_score = 0.4;
|
|
static const float min_bbox_area = 0.1 * 0.1;
|
|
static const bool batch_poses = true;
|
|
|
|
static const std::map<std::string, bool> whole_body = {
|
|
{"foots", true},
|
|
{"face", true},
|
|
{"hands", true},
|
|
};
|
|
|
|
// =================================================================================================
|
|
// =================================================================================================
|
|
|
|
class Rpt2DWrapperNode : public rclcpp::Node
|
|
{
|
|
public:
|
|
Rpt2DWrapperNode(const std::string &node_name)
|
|
: Node(node_name)
|
|
{
|
|
this->stop_flag = false;
|
|
this->last_input_image = cv::Mat();
|
|
this->last_input_time = 0.0;
|
|
|
|
// QoS
|
|
rclcpp::QoS qos_profile(1);
|
|
qos_profile.reliable();
|
|
qos_profile.keep_last(1);
|
|
|
|
// Setup subscriber
|
|
image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
|
|
img_input_topic, qos_profile,
|
|
std::bind(&Rpt2DWrapperNode::callbackImages, this, std::placeholders::_1));
|
|
|
|
// Setup publisher
|
|
pose_pub_ = this->create_publisher<std_msgs::msg::String>(pose_out_topic, qos_profile);
|
|
|
|
// Load model
|
|
bool use_wb = utils_pipeline::use_whole_body(whole_body);
|
|
this->kpt_model = std::make_unique<utils_2d_pose::PosePredictor>(
|
|
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_;
|
|
|
|
// 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
|
|
std::thread model_thread;
|
|
std::mutex mutex;
|
|
std::atomic<bool> stop_flag;
|
|
|
|
cv::Mat last_input_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)
|
|
{
|
|
std_msgs::msg::String msg;
|
|
msg.data = data.dump();
|
|
pose_pub_->publish(msg);
|
|
}
|
|
};
|
|
|
|
// =================================================================================================
|
|
|
|
void Rpt2DWrapperNode::callbackImages(const sensor_msgs::msg::Image::SharedPtr msg)
|
|
{
|
|
// Load or convert image to Bayer format
|
|
cv::Mat bayer_image;
|
|
try
|
|
{
|
|
if (msg->encoding == "mono8")
|
|
{
|
|
cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, msg->encoding);
|
|
bayer_image = cv_ptr->image;
|
|
}
|
|
else if (msg->encoding == "bayer_rggb8")
|
|
{
|
|
cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, msg->encoding);
|
|
bayer_image = cv_ptr->image;
|
|
}
|
|
else if (msg->encoding == "rgb8")
|
|
{
|
|
cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, "rgb8");
|
|
cv::Mat color_image = cv_ptr->image;
|
|
bayer_image = utils_pipeline::rgb2bayer(color_image);
|
|
}
|
|
else
|
|
{
|
|
throw std::runtime_error("Unknown image encoding: " + msg->encoding);
|
|
}
|
|
}
|
|
catch (const std::exception &e)
|
|
{
|
|
RCLCPP_ERROR(this->get_logger(), "cv_bridge exception: %s", e.what());
|
|
return;
|
|
}
|
|
|
|
// Get timestamp
|
|
double time_stamp = msg->header.stamp.sec + msg->header.stamp.nanosec / 1.0e9;
|
|
|
|
// Store in member variables with lock
|
|
{
|
|
std::lock_guard<std::mutex> lk(mutex);
|
|
this->last_input_image = std::move(bayer_image);
|
|
this->last_input_time = time_stamp;
|
|
}
|
|
}
|
|
|
|
// =================================================================================================
|
|
|
|
void Rpt2DWrapperNode::callbackModel()
|
|
{
|
|
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);
|
|
std::vector<cv::Mat> images_2d;
|
|
images_2d.push_back(rgb_image);
|
|
|
|
// Predict 2D poses
|
|
auto poses_2d_all = kpt_model->predict(images_2d);
|
|
auto poses_2d_upd = utils_pipeline::update_keypoints(
|
|
poses_2d_all, joint_names_2d, whole_body);
|
|
auto &poses_2d = poses_2d_upd[0];
|
|
|
|
// Drop persons with no joints
|
|
std::vector<std::vector<std::array<float, 3>>> valid_poses;
|
|
for (auto &person : poses_2d)
|
|
{
|
|
float sum_conf = 0.0;
|
|
for (auto &kp : person)
|
|
{
|
|
sum_conf += kp[2];
|
|
}
|
|
if (sum_conf > 0.0)
|
|
{
|
|
valid_poses.push_back(person);
|
|
}
|
|
}
|
|
|
|
// Round poses to 3 decimal places
|
|
for (auto &person : valid_poses)
|
|
{
|
|
for (auto &kp : person)
|
|
{
|
|
kp[0] = std::round(kp[0] * 1000.0) / 1000.0;
|
|
kp[1] = std::round(kp[1] * 1000.0) / 1000.0;
|
|
kp[2] = std::round(kp[2] * 1000.0) / 1000.0;
|
|
}
|
|
}
|
|
|
|
// Build JSON
|
|
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;
|
|
}
|
|
|
|
// =================================================================================================
|
|
// =================================================================================================
|
|
|
|
int main(int argc, char **argv)
|
|
{
|
|
rclcpp::init(argc, argv);
|
|
|
|
auto node = std::make_shared<Rpt2DWrapperNode>("rpt2D_wrapper");
|
|
rclcpp::spin(node);
|
|
|
|
rclcpp::shutdown();
|
|
return 0;
|
|
}
|