Journal Article

NMS算法详解:从原理到C++/Java实现

深入讲解NMS(非极大值抑制)算法原理,附带C++和Java完整实现代码,适合目标检测开发者阅读。

8 min2 views

什么是NMS

NMS(Non-Maximum Suppression,非极大值抑制)是目标检测领域中最常用的后处理算法之一。它的作用是消除冗余的检测框,保留最准确的检测结果。

算法原理

  1. 对所有检测框按置信度分数从高到低排序
  2. 选择置信度最高的检测框,将其加入最终结果
  3. 计算该框与其余所有框的IoU(Intersection over Union)
  4. 删除IoU超过阈值的框
  5. 重复步骤2-4,直到所有框都被处理

IoU计算

IoU(Intersection over Union)是两个框的交集面积与并集面积的比值:

code
IoU = 交集面积 / 并集面积

C++ 实现

cpp
#include <vector>
#include <algorithm>

struct BBox {
    float x1, y1, x2, y2;
    float score;
};

float computeIoU(const BBox& a, const BBox& b) {
    float interX1 = std::max(a.x1, b.x1);
    float interY1 = std::max(a.y1, b.y1);
    float interX2 = std::min(a.x2, b.x2);
    float interY2 = std::min(a.y2, b.y2);
    
    float interArea = std::max(0.0f, interX2 - interX1) * std::max(0.0f, interY2 - interY1);
    float areaA = (a.x2 - a.x1) * (a.y2 - a.y1);
    float areaB = (b.x2 - b.x1) * (b.y2 - b.y1);
    
    return interArea / (areaA + areaB - interArea);
}

std::vector<BBox> nms(std::vector<BBox>& boxes, float threshold) {
    std::sort(boxes.begin(), boxes.end(), [](const BBox& a, const BBox& b) {
        return a.score > b.score;
    });
    
    std::vector<BBox> result;
    std::vector<bool> suppressed(boxes.size(), false);
    
    for (size_t i = 0; i < boxes.size(); ++i) {
        if (suppressed[i]) continue;
        result.push_back(boxes[i]);
        for (size_t j = i + 1; j < boxes.size(); ++j) {
            if (suppressed[j]) continue;
            float iou = computeIoU(boxes[i], boxes[j]);
            if (iou > threshold) {
                suppressed[j] = true;
            }
        }
    }
    return result;
}

Java 实现

java
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

public class NMS {
    static class BBox {
        float x1, y1, x2, y2;
        float score;
        public BBox(float x1, float y1, float x2, float y2, float score) {
            this.x1 = x1; this.y1 = y1;
            this.x2 = x2; this.y2 = y2;
            this.score = score;
        }
    }
    
    public static float computeIoU(BBox a, BBox b) {
        float interX1 = Math.max(a.x1, b.x1);
        float interY1 = Math.max(a.y1, b.y1);
        float interX2 = Math.min(a.x2, b.x2);
        float interY2 = Math.min(a.y2, b.y2);
        float interArea = Math.max(0, interX2 - interX1) * Math.max(0, interY2 - interY1);
        float areaA = (a.x2 - a.x1) * (a.y2 - a.y1);
        float areaB = (b.x2 - b.x1) * (b.y2 - b.y1);
        return interArea / (areaA + areaB - interArea);
    }
    
    public static List<BBox> nms(List<BBox> boxes, float threshold) {
        Collections.sort(boxes, (a, b) -> Float.compare(b.score, a.score));
        List<BBox> result = new ArrayList<>();
        boolean[] suppressed = new boolean[boxes.size()];
        for (int i = 0; i < boxes.size(); i++) {
            if (suppressed[i]) continue;
            result.add(boxes.get(i));
            for (int j = i + 1; j < boxes.size(); j++) {
                if (suppressed[j]) continue;
                float iou = computeIoU(boxes.get(i), boxes.get(j));
                if (iou > threshold) suppressed[j] = true;
            }
        }
        return result;
    }
}

应用场景

  • YOLO系列:YOLOv3、YOLOv4、YOLOv5等
  • SSD:Single Shot MultiBox Detector
  • Faster R-CNN:Region-based CNN
  • 目标跟踪:多目标跟踪中的检测框筛选

参数选择

  • IoU阈值:通常设置为0.5或0.3
  • 置信度阈值:根据实际需求调整

总结

NMS是目标检测流程中不可或缺的一环,理解其原理并能够手动实现对于深入理解目标检测算法非常有帮助。