yolo roi binary box find light

This commit is contained in:
cyy_mac
2026-03-26 03:15:10 +08:00
parent bda0306ce0
commit 550414dd3b
2 changed files with 59 additions and 45 deletions

View File

@@ -146,7 +146,7 @@ private:
std::vector<fyt::auto_aim::Armor> processROIs( std::vector<fyt::auto_aim::Armor> processROIs(
const cv::Mat& input, const cv::Mat& input,
const cv::Mat& gray_img, const cv::Mat& gray_img,
const cv::Mat& binary_img); const std::vector<cv::Rect>& rois);
std::unique_ptr<YoloTensorRT> tensorrt_; std::unique_ptr<YoloTensorRT> tensorrt_;
std::vector<YoloObject> yolo_objects_; std::vector<YoloObject> yolo_objects_;

View File

@@ -68,14 +68,6 @@ std::vector<fyt::auto_aim::Armor> Detector::detect(const cv::Mat& input) {
cv::Mat gray_img; cv::Mat gray_img;
cv::cvtColor(input, gray_img, cv::COLOR_BGR2GRAY); cv::cvtColor(input, gray_img, cv::COLOR_BGR2GRAY);
cv::Mat full_binary_img;
cv::threshold(gray_img, full_binary_img, binary_thres, 255, cv::THRESH_BINARY);
// Store binary image for debug visualization
if (enable_debug) {
debug_binary_img = full_binary_img.clone();
}
uint64_t current_frame = frame_counter_.fetch_add(1, std::memory_order_relaxed); uint64_t current_frame = frame_counter_.fetch_add(1, std::memory_order_relaxed);
// Update YOLO objects based on mode // Update YOLO objects based on mode
@@ -102,19 +94,7 @@ std::vector<fyt::auto_aim::Armor> Detector::detect(const cv::Mat& input) {
return armors_; return armors_;
} }
// Process ROIs to detect armors // Calculate ROIs from YOLO objects first
armors_ = processROIs(input, gray_img, full_binary_img);
return armors_;
}
std::vector<fyt::auto_aim::Armor> Detector::processROIs(
const cv::Mat& input,
const cv::Mat& gray_img,
const cv::Mat& binary_img) {
std::vector<fyt::auto_aim::Armor> result_armors;
// 1. Calculate all YOLO ROIs for filtering
std::vector<cv::Rect> rois; std::vector<cv::Rect> rois;
for (const auto& yolo_obj : yolo_objects_) { for (const auto& yolo_obj : yolo_objects_) {
cv::Rect roi = expandBBox(yolo_obj.rect, input.size(), params.roi_expand_pixel); cv::Rect roi = expandBBox(yolo_obj.rect, input.size(), params.roi_expand_pixel);
@@ -124,31 +104,50 @@ std::vector<fyt::auto_aim::Armor> Detector::processROIs(
} }
} }
// 2. Find lights on the FULL image natively (exact armor_detector logic) if (rois.empty()) {
cv::Rect full_roi(0, 0, input.cols, input.rows); return armors_;
auto all_lights = findLightsInROI(input, gray_img, binary_img, full_roi); }
// 3. Filter lights: only keep those inside YOLO ROIs // Store binary image for debug visualization (crop to first ROI)
std::vector<fyt::auto_aim::Light> filtered_lights; if (enable_debug) {
for (const auto& light : all_lights) { const auto& roi = debug_rois_.front();
bool in_roi = false; cv::Mat roi_gray = gray_img(roi);
cv::threshold(roi_gray, debug_binary_img, binary_thres, 255, cv::THRESH_BINARY);
}
// Process ROIs to detect armors (only in ROI regions)
armors_ = processROIs(input, gray_img, rois);
return armors_;
}
std::vector<fyt::auto_aim::Armor> Detector::processROIs(
const cv::Mat& input,
const cv::Mat& gray_img,
const std::vector<cv::Rect>& rois) {
std::vector<fyt::auto_aim::Armor> result_armors;
// Find lights in each ROI region (with local binary for each ROI)
std::vector<fyt::auto_aim::Light> all_lights;
for (const auto& roi : rois) { for (const auto& roi : rois) {
if (roi.contains(light.center)) { // Crop ROI from gray image
in_roi = true; cv::Mat roi_gray = gray_img(roi);
break;
} // Apply binary threshold within ROI
} cv::Mat roi_binary;
if (in_roi) { cv::threshold(roi_gray, roi_binary, binary_thres, 255, cv::THRESH_BINARY);
filtered_lights.push_back(light);
} // Find lights in this ROI
auto lights = findLightsInROI(input, gray_img, roi_binary, roi);
all_lights.insert(all_lights.end(), lights.begin(), lights.end());
} }
if (filtered_lights.empty()) { if (all_lights.empty()) {
return result_armors; return result_armors;
} }
// 4. Match lights to form armors natively in global coordinates // 4. Match lights to form armors in global coordinates
auto armors = matchLightsInROI(filtered_lights); auto armors = matchLightsInROI(all_lights);
// 5. Extract numbers, classify and refine corners // 5. Extract numbers, classify and refine corners
if (classifier != nullptr) { if (classifier != nullptr) {
@@ -243,17 +242,29 @@ std::vector<fyt::auto_aim::Light> Detector::findLightsInROI(
light = fyt::auto_aim::Light(contour); light = fyt::auto_aim::Light(contour);
} }
// Transform light coordinates from ROI-local to global
light.center += cv::Point2f(roi.x, roi.y);
light.top += cv::Point2f(roi.x, roi.y);
light.bottom += cv::Point2f(roi.x, roi.y);
// Also update the RotatedRect center
light.center = light.center; // Already updated above
if (isLight(light)) { if (isLight(light)) {
// Color detection within ROI // Color detection within ROI
int sum_r = 0; int sum_r = 0;
int sum_b = 0; int sum_b = 0;
int sample_count = 0; int sample_count = 0;
// ROI offset for coordinate transformation
const int roi_x = roi.x;
const int roi_y = roi.y;
if (light_params.use_fit_line && !points.empty()) { if (light_params.use_fit_line && !points.empty()) {
auto b_rect = cv::boundingRect(contour); auto b_rect = cv::boundingRect(contour);
for (const auto& point : points) { for (const auto& point : points) {
const int x = point.x + b_rect.x; // Add both b_rect offset (relative to ROI) and ROI offset (global)
const int y = point.y + b_rect.y; const int x = point.x + b_rect.x + roi_x;
const int y = point.y + b_rect.y + roi_y;
const auto& pixel = rgb_img.at<cv::Vec3b>(y, x); const auto& pixel = rgb_img.at<cv::Vec3b>(y, x);
sum_b += pixel[0]; sum_b += pixel[0];
sum_r += pixel[2]; sum_r += pixel[2];
@@ -261,7 +272,10 @@ std::vector<fyt::auto_aim::Light> Detector::findLightsInROI(
} }
} else { } else {
for (const auto& point : contour) { for (const auto& point : contour) {
const auto& pixel = rgb_img.at<cv::Vec3b>(point.y, point.x); // Add ROI offset for global coordinates
const int x = point.x + roi_x;
const int y = point.y + roi_y;
const auto& pixel = rgb_img.at<cv::Vec3b>(y, x);
sum_b += pixel[0]; sum_b += pixel[0];
sum_r += pixel[2]; sum_r += pixel[2];
++sample_count; ++sample_count;