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)
while order.size > 0:
index = order[0]
res_bboxes.append(bboxes[index])
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);
}
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;
}