参考上海交通大学代码进行修改

Catalyst-MDVS2/inc/ArmorDetector.h

@@ -0,0 +1,65 @@
+#ifndef ARMORDETECTOR_H
+#define ARMORDETECTOR_H
+
+#include <opencv2/opencv.hpp>
+#include <vector>
+#include <string>
+#include "config.h"
+
+struct LightBar {
+    cv::Point2f center;
+    cv::Size2f size;
+    float angle;          // in degrees
+    float angle_rad;      // in radians
+    float area;
+    std::vector<cv::Point2f> center_line;  // [end1, end2]
+    float center_line_length;
+    std::vector<cv::Point2f> box;
+};
+
+struct ArmorPlate {
+    std::string color;
+    cv::Point2f center;
+    double confidence;
+    std::pair<LightBar, LightBar> pair;
+    std::vector<cv::Point2f> corners_2d;  // Can be computed later for 3D pose
+
+    // Function to get bounding rectangle of the armor plate
+    cv::Rect2d getBoundingRect() const {
+        if (corners_2d.size() >= 4) {
+            cv::Point2f min_pt = corners_2d[0];
+            cv::Point2f max_pt = corners_2d[0];
+            for (const auto& pt : corners_2d) {
+                min_pt.x = std::min(min_pt.x, pt.x);
+                min_pt.y = std::min(min_pt.y, pt.y);
+                max_pt.x = std::max(max_pt.x, pt.x);
+                max_pt.y = std::max(max_pt.y, pt.y);
+            }
+            return cv::Rect2d(min_pt.x, min_pt.y, max_pt.x - min_pt.x, max_pt.y - min_pt.y);
+        }
+        // Fallback: use center and a fixed size
+        return cv::Rect2d(center.x - 30, center.y - 15, 60, 30);
+    }
+};
+
+class ArmorDetector {
+public:
+    ArmorDetector();
+    
+    // Main detection function
+    void detect(const cv::Mat& mask, const std::string& target_color, 
+                std::vector<LightBar>& valid_light_bars, 
+                std::vector<ArmorPlate>& armor_plates);
+
+private:
+    float angle_threshold_rad;
+    
+    // Helper methods
+    std::vector<LightBar> extract_initial_light_bars(const cv::Mat& mask);
+    std::vector<std::vector<cv::Point2f>> merge_nearby_light_bars(const std::vector<LightBar>& initial_light_bars);
+    std::vector<LightBar> filter_valid_light_bars(const std::vector<std::vector<cv::Point2f>>& processed_boxes);
+    std::vector<ArmorPlate> pair_light_bars_to_armor(const std::vector<LightBar>& light_bars, const std::string& target_color);
+    double calculate_iou(const cv::Rect& b1, const cv::Rect& b2);
+};
+
+#endif // ARMORDETECTOR_H