NMS和soft-NMS

本文深入解析了Non-Maximum Suppression (NMS)算法,一种在目标检测中广泛使用的技术,旨在去除重复的边界框,提升检测精度。文章提供了NMS的C++和Python实现代码,并介绍了soft-NMS这一优化变体。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

   NMS在目标检测中非常常用,核心原理很容易理解,2017年开始相继出现了soft-NMS等等NMS的变体,优化的核心都是尽可能去除重复框的同时尽可能保证不漏检,以下是NMS的C++实现

static void sort(int n, const float* x, int* indices)  
{  
// 排序函数(降序排序),排序后进行交换的是indices中的数据  
// n:排序总数// x:带排序数// indices:初始为0~n-1数目   
  
    int i, j;  
    for (i = 0; i < n; i++)  
        for (j = i + 1; j < n; j++)  
        {  
            if (x[indices[j]] > x[indices[i]])  
            {  
                //float x_tmp = x[i];  
                int index_tmp = indices[i];  
                //x[i] = x[j];  
                indices[i] = indices[j];  
                //x[j] = x_tmp;  
                indices[j] = index_tmp;  
            }  
        }  
}


int nonMaximumSuppression(int numBoxes, const CvPoint *points,  
                          const CvPoint *oppositePoints, const float *score,  
                          float overlapThreshold,  
                          int *numBoxesOut, CvPoint **pointsOut,  
                          CvPoint **oppositePointsOut, float **scoreOut)  
{  
  
// numBoxes:窗口数目// points:窗口左上角坐标点// oppositePoints:窗口右下角坐标点  
// score:窗口得分// overlapThreshold:重叠阈值控制// numBoxesOut:输出窗口数目  
// pointsOut:输出窗口左上角坐标点// oppositePoints:输出窗口右下角坐标点  
// scoreOut:输出窗口得分  
    int i, j, index;  
    float* box_area = (float*)malloc(numBoxes * sizeof(float));    // 定义窗口面积变量并分配空间   
    int* indices = (int*)malloc(numBoxes * sizeof(int));          // 定义窗口索引并分配空间   
    int* is_suppressed = (int*)malloc(numBoxes * sizeof(int));    // 定义是否抑制表标志并分配空间   
    // 初始化indices、is_supperssed、box_area信息   
    for (i = 0; i < numBoxes; i++)  
    {  
        indices[i] = i;  
        is_suppressed[i] = 0;  
        box_area[i] = (float)( (oppositePoints[i].x - points[i].x + 1) *  
                                (oppositePoints[i].y - points[i].y + 1));  
    }  
    // 对输入窗口按照分数比值进行排序,排序后的编号放在indices中   
    sort(numBoxes, score, indices);  
    for (i = 0; i < numBoxes; i++)                // 循环所有窗口   
    {  
        if (!is_suppressed[indices[i]])           // 判断窗口是否被抑制   
        {  
            for (j = i + 1; j < numBoxes; j++)    // 循环当前窗口之后的窗口   
            {  
                if (!is_suppressed[indices[j]])   // 判断窗口是否被抑制   
                {  
                    int x1max = max(points[indices[i]].x, points[indices[j]].x);                     // 求两个窗口左上角x坐标最大值   
                    int x2min = min(oppositePoints[indices[i]].x, oppositePoints[indices[j]].x);     // 求两个窗口右下角x坐标最小值   
                    int y1max = max(points[indices[i]].y, points[indices[j]].y);                     // 求两个窗口左上角y坐标最大值   
                    int y2min = min(oppositePoints[indices[i]].y, oppositePoints[indices[j]].y);     // 求两个窗口右下角y坐标最小值   
                    int overlapWidth = x2min - x1max + 1;            // 计算两矩形重叠的宽度   
                    int overlapHeight = y2min - y1max + 1;           // 计算两矩形重叠的高度   
                    if (overlapWidth > 0 && overlapHeight > 0)  
                    {  
                        float overlapPart = (overlapWidth * overlapHeight) / box_area[indices[j]];    // 计算重叠的比率   
                        if (overlapPart > overlapThreshold)          // 判断重叠比率是否超过重叠阈值   
                        {  
                            is_suppressed[indices[j]] = 1;           // 将窗口j标记为抑制   
                        }  
                    }  
                }  
            }  
        }  
    }  
  
    *numBoxesOut = 0;    // 初始化输出窗口数目0   
    for (i = 0; i < numBoxes; i++)  
    {  
        if (!is_suppressed[i]) (*numBoxesOut)++;    // 统计输出窗口数目   
    }  
  
    *pointsOut = (CvPoint *)malloc((*numBoxesOut) * sizeof(CvPoint));           // 分配输出窗口左上角坐标空间   
    *oppositePointsOut = (CvPoint *)malloc((*numBoxesOut) * sizeof(CvPoint));   // 分配输出窗口右下角坐标空间   
    *scoreOut = (float *)malloc((*numBoxesOut) * sizeof(float));                // 分配输出窗口得分空间   
    index = 0;  
    for (i = 0; i < numBoxes; i++)                  // 遍历所有输入窗口   
    {  
        if (!is_suppressed[indices[i]])             // 将未发生抑制的窗口信息保存到输出信息中   
        {  
            (*pointsOut)[index].x = points[indices[i]].x;  
            (*pointsOut)[index].y = points[indices[i]].y;  
            (*oppositePointsOut)[index].x = oppositePoints[indices[i]].x;  
            (*oppositePointsOut)[index].y = oppositePoints[indices[i]].y;  
            (*scoreOut)[index] = score[indices[i]];  
            index++;  
        }  
  
    }  
  
    free(indices);          // 释放indices空间   
    free(box_area);         // 释放box_area空间   
    free(is_suppressed);    // 释放is_suppressed空间   
  
    return LATENT_SVM_OK;  
}

 

