目标检测中NMS优化算法对比：Soft-NMS与DIoU-NMS实战

管老太

1. 项目概述

在计算机视觉领域，目标检测一直是最基础也最具挑战性的任务之一。作为目标检测流程中的关键后处理步骤，非极大值抑制(NMS)算法直接影响着最终检测结果的准确性和完整性。我在多个工业检测项目中深刻体会到，传统NMS算法在面对密集目标场景时的表现往往不尽如人意，这也是促使我深入研究NMS优化算法的直接原因。

本次实战对比测试基于YOLOv5框架，重点分析两种主流NMS优化算法——Soft-NMS和DIoU-NMS。通过完整的实验设计和详实的数据对比，我将带大家深入了解这两种算法的实现原理、代码适配方法以及各自的适用场景。无论你是刚入门目标检测的新手，还是正在为项目选择合适NMS算法的工程师，这篇文章都能提供实用的参考价值。

2. 核心算法原理解析

2.1 传统NMS的局限性

传统NMS的工作原理可以用一个简单的场景来理解：假设你在人群中寻找朋友，传统NMS就像是一个严格的筛选器，只要看到两个相似的人站得很近，就会直接认定其中一个是"假朋友"并将其排除。这种简单粗暴的处理方式在实际应用中会带来两个主要问题：

首先，当目标密集排列时（比如超市货架上的商品），高置信度的检测框会"霸道"地抑制掉周围真实的检测结果。我在一个药品包装检测项目中就遇到过这种情况，传统NMS导致约15%的真实药品包装被错误抑制。

其次，传统NMS只考虑IoU（交并比）指标，完全忽略了检测框之间的位置关系。这就像是用一把尺子只测量两个物体的重叠面积，而不关心它们中心点的距离。在实际场景中，这种处理方式会导致检测框定位精度下降。

2.2 Soft-NMS算法原理

Soft-NMS的核心思想可以用"温柔拒绝"来形象描述。不同于传统NMS的"一刀切"做法，Soft-NMS采用了一种更柔和的抑制策略：

当检测到重叠框时，不是直接将其置信度置零，而是根据重叠程度对其进行降权处理。具体来说，算法提供了两种降权方式：

线性加权：score = score * (1 - iou)
高斯加权：score = score * e^(-iou^2/σ)

我在实现时发现，高斯加权方式通常能获得更平滑的抑制效果，特别是在目标密集但重叠程度不一的场景中。通过调整σ参数，可以灵活控制抑制的强度，这在不同的应用场景中非常实用。

2.3 DIoU-NMS算法原理

DIoU-NMS的改进思路则更加"全面"。它不仅考虑了检测框之间的重叠面积，还引入了中心点距离和框的大小信息：

DIoU = IoU - ρ²(b,b^gt)/c²

其中ρ表示两个框中心点的欧式距离，c是最小外接矩形的对角线长度。这种设计使得DIoU-NMS能够更准确地评估检测框之间的空间关系。

在一个车辆检测项目中，我发现DIoU-NMS对于解决"停车场上相邻车辆"的检测问题特别有效。传统NMS在这里经常会把相邻车辆误认为同一辆车，而DIoU-NMS则能更好地保持各自的独立性。

3. 实验环境搭建与代码实现

3.1 实验环境配置

为了确保实验结果的可比性，我建议使用以下标准配置：

硬件：至少配备NVIDIA GTX 1660以上的GPU
软件：
- Python 3.8+
- PyTorch 1.7+
- CUDA 11.0
- YOLOv5最新稳定版

注意：不同版本的PyTorch可能会在NMS计算上存在细微差异，建议统一使用相同版本进行对比实验。

3.2 代码适配要点

在YOLOv5中实现NMS算法替换需要修改utils/general.py文件中的non_max_suppression函数。以下是关键修改步骤：

python复制# Soft-NMS实现核心代码
def soft_nms(dets, scores, iou_thres, sigma=0.5, score_threshold=0.25):
    """
    dets: 检测框坐标 [x1,y1,x2,y2]
    scores: 置信度分数
    iou_thres: IoU阈值
    sigma: 高斯函数参数
    score_threshold: 最终得分阈值
    """
    N = dets.shape[0]
    indexes = np.arange(N)
    
    for i in range(N):
        max_pos = i
        max_score = scores[i]
        
        pos = i + 1
        while pos < N:
            if scores[pos] > max_score:
                max_score = scores[pos]
                max_pos = pos
            pos += 1
            
        dets[[i, max_pos]] = dets[[max_pos, i]]
        scores[[i, max_pos]] = scores[[max_pos, i]]
        indexes[[i, max_pos]] = indexes[[max_pos, i]]
        
        pos = i + 1
        while pos < N:
            iou = bbox_iou(dets[i], dets[pos])
            # 高斯加权
            scores[pos] = scores[pos] * np.exp(-(iou * iou) / sigma)
            
            if scores[pos] < score_threshold:
                dets[[pos, N-1]] = dets[[N-1, pos]]
                scores[[pos, N-1]] = scores[[N-1, pos]]
                indexes[[pos, N-1]] = indexes[[N-1, pos]]
                N -= 1
                pos -= 1
            pos += 1
    
    keep = indexes[:N]
    return keep