272 lines
9.8 KiB
C++
272 lines
9.8 KiB
C++
#include "ArmorDetector.h"
|
|
#include <cmath>
|
|
#include <algorithm>
|
|
|
|
ArmorDetector::ArmorDetector() {
|
|
angle_threshold_rad = HORIZONTAL_ANGLE_THRESHOLD_RAD;
|
|
}
|
|
|
|
void ArmorDetector::detect(const cv::Mat& mask, const std::string& target_color,
|
|
std::vector<LightBar>& valid_light_bars,
|
|
std::vector<ArmorPlate>& armor_plates) {
|
|
std::vector<LightBar> initial_light_bars = extract_initial_light_bars(mask);
|
|
if (initial_light_bars.empty()) {
|
|
valid_light_bars.clear();
|
|
armor_plates.clear();
|
|
return;
|
|
}
|
|
|
|
std::vector<std::vector<cv::Point2f>> processed_boxes = merge_nearby_light_bars(initial_light_bars);
|
|
if (processed_boxes.empty()) {
|
|
valid_light_bars.clear();
|
|
armor_plates.clear();
|
|
return;
|
|
}
|
|
|
|
valid_light_bars = filter_valid_light_bars(processed_boxes);
|
|
if (valid_light_bars.size() < 2) {
|
|
armor_plates.clear();
|
|
return;
|
|
}
|
|
|
|
armor_plates = pair_light_bars_to_armor(valid_light_bars, target_color);
|
|
if (armor_plates.empty()) {
|
|
return;
|
|
}
|
|
}
|
|
|
|
std::vector<LightBar> ArmorDetector::extract_initial_light_bars(const cv::Mat& mask) {
|
|
std::vector<std::vector<cv::Point>> contours;
|
|
std::vector<cv::Vec4i> hierarchy;
|
|
cv::findContours(mask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
|
|
|
|
std::vector<LightBar> initial_light_bars;
|
|
|
|
for (size_t i = 0; i < contours.size(); i++) {
|
|
if (contours[i].size() < 5) continue; // Need at least 5 points to fit ellipse
|
|
|
|
cv::RotatedRect rect = cv::minAreaRect(contours[i]);
|
|
cv::Point2f vertices[4];
|
|
rect.points(vertices);
|
|
|
|
std::vector<cv::Point2f> box;
|
|
for (int j = 0; j < 4; j++) {
|
|
box.push_back(vertices[j]);
|
|
}
|
|
|
|
float contour_area = cv::contourArea(contours[i]);
|
|
float rect_area = rect.size.width * rect.size.height;
|
|
|
|
// Filter: sufficient area + high fill ratio
|
|
if (rect_area > 30 && rect_area > 0 && (contour_area / rect_area) > 0.4) {
|
|
LightBar light_bar;
|
|
light_bar.center = rect.center;
|
|
light_bar.size = rect.size;
|
|
light_bar.angle = rect.angle;
|
|
light_bar.area = rect_area;
|
|
light_bar.box = box;
|
|
|
|
initial_light_bars.push_back(light_bar);
|
|
}
|
|
}
|
|
|
|
return initial_light_bars;
|
|
}
|
|
|
|
std::vector<std::vector<cv::Point2f>> ArmorDetector::merge_nearby_light_bars(const std::vector<LightBar>& initial_light_bars) {
|
|
std::vector<std::vector<cv::Point2f>> processed_boxes;
|
|
std::vector<bool> visited(initial_light_bars.size(), false);
|
|
|
|
for (size_t i = 0; i < initial_light_bars.size(); i++) {
|
|
if (visited[i]) continue;
|
|
|
|
std::vector<int> nearby_indices = {static_cast<int>(i)};
|
|
cv::Point2f center1 = initial_light_bars[i].center;
|
|
|
|
for (size_t j = i + 1; j < initial_light_bars.size(); j++) {
|
|
if (visited[j]) continue;
|
|
|
|
cv::Point2f center2 = initial_light_bars[j].center;
|
|
float distance = std::sqrt(std::pow(center2.x - center1.x, 2) + std::pow(center2.y - center1.y, 2));
|
|
|
|
// Distance + IOU filtering
|
|
if (distance < NEARBY_LIGHT_BAR_THRESHOLD) {
|
|
cv::Rect b1 = cv::boundingRect(initial_light_bars[i].box);
|
|
cv::Rect b2 = cv::boundingRect(initial_light_bars[j].box);
|
|
|
|
double iou = calculate_iou(b1, b2);
|
|
if (iou > LIGHT_BAR_IOU_THRESHOLD) {
|
|
nearby_indices.push_back(static_cast<int>(j));
|
|
visited[j] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Merge nearby light bar vertices to generate new rectangle
|
|
std::vector<cv::Point2f> all_points;
|
|
for (int k : nearby_indices) {
|
|
for (const auto& point : initial_light_bars[k].box) {
|
|
all_points.push_back(point);
|
|
}
|
|
}
|
|
|
|
if (!all_points.empty()) {
|
|
cv::RotatedRect merged_rect = cv::minAreaRect(all_points);
|
|
cv::Point2f vertices[4];
|
|
merged_rect.points(vertices);
|
|
|
|
std::vector<cv::Point2f> merged_box;
|
|
for (int k = 0; k < 4; k++) {
|
|
merged_box.push_back(vertices[k]);
|
|
}
|
|
|
|
processed_boxes.push_back(merged_box);
|
|
visited[i] = true;
|
|
}
|
|
}
|
|
|
|
return processed_boxes;
|
|
}
|
|
|
|
std::vector<LightBar> ArmorDetector::filter_valid_light_bars(const std::vector<std::vector<cv::Point2f>>& processed_boxes) {
|
|
std::vector<LightBar> valid_light_bars;
|
|
|
|
for (const auto& box : processed_boxes) {
|
|
if (box.size() < 4) continue;
|
|
|
|
cv::RotatedRect rect = cv::minAreaRect(box);
|
|
cv::Point2f center = rect.center;
|
|
float width = rect.size.width;
|
|
float height = rect.size.height;
|
|
|
|
float length = width > height ? width : height;
|
|
float bar_width = width > height ? height : width;
|
|
|
|
float main_angle = rect.angle;
|
|
if (width <= height) {
|
|
main_angle += 90.0;
|
|
}
|
|
|
|
// Normalize angle to [-90, 90]
|
|
main_angle = std::fmod(main_angle, 180);
|
|
if (main_angle > 90) main_angle -= 180;
|
|
if (main_angle < -90) main_angle += 180;
|
|
|
|
float main_angle_rad = main_angle * CV_PI / 180.0f;
|
|
|
|
cv::Point2f end1(
|
|
center.x + (length / 2) * std::cos(main_angle_rad),
|
|
center.y + (length / 2) * std::sin(main_angle_rad)
|
|
);
|
|
cv::Point2f end2(
|
|
center.x - (length / 2) * std::cos(main_angle_rad),
|
|
center.y - (length / 2) * std::sin(main_angle_rad)
|
|
);
|
|
|
|
float rect_area = length * bar_width;
|
|
float contour_area = cv::contourArea(box);
|
|
float area_ratio = rect_area > 0 ? contour_area / rect_area : 0;
|
|
|
|
// Filter: sufficient area + high fill ratio + non-horizontal
|
|
if (rect_area > 40 && area_ratio > 0.4 && std::abs(main_angle_rad) > angle_threshold_rad) {
|
|
LightBar light_bar;
|
|
light_bar.center = center;
|
|
light_bar.size = cv::Size2f(length, bar_width);
|
|
light_bar.angle = main_angle;
|
|
light_bar.angle_rad = main_angle_rad;
|
|
light_bar.area = rect_area;
|
|
light_bar.center_line = {end1, end2};
|
|
light_bar.center_line_length = length;
|
|
light_bar.box = box;
|
|
|
|
valid_light_bars.push_back(light_bar);
|
|
}
|
|
}
|
|
|
|
return valid_light_bars;
|
|
}
|
|
|
|
std::vector<ArmorPlate> ArmorDetector::pair_light_bars_to_armor(const std::vector<LightBar>& light_bars, const std::string& target_color) {
|
|
std::vector<ArmorPlate> armor_plates;
|
|
std::vector<bool> used(light_bars.size(), false);
|
|
|
|
for (size_t i = 0; i < light_bars.size(); i++) {
|
|
if (used[i]) continue;
|
|
const LightBar& lb1 = light_bars[i];
|
|
|
|
for (size_t j = i + 1; j < light_bars.size(); j++) {
|
|
if (used[j]) continue;
|
|
const LightBar& lb2 = light_bars[j];
|
|
|
|
// Angle difference filtering
|
|
float angle_diff = std::abs(lb1.angle_rad - lb2.angle_rad);
|
|
float angle_diff_2 = 2.0f * CV_PI - angle_diff;
|
|
angle_diff = std::min(angle_diff, angle_diff_2);
|
|
if (angle_diff > ARMOR_ANGLE_DIFF_THRESHOLD) {
|
|
continue;
|
|
}
|
|
|
|
// Distance ratio filtering
|
|
float dx = lb2.center.x - lb1.center.x;
|
|
float dy = lb2.center.y - lb1.center.y;
|
|
float distance = std::sqrt(dx * dx + dy * dy);
|
|
float avg_length = (lb1.center_line_length + lb2.center_line_length) / 2;
|
|
float distance_ratio = avg_length > 0 ? distance / avg_length : 0;
|
|
if (distance_ratio <= ARMOR_DISTANCE_RATIO_MIN || distance_ratio >= ARMOR_DISTANCE_RATIO_MAX) {
|
|
continue;
|
|
}
|
|
|
|
// Length difference filtering
|
|
float length_diff_ratio = avg_length > 0 ?
|
|
std::abs(lb1.center_line_length - lb2.center_line_length) / avg_length : 0;
|
|
if (length_diff_ratio > ARMOR_LENGTH_DIFF_RATIO) {
|
|
continue;
|
|
}
|
|
|
|
cv::Point2f armor_center(
|
|
(lb1.center.x + lb2.center.x) / 2,
|
|
(lb1.center.y + lb2.center.y) / 2
|
|
);
|
|
|
|
double confidence = (lb1.area + lb2.area) / (distance + 1);
|
|
|
|
ArmorPlate armor_plate;
|
|
armor_plate.color = target_color;
|
|
armor_plate.center = armor_center;
|
|
armor_plate.confidence = confidence;
|
|
armor_plate.pair = std::make_pair(lb1, lb2);
|
|
armor_plate.corners_2d = std::vector<cv::Point2f>(); // Will be computed later if needed
|
|
|
|
armor_plates.push_back(armor_plate);
|
|
|
|
used[i] = true;
|
|
used[j] = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Sort by confidence in descending order
|
|
std::sort(armor_plates.begin(), armor_plates.end(),
|
|
[](const ArmorPlate& a, const ArmorPlate& b) {
|
|
return a.confidence > b.confidence;
|
|
});
|
|
|
|
return armor_plates;
|
|
}
|
|
|
|
double ArmorDetector::calculate_iou(const cv::Rect& b1, const cv::Rect& b2) {
|
|
cv::Point2f inter_tl(std::max(b1.x, b2.x), std::max(b1.y, b2.y));
|
|
cv::Point2f inter_br(std::min(b1.x + b1.width, b2.x + b2.width),
|
|
std::min(b1.y + b1.height, b2.y + b2.height));
|
|
|
|
if (inter_br.x <= inter_tl.x || inter_br.y <= inter_tl.y) {
|
|
return 0.0;
|
|
}
|
|
|
|
float inter_area = (inter_br.x - inter_tl.x) * (inter_br.y - inter_tl.y);
|
|
float b1_area = b1.width * b1.height;
|
|
float b2_area = b2.width * b2.height;
|
|
float union_area = b1_area + b2_area - inter_area;
|
|
|
|
return union_area > 0 ? inter_area / union_area : 0.0;
|
|
} |