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

本文主要是介绍面试题-手撕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;
}