hongshaorou/YOPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial Commit (tested training, testing, and TRT conversion)

Browse Source
This commit is contained in:
Lu Junjie
2024-10-20 17:01:07 +08:00
parent 86d2f311f8
commit 5738088bae
221 changed files with 59249 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

293
flightlib/src/ros_nodes/flight_pilot.cpp Normal file
View File

@@ -0,0 +1,293 @@
#include "flightlib/ros_nodes/flight_pilot.hpp"
namespace flightros {
FlightPilot::FlightPilot(const ros::NodeHandle& nh, const ros::NodeHandle& pnh)
: nh_(nh), pnh_(pnh), scene_id_(4), unity_ready_(false), unity_render_(false), receive_id_(0), frameID(0), main_loop_freq_(50.0) {
// quad initialization
quad_ptr_ = std::make_shared<Quadrotor>();
// load parameters
std::string cfg_path = getenv("FLIGHTMARE_PATH") + std::string("/flightlib/configs/quadrotor_ros.yaml");
YAML::Node cfg_ = YAML::LoadFile(cfg_path);
loadParams(cfg_);
configCamera(cfg_);
quad_ptr_->setSize(quad_size_);
// initialization
quad_state_.setZero();
if (scene_id_ == UnityScene::NATUREFOREST) {
quad_state_.x(QS::POSX) = 51; // 41-61
quad_state_.x(QS::POSY) = 108; // 98-118
quad_state_.x(QS::POSZ) = 34;
}
quad_ptr_->reset(quad_state_);
sgm_.reset(new sgm_gpu::SgmGpu(width_, height_));
// initialize subscriber and publisher
left_img_pub = nh_.advertise<sensor_msgs::Image>("RGB_image", 1);
stereo_pub = nh_.advertise<sensor_msgs::Image>("stereo_image", 1);
depth_pub = nh_.advertise<sensor_msgs::Image>("depth_image", 1);
cam_info_pub = nh_.advertise<sensor_msgs::CameraInfo>("camera_info", 1);
state_est_sub_ = nh_.subscribe(odom_topic_, 1, &FlightPilot::poseCallback, this);
spawn_tree_sub_ = nh_.subscribe("/spawn_tree", 1, &FlightPilot::spawnTreeCallback, this);
clear_tree_sub_ = nh_.subscribe("/clear_tree", 1, &FlightPilot::clearTreeCallback, this);
save_pc_sub_ = nh_.subscribe("/save_pc", 1, &FlightPilot::savePointcloudCallback, this);
timestamp = ros::Time::now();
// connect unity and setup unity
setUnity(unity_render_);
connectUnity();
if (!unity_ready_)
ROS_ERROR("Start the gazebo and unity first!");
spawnTreesAndSavePointCloud();
timer_main_loop_ = nh_.createTimer(ros::Rate(main_loop_freq_), &FlightPilot::mainLoopCallback, this);
}
FlightPilot::~FlightPilot() { disconnectUnity(); }
void FlightPilot::spawnTreeCallback(const std_msgs::Empty::ConstPtr& msg) {
if (!unity_ready_ || unity_bridge_ptr_ == nullptr)
return;
unity_bridge_ptr_->spawnTrees(bounding_box_, bounding_box_origin_, avg_tree_spacing_);
}
void FlightPilot::clearTreeCallback(const std_msgs::Empty::ConstPtr& msg) { unity_bridge_ptr_->rmTrees(); }
// If the point cloud is not saved when init, it also can be saved by this callback.
void FlightPilot::savePointcloudCallback(const std_msgs::Empty::ConstPtr& msg) {
Vector<3> min_corner = bounding_box_origin_ - 0.5 * bounding_box_;
Vector<3> max_corner = bounding_box_origin_ + 0.5 * bounding_box_;
unity_bridge_ptr_->generatePointcloud(min_corner, max_corner, ply_id_, ply_path_, scene_id_, pointcloud_resolution_);
}
void FlightPilot::poseCallback(const nav_msgs::Odometry::ConstPtr& msg) {
quad_state_.x[QS::POSX] = (Scalar)msg->pose.pose.position.x;
quad_state_.x[QS::POSY] = (Scalar)msg->pose.pose.position.y;
quad_state_.x[QS::POSZ] = (Scalar)msg->pose.pose.position.z;
quad_state_.x[QS::ATTW] = (Scalar)msg->pose.pose.orientation.w;
quad_state_.x[QS::ATTX] = (Scalar)msg->pose.pose.orientation.x;
quad_state_.x[QS::ATTY] = (Scalar)msg->pose.pose.orientation.y;
quad_state_.x[QS::ATTZ] = (Scalar)msg->pose.pose.orientation.z;
quad_ptr_->setState(quad_state_);
timestamp = msg->header.stamp;
}
void FlightPilot::mainLoopCallback(const ros::TimerEvent& event) {
if (!unity_render_ || !unity_ready_)
return;
frameID++;
ros::Time timestamp_ = timestamp;
unity_bridge_ptr_->getRender(frameID); // 1ms
FrameID frameID_rt;
unity_bridge_ptr_->handleOutput(frameID_rt); // 30ms
while (frameID != frameID_rt)
unity_bridge_ptr_->handleOutput(frameID_rt);
cv::Mat left_img, right_img, depth_img;
// publish RGB image
rgb_camera_left->getRGBImage(left_img);
sensor_msgs::ImagePtr img_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", left_img).toImageMsg();
img_msg->header.stamp = timestamp_;
left_img_pub.publish(img_msg);
// publish camera info
int width = rgb_camera_left->getWidth();
int hight = rgb_camera_left->getHeight();
float fov = rgb_camera_left->getFOV();
float cx = width / 2.0;
float cy = hight / 2.0;
float fx = cy / tan(0.5 * fov * M_PI / 180.0); // 他这个特殊，fov是竖直方向的
float fy = fx;
boost::array<double, 9> K = {fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0};
sensor_msgs::CameraInfo cam_msg;
cam_msg.height = hight;
cam_msg.width = width;
cam_msg.K = K;
cam_info_pub.publish(cam_msg);
// publish depth image
if (use_depth) {
rgb_camera_left->getDepthMap(depth_img);
depth_img = depth_img * 1000.0;
depth_img = cv::min(depth_img, 20.0);
sensor_msgs::ImagePtr depth_msg = cv_bridge::CvImage(std_msgs::Header(), "32FC1", depth_img).toImageMsg();
depth_msg->header.stamp = timestamp_;
depth_pub.publish(depth_msg);
}
// publish stereo image
if (use_stereo) {
rgb_camera_right->getRGBImage(right_img);
cv::Mat stereo_(height_, width_, CV_32FC1);
computeDepthImage(left_img, right_img, &stereo_);
if (use_depth) {
// Complete the NaN values in the depth map, as the RealSense performs better than SGM.
cv::Mat mask, mask1, mask2;
cv::compare(stereo_, 0, mask1, cv::CMP_EQ); // 将 A 中为 0 的位置置为 255，其余位置置为 0
cv::compare(stereo_, 20, mask2, cv::CMP_GT); // 将 A 中大于 20 的位置置为 255，其余位置置为 0
mask = mask1 | mask2; // 将两个掩码进行逻辑或操作
depth_img.copyTo(stereo_, mask); // 将 B 中 mask 为 255 的位置的值复制到 A 中
}
sensor_msgs::ImagePtr stereo_msg = cv_bridge::CvImage(std_msgs::Header(), "32FC1", stereo_).toImageMsg();
stereo_msg->header.stamp = timestamp_;
stereo_pub.publish(stereo_msg);
}
}
bool FlightPilot::setUnity(const bool render) {
unity_render_ = render;
if (unity_render_ && unity_bridge_ptr_ == nullptr) {
// create unity bridge
unity_bridge_ptr_ = UnityBridge::getInstance();
unity_bridge_ptr_->addQuadrotor(quad_ptr_);
ROS_INFO("[%s] Unity Bridge is created.", pnh_.getNamespace().c_str());
}
return true;
}
bool FlightPilot::spawnTreesAndSavePointCloud() {
if (!unity_ready_ || unity_bridge_ptr_ == nullptr)
return false;
unity_bridge_ptr_->spawnTrees(bounding_box_, bounding_box_origin_, avg_tree_spacing_);
// Saving point cloud during the testing is much time-consuming, but can be used for evaluation.
// The following can be commented if evaluation is not needed.
Vector<3> min_corner = bounding_box_origin_ - 0.5 * bounding_box_;
Vector<3> max_corner = bounding_box_origin_ + 0.5 * bounding_box_;
unity_bridge_ptr_->generatePointcloud(min_corner, max_corner, ply_id_, ply_path_, scene_id_, pointcloud_resolution_);
return true;
}
bool FlightPilot::connectUnity() {
if (!unity_render_ || unity_bridge_ptr_ == nullptr)
return false;
unity_ready_ = unity_bridge_ptr_->connectUnity(scene_id_);
return unity_ready_;
}
bool FlightPilot::disconnectUnity() {
if (unity_render_ && unity_bridge_ptr_ != nullptr)
;
unity_bridge_ptr_->disconnectUnity();
unity_ready_ = false;
}
bool FlightPilot::loadParams(const YAML::Node& cfg) {
// ros
main_loop_freq_ = cfg["main_loop_freq"].as<int>();
odom_topic_ = cfg["odom_topic"].as<std::string>();
// camera
width_ = cfg["rgb_camera_left"]["width"].as<int>();
height_ = cfg["rgb_camera_left"]["height"].as<int>();
fov_ = cfg["rgb_camera_left"]["fov"].as<Scalar>();
use_depth = cfg["rgb_camera_left"]["enable_depth"].as<bool>();
use_stereo = cfg["rgb_camera_right"]["on"].as<bool>();
// scence
scene_id_ = cfg["scene_id"].as<int>();
unity_render_ = cfg["unity_render"].as<bool>();
Scalar quad_size_i = cfg["quad_size"].as<Scalar>();
quad_size_ = Vector<3>(quad_size_i, quad_size_i, quad_size_i);
avg_tree_spacing_ = cfg["unity"]["avg_tree_spacing"].as<Scalar>();
for (int i = 0; i < 3; ++i) {
bounding_box_(i) = cfg["unity"]["bounding_box"][i].as<Scalar>();
bounding_box_origin_(i) = cfg["unity"]["bounding_box_origin"][i].as<Scalar>();
}
pointcloud_resolution_ = cfg["unity"]["pointcloud_resolution"].as<Scalar>();
ply_path_ = getenv("FLIGHTMARE_PATH") + cfg["ply_path"].as<std::string>();
if (!std::filesystem::exists(ply_path_)) {
std::filesystem::create_directories(ply_path_);
std::cout << "Directory created: " << ply_path_ << std::endl;
}
return true;
}
void FlightPilot::computeDepthImage(const cv::Mat& left_frame, const cv::Mat& right_frame, cv::Mat* const depth) {
cv::Mat disparity(height_, width_, CV_8UC1);
sgm_->computeDisparity(left_frame, right_frame, disparity);
disparity.convertTo(disparity, CV_32FC1);
// compute depth from disparity
float f = (width_ / 2.0) / std::tan((M_PI * (fov_ / 180.0)) / 2.0);
// depth = stereo_baseline_ * f / disparity
for (int r = 0; r < height_; ++r) {
for (int c = 0; c < width_; ++c) {
if (disparity.at<float>(r, c) == 0.0f) {
depth->at<float>(r, c) = 0.0f;
} else if (disparity.at<float>(r, c) == 255.0f) {
depth->at<float>(r, c) = 0.0f;
} else {
depth->at<float>(r, c) = static_cast<float>(stereo_baseline_) * f / disparity.at<float>(r, c);
}
}
}
}
bool FlightPilot::configCamera(const YAML::Node& cfg) {
if (!cfg["rgb_camera_left"] || !cfg["rgb_camera_right"]) {
ROS_ERROR("Cannot config stereo Camera");
return false;
}
// create left camera --------------------------------------------
rgb_camera_left = std::make_shared<RGBCamera>();
// load camera settings
std::vector<Scalar> t_BC_vec = cfg["rgb_camera_left"]["t_BC"].as<std::vector<Scalar>>();
std::vector<Scalar> r_BC_vec = cfg["rgb_camera_left"]["r_BC"].as<std::vector<Scalar>>();
Vector<3> t_BC(t_BC_vec.data());
Matrix<3, 3> r_BC = (AngleAxis(r_BC_vec[0] * M_PI / 180.0, Vector<3>::UnitX()) * AngleAxis(r_BC_vec[1] * M_PI / 180.0, Vector<3>::UnitY()) *
AngleAxis(r_BC_vec[2] * M_PI / 180.0, Vector<3>::UnitZ()))
.toRotationMatrix(); // the rotation order has been verified
// Convert the horizontal FOV (usually used) to vertical FOV (flightmare).
Scalar rgb_fov_deg_ = cfg["rgb_camera_left"]["fov"].as<Scalar>();
double hor_fov_radians = (M_PI * (rgb_fov_deg_ / 180.0));
Scalar img_rows_ = cfg["rgb_camera_left"]["height"].as<Scalar>();
Scalar img_cols_ = cfg["rgb_camera_left"]["width"].as<Scalar>();
double flightmare_fov = 2. * std::atan(std::tan(hor_fov_radians / 2) * img_rows_ / img_cols_);
flightmare_fov = (flightmare_fov / M_PI) * 180.0;
rgb_camera_left->setFOV(flightmare_fov);
rgb_camera_left->setWidth(cfg["rgb_camera_left"]["width"].as<int>());
rgb_camera_left->setHeight(cfg["rgb_camera_left"]["height"].as<int>());
rgb_camera_left->setRelPose(t_BC, r_BC);
rgb_camera_left->enableOpticalFlow(cfg["rgb_camera_left"]["enable_opticalflow"].as<bool>());
rgb_camera_left->enableSegmentation(cfg["rgb_camera_left"]["enable_segmentation"].as<bool>());
rgb_camera_left->enableDepth(cfg["rgb_camera_left"]["enable_depth"].as<bool>());
// add camera to the quadrotor
quad_ptr_->addRGBCamera(rgb_camera_left);
// create right camera --------------------------------------------
if (use_stereo) {
rgb_camera_right = std::make_shared<RGBCamera>();
// load camera settings
std::vector<Scalar> t_BC_vec_r = cfg["rgb_camera_right"]["t_BC"].as<std::vector<Scalar>>();
std::vector<Scalar> r_BC_vec_r = cfg["rgb_camera_right"]["r_BC"].as<std::vector<Scalar>>();
Vector<3> t_BC_r(t_BC_vec_r.data());
Matrix<3, 3> r_BC_r = (AngleAxis(r_BC_vec[0] * M_PI / 180.0, Vector<3>::UnitX()) * AngleAxis(r_BC_vec[1] * M_PI / 180.0, Vector<3>::UnitY()) *
AngleAxis(r_BC_vec[2] * M_PI / 180.0, Vector<3>::UnitZ()))
.toRotationMatrix();
rgb_camera_right->setFOV(flightmare_fov);
rgb_camera_right->setWidth(cfg["rgb_camera_left"]["width"].as<int>());
rgb_camera_right->setHeight(cfg["rgb_camera_left"]["height"].as<int>());
rgb_camera_right->setRelPose(t_BC_r, r_BC_r);
rgb_camera_right->enableOpticalFlow(false);
rgb_camera_right->enableSegmentation(false);
rgb_camera_right->enableDepth(false);
// add camera to the quadrotor
quad_ptr_->addRGBCamera(rgb_camera_right);
stereo_baseline_ = fabs(t_BC(1) - t_BC_r(1));
}
return true;
}
} // namespace flightros

