为了在C++中实现非极大值抑制(NMS),我们需要以下步骤:

  1. 定义边界框和置信度的结构。
  2. 实现计算IoU的函数。
  3. 实现NMS的函数。

下面是一个完整的C++代码示例:

完整的C++代码实现NMS

#include <algorithm> // std::max, std::min, std::sort
#include <vector>    // std::vector
#include <iostream>  // std::cout

struct Box {
    float x, y, w, h; // x, y 为左上角坐标,w, h 为宽度和高度
    float confidence; // 置信度
};

// 计算两个边界框的 IoU
float iou(const Box& box1, const Box& box2) {
    float xi1 = std::max(box1.x, box2.x);
    float yi1 = std::max(box1.y, box2.y);
    float xi2 = std::min(box1.x + box1.w, box2.x + box2.w);
    float yi2 = std::min(box1.y + box1.h, box2.y + box2.h);
    
    float inter_width = std::max(0.0f, xi2 - xi1);
    float inter_height = std::max(0.0f, yi2 - yi1);
    float inter_area = inter_width * inter_height;
    
    float box1_area = box1.w * box1.h;
    float box2_area = box2.w * box2.h;
    
    float union_area = box1_area + box2_area - inter_area;
    
    return inter_area / union_area;
}

// 非极大值抑制 (NMS)
std::vector<Box> non_max_suppression(const std::vector<Box>& boxes, float iou_threshold) {
    // 按置信度降序排序
    std::vector<Box> sorted_boxes = boxes;
    std::sort(sorted_boxes.begin(), sorted_boxes.end(), [](const Box& a, const Box& b) {
        return a.confidence > b.confidence;
    });
    
    std::vector<Box> selected_boxes;
    
    std::vector<bool> suppressed(sorted_boxes.size(), false);
    
    for (size_t i = 0; i < sorted_boxes.size(); ++i) {
        if (suppressed[i]) continue;
        const Box& boxA = sorted_boxes[i];
        selected_boxes.push_back(boxA);
        
        for (size_t j = i + 1; j < sorted_boxes.size(); ++j) {
            if (suppressed[j]) continue;
            const Box& boxB = sorted_boxes[j];
            if (iou(boxA, boxB) > iou_threshold) {
                suppressed[j] = true;
            }
        }
    }
    
    return selected_boxes;
}

int main() {
    std::vector<Box> boxes = {
        {30, 40, 50, 60, 0.9},
        {32, 42, 48, 58, 0.8},
        {28, 38, 52, 62, 0.75},
        {100, 200, 30, 40, 0.6},
        {29, 39, 51, 61, 0.4}
    };
    
    float iou_threshold = 0.5;
    std::vector<Box> result = non_max_suppression(boxes, iou_threshold);
    
    std::cout << "Selected boxes after NMS:" << std::endl;
    for (const auto& box : result) {
        std::cout << "Box: (" << box.x << ", " << box.y << ", " << box.w << ", " << box.h 
                  << "), Confidence: " << box.confidence << std::endl;
    }
    
    return 0;
}

解释

  1. Box结构体

    • Box结构体包含左上角坐标(x, y)、宽度和高度(w, h)、以及置信度(confidence)。
  2. 计算IoU的函数

    • iou函数计算两个边界框之间的交集和并集,然后返回IoU。
  3. NMS函数

    • non_max_suppression函数首先按置信度对边界框进行排序。
    • 然后遍历排序后的边界框,如果当前边界框没有被抑制,选择它并将其添加到最终结果中。
    • 接着,计算当前边界框与其余边界框的IoU,如果IoU超过阈值,则将这些边界框标记为已抑制。
  4. 主函数测试

    • 定义了一组边界框和置信度,并调用non_max_suppression函数进行NMS处理。
    • 最后输出经过NMS处理后的边界框。

通过以上步骤,你可以在C++中实现非极大值抑制(NMS),用于目标检测中的后处理。

更多推荐