diff --git a/src/rm_auto_aim/armor_detector/include/armor_detector/armor_reproject.hpp b/src/rm_auto_aim/armor_detector/include/armor_detector/armor_reproject.hpp new file mode 100644 index 0000000..cc0670f --- /dev/null +++ b/src/rm_auto_aim/armor_detector/include/armor_detector/armor_reproject.hpp @@ -0,0 +1,205 @@ +// Copyright (C) FYT Vision Group. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef ARMOR_DETECTOR_ARMOR_REPROJECT_HPP_ +#define ARMOR_DETECTOR_ARMOR_REPROJECT_HPP_ + +// std +#include +// third party +#include +#include +#include +#include +#include +// ros msgs +#include +#include +// project +#include "armor_detector/types.hpp" + +namespace fyt::auto_aim { + +/** + * @brief Get armor corner points in armor coordinate system + * @param is_big true for large armor, false for small armor + * @return 4 corner points in clockwise order starting from bottom-left + * + * Armor coordinate system: + * x-axis: points out of the armor face (normal direction) + * y-axis: along the width, positive to the right when facing armor + * z-axis: along the height, positive upward + * + * Corner order (clockwise from bottom-left when facing armor): + * 0: bottom-left (y = +w/2, z = -h/2) + * 1: top-left (y = +w/2, z = +h/2) + * 2: top-right (y = -w/2, z = +h/2) + * 3: bottom-right (y = -w/2, z = -h/2) + */ +static inline std::vector getArmorCornerPoints(bool is_big) { + double w = is_big ? LARGE_ARMOR_WIDTH : SMALL_ARMOR_WIDTH; + double h = is_big ? LARGE_ARMOR_HEIGHT : SMALL_ARMOR_HEIGHT; + return std::vector{ + cv::Point3f(0, w / 2, -h / 2), // 0: bottom-left + cv::Point3f(0, w / 2, h / 2), // 1: top-left + cv::Point3f(0, -w / 2, h / 2), // 2: top-right + cv::Point3f(0, -w / 2, -h / 2) // 3: bottom-right + }; +} + +/** + * @brief Convert geometry_msgs::Pose to rotation matrix and translation vector + */ +static inline void poseToRotTvec( + const geometry_msgs::msg::Pose& pose, + cv::Matx33d& rot_mat, + cv::Vec3d& tvec) { + tf2::Quaternion q( + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w); + tf2::Matrix3x3 rot_mat_tf(q); + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + rot_mat(i, j) = rot_mat_tf[i][j]; + } + } + tvec = cv::Vec3d(pose.position.x, pose.position.y, pose.position.z); +} + +/** + * @brief Reproject armor pose to 2D image points + * @param pose Armor pose in camera coordinate system + * @param K Camera intrinsic matrix + * @param dist_coeffs Camera distortion coefficients + * @param is_big Whether the armor is large or small + * @return 4 corner points in image coordinates + */ +static inline std::vector reprojectArmorToImage( + const geometry_msgs::msg::Pose& pose, + const cv::Mat& K, + const cv::Mat& dist_coeffs, + bool is_big) { + cv::Matx33d rot_mat; + cv::Vec3d tvec; + poseToRotTvec(pose, rot_mat, tvec); + + auto corners_armor = getArmorCornerPoints(is_big); + + // Transform corners from armor frame to camera frame + std::vector corners_cam; + for (const auto& p_armor : corners_armor) { + cv::Point3d p_cam; + p_cam.x = rot_mat(0, 0) * p_armor.x + rot_mat(0, 1) * p_armor.y + rot_mat(0, 2) * p_armor.z + tvec[0]; + p_cam.y = rot_mat(1, 0) * p_armor.x + rot_mat(1, 1) * p_armor.y + rot_mat(1, 2) * p_armor.z + tvec[1]; + p_cam.z = rot_mat(2, 0) * p_armor.x + rot_mat(2, 1) * p_armor.y + rot_mat(2, 2) * p_armor.z + tvec[2]; + corners_cam.push_back(p_cam); + } + + // Project to image plane + cv::Mat rvec = cv::Mat::zeros(3, 1, CV_64F); + std::vector corners_2d; + cv::projectPoints(corners_cam, rvec, tvec, K, dist_coeffs, corners_2d); + return corners_2d; +} + +/** + * @brief Draw armor quadrilateral on image + * @param img Image to draw on + * @param corners 4 corner points in image coordinates + * @param color Line color + * @param thickness Line thickness + */ +static inline void drawArmorQuad( + cv::Mat& img, + const std::vector& corners, + const cv::Scalar& color = cv::Scalar(0, 255, 0), + int thickness = 2) { + if (corners.size() != 4) return; + + // Draw 4 edges + for (int i = 0; i < 4; i++) { + cv::line(img, corners[i], corners[(i + 1) % 4], color, thickness); + } + + // Draw center point + cv::Point2f center; + for (const auto& p : corners) { + center += p; + } + center /= 4; + cv::circle(img, center, 3, color, -1); +} + +/** + * @brief Draw armor with number label on image + * @param img Image to draw on + * @param armor Armor message containing pose and number + * @param K Camera intrinsic matrix + * @param dist_coeffs Camera distortion coefficients + * @param color Line color + * @param thickness Line thickness + */ +static inline void drawArmorFromMsg( + cv::Mat& img, + const rm_interfaces::msg::Armor& armor, + const cv::Mat& K, + const cv::Mat& dist_coeffs, + const cv::Scalar& color = cv::Scalar(0, 255, 0), + int thickness = 2) { + bool is_big = (armor.type == "large"); + auto corners = reprojectArmorToImage(armor.pose, K, dist_coeffs, is_big); + drawArmorQuad(img, corners, color, thickness); + + // Draw number label near the armor + if (!armor.number.empty() && corners.size() == 4) { + cv::putText( + img, + armor.number, + corners[0], + cv::FONT_HERSHEY_SIMPLEX, + 0.5, + color, + 1); + } +} + +/** + * @brief Build camera matrix from CameraInfo message + */ +static inline cv::Mat buildCameraMatrix(const sensor_msgs::msg::CameraInfo& cam_info) { + cv::Mat K = cv::Mat::eye(3, 3, CV_64F); + K.at(0, 0) = cam_info.k[0]; + K.at(0, 2) = cam_info.k[2]; + K.at(1, 1) = cam_info.k[4]; + K.at(1, 2) = cam_info.k[5]; + K.at(2, 2) = 1.0; + return K; +} + +/** + * @brief Build distortion coefficient vector from CameraInfo message + */ +static inline cv::Mat buildDistCoeffs(const sensor_msgs::msg::CameraInfo& cam_info) { + // CameraInfo uses 5 distortion coefficients [k1, k2, p1, p2, k3] + cv::Mat D = cv::Mat::zeros(1, 5, CV_64F); + for (int i = 0; i < 5 && i < static_cast(cam_info.d.size()); i++) { + D.at(0, i) = cam_info.d[i]; + } + return D; +} + +} // namespace fyt::auto_aim +#endif // ARMOR_DETECTOR_ARMOR_REPROJECT_HPP_ diff --git a/src/rm_auto_aim/armor_detector/src/armor_detector_node.cpp b/src/rm_auto_aim/armor_detector/src/armor_detector_node.cpp index 1b81a67..7444cd7 100644 --- a/src/rm_auto_aim/armor_detector/src/armor_detector_node.cpp +++ b/src/rm_auto_aim/armor_detector/src/armor_detector_node.cpp @@ -48,6 +48,7 @@ #include // project #include "armor_detector/armor_detector_node.hpp" +#include "armor_detector/armor_reproject.hpp" #include "armor_detector/ba_solver.hpp" #include "armor_detector/types.hpp" #include "rm_utils/assert.hpp" @@ -337,6 +338,33 @@ void ArmorDetectorNode::processFrame(PendingFrame frame) { } const double marker_ms = durationMs(std::chrono::steady_clock::now() - stage_start); + // Publishing armor reprojection on result image + stage_start = std::chrono::steady_clock::now(); + { + std::lock_guard lock(processing_mutex_); + if (debug_enabled && debug_result_img_ && debug_result_pub_active_ && + armor_pose_estimator_ != nullptr && cam_info_ != nullptr) { + // Re-convert image from ROS msg for reprojection drawing + auto debug_img = cv_bridge::toCvShare(frame.img_msg, "rgb8")->image; + if (!debug_img.empty()) { + cv::Mat K = buildCameraMatrix(*cam_info_); + cv::Mat dist_coeffs = buildDistCoeffs(*cam_info_); + + // Draw reprojection for each armor with pose + for (const auto& armor : armors_msg.armors) { + if (armor.pose.position.x > 0) { // Valid pose check + drawArmorFromMsg(debug_img, armor, K, dist_coeffs, cv::Scalar(255, 0, 255), 2); + } + } + + result_img_pub_.publish( + cv_bridge::CvImage(frame.img_msg->header, "rgb8", debug_img).toImageMsg()); + } + } + } + const double reproject_ms = durationMs(std::chrono::steady_clock::now() - stage_start); + (void)reproject_ms; // Avoid unused warning if debug is off + // Publishing detected armors stage_start = std::chrono::steady_clock::now(); armors_pub_->publish(armors_msg);