13
flightlib/src/ros_nodes/flight_pilot_node.cpp Normal file
View File

@@ -0,0 +1,13 @@
#include <ros/ros.h>
#include "flightlib/ros_nodes/flight_pilot.hpp"
int main(int argc, char** argv) {
ros::init(argc, argv, "flight_pilot");
flightros::FlightPilot pilot(ros::NodeHandle(), ros::NodeHandle("~"));
// spin the ros
ros::spin();
return 0;
}

101
flightlib/src/ros_nodes/map_visual_node.cpp Normal file
View File

@@ -0,0 +1,101 @@
// Publish the surrounding point cloud map based on the drone's position for visualization.
#include <nav_msgs/Odometry.h>
#include <pcl/common/common.h>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl_ros/point_cloud.h>
#include <ros/ros.h>
#include <tf/transform_broadcaster.h>
#include <time.h>
#include <yaml-cpp/yaml.h>
#include <Eigen/Core>
#include "visualization_msgs/Marker.h"
namespace map_visual {
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
void pcl_input() {
std::string cfg_path = getenv("FLIGHTMARE_PATH") + std::string("/flightlib/configs/quadrotor_ros.yaml");
YAML::Node cfg_ = YAML::LoadFile(cfg_path);
std::string ply_path_ = getenv("FLIGHTMARE_PATH") + cfg_["ply_path"].as<std::string>() + "pointcloud-0.ply";
pcl::io::loadPLYFile<pcl::PointXYZ>(ply_path_, *cloud);
std::cout << "size of pointcloud: " << cloud->points.size() << std::endl;
}
void odom_cb(const nav_msgs::Odometry::ConstPtr odom_msg, ros::Publisher* local_map_pub, ros::Publisher* mesh_pub,
tf::TransformBroadcaster* uav_tf_br) {
// 1. publish tf
tf::Transform transform;
transform.setOrigin(tf::Vector3(odom_msg->pose.pose.position.x, odom_msg->pose.pose.position.y, odom_msg->pose.pose.position.z));
tf::Quaternion q(0, 0, 0, 1);
// tf::Quaternion q(odom_msg->pose.pose.orientation.x, odom_msg->pose.pose.orientation.y,
// odom_msg->pose.pose.orientation.z, odom_msg->pose.pose.orientation.w);
transform.setRotation(q);
uav_tf_br->sendTransform(tf::StampedTransform(transform, ros::Time::now(), "world", "uav"));
// 2. publish map
Eigen::Vector3f pose_cur(odom_msg->pose.pose.position.x, odom_msg->pose.pose.position.y, odom_msg->pose.pose.position.z);
Eigen::Vector3f local_map_half_size(8, 8, 1);
pcl::PointCloud<pcl::PointXYZ>::Ptr local_map(new pcl::PointCloud<pcl::PointXYZ>);
for (auto& x : cloud->points) {
if (x.z > pose_cur(2) - 6 && x.z < pose_cur(2) + local_map_half_size(2) && x.x > pose_cur(0) - local_map_half_size(0) &&
x.x < pose_cur(0) + local_map_half_size(0) && x.y > pose_cur(1) - local_map_half_size(1) && x.y < pose_cur(1) + local_map_half_size(1)) {
local_map->points.push_back(x);
}
}
pcl_conversions::toPCL(ros::Time::now(), local_map->header.stamp);
local_map->header.frame_id = "world";
local_map_pub->publish(local_map);
// 3. publish UAV model
std::string mesh_resource = std::string("file://") + getenv("FLIGHTMARE_PATH") + std::string("/flightlib/src/ros_nodes/model/uav.dae");
visualization_msgs::Marker meshROS;
meshROS.header.frame_id = "world";
meshROS.header.stamp = ros::Time::now();
meshROS.ns = "mesh";
meshROS.id = 0;
meshROS.type = visualization_msgs::Marker::MESH_RESOURCE;
meshROS.action = visualization_msgs::Marker::ADD;
meshROS.pose.position.x = odom_msg->pose.pose.position.x - 0.2;
meshROS.pose.position.y = odom_msg->pose.pose.position.y;
meshROS.pose.position.z = odom_msg->pose.pose.position.z;
meshROS.pose.orientation.w = odom_msg->pose.pose.orientation.w;
meshROS.pose.orientation.x = odom_msg->pose.pose.orientation.x;
meshROS.pose.orientation.y = odom_msg->pose.pose.orientation.y;
meshROS.pose.orientation.z = odom_msg->pose.pose.orientation.z;
float scale = 2;
meshROS.scale.x = scale;
meshROS.scale.y = scale;
meshROS.scale.z = scale;
meshROS.color.a = 1;
meshROS.color.r = 1;
meshROS.color.g = 1;
meshROS.color.b = 1;
meshROS.mesh_resource = mesh_resource;
meshROS.mesh_use_embedded_materials = true;
mesh_pub->publish(meshROS);
}
} // namespace map_visual
using namespace map_visual;
int main(int argc, char** argv) {
pcl_input();
ros::init(argc, argv, "map_visual");
ros::NodeHandle nh;
tf::TransformBroadcaster uav_tf_br;
ros::Publisher local_map_pub = nh.advertise<pcl::PointCloud<pcl::PointXYZ>>("/local_map", 1);
ros::Publisher mesh_pub = nh.advertise<visualization_msgs::Marker>("/uav_mesh", 1);
ros::Subscriber odom_sub =
nh.subscribe<nav_msgs::Odometry>("/juliett/ground_truth/odom", 1, boost::bind(&odom_cb, _1, &local_map_pub, &mesh_pub, &uav_tf_br));
std::cout << "Map visual node OK!" << std::endl;
ros::spin();
}

