c++实现probiou nms函数

使用opencv实现

#include <opencv2/opencv.hpp>
#include <vector>
#include <algorithm>

// 计算旋转矩形的协方差矩阵参数 (a, b, c)
void getCovarianceParams(const cv::RotatedRect& obb, double& a, double& b, double& c) {
    const double w = obb.size.width;
    const double h = obb.size.height;
    const double angle_rad = obb.angle * CV_PI / 180.0;
    
    const double cos_theta = std::cos(angle_rad);
    const double sin_theta = std::sin(angle_rad);
    
    a = (w*w*cos_theta*cos_theta + h*h*sin_theta*sin_theta) / 12.0;
    b = (w*w*sin_theta*sin_theta + h*h*cos_theta*cos_theta) / 12.0;
    c = (w*w - h*h) * sin_theta * cos_theta / 12.0;
}

// 计算两个旋转框的ProbIoU
double calculateProbIoU(const cv::RotatedRect& obb1, const cv::RotatedRect& obb2, double eps = 1e-7) {
    cv::Point2f center1 = obb1.center;
    cv::Point2f center2 = obb2.center;
    double x1 = center1.x, y1 = center1.y;
    double x2 = center2.x, y2 = center2.y;

    double a1, b1, c1, a2, b2, c2;
    getCovarianceParams(obb1, a1, b1, c1);
    getCovarianceParams(obb2, a2, b2, c2);

    double t1 = ((a1 + a2) * std::pow(y1 - y2, 2) + (b1 + b2) * std::pow(x1 - x2, 2)) / 
               ((a1 + a2) * (b1 + b2) - std::pow(c1 + c2, 2) + eps) * 0.25;
    double t2 = ((c1 + c2) * (x2 - x1) * (y1 - y2)) / 
               ((a1 + a2) * (b1 + b2) - std::pow(c1 + c2, 2) + eps) * 0.5;
    double t3 = std::log(((a1 + a2) * (b1 + b2) - std::pow(c1 + c2, 2)) / 
               (4 * std::sqrt(std::max(a1*b1 - c1*c1, 0.0) * std::max(a2*b2 - c2*c2, 0.0)) + eps) + eps) * 0.5;

    double bd = std::clamp(t1 + t2 + t3, eps, 100.0);
    double hd = std::sqrt(1.0 - std::exp(-bd) + eps);
    return 1.0 - hd;
}

// 基于ProbIoU的旋转框NMS
void NMSBoxesRotated(const std::vector<cv::RotatedRect>& boxes, 
                    const std::vector<float>& scores,
                    float score_threshold, 
                    float nms_threshold,
                    std::vector<int>& indices,
                    float eps = 1e-7) {
    
    // 1. 过滤低分框
    std::vector<int> valid_indices;
    for (size_t i = 0; i < scores.size(); ++i) {
        if (scores[i] >= score_threshold) {
            valid_indices.push_back(i);
        }
    }

    // 2. 按分数降序排序
    std::sort(valid_indices.begin(), valid_indices.end(),
        [&scores](int lhs, int rhs) { return scores[lhs] > scores[rhs]; });

    // 3. 进行ProbIoU NMS
    std::vector<bool> is_suppressed(valid_indices.size(), false);
    for (size_t i = 0; i < valid_indices.size(); ++i) {
        if (is_suppressed[i]) continue;
        
        indices.push_back(valid_indices[i]);
        
        for (size_t j = i + 1; j < valid_indices.size(); ++j) {
            if (is_suppressed[j]) continue;
            
            double piou = calculateProbIoU(boxes[valid_indices[i]], 
                                          boxes[valid_indices[j]], 
                                          eps);
            if (piou > nms_threshold) {
                is_suppressed[j] = true;
            }
        }
    }
}

// 使用示例
int main() {
    // 模拟输入数据
    std::vector<cv::RotatedRect> boxes = {
        cv::RotatedRect(cv::Point2f(100, 100), cv::Size2f(50, 30), 45),
        cv::RotatedRect(cv::Point2f(110, 110), cv::Size2f(55, 35), 40),
        cv::RotatedRect(cv::Point2f(200, 200), cv::Size2f(60, 40), 30)
    };
    std::vector<float> scores = {0.9f, 0.8f, 0.95f};
    
    // 调用NMS
    std::vector<int> indices;
    NMSBoxesRotated(boxes, scores, 0.5f, 0.4f, indices);
    
    // 输出结果
    for (int idx : indices) {
        std::cout << "保留框索引: " << idx 
                  << ", 分数: " << scores[idx] 
                  << ", 中心点: " << boxes[idx].center 
                  << std::endl;
    }
    
    return 0;
}