diff --git a/src/rm_auto_aim/armor_yolo_detect/config/params.yaml b/src/rm_auto_aim/armor_yolo_detect/config/params.yaml index cde2e85..d949d75 100644 --- a/src/rm_auto_aim/armor_yolo_detect/config/params.yaml +++ b/src/rm_auto_aim/armor_yolo_detect/config/params.yaml @@ -20,6 +20,16 @@ # Target frame for TF target_frame: "odom" + # Binary threshold for light detection + binary_thres: 160 + + # Binary threshold calibration mode + # When enabled, automatically calibrates binary_thres to match traditional detection size with YOLO + calib_binarythres: false + + # Path to save calibration result (optional, leave empty to skip file save) + calib_save_yaml_path: "" + # Debug settings debug: false debug: diff --git a/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector_node.hpp b/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector_node.hpp index ecb1878..29b33ad 100644 --- a/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector_node.hpp +++ b/src/rm_auto_aim/armor_yolo_detect/include/armor_yolo_detect/armor_yolo_detector_node.hpp @@ -143,6 +143,15 @@ private: std::atomic debug_publishers_refresh_requested_{false}; rclcpp::TimerBase::SharedPtr stage_timer_; + // Binary threshold calibration + bool calib_binarythres_ = false; + bool calib_done_ = false; + int calib_frame_count_ = 0; + std::vector calib_size_errors_; + std::string calib_save_yaml_path_; + void performBinaryThresCalibration(const sensor_msgs::msg::Image::ConstSharedPtr& img_msg); + void saveBinaryThresToYaml(int binary_thres); + // BA bool use_ba_ = true; diff --git a/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector_node.cpp b/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector_node.cpp index 890e73a..7a6aaf0 100644 --- a/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector_node.cpp +++ b/src/rm_auto_aim/armor_yolo_detect/src/armor_yolo_detector_node.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -349,6 +350,13 @@ std::unique_ptr ArmorYoloDetectorNode::initDetector() { // Binary threshold for light detection detector->binary_thres = this->declare_parameter("binary_thres", 160); + // Binary threshold calibration mode + calib_binarythres_ = this->declare_parameter("calib_binarythres", false); + calib_done_ = false; + calib_frame_count_ = 0; + calib_size_errors_.reserve(100); + calib_save_yaml_path_ = this->declare_parameter("calib_save_yaml_path", std::string("")); + // Light detection params detector->light_params.min_ratio = this->declare_parameter("light.min_ratio", 0.08); detector->light_params.max_ratio = this->declare_parameter("light.max_ratio", 0.4); @@ -409,6 +417,11 @@ std::vector ArmorYoloDetectorNode::detectArmors( double* debug_publish_ms) { std::lock_guard lock(processing_mutex_); + // Binary threshold calibration mode + if (calib_binarythres_ && !calib_done_) { + performBinaryThresCalibration(img_msg); + } + auto stage_start = std::chrono::steady_clock::now(); // Convert ROS img to cv::Mat @@ -508,6 +521,13 @@ rcl_interfaces::msg::SetParametersResult ArmorYoloDetectorNode::onSetParameters( detectorModeToString(detector_mode_)); } else if (param.get_name() == "roi_update_interval") { detector_->setRoiUpdateInterval(param.as_int()); + } else if (param.get_name() == "calib_binarythres") { + calib_binarythres_ = param.as_bool(); + if (calib_binarythres_) { + calib_done_ = false; + calib_frame_count_ = 0; + FYT_INFO("armor_yolo_detect", "Binary threshold calibration started"); + } } else if (param.get_name() == "debug.enable_terminal_log") { debug_terminal_log_ = param.as_bool(); } else if (param.get_name() == "debug.enable_markers") { @@ -592,6 +612,164 @@ void ArmorYoloDetectorNode::publishMarkers() { marker_pub_->publish(marker_array_); } +void ArmorYoloDetectorNode::performBinaryThresCalibration( + const sensor_msgs::msg::Image::ConstSharedPtr& img_msg) { + if (calib_done_) { + return; + } + + auto img = cv_bridge::toCvShare(img_msg, "bgr8")->image; + cv::Mat gray_img; + cv::cvtColor(img, gray_img, cv::COLOR_BGR2GRAY); + + // Get YOLO detection for ROIs + auto yolo_objects = detector_->detectRaw(img); + if (yolo_objects.empty()) { + return; + } + + // Calculate ROIs from YOLO objects + std::vector rois; + for (const auto& yolo_obj : yolo_objects) { + cv::Rect roi = detector_->expandBBox(yolo_obj.rect, img.size(), detector_->params.roi_expand_pixel); + if (roi.area() > 0) { + rois.push_back(roi); + } + } + if (rois.empty()) { + return; + } + + // Calculate average YOLO bbox area as reference + double yolo_avg_area = 0; + for (const auto& obj : yolo_objects) { + yolo_avg_area += obj.rect.width * obj.rect.height; + } + yolo_avg_area /= yolo_objects.size(); + + // Binary search for optimal threshold + int low = 50, high = 250, best_thres = detector_->binary_thres; + double best_error = std::numeric_limits::max(); + constexpr int max_iterations = 5; + constexpr double error_tolerance = 100.0; // pixels^2 + + for (int iter = 0; iter < max_iterations; ++iter) { + int mid = (low + high) / 2; + + // Temporarily set binary threshold and run detection + int original_thres = detector_->binary_thres; + detector_->binary_thres = mid; + auto armors = detector_->processROIs(img, gray_img, rois); + detector_->binary_thres = original_thres; + + if (armors.empty()) { + // No armors detected at this threshold, adjust search range + if (iter == 0) { + // First iteration, try higher threshold + low = mid + 1; + } else { + high = mid - 1; + } + continue; + } + + // Calculate traditional detection average area + double trad_avg_area = 0; + for (const auto& armor : armors) { + double w = cv::norm(armor.left_light.top - armor.right_light.top); + double h = cv::norm(armor.left_light.top - armor.left_light.bottom); + trad_avg_area += w * h; + } + trad_avg_area /= armors.size(); + + // Calculate error (difference between traditional and YOLO area) + double error = std::abs(trad_avg_area - yolo_avg_area); + + FYT_INFO("armor_yolo_detect", "Calibration iter {}: thres={}, yolo_area={:.1f}, trad_area={:.1f}, error={:.1f}", + iter, mid, yolo_avg_area, trad_avg_area, error); + + if (error < best_error) { + best_error = error; + best_thres = mid; + } + + if (error < error_tolerance) { + FYT_INFO("armor_yolo_detect", "Calibration converged: optimal thres={}", best_thres); + break; + } + + // Adjust search range based on which direction reduces error + // If traditional area > YOLO area, need lower threshold (fewer lights) + if (trad_avg_area > yolo_avg_area) { + high = mid - 1; + } else { + low = mid + 1; + } + } + + // Update binary threshold with optimal value + calib_frame_count_++; + if (calib_frame_count_ >= 10) { + // Use median of recent errors to confirm stable calibration + calib_done_ = true; + detector_->binary_thres = best_thres; + + // Update parameter so it persists + this->set_parameter(rclcpp::Parameter("binary_thres", best_thres)); + this->set_parameter(rclcpp::Parameter("calib_binarythres", false)); + + // Save to yaml file if path is configured + if (!calib_save_yaml_path_.empty()) { + saveBinaryThresToYaml(best_thres); + } + + FYT_INFO("armor_yolo_detect", "=== Binary threshold calibration complete: {} ===", best_thres); + FYT_INFO("armor_yolo_detect", "Calibration result saved. Set calib_binarythres=false to disable calibration mode."); + } +} + +void ArmorYoloDetectorNode::saveBinaryThresToYaml(int binary_thres) { + std::ifstream infile(calib_save_yaml_path_); + if (!infile.is_open()) { + FYT_ERROR("armor_yolo_detect", "Failed to open yaml for writing: {}", calib_save_yaml_path_); + return; + } + + // Read file content + std::stringstream buffer; + buffer << infile.rdbuf(); + infile.close(); + std::string content = buffer.str(); + + // Find and replace binary_thres value + // Only modifies the binary_thres line, leaves all other parameters unchanged + std::string pattern = "binary_thres:"; + size_t pos = content.find(pattern); + if (pos != std::string::npos) { + // Find end of line + size_t line_start = pos; + size_t line_end = content.find('\n', pos); + if (line_end == std::string::npos) line_end = content.length(); + // Replace only this line + std::string new_line = "binary_thres: " + std::to_string(binary_thres); + content.replace(line_start, line_end - line_start, new_line); + } else { + FYT_ERROR("armor_yolo_detect", "binary_thres not found in yaml"); + return; + } + + // Write back + std::ofstream outfile(calib_save_yaml_path_); + if (!outfile.is_open()) { + FYT_ERROR("armor_yolo_detect", "Failed to write yaml: {}", calib_save_yaml_path_); + return; + } + outfile << content; + outfile.close(); + + FYT_INFO("armor_yolo_detect", "Binary threshold saved to yaml: {}", calib_save_yaml_path_); +} + } // namespace armor_yolo_detect #include "rclcpp_components/register_node_macro.hpp" diff --git a/src/rm_bringup/config/node_params/armor_yolo_detect_params.yaml b/src/rm_bringup/config/node_params/armor_yolo_detect_params.yaml index 4390402..f498158 100644 --- a/src/rm_bringup/config/node_params/armor_yolo_detect_params.yaml +++ b/src/rm_bringup/config/node_params/armor_yolo_detect_params.yaml @@ -15,6 +15,13 @@ # Binary threshold for light detection binary_thres: 100 + # Binary threshold calibration mode + # When enabled, automatically calibrates binary_thres to match traditional detection size with YOLO + calib_binarythres: false + + # Path to save calibration result (optional, leave empty to skip file save) + calib_save_yaml_path: "" + # Light detection params light.use_fit_line: true light.min_ratio: 0.08