65 lines
2.2 KiB
C++
65 lines
2.2 KiB
C++
#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
|