以下是python版本的实现

import numpy as np  
  
def py_cpu_nms(dets, thresh):  
    """Pure Python NMS baseline."""  
    x1 = dets[:, 0]  
    y1 = dets[:, 1]  
    x2 = dets[:, 2]  
    y2 = dets[:, 3]  
    scores = dets[:, 4]  #bbox打分
  
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)  
#打分从大到小排列,取index  
    order = scores.argsort()[::-1]  
#keep为最后保留的边框  
    keep = []  
    while order.size > 0:  
#order[0]是当前分数最大的窗口,肯定保留  
        i = order[0]  
        keep.append(i)  
#计算窗口i与其他所有窗口的交叠部分的面积
        xx1 = np.maximum(x1[i], x1[order[1:]])  
        yy1 = np.maximum(y1[i], y1[order[1:]])  
        xx2 = np.minimum(x2[i], x2[order[1:]])  
        yy2 = np.minimum(y2[i], y2[order[1:]])  
  
        w = np.maximum(0.0, xx2 - xx1 + 1)  
        h = np.maximum(0.0, yy2 - yy1 + 1)  
        inter = w * h  
#交/并得到iou值  
        ovr = inter / (areas[i] + areas[order[1:]] - inter)  
#inds为所有与窗口i的iou值小于threshold值的窗口的index,其他窗口此次都被窗口i吸收  
        inds = np.where(ovr <= thresh)[0]  
#order里面只保留与窗口i交叠面积小于threshold的那些窗口,由于ovr长度比order长度少1(不包含i),所以inds+1对应到保留的窗口
        order = order[inds + 1]  
  
    return keep

 

 

以下是soft-NMS的具体实现,由于score每次都需要改变,因为取消了排序的过程,两次循环中,先寻找当前score的最大值,和外层循环位置进行交换,再在第二层循环中对后续的框进行score衰减

