手撕NMS及其变种

NMS Python版

def NMS(boxes,scores, thresholds):
    x1 = boxes[:,0]
    y1 = boxes[:,1]
    x2 = boxes[:,2]
    y2 = boxes[:,3]
    areas = (x2-x1)*(y2-y1)
    _,order = scores.sort(0,descending=True)
    keep = []
    while order.numel() > 0:
        i = order[0]
        keep.append(i)
        if order.numel() == 1:
            break
        xx1 = x1[order[1:]].clamp(min=x1[i])
        yy1 = y1[order[1:]].clamp(min=y1[i])
        xx2 = x2[order[1:]].clamp(max=x2[i])
        yy2 = y2[order[1:]].clamp(max=y2[i])
        w = (xx2-xx1).clamp(min=0)
        h = (yy2-yy1).clamp(min=0)
        inter = w*h
        ovr = inter/(areas[i] + areas[order[1:]] - inter)
        ids = (ovr<=thresholds).nonzero().squeeze()
        if ids.numel() == 0:
            break
        order = order[ids+1]
    return torch.LongTensor(keep)

现有工具包：torchvision.ops.nms(boxes, scores, iou_threshold)

 多类别NMS

#实现源码
max_coordinate = boxes.max()
offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
boxes_for_nms = boxes + offsets[:, None]
keep = nms(boxes_for_nms, scores, iou_threshold)
return keep

#使用方法
torchvision.ops.boxes.batched_nms(boxes, scores, classes, nms_thresh)

 NMS C++版本

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

// 边界框结构
struct BoundingBox {
    float x, y, w, h, confidence;

    // 计算边界框的面积
    float area() const {
        return w * h;
    }
};

// 计算两个边界框之间的IoU
float iou(const BoundingBox& a, const BoundingBox& b) {
    const float area_a = a.area();
    const float area_b = b.area();

    // 计算重叠区域的坐标范围
    const float x1 = std::max(a.x, b.x);
    const float y1 = std::max(a.y, b.y);
    const float x2 = std::min(a.x + a.w, b.x + b.w);
    const float y2 = std::min(a.y + a.h, b.y + b.h);

    // 计算重叠区域的面积
    const float intersection_area = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);

    // 计算并集区域的面积
    const float union_area = area_a + area_b - intersection_area;

    // 计算IoU
    return union_area > 0 ? intersection_area / union_area : 0;
}

// 非极大值抑制函数
std::vector<BoundingBox> nms(std::vector<BoundingBox>& boxes, float threshold) {
    // 根据置信度排序
    std::sort(boxes.begin(), boxes.end(),
              [](const BoundingBox& a, const BoundingBox& b) { return a.confidence > b.confidence; });

    std::vector<BoundingBox> result;
    for (size_t i = 0; i < boxes.size(); ++i) {
        bool keep = true;
        for (size_t j = 0; j < result.size(); ++j) {
            if (iou(boxes[i], result[j]) > threshold) {
                keep = false;
                break;
            }
        }
        if (keep) {
            result.push_back(boxes[i]);
        }
    }
    return result;
}

int main() {
    std::vector<BoundingBox> boxes;

    // 添加一些示例边界框数据
    boxes.push_back({10, 10, 20, 20, 0.9});
    boxes.push_back({15, 15, 25, 25, 0.8});
    boxes.push_back({30, 30, 20, 20, 0.7});
    boxes.push_back({40, 40, 15, 15, 0.85});

    // 运行非极大值抑制算法，保留不重叠的边界框
    float threshold = 0.5;
    std::vector<BoundingBox> result = nms(boxes, threshold);

    // 输出保留下来的边界框
    std::cout << "保留的边界框：" << std::endl;
    for (const BoundingBox& box : result) {
        std::cout << "x: " << box.x << ", y: " << box.y << ", w: " << box.w << ", h: " << box.h << ", confidence: " << box.confidence << std::endl;
    }

    return 0;
}