面试题-手撕NMS(非极大值抑制)

非极大值抑制（Non-Maximum Suppression，NMS）是一种常用于目标检测和计算机视觉中的算法，用于去除重叠的边界框，保留最可能是真实目标的边界框。

其核心就是对一组检测框，找出其中属于同一个类别且分数最高的那个框，然后把和这个框的IOU值大于阈值的那些框都删掉。

在NMS中，其实用到了计算IOU的方法，可以参考：面试题-手撕IOU计算

下面是代码：

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

struct BoundingBox {
    float x1, y1, x2, y2;
    float score;
    int category;
};

bool compareScores(const BoundingBox& a, const BoundingBox& b) {
    return a.score > b.score; // 按照置信度分数降序排序
}

float intersectionArea(const BoundingBox& rect1, const BoundingBox& rect2) {
    float overlapWidth = std::max(0.0f, std::min(rect1.x2, rect2.x2) - std::max(rect1.x1, rect2.x1));
    float overlapHeight = std::max(0.0f, std::min(rect1.y2, rect2.y2) - std::max(rect1.y1, rect2.y1));

    return overlapWidth * overlapHeight;
}

float calculateIOU(const BoundingBox& rect1, const BoundingBox& rect2) {
    float area1 = (rect1.x2 - rect1.x1) * (rect1.y2 - rect1.y1);
    float area2 = (rect2.x2 - rect2.x1) * (rect2.y2 - rect2.y1);

    float intersection = intersectionArea(rect1, rect2);
    float unionArea = area1 + area2 - intersection;

    return intersection / unionArea;
}

std::vector<int> nms(const std::vector<BoundingBox>& bboxes, float threshold) {
    std::vector<int> keep;
    std::vector<bool> suppressed(bboxes.size(), false);

    std::vector<BoundingBox> sortedBBoxes = bboxes;
    std::sort(sortedBBoxes.begin(), sortedBBoxes.end(), compareScores);

    for (size_t i = 0; i < sortedBBoxes.size(); ++i) {
        if (suppressed[i]) continue;

        keep.push_back(i);

        for (size_t j = i + 1; j < sortedBBoxes.size(); ++j) {
            if (!suppressed[j] && sortedBBoxes[i].category == sortedBBoxes[j].category) {
                float iou = calculateIOU(sortedBBoxes[i], sortedBBoxes[j]);
                if (iou > threshold) {
                    suppressed[j] = true; // 标记重叠的边界框为已抑制
                }
            }
        }
    }

    return keep;
}

int main() {
    std::vector<BoundingBox> bboxes = {
        {10, 10, 50, 50, 0.9, 1},
        {20, 20, 60, 60, 0.85, 2},
        {30, 30, 70, 70, 0.95, 1},
        {40, 40, 80, 80, 0.75, 2}
    };

    float threshold = 0.5;
    std::vector<int> indices = nms(bboxes, threshold);

    std::cout << "Indices to keep after NMS: ";
    for (auto idx : indices) {
        std::cout << idx << " ";
    }
    std::cout << std::endl;

    return 0;
}