def cpu_soft_nms(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0):
    cdef unsigned int N = boxes.shape[0]
    cdef float iw, ih, box_area
    cdef float ua
    cdef int pos = 0
    cdef float maxscore = 0
    cdef int maxpos = 0
    cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov

    for i in range(N):
        maxscore = boxes[i, 4]
        maxpos = i

        tx1 = boxes[i,0]
        ty1 = boxes[i,1]
        tx2 = boxes[i,2]
        ty2 = boxes[i,3]
        ts = boxes[i,4]

        pos = i + 1
    # get max box
        while pos < N:
            if maxscore < boxes[pos, 4]:
                maxscore = boxes[pos, 4]
                maxpos = pos
            pos = pos + 1

    # add max box as a detection 
        boxes[i,0] = boxes[maxpos,0]
        boxes[i,1] = boxes[maxpos,1]
        boxes[i,2] = boxes[maxpos,2]
        boxes[i,3] = boxes[maxpos,3]
        boxes[i,4] = boxes[maxpos,4]

    # swap ith box with position of max box
        boxes[maxpos,0] = tx1
        boxes[maxpos,1] = ty1
        boxes[maxpos,2] = tx2
        boxes[maxpos,3] = ty2
        boxes[maxpos,4] = ts

        tx1 = boxes[i,0]
        ty1 = boxes[i,1]
        tx2 = boxes[i,2]
        ty2 = boxes[i,3]
        ts = boxes[i,4]

        pos = i + 1
    # NMS iterations, note that N changes if detection boxes fall below threshold
        while pos < N:
            x1 = boxes[pos, 0]
            y1 = boxes[pos, 1]
            x2 = boxes[pos, 2]
            y2 = boxes[pos, 3]
            s = boxes[pos, 4]

            area = (x2 - x1 + 1) * (y2 - y1 + 1)
            iw = (min(tx2, x2) - max(tx1, x1) + 1)
            if iw > 0:
                ih = (min(ty2, y2) - max(ty1, y1) + 1)
                if ih > 0:
                    ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih)
                    ov = iw * ih / ua #iou between max box and detection box

                    if method == 1: # linear
                        if ov > Nt: 
                            weight = 1 - ov
                        else:
                            weight = 1
                    elif method == 2: # gaussian
                        weight = np.exp(-(ov * ov)/sigma)
                    else: # original NMS
                        if ov > Nt: 
                            weight = 0
                        else:
                            weight = 1

                    boxes[pos, 4] = weight*boxes[pos, 4]

            # if box score falls below threshold, discard the box by swapping with last box
            # update N
                    if boxes[pos, 4] < threshold:
                        boxes[pos,0] = boxes[N-1, 0]
                        boxes[pos,1] = boxes[N-1, 1]
                        boxes[pos,2] = boxes[N-1, 2]
                        boxes[pos,3] = boxes[N-1, 3]
                        boxes[pos,4] = boxes[N-1, 4]
                        N = N - 1
                        pos = pos - 1

            pos = pos + 1

    keep = [i for i in range(N)]
    return keep

 

