【NMS，非极大值抑制】Python和C++的实现

1、Python实现

• 代码如下：

def nms(bboxes, threshold):
    """

    :param bboxes: [N, 5]的ndarray，左上角和右下角的坐标以及置信度分数score
    :param threshold: IoU阈值
    :return:
    """
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    score = bboxes[:, 4]
    order = score.argsort()[::-1]
    res_bboxes = []

    area = (x2 - x1) * (y2 - y1)
	# 注意获取order的元素个数只能用order.size或者np.size(size)，不能用order.size()
    while order.size > 0:  
        index = order[0]
        res_bboxes.append(bboxes[index])

        # 计算出第一个bbox与其余bboxes的相交区域的坐标
        x11 = np.maximum(x1[index], x1[order[1:]])
        y11 = np.maximum(y1[index], y1[order[1:]])
        x22 = np.minimum(x2[index], x2[order[1:]])
        y22 = np.minimum(y2[index], y2[order[1:]])

        w = np.maximum(0.0, x22 - x11)
        h = np.maximum(0.0, y22 - y11)
        inter = w * h
        iou = inter / (area[index] + area[order[1:]] - inter)

        retain_id = np.where(iou <= threshold)[0]
        order = order[retain_id + 1]

    return res_bboxes


if __name__ == '__main__':
    bounding = np.array([[187, 82, 337, 317, 0.9], [150, 67, 305, 282, 0.75], [246, 121, 368, 304, 0.8]])
    threshold = 0.4
    picked_boxes = nms(bounding, threshold)
    print('阈值threshold为:', threshold)
    print('最终bbox列表：', picked_boxes)

2、C++实现

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

// 定义矩形框结构体
struct BoundingBox {
    float x1, y1, x2, y2;
    float score;
};

bool compareBbox(const BoundingBox& bbox1, const BoundingBox& bbox2) {
    return bbox1.score > bbox2.score;
}

// 计算矩形框的面积
float bboxArea(const BoundingBox& bbox) {
    return (bbox.x2 - bbox.x1) * (bbox.y2 - bbox.y1);
}

// 计算两个矩形框的交并比（Intersection over Union，IoU）
float bboxIoU(const BoundingBox& bbox1, const BoundingBox& bbox2) {
    // 计算两个矩形框的相交部分的坐标
    float x1 = max(bbox1.x1, bbox2.x1);
    float y1 = max(bbox1.y1, bbox2.y1);
    float x2 = min(bbox1.x2, bbox2.x2);
    float y2 = min(bbox1.y2, bbox2.y2);

    // 计算相交部分的面积
    float interArea = max(0.0f, x2 - x1) * max(0.0f, y2 - y1);

    // 计算并集的面积
    float unionArea = bboxArea(bbox1) + bboxArea(bbox2) - interArea;

    // 计算交并比
    return interArea / unionArea;
}

// 非极大值抑制
vector<BoundingBox> nms(const vector<BoundingBox>& bboxes, float threshold) {
    vector<BoundingBox> selected;

    // 按照矩形框的得分进行排序
    vector<BoundingBox> sortedBboxes(bboxes);
    sort(sortedBboxes.begin(), sortedBboxes.end(), compareBbox);

    // 遍历排序后的矩形框
    while (!sortedBboxes.empty()) {
        const BoundingBox& current = sortedBboxes[0];
        selected.push_back(current);

        // 从列表中删除当前矩形框
        sortedBboxes.erase(sortedBboxes.begin());

        // 计算当前矩形框与其他剩余矩形框的交并比
        vector<BoundingBox>::iterator iter = sortedBboxes.begin();
        while (iter != sortedBboxes.end()) {
            const BoundingBox& bbox = *iter;
            float iou = bboxIoU(current, bbox);

            // 如果交并比大于阈值，则删除该矩形框
            if (iou > threshold) {
                iter = sortedBboxes.erase(iter);
            } else {
                ++iter;
            }
        }
    }

    return selected;
}

int main() {
    // 示例使用的原始矩形框列表
    vector<BoundingBox> bboxes = {
        {10, 20, 50, 60, 0.9f},
        {30, 35, 70, 80, 0.8f},
        {15, 25, 45, 55, 0.7f},
        {55, 60, 90, 100, 0.95f},
        {80, 85, 120, 130, 0.75f}
    };

    // 设置非极大值抑制的阈值
    float threshold = 0.5f;

    // 应用非极大值抑制
    vector<BoundingBox> selectedBboxes = nms(bboxes, threshold);

    // 输出保留下来的矩形框
    for (const BoundingBox& bbox : selectedBboxes) {
        cout << "Selected BBox: (" << bbox.x1 << ", " << bbox.y1 << ", " << bbox.x2 << ", " << bbox.y2 << "), Score: " << bbox.score << endl;
    }

    return 0;
}