10597
flightlib/src/ros_nodes/model/uav.dae Executable file
View File

File diff suppressed because one or more lines are too long

145
flightlib/src/ros_nodes/traj_eval_node.cpp Normal file
View File

@@ -0,0 +1,145 @@
// A very rough evaluation: Input the map point cloud, receive odometry, and calculate: execution time,
// trajectory length, distance to the nearest obstacle, and dynamic cost (logging method to be optimized).
#include <nav_msgs/Odometry.h>
#include <pcl/common/common.h>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl/search/kdtree.h>
#include <ros/ros.h>
#include <time.h>
#include <Eigen/Core>
#include "flightlib/controller/ctrl_ref.h"
namespace traj_eval {
float start_record = -39;
float finish_record = 18;
// eval
std::ofstream log_, log_x, ctrl_log;
Eigen::Vector3f pose_last;
Eigen::Vector3f acc_last;
ros::Time start, end;
bool odom_init = false;
bool odom_finish = false;
bool first_cmd = true;
float length_ = 0;
float dist_ = 0;
float ctrl_cost_ = 0;
int num_ = 0;
// map
pcl::search::KdTree<pcl::PointXYZ> kdtree;
float resolution = 0.2;
Eigen::Vector3i m_size;
Eigen::Vector3i m_origin;
void map_input() {
std::string ply_path_ = getenv("FLIGHTMARE_PATH") + std::string("/flightrender/RPG_Flightmare/pointcloud_data/pointcloud-0.ply");
pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
pcl::io::loadPLYFile<pcl::PointXYZ>(ply_path_, *cloud);
pcl::PointXYZ min, max;
pcl::getMinMax3D(*cloud, min, max);
m_size(0) = ceil((max.x - min.x) / resolution);
m_size(1) = ceil((max.y - min.y) / resolution);
m_size(2) = ceil((max.z - min.z + 0.01) / resolution);
m_origin = Eigen::Vector3i(min.x, min.y, min.z);
kdtree.setInputCloud(cloud);
}
int to_id(int x, int y, int z) { return x * m_size(1) * m_size(2) + y * m_size(2) + z; }
float distance_at(Eigen::Vector3f pose) {
pcl::PointXYZ drone_;
drone_.x = pose(0);
drone_.y = pose(1);
drone_.z = pose(2);
int K = 1;
std::vector<int> indices(K);
std::vector<float> distances(K); // 存储近邻点对应距离的平方
kdtree.nearestKSearch(drone_, K, indices, distances);
return std::sqrt(distances[0]);
}
void odom_cb(const nav_msgs::Odometry::Ptr odom_msg) {
if (!odom_init && odom_msg->pose.pose.position.x > start_record) {
odom_init = true;
pose_last = Eigen::Vector3f(odom_msg->pose.pose.position.x, odom_msg->pose.pose.position.y, odom_msg->pose.pose.position.z);
start = ros::Time::now();
std::cout << "start!" << std::endl;
return;
}
if (!odom_init || odom_finish)
return;
Eigen::Vector3f pose_cur(odom_msg->pose.pose.position.x, odom_msg->pose.pose.position.y, odom_msg->pose.pose.position.z);
length_ += (pose_cur - pose_last).norm();
dist_ += distance_at(pose_cur);
log_ << distance_at(pose_cur) << std::endl;
log_x << pose_cur(0) << std::endl;
pose_last = pose_cur;
num_++;
if (odom_msg->pose.pose.position.x > finish_record) {
odom_finish = true;
end = ros::Time::now();
std::cout << "finish! \n time:" << (end - start).toSec() << " s,\nlength:" << length_ << " m,\ndist:" << dist_ / num_
<< " m,\nctrl cost:" << ctrl_cost_ << " m2/s5" << std::endl;
log_.close();
log_x.close();
ctrl_log.close();
}
}
void ctrl_cb(const quad_pos_ctrl::ctrl_ref& ctrl_msg) {
if (!odom_init || odom_finish)
return;
if (first_cmd) {
first_cmd = false;
acc_last = Eigen::Vector3f(ctrl_msg.acc_ref[0], -ctrl_msg.acc_ref[1], -ctrl_msg.acc_ref[2]);
return;
}
Eigen::Vector3f cur_acc(ctrl_msg.acc_ref[0], -ctrl_msg.acc_ref[1], -ctrl_msg.acc_ref[2]);
Eigen::Vector3f d_acc = (acc_last - cur_acc) / 0.02;
float acc_norm2 = d_acc.dot(d_acc);
ctrl_cost_ += 0.02 * acc_norm2;
acc_last = cur_acc;
ctrl_log << ctrl_msg.pos_ref[0] << ',';
ctrl_log << ctrl_msg.pos_ref[1] << ',';
ctrl_log << ctrl_msg.pos_ref[2] << ',';
ctrl_log << ctrl_msg.vel_ref[0] << ',';
ctrl_log << ctrl_msg.vel_ref[1] << ',';
ctrl_log << ctrl_msg.vel_ref[2] << ',';
ctrl_log << ctrl_msg.acc_ref[0] << ',';
ctrl_log << ctrl_msg.acc_ref[1] << ',';
ctrl_log << ctrl_msg.acc_ref[2] << std::endl;
}
} // namespace traj_eval
using namespace traj_eval;
int main(int argc, char** argv) {
map_input();
// Move to the same log file.
std::string log_file = getenv("FLIGHTMARE_PATH") + std::string("/run/utils/dist.csv");
std::cout << "log path:" << log_file << std::endl;
log_.open(log_file.c_str(), std::ios::out);
std::string log_file2 = getenv("FLIGHTMARE_PATH") + std::string("/run/utils/dist_x.csv");
log_x.open(log_file2.c_str(), std::ios::out);
std::string log_file3 = getenv("FLIGHTMARE_PATH") + std::string("/run/utils/ctrl_log.csv");
ctrl_log.open(log_file3.c_str(), std::ios::out);
ros::init(argc, argv, "traj_eval");
ros::NodeHandle nh;
ros::Subscriber odom_sub = nh.subscribe("/juliett/ground_truth/odom", 1, odom_cb);
ros::Subscriber ctrl_sub = nh.subscribe("/juliett_pos_ctrl_node/controller/ctrl_ref", 1, ctrl_cb);
ros::spin();
}

