增加重投影

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 <vector>
+// third party
+#include <opencv2/core.hpp>
+#include <opencv2/core/types.hpp>
+#include <opencv2/imgproc.hpp>
+#include <tf2/LinearMath/Quaternion.h>
+#include <tf2/convert.h>
+// ros msgs
+#include <geometry_msgs/msg/pose.hpp>
+#include <rm_interfaces/msg/armor.hpp>
+// 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<cv::Point3f> 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>{
+    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<cv::Point2f> 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<cv::Point3f> 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<cv::Point2f> 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<cv::Point2f>& 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<double>(0, 0) = cam_info.k[0];
+  K.at<double>(0, 2) = cam_info.k[2];
+  K.at<double>(1, 1) = cam_info.k[4];
+  K.at<double>(1, 2) = cam_info.k[5];
+  K.at<double>(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<int>(cam_info.d.size()); i++) {
+    D.at<double>(0, i) = cam_info.d[i];
+  }
+  return D;
+}
+
+}  // namespace fyt::auto_aim
+#endif  // ARMOR_DETECTOR_ARMOR_REPROJECT_HPP_

src/rm_auto_aim/armor_detector/src/armor_detector_node.cpp

@@ -48,6 +48,7 @@
 #include <opencv2/imgproc.hpp>
 // 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<std::mutex> 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);