<think>嗯,用户想在YOLOv8中用Soft-NMS或DIoU-NMS替换传统的NMS,以减少重叠目标的漏检。首先,我需要回忆一下YOLOv8的结构,尤其是后处理部分。YOLOv8使用的是Ultralytics的框架,可能在后处理中调用了NMS的实现。传统的NMS是直接根据IOU阈值来抑制重叠的框,而Soft-NMSDIoU-NMS则是改进的方法,可能需要在后处理代码中进行修改。 用户需要详细的步骤代码,所以我要先确定YOLOv8中NMS的具体位置。通常,检测模型的后处理会在检测头的输出后执行NMS。在YOLOv8的代码库中,可能有一个名为non_max_suppression的函数,负责处理这个过程。我需要找到这个函数的定义位置,通常在utils/ops.py或者类似的路径下。 接下来,我需要考虑如何修改这个函数,替换成Soft-NMS或DIoU-NMS。首先,要确认现有的NMS实现是否支持替换,或者是否需要重写一个新的函数。比如,PyTorch自带的NMS函数是torchvision.ops.nms,但Soft-NMSDIoU-NMS可能需要自己实现或者找到现有的实现代码。 然后,我需要编写Soft-NMSDIoU-NMS的代码。对于Soft-NMS,传统的做法是根据IOU的值来降低其他框的分数,而不是直接移除。而DIoU-NMS则是使用DIoU距离来替代传统的IOU,考虑中心点距离,可能更适合遮挡情况。需要找到这两种方法的PyTorch实现,或者自己实现。 另外,YOLOv8的non_max_suppression函数可能有一些参数需要调整,比如是否使用传统的NMS或者其他的变种。可能需要添加参数来控制使用哪种NMS方法,或者在函数内部进行条件判断。 用户可能需要修改调用NMS的地方,比如在模型预测后的后处理步骤中,传入不同的参数来选择Soft-NMS或DIoU-NMS。同时,还需要注意输入输出的格式是否兼容,比如分数、框的坐标格式等。 此外,还需要测试修改后的NMS效果,确保没有引入错误,并且实际减少了重叠目标的漏检。可能还需要调整相关的超参数,比如Soft-NMS中的sigma或者DIoU-NMS中的beta值。 现在,我需要整理这些步骤,给出具体的代码修改示例。比如,找到non_max_suppression函数,替换其中的NMS实现,添加Soft-NMSDIoU-NMS的选项。然后,在模型调用时指定使用哪种NMS方法。 可能还需要处理一些兼容性问题,比如确保输入张量的设备类型(CPU/GPU)正确,以及梯度计算是否会影响NMS过程(不过在推理阶段通常不需要梯度)。另外,要注意不同版本的库之间的差异,比如torchvision的版本是否支持某些函数。 最后,需要给出用户如何验证修改是否生效的方法,比如运行检测并观察重叠目标的检测情况,或者使用评估指标来量化改进效果。</think>在YOLOv8中替换传统NMSSoft-NMS或DIoU-NMS,需要修改后处理阶段的非极大值抑制实现。以下是分步骤解决方案代码实现: --- ### 一、修改原理 1. YOLOv8默认使用`torchvision.ops.batched_nms`实现传统NMS 2. 需要重写`non_max_suppression`函数,加入新的NMS方法选择逻辑 3. 主要修改文件:`utils/ops.py`中的`non_max_suppression`函数 --- ### 二、具体步骤 & 代码实现 #### 步骤1:添加DIoU计算函数 ```python # 在utils/metrics.py中添加 def box_diou(box1, box2): """ 计算DIoU box1: (N, 4) [x1,y1,x2,y2] box2: (M, 4) """ # 计算交并比 inter = (torch.min(box1[:, None, 2], box2[:, 2]) - torch.max(box1[:, None, 0], box2[:, 0])).clamp(0) * \ (torch.min(box1[:, None, 3], box2[:, 3]) - torch.max(box1[:, None, 1], box2[:, 1])).clamp(0) area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1]) area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1]) union = area1[:, None] + area2 - inter iou = inter / union # 计算中心点距离 center1 = torch.stack([(box1[:, 0] + box1[:, 2])/2, (box1[:, 1] + box1[:, 3])/2], dim=1) center2 = torch.stack([(box2[:, 0] + box2[:, 2])/2, (box2[:, 1] + box2[:, 3])/2], dim=1) center_dist = torch.cdist(center1, center2, p=2) # 欧氏距离 # 计算最小闭包矩形对角线长度 enclose_wh = torch.stack([ torch.max(box1[:, None, 2], box2[:, 2]) - torch.min(box1[:, None, 0], box2[:, 0]), torch.max(box1[:, None, 3], box2[:, 3]) - torch.min(box1[:, None, 1], box2[:, 1]) ], dim=2) enclose_diag = torch.norm(enclose_wh, p=2, dim=2) # DIoU计算 diou = iou - (center_dist ** 2) / (enclose_diag ** 2 + 1e-7) return diou ``` #### 步骤2:实现Soft-NMSDIoU-NMS ```python # 修改utils/ops.py中的non_max_suppression函数 def non_max_suppression( prediction, nms_method='standard', # 添加方法选择参数 ... # 其他原有参数 ): # ... 原有代码保持不变 # 替换NMS核心代码部分 if nms_method == 'standard': # 原版NMS i = torchvision.ops.nms(boxes, scores, iou_thres) elif nms_method == 'soft': # Soft-NMS实现 i = soft_nms(boxes, scores, iou_thres, sigma=0.5) elif nms_method == 'diou': # DIoU-NMS实现 i = diou_nms(boxes, scores, iou_thres) # ... 后续代码保持不变 # 新增Soft-NMS实现 def soft_nms(boxes, scores, iou_thres, sigma=0.5): """ Soft-NMS实现 boxes: (N, 4) [x1,y1,x2,y2] scores: (N,) """ keep = [] idxs = scores.argsort(descending=True) while idxs.numel() > 0: # 保留当前最高分框 keep.append(idxs[0]) if idxs.numel() == 1: break # 计算IoU ious = torchvision.ops.box_iou(boxes[idxs[0]].unsqueeze(0), boxes[idxs[1:]]).squeeze() # 分数衰减 scores[idxs[1:]] *= torch.exp(-(ious ** 2) / sigma) # 过滤低分框 mask = scores[idxs[1:]] > 0.01 # 保留分数阈值 idxs = idxs[1:][mask] return torch.tensor(keep, device=boxes.device) # 新增DIoU-NMS实现 def diou_nms(boxes, scores, iou_thres): """ DIoU-NMS实现 boxes: (N, 4) [x1,y1,x2,y2] scores: (N,) """ keep = [] idxs = scores.argsort(descending=True) while idxs.numel() > 0: keep.append(idxs[0]) if idxs.numel() == 1: break # 计算DIoU dious = box_diou(boxes[idxs[0]].unsqueeze(0), boxes[idxs[1:]]) # DIoU阈值过滤 mask = dious < iou_thres idxs = idxs[1:][mask.squeeze()] return torch.tensor(keep, device=boxes.device) ``` #### 步骤3:调用时指定NMS方法 ```python from ultralytics import YOLO model = YOLO('yolov8n.pt') results = model.predict( source='path/to/image.jpg', nms_method='diou', # 指定NMS方法 iou_thres=0.5 # 调整阈值 ) ``` --- ### 三、关键参数调整建议 1. **Soft-NMS参数**: - `sigma`:控制分数衰减速度(默认0.5) - 建议调整范围:0.3-0.7 2. **DIoU-NMS参数**: - `iou_thres`:可适当降低(0.4-0.6) - 结合DIoU特性,可增加中心距离权重 --- ### 四、验证效果 ```python # 验证不同NMS方法的mAP from ultralytics.yolo.engine.validator import BaseValidator class CustomValidator(BaseValidator): def __init__(self, nms_method='standard'): super().__init__() self.nms_method = nms_method def _get_dataloader(self): # 原有代码保持不变... def _process_batch(self, detections, gt_bboxes): # 修改NMS方法调用 detections = non_max_suppression( detections, nms_method=self.nms_method, ... ) # 后续评估代码保持不变... # 测试不同NMS方法 for method in ['standard', 'soft', 'diou']: validator = CustomValidator(nms_method=method) metrics = validator(model) print(f"{method} mAP50: {metrics.box.map50}") ``` --- ### 五、注意事项 1. 训练阶段不需要修改NMS,仅影响推理阶段 2. 不同场景适用性: - **Soft-NMS**:适合密集小目标场景 - **DIoU-NMS**:适合遮挡严重的场景 3. 计算开销排序:`DIoU-NMS > Soft-NMS > Standard NMS` 4. 建议配合调整`conf_thres`(置信度阈值)进行优化 通过上述修改,可以有效改善重叠目标的漏检问题。实际应用中需要根据具体场景进行参数调优。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值