347
flightlib/src/ros_nodes/yopo_planner_node.cpp Normal file
View File

@@ -0,0 +1,347 @@
#include <nav_msgs/Odometry.h>
#include <pcl/point_types.h>
#include <pcl_ros/point_cloud.h>
#include <ros/ros.h>
#include <std_msgs/Float32MultiArray.h>
#include <yaml-cpp/yaml.h>
#include <Eigen/Core>
#include "flightlib/controller/PositionCommand.h"
#include "flightlib/controller/ctrl_ref.h"
#include "flightlib/grad_traj_optimization/opt_utile.h"
#include "flightlib/grad_traj_optimization/traj_optimization_bridge.h"
namespace yopo_net {
nav_msgs::Odometry odom_msg;
quad_pos_ctrl::ctrl_ref ctrl_ref_last;
quadrotor_msgs::PositionCommand pos_cmd_last;
bool odom_init = false;
bool odom_ref_init = false;
bool yopo_init = false;
bool done = false;
float traj_time = 2.0;
float sample_t = 0.0;
float last_yaw_ = 0; // NWU
TrajOptimizationBridge* traj_opt_bridge;
TrajOptimizationBridge* traj_opt_bridge_for_vis;
std::vector<std::pair<Eigen::Vector3d, Eigen::Vector3d>> lattice_nodes;
Eigen::Vector3d goal_(100, 0, 2);
Eigen::Quaterniond quat_(1, 0, 0, 0);
Eigen::Vector3d last_pos_(0, 0, 0), last_vel_(0, 0, 0), last_acc_(0, 0, 0);
ros::Publisher trajs_visual_pub, best_traj_visual_pub, state_ref_pub, ctrl_pub, mpc_ctrl_pub, so3_ctrl_pub, lattice_trajs_visual_pub;
ros::Subscriber odom_sub, odom_ref_sub, yopo_best_sub, yopo_all_sub, goal_sub;
void odom_cb(const nav_msgs::Odometry::Ptr msg) {
odom_msg = *msg;
odom_init = true;
quat_.w() = odom_msg.pose.pose.orientation.w;
quat_.x() = odom_msg.pose.pose.orientation.x;
quat_.y() = odom_msg.pose.pose.orientation.y;
quat_.z() = odom_msg.pose.pose.orientation.z;
if (!odom_ref_init) {
last_pos_ << odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z;
last_vel_ << 0, 0, 0;
last_acc_ << 0, 0, 0;
}
// check if reach the goal
Eigen::Vector3d dist(odom_msg.pose.pose.position.x - goal_(0), odom_msg.pose.pose.position.y - goal_(1),
odom_msg.pose.pose.position.z - goal_(2));
if (dist.norm() < 4) {
if (!done)
printf("Done!\n");
done = true;
}
}
void goal_cb(const std_msgs::Float32MultiArray::Ptr msg) {
Eigen::Vector3d last_goal = goal_;
goal_(0) = msg->data[0];
goal_(1) = msg->data[1];
goal_(2) = msg->data[2];
if (last_goal != goal_)
done = false;
}
void traj_to_pcl(TrajOptimizationBridge* traj_opt_bridge_input, pcl::PointCloud<pcl::PointXYZI>::Ptr cloud, double cost = 0.0) {
for (float dt = 0.0; dt < traj_time; dt = dt + 0.05) {
Eigen::Vector3d next_pos, next_vel, next_acc;
traj_opt_bridge_input->getNextState(next_pos, next_vel, next_acc, dt);
pcl::PointXYZI clrP;
clrP.x = next_pos(0);
clrP.y = next_pos(1);
clrP.z = next_pos(2);
clrP.intensity = cost;
cloud->points.push_back(clrP);
}
}
void yopo_cb(const std_msgs::Float32MultiArray::ConstPtr msg) {
if (!odom_init)
return;
std::vector<double> endstate;
for (int i = 0; i < msg->data.size(); i++) {
endstate.push_back(static_cast<double>(msg->data[i]));
}
traj_opt_bridge->setState(last_pos_.cast<double>(), quat_.cast<double>(), last_vel_.cast<double>(), last_acc_.cast<double>());
traj_opt_bridge->solveBVP(endstate);
// int milliseconds_yopo_pub = msg->layout.data_offset; // 预测时采用的状态和图像的时间戳(ms)
// int milliseconds_now = uint64_t(ros::Time::now().toSec() * 1000) % 1000000;
// double delta_t = double(milliseconds_now - milliseconds_yopo_pub) / 1000;
// std::cout<<"yopo开始预测到当前时间: "<<delta_t<<std::endl;
sample_t = 0.0; // = delta_t
pcl::PointCloud<pcl::PointXYZI>::Ptr best_traj_cld(new pcl::PointCloud<pcl::PointXYZI>);
traj_to_pcl(traj_opt_bridge, best_traj_cld);
pcl_conversions::toPCL(ros::Time::now(), best_traj_cld->header.stamp); // for test
best_traj_cld->header.frame_id = "world";
best_traj_visual_pub.publish(best_traj_cld);
yopo_init = true;
}
void trajs_vis_cb(const std_msgs::Float32MultiArray::ConstPtr msg) {
if (!odom_init)
return;
// ---------------- visualization of all trajs --------------------
std::vector<std::vector<double>> endstates_b;
endstates_b.resize(msg->layout.dim[0].size);
std::vector<double> scores;
for (int i = 0; i < msg->layout.dim[0].size; i++) {
for (int j = 0; j < msg->layout.dim[1].size - 1; j++) {
endstates_b[i].push_back(static_cast<double>(msg->data[i * msg->layout.dim[1].size + j]));
}
scores.push_back(static_cast<double>(msg->data[(i + 1) * msg->layout.dim[1].size - 1]));
}
traj_opt_bridge_for_vis->setState(last_pos_.cast<double>(), quat_.cast<double>(), last_vel_.cast<double>(), last_acc_.cast<double>());
pcl::PointCloud<pcl::PointXYZI>::Ptr trajs_cld(new pcl::PointCloud<pcl::PointXYZI>);
for (size_t i = 0; i < endstates_b.size(); i++) {
traj_opt_bridge_for_vis->solveBVP(endstates_b[i]);
traj_to_pcl(traj_opt_bridge_for_vis, trajs_cld, scores[i]);
}
pcl_conversions::toPCL(ros::Time::now(), trajs_cld->header.stamp);
trajs_cld->header.frame_id = "world";
trajs_visual_pub.publish(trajs_cld);
// ---------------- visualization of lattice ------------------------
pcl::PointCloud<pcl::PointXYZI>::Ptr lattice_trajs_cld(new pcl::PointCloud<pcl::PointXYZI>);
Eigen::Vector3d pos_1(0.0, 0.0, 0.0), vel_1(0.0, 0.0, 0.0), acc_1(0.0, 0.0, 0.0);
for (size_t i = 0; i < lattice_nodes.size(); i++) {
pos_1 = lattice_nodes[i].first;
vel_1 = lattice_nodes[i].second;
std::vector<double> endstate_lattice = {pos_1(0), vel_1(0), acc_1(0), pos_1(1), vel_1(1), acc_1(1), pos_1(2), vel_1(2), acc_1(2)};
traj_opt_bridge_for_vis->solveBVP(endstate_lattice);
traj_to_pcl(traj_opt_bridge_for_vis, lattice_trajs_cld);
}
pcl_conversions::toPCL(ros::Time::now(), lattice_trajs_cld->header.stamp);
lattice_trajs_cld->header.frame_id = "world";
lattice_trajs_visual_pub.publish(lattice_trajs_cld);
}
std::pair<float, float> calculate_yaw(float sample_t, float dt) {
constexpr float PI = 3.1415926;
constexpr float YAW_DOT_MAX_PER_SEC = 0.3 * PI;
std::pair<float, float> yaw_yawdot(0, 0);
float yaw = 0;
float yawdot = 0;
// dir of vel
Eigen::Vector3d nxt_p, nxt_v, nxt_a;
traj_opt_bridge->getNextState(nxt_p, nxt_v, nxt_a, sample_t);
Eigen::Vector3d dir_vel = nxt_v / nxt_v.norm();
// dir of goal
Eigen::Vector3d dir_goal(goal_(0) - nxt_p(0), goal_(1) - nxt_p(1), goal_(2) - nxt_p(2));
float goal_dist = dir_goal.norm();
dir_goal = dir_goal / goal_dist;
// or just dir_des = dir_vel
Eigen::Vector3d dir_des = dir_vel + dir_goal;
float yaw_temp = goal_dist > 0.2 ? atan2(dir_des(1), dir_des(0)) : last_yaw_;
float max_yaw_change = YAW_DOT_MAX_PER_SEC * dt;
if (yaw_temp - last_yaw_ > PI) {
if (yaw_temp - last_yaw_ - 2 * PI < -max_yaw_change) {
yaw = last_yaw_ - max_yaw_change;
if (yaw < -PI)
yaw += 2 * PI;
yawdot = -YAW_DOT_MAX_PER_SEC;
} else {
yaw = yaw_temp;
if (yaw - last_yaw_ > PI)
yawdot = -YAW_DOT_MAX_PER_SEC;
else
yawdot = (yaw_temp - last_yaw_) / dt;
}
} else if (yaw_temp - last_yaw_ < -PI) {
if (yaw_temp - last_yaw_ + 2 * PI > max_yaw_change) {
yaw = last_yaw_ + max_yaw_change;
if (yaw > PI)
yaw -= 2 * PI;
yawdot = YAW_DOT_MAX_PER_SEC;
} else {
yaw = yaw_temp;
if (yaw - last_yaw_ < -PI)
yawdot = YAW_DOT_MAX_PER_SEC;
else
yawdot = (yaw_temp - last_yaw_) / dt;
}
} else {
if (yaw_temp - last_yaw_ < -max_yaw_change) {
yaw = last_yaw_ - max_yaw_change;
if (yaw < -PI)
yaw += 2 * PI;
yawdot = -YAW_DOT_MAX_PER_SEC;
} else if (yaw_temp - last_yaw_ > max_yaw_change) {
yaw = last_yaw_ + max_yaw_change;
if (yaw > PI)
yaw -= 2 * PI;
yawdot = YAW_DOT_MAX_PER_SEC;
} else {
yaw = yaw_temp;
if (yaw - last_yaw_ > PI)
yawdot = -YAW_DOT_MAX_PER_SEC;
else if (yaw - last_yaw_ < -PI)
yawdot = YAW_DOT_MAX_PER_SEC;
else
yawdot = (yaw_temp - last_yaw_) / dt;
}
}
last_yaw_ = yaw;
yaw_yawdot.first = yaw;
yaw_yawdot.second = yawdot;
return yaw_yawdot;
}
void ref_pub_cb(const ros::TimerEvent&) {
if (!yopo_init)
return;
if (done) {
// single state control for smoother performance
ctrl_ref_last.header.stamp = ros::Time::now();
ctrl_ref_last.vel_ref = {0, 0, 0};
ctrl_ref_last.acc_ref = {0, 0, 0};
ctrl_ref_last.ref_mask = 1;
ctrl_pub.publish(ctrl_ref_last);
// un-smooth, just for simpler demonstration
pos_cmd_last.header.stamp = ros::Time::now();
pos_cmd_last.velocity.x = 0.95 * pos_cmd_last.velocity.x;
pos_cmd_last.velocity.y = 0.95 * pos_cmd_last.velocity.y;
pos_cmd_last.velocity.z = 0.95 * pos_cmd_last.velocity.z;
pos_cmd_last.acceleration.x = 0.95 * pos_cmd_last.acceleration.x;
pos_cmd_last.acceleration.y = 0.95 * pos_cmd_last.acceleration.y;
pos_cmd_last.acceleration.z = 0.95 * pos_cmd_last.acceleration.z;
pos_cmd_last.yaw_dot = 0.95 * pos_cmd_last.yaw_dot;
so3_ctrl_pub.publish(pos_cmd_last);
return;
}
sample_t += 0.02;
Eigen::Vector3d desired_p, desired_v, desired_a;
traj_opt_bridge->getNextState(desired_p, desired_v, desired_a, sample_t);
std::pair<float, float> yaw_yawdot(0, 0);
yaw_yawdot = calculate_yaw(sample_t, 0.02);
// Realworld & our PID Controller
quad_pos_ctrl::ctrl_ref ctrl_msg;
ctrl_msg.header.stamp = ros::Time::now();
ctrl_msg.pos_ref = {desired_p(0), -desired_p(1), -desired_p(2)};
ctrl_msg.vel_ref = {desired_v(0), -desired_v(1), -desired_v(2)};
ctrl_msg.acc_ref = {desired_a(0), -desired_a(1), -desired_a(2)};
ctrl_msg.yaw_ref = -yaw_yawdot.first;
ctrl_msg.ref_mask = 7;
ctrl_ref_last = ctrl_msg;
ctrl_pub.publish(ctrl_msg);
// SO3 Simulator & SO3 Controller
quadrotor_msgs::PositionCommand cmd;
cmd.header.frame_id = "world";
cmd.header.stamp = ros::Time::now();
cmd.trajectory_flag = quadrotor_msgs::PositionCommand::TRAJECTORY_STATUS_READY;
cmd.position.x = desired_p(0);
cmd.position.y = desired_p(1);
cmd.position.z = desired_p(2);
cmd.velocity.x = desired_v(0);
cmd.velocity.y = desired_v(1);
cmd.velocity.z = desired_v(2);
cmd.acceleration.x = desired_a(0);
cmd.acceleration.y = desired_a(1);
cmd.acceleration.z = desired_a(2);
cmd.yaw = yaw_yawdot.first;
cmd.yaw_dot = yaw_yawdot.second;
pos_cmd_last = cmd;
so3_ctrl_pub.publish(cmd);
// update the desire state for next planning
last_pos_ = desired_p;
last_vel_ = desired_v;
last_acc_ = desired_a;
// for reference of yopo network
nav_msgs::Odometry odom_;
odom_.header.stamp = ros::Time::now();
odom_.pose.pose.position.x = desired_p(0);
odom_.pose.pose.position.y = desired_p(1);
odom_.pose.pose.position.z = desired_p(2);
odom_.twist.twist.linear.x = desired_v(0);
odom_.twist.twist.linear.y = desired_v(1);
odom_.twist.twist.linear.z = desired_v(2);
odom_.twist.twist.angular.x = desired_a(0);
odom_.twist.twist.angular.y = desired_a(1);
odom_.twist.twist.angular.z = desired_a(2);
state_ref_pub.publish(odom_);
odom_ref_init = true;
}
} // namespace yopo_net
using namespace yopo_net;
int main(int argc, char** argv) {
ros::init(argc, argv, "yopo_test");
ros::NodeHandle nh;
string cfg_path = getenv("FLIGHTMARE_PATH") + std::string("/flightlib/configs/traj_opt.yaml");
YAML::Node cfg_ = YAML::LoadFile(cfg_path);
traj_time = 2 * cfg_["radio_range"].as<double>() / cfg_["vel_max"].as<double>();
int horizon_num = cfg_["horizon_num"].as<int>();
int vertical_num = cfg_["vertical_num"].as<int>();
int vel_num = cfg_["vel_num"].as<int>();
int radio_num = cfg_["radio_num"].as<int>();
double horizon_fov = cfg_["horizon_camera_fov"].as<double>() * (horizon_num - 1) / horizon_num;
double vertical_fov = cfg_["vertical_camera_fov"].as<double>() * (vertical_num - 1) / vertical_num;
double vel_fov = cfg_["vel_fov"].as<double>();
double radio_range = cfg_["radio_range"].as<double>();
double vel_prefile = cfg_["vel_prefile"].as<double>();
getLatticeGuiding(lattice_nodes, horizon_num, vertical_num, radio_num, vel_num, horizon_fov, vertical_fov, radio_range, vel_fov, vel_prefile);
traj_opt_bridge = new TrajOptimizationBridge();
traj_opt_bridge_for_vis = new TrajOptimizationBridge();
lattice_trajs_visual_pub = nh.advertise<pcl::PointCloud<pcl::PointXYZI>>("/yopo_net/lattice_trajs_visual", 1);
trajs_visual_pub = nh.advertise<pcl::PointCloud<pcl::PointXYZI>>("/yopo_net/trajs_visual", 1);
best_traj_visual_pub = nh.advertise<pcl::PointCloud<pcl::PointXYZI>>("/yopo_net/best_traj_visual", 1);
state_ref_pub = nh.advertise<nav_msgs::Odometry>("/juliett/state_ref/odom", 10);
// our PID Controller (realworld) & SO3 Controller (simulation)
ctrl_pub = nh.advertise<quad_pos_ctrl::ctrl_ref>("/juliett_pos_ctrl_node/controller/ctrl_ref", 10);
so3_ctrl_pub = nh.advertise<quadrotor_msgs::PositionCommand>("/so3_control/pos_cmd", 10);
odom_sub = nh.subscribe("/juliett/ground_truth/odom", 1, yopo_net::odom_cb, ros::TransportHints().tcpNoDelay());
yopo_best_sub = nh.subscribe("/yopo_net/pred_endstate", 1, yopo_net::yopo_cb, ros::TransportHints().tcpNoDelay());
yopo_all_sub = nh.subscribe("/yopo_net/pred_endstates", 1, trajs_vis_cb);
goal_sub = nh.subscribe("/yopo_net/goal", 1, goal_cb);
ros::Timer ref_timer = nh.createTimer(ros::Duration(0.02), ref_pub_cb);
std::cout << "YOPO Planner Node OK!" << std::endl;
ros::spin();
}