From 550414dd3bb55eb2147b90f78c38a84eafe584d1 Mon Sep 17 00:00:00 2001 From: cyy_mac Date: Thu, 26 Mar 2026 03:15:10 +0800 Subject: [PATCH] yolo roi binary box find light --- .../armor_yolo_detect/armor_yolo_detector.hpp | 2 +- .../src/armor_yolo_detector.cpp | 102 ++++++++++-------- 2 files changed, 59 insertions(+), 45 deletions(-) diff --git a/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector.hpp b/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector.hpp index e90af7e..4fca9e6 100644 --- a/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector.hpp +++ b/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector.hpp @@ -146,7 +146,7 @@ private: std::vector processROIs( const cv::Mat& input, const cv::Mat& gray_img, - const cv::Mat& binary_img); + const std::vector& rois); std::unique_ptr tensorrt_; std::vector yolo_objects_; diff --git a/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector.cpp b/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector.cpp index 0b5512e..e3cbb11 100644 --- a/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector.cpp +++ b/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector.cpp @@ -68,14 +68,6 @@ std::vector Detector::detect(const cv::Mat& input) { cv::Mat gray_img; 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); // Update YOLO objects based on mode @@ -102,19 +94,7 @@ std::vector Detector::detect(const cv::Mat& input) { return armors_; } - // Process ROIs to detect armors - armors_ = processROIs(input, gray_img, full_binary_img); - - return armors_; -} - -std::vector Detector::processROIs( - const cv::Mat& input, - const cv::Mat& gray_img, - const cv::Mat& binary_img) { - std::vector result_armors; - - // 1. Calculate all YOLO ROIs for filtering + // Calculate ROIs from YOLO objects first std::vector rois; for (const auto& yolo_obj : yolo_objects_) { cv::Rect roi = expandBBox(yolo_obj.rect, input.size(), params.roi_expand_pixel); @@ -124,31 +104,50 @@ std::vector Detector::processROIs( } } - // 2. Find lights on the FULL image natively (exact armor_detector logic) - cv::Rect full_roi(0, 0, input.cols, input.rows); - auto all_lights = findLightsInROI(input, gray_img, binary_img, full_roi); - - // 3. Filter lights: only keep those inside YOLO ROIs - std::vector filtered_lights; - for (const auto& light : all_lights) { - bool in_roi = false; - for (const auto& roi : rois) { - if (roi.contains(light.center)) { - in_roi = true; - break; - } - } - if (in_roi) { - filtered_lights.push_back(light); - } + if (rois.empty()) { + return armors_; } - if (filtered_lights.empty()) { + // Store binary image for debug visualization (crop to first ROI) + if (enable_debug) { + const auto& roi = debug_rois_.front(); + 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 Detector::processROIs( + const cv::Mat& input, + const cv::Mat& gray_img, + const std::vector& rois) { + std::vector result_armors; + + // Find lights in each ROI region (with local binary for each ROI) + std::vector all_lights; + for (const auto& roi : rois) { + // Crop ROI from gray image + cv::Mat roi_gray = gray_img(roi); + + // Apply binary threshold within ROI + cv::Mat roi_binary; + cv::threshold(roi_gray, roi_binary, binary_thres, 255, cv::THRESH_BINARY); + + // 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 (all_lights.empty()) { return result_armors; } - // 4. Match lights to form armors natively in global coordinates - auto armors = matchLightsInROI(filtered_lights); + // 4. Match lights to form armors in global coordinates + auto armors = matchLightsInROI(all_lights); // 5. Extract numbers, classify and refine corners if (classifier != nullptr) { @@ -243,17 +242,29 @@ std::vector Detector::findLightsInROI( 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)) { // Color detection within ROI int sum_r = 0; int sum_b = 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()) { auto b_rect = cv::boundingRect(contour); for (const auto& point : points) { - const int x = point.x + b_rect.x; - const int y = point.y + b_rect.y; + // Add both b_rect offset (relative to ROI) and ROI offset (global) + 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(y, x); sum_b += pixel[0]; sum_r += pixel[2]; @@ -261,7 +272,10 @@ std::vector Detector::findLightsInROI( } } else { for (const auto& point : contour) { - const auto& pixel = rgb_img.at(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(y, x); sum_b += pixel[0]; sum_r += pixel[2]; ++sample_count;