NMS和Soft-NMS的原理和Pytorch代码实现

5 篇文章 0 订阅

NMS都不会,做什么Detection!

Non-maximum suppression(非极大值抑制)算法

NMS原理:

  1. 首先得出所有的预测框集合B、 对应框的得分Scores, NMS(IoU)阈值T;

  2. 定义存放侯选框的集合H(初始为Null), 对Scores排序选出得分最大的框为maxBox
    maxBox从集合B中移到集合H中,集合B中没有maxBox框了;

  3. 计算maxBoxB中剩余的所有框的IoU, 将IoU大于T的从B中删除(认为和maxBox重叠了);

  4. 重复2~3步骤,直到集合BNull, 集合H中存放的框就是NMS处理的结果;

    重复步骤是:
    (1)对集合B中剩余框对应的得分进行排序, 选出最大得分的框maxBox,并从集合B中移到集合H中。

    (2) 计算这个得分最大的框maxBox和集合B中框的IoU阈值,将大于IoU阈值的框从B中删除。

NMS代码实现

1. Pytorch代码实现
from torch import Tensor
import torch


def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by its
    (x1, y1, x2, y2) coordinates.

    Arguments:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format

    Returns:
        area (Tensor[N]): area for each box
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """
    Return intersection-over-union (Jaccard index) of boxes.

    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.

    Arguments:
        boxes1 (Tensor[N, 4])
        boxes2 (Tensor[M, 4])

    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
    """
    area1 = box_area(boxes1)  # 每个框的面积 (N,)
    area2 = box_area(boxes2)  # (M,)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] # N中一个和M个比较; 所以由N,M 个
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]  #小于0的为0  clamp 钳;夹钳;
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]  

    iou = inter / (area1[:, None] + area2 - inter)
    return iou  # NxM, boxes1中每个框和boxes2中每个框的IoU值;


def nms(boxes: Tensor, scores: Tensor, iou_threshold: float):
    """
    :param boxes: [N, 4], 此处传进来的框,是经过筛选(NMS之前选取过得分TopK)之后, 在传入之前处理好的;
    :param scores: [N]
    :param iou_threshold: 0.7
    :return:
    """
    keep = []  # 最终保留的结果, 在boxes中对应的索引;
    idxs = scores.argsort()  # 值从小到大的 索引

    while idxs.numel() > 0:  # 循环直到null; numel(): 数组元素个数
        # 得分最大框对应的索引, 以及对应的坐标
        max_score_index = idxs[-1]
        max_score_box = boxes[max_score_index][None, :]  # [1, 4]
        keep.append(max_score_index)
        if idxs.size(0) == 1:  # 就剩余一个框了;
            break
        idxs = idxs[:-1]  # 将得分最大框 从索引中删除; 剩余索引对应的框 和 得分最大框 计算IoU;
        other_boxes = boxes[idxs]  # [?, 4]
        ious = box_iou(max_score_box, other_boxes)  # 一个框和其余框比较 1XM
        idxs = idxs[ious[0] <= iou_threshold]

    keep = idxs.new(keep)  # Tensor
    return keep
2. Pytorch代码实现
import torch
def nms(boxes, scores, overlap=0.7, top_k=200):
    """
    输入:
        boxes: 存储一个图片的所有预测框。[num_positive,4].
        scores:置信度。如果为多分类则需要将nms函数套在一个循环内。[num_positive].
        overlap: nms抑制时iou的阈值.
        top_k: 先选取置信度前top_k个框再进行nms.
    返回:
        nms后剩余预测框的索引.
    """

    keep = scores.new(scores.size(0)).zero_().long()
    # 保存留下来的box的索引 [num_positive]
    # 函数new(): 构建一个有相同数据类型的tensor

    # 如果输入box为空则返回空Tensor
    if boxes.numel() == 0:
        return keep

    x1 = boxes[:, 0]  # x1 坐标
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = torch.mul(x2 - x1, y2 - y1)  # 并行化计算所有框的面积
    v, idx = scores.sort(0)  # 升序排序
    idx = idx[-top_k:]  # 前top-k的索引,从小到大
    xx1 = boxes.new()
    yy1 = boxes.new()
    xx2 = boxes.new()  # new() 无参数,创建 相同类型的空值;
    yy2 = boxes.new()
    w = boxes.new()
    h = boxes.new()

    count = 0
    while idx.numel() > 0:
        i = idx[-1]  # 目前最大score对应的索引  # 选取得分最大的框索引;
        keep[count] = i  # 存储在keep中
        count += 1
        if idx.size(0) == 1:  # 跳出循环条件:box被筛选完了
            break
        idx = idx[:-1]  # 去掉最后一个

        # 剩下boxes的信息存储在xx,yy中
        torch.index_select(x1, 0, idx, out=xx1)  # 从x1中再维度0选取索引为idx 数据 输出到xx1中;
        torch.index_select(y1, 0, idx, out=yy1)  # torch.index_select() # 从tensor中按指定维度和索引 取值;
        torch.index_select(x2, 0, idx, out=xx2)
        torch.index_select(y2, 0, idx, out=yy2)

        # 计算当前最大置信框与其他剩余框的交集,不知道clamp的同学确实容易被误导
        xx1 = torch.clamp(xx1, min=x1[i])  # max(x1,xx1)  # x1 y1 的最大值
        yy1 = torch.clamp(yy1, min=y1[i])  # max(y1,yy1)
        xx2 = torch.clamp(xx2, max=x2[i])  # min(x2,xx2)  # x2 x3 最小值;
        yy2 = torch.clamp(yy2, max=y2[i])  # min(y2,yy2)
        w.resize_as_(xx2)
        h.resize_as_(yy2)
        w = xx2 - xx1  # w=min(x2,xx2)−max(x1,xx1)
        h = yy2 - yy1  # h=min(y2,yy2)−max(y1,yy1)
        w = torch.clamp(w, min=0.0)  # max(w,0)
        h = torch.clamp(h, min=0.0)  # max(h,0)
        inter = w * h

        # 计算当前最大置信框与其他剩余框的IOU
        # IoU = i / (area(a) + area(b) - i)
        rem_areas = torch.index_select(area, 0, idx)  # 剩余的框的面积
        union = rem_areas + area[i] - inter  # 并集
        IoU = inter / union  # 计算iou

        # 选出IoU <= overlap的boxes(注意le函数的使用)
        idx = idx[IoU.le(overlap)]  # le: 小于等于 返回的bool , 去除大于overlap的值;
    return keep, count

参考自:链接

3. Numpy代码实现
import numpy as np
from numpy import array

def box_area(boxes :array):
    """
    :param boxes: [N, 4]
    :return: [N]
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])

def box_iou(box1 :array, box2: array):
    """
    :param box1: [N, 4]
    :param box2: [M, 4]
    :return: [N, M]
    """
    area1 = box_area(box1)  # N
    area2 = box_area(box2)  # M
    # broadcasting, 两个数组各维度大小 从后往前对比一致, 或者 有一维度值为1;
    lt = np.maximum(box1[:, np.newaxis, :2], box2[:, :2])
    rb = np.minimum(box1[:, np.newaxis, 2:], box2[:, 2:])
    wh = rb - lt
    wh = np.maximum(0, wh) # [N, M, 2]
    inter = wh[:, :, 0] * wh[:, :, 1]
    iou = inter / (area1[:, np.newaxis] + area2 - inter)
    return iou  # NxM

def numpy_nms(boxes :array, scores :array, iou_threshold :float):

    idxs = scores.argsort()  # 按分数 降序排列的索引 [N]
    keep = []
    while idxs.size > 0:  # 统计数组中元素的个数
        max_score_index = idxs[-1]
        max_score_box = boxes[max_score_index][None, :]
        keep.append(max_score_index)

        if idxs.size == 1:
            break
        idxs = idxs[:-1]  # 将得分最大框 从索引中删除; 剩余索引对应的框 和 得分最大框 计算IoU;
        other_boxes = boxes[idxs]  # [?, 4]
        ious = box_iou(max_score_box, other_boxes)  # 一个框和其余框比较 1XM
        idxs = idxs[ious[0] <= iou_threshold]

    keep = np.array(keep)  # Tensor
    return keep

Soft NMS原理和代码

Soft-NMS的原理

原理步骤:

  1. 首先得出所有的预测框集合B、 对应框的得分Scores, NMS(IoU)阈值T;

  2. 定义存放侯选框的集合H(初始为Null),对Scores排序选出得分最大的框为maxBox, 将maxBox从集合B中移到集合H中,集合B中没有maxBox框了;

  3. 计算maxBoxB中剩余的所有框的IoU, 将IoU大于T的框的得分 按某种方式 降低(不删除了),

  4. 重复2~3步骤,直到集合BNull, 集合H中存放的框就是Soft-NMS处理的结果;
    重复步骤是:
    (1)对集合B中剩余框对应的得分进行排序(因为分数变化,必须排序), 选出最大得分的框maxBox,并从集合B中移到集合H中。

    (2) 计算这个得分最大的框maxBox和集合B中框的IoU阈值,将大于IoU阈值的框对应的得分降低。

  5. Soft-NMS返回的结果是 框以及框对应的得分(得分是Soft-NMS抑制后的),说白了,就是抑制了框对应的得分, 使用时需要一个得分阈值。
    在这里插入图片描述

Soft-NMS权重函数的形式

M M M表示得分最大框, b i b_i bi是除去得分最大框后剩余的每个框。
原来的NMS可以描述如下:将IoU大于阈值的窗口的得分全部置为0。
s i = { s i , i o u ( M , b i ) < T 0 , i o u ( M , b i ) > = T s_i = \left\{ \begin{array}{lr}s_i, iou(M,b_i) < T \\ 0, iou(M,b_i) >= T \end{array} \right. si={si,iou(M,bi)<T0,iou(M,bi)>=T

(1) 线性加权抑制得分
s i = { s i , i o u ( M , b i ) < T s i ( 1 − i o u ( M , b i ) ) , i o u ( M , b i ) > = T s_i = \left\{ \begin{array}{lr}s_i, iou(M,b_i) < T \\ s_i(1 - iou(M, b_i)), iou(M,b_i) > =T \end{array} \right. si={si,iou(M,bi)<Tsi(1iou(M,bi)),iou(M,bi)>=T
(2) 高斯加权抑制得分
s i = s i e − i o u ( M , b i ) 2 σ , ∀ b i ∉ D s_i = s_ie^{-\frac{iou(M, b_i)^2}{\sigma}}, \forall b_i \notin D si=sieσiou(M,bi)2,bi/D

NMS-Soft 实现代码

Pytorch 代码 未验证 没有合适比较代码【欢迎指正错误】

from torch import Tensor
import torch


def box_area(boxes: Tensor) -> Tensor:
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    area1 = box_area(boxes1)  # 每个框的面积 (N,)
    area2 = box_area(boxes2)  # (M,)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] # N中一个和M个比较; 所以由N,M 个
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]  # 删除面积小于0 不相交的  clamp 钳;夹钳;
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]  # 切片的用法 相乘维度减1
    iou = inter / (area1[:, None] + area2 - inter)
    return iou  # NxM, boxes1中每个框和boxes2中每个框的IoU值;


def soft_nms(boxes: Tensor, scores: Tensor, soft_threshold=0.01,  iou_threshold=0.7, weight_method=2, sigma=0.5):
    """
    :param boxes: [N, 4], 此处传进来的框,是经过筛选(选取的得分TopK)之后的
    :param scores: [N]
    :param iou_threshold: 0.7
    :param soft_threshold soft nms 过滤掉得分太低的框 (手动设置)
    :param weight_method 权重方法 1. 线性 2. 高斯
    :return:
    """
    keep = []
    idxs = scores.argsort()
    while idxs.numel() > 0:  # 循环直到null; numel(): 数组元素个数
        # 由于scores得分会改变,所以每次都要重新排序,获取得分最大值
        idxs = scores.argsort()  # 评分排序
        if idxs.size(0) == 1:  # 就剩余一个框了;
            keep.append(idxs[-1])  
            break
        keep_len = len(keep)
        max_score_index = idxs[-(keep_len + 1)]
        max_score_box = boxes[max_score_index][None, :]  # [1, 4]
        idxs = idxs[:-(keep_len + 1)]
        other_boxes = boxes[idxs]  # [?, 4]
        keep.append(max_score_index)  # 位置不能边
        ious = box_iou(max_score_box, other_boxes)  # 一个框和其余框比较 1XM
        # Soft NMS 处理, 和 得分最大框 IOU大于阈值的框, 进行得分抑制
        if weight_method == 1:   # 线性抑制  # 整个过程 只修改分数
            ge_threshod_bool = ious[0] >= iou_threshold
            ge_threshod_idxs = idxs[ge_threshod_bool]
            scores[ge_threshod_idxs] *= (1. - ious[0][ge_threshod_bool])  # 小于IoU阈值的不变
            # idxs = idxs[scores[idxs] >= soft_threshold]  # 小于soft_threshold删除, 经过抑制后 阈值会越来越小;
        elif weight_method == 2:  # 高斯抑制, 不管大不大于阈值,都计算权重
            scores[idxs] *= torch.exp(-(ious[0] * ious[0]) / sigma) # 权重(0, 1]
            # idxs = idxs[scores[idxs] >= soft_threshold]
        # else:  # NMS
        #     idxs = idxs[ious[0] <= iou_threshold]

    # keep = scores[scores > soft_threshold].int()
    keep = idxs.new(keep)  # Tensor
    keep = keep[scores[keep] > soft_threshold]  # 最后处理阈值
    boxes = boxes[keep]  # 保留下来的框
    scores = scores[keep]  # soft nms抑制后得分
    return boxes, scores

对应代码和测试在NMS_SoftNMS项目,欢迎指正!

  • 10
    点赞
  • 41
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
### 回答1: soft-nms pytorch代码是一个基于PyTorch框架实现的软非极大值抑制算法的源代码。该算法可以在目标检测中用于抑制重叠的检测框,从而提高检测的准确性和效率。该源代码提供了一个简单易用的接口,可以方便地集成到现有的目标检测系统中。 ### 回答2: Soft-NMS是一种针对目标检测中非极大值抑制(Non-Maximum Suppression, NMS)的改进算法。Soft-NMS的目的是尽可能保留更多的检测结果,其中,保留哪些检测结果以及去除哪些检测结果是通过计算每个检测框和其他检测框的IoU(Intersection over Union)来决定的。 PyTorch是一个流行的深度学习框架,在PyTorch实现Soft-NMS非常简单。下面是PyTorch官方提供的Soft-NMS代码: ```python def soft_nms(dets, sigma=0.5, Nt=0.3, threshold=0.001, method=2): """ :param dets: [[x1, y1, x2, y2, score], ...] :param sigma: :param Nt: :param threshold: :param method: :return: indexes to keep """ N = dets.shape[0] indexes = np.arange(N) for i in range(N): pos = i + 1 if i != N - 1: maximum = np.argmax(dets[pos:, 4]) if dets[i, 4] < dets[maximum + pos, 4]: dets[[i, maximum + pos], :] = dets[[maximum + pos, i], :] indexes[[i, maximum + pos]] = indexes[[maximum + pos, i]] t_bbox, t_score = dets[i, :4], dets[i, 4] if method == 1: weight = np.exp(-(np.power((dets[i + 1:, :4] - t_bbox), 2).sum(1) / sigma)) else: weight = np.zeros((N - i - 1,)) _idx = np.where(dets[i + 1:, 4] >= Nt)[0] + i + 1 if len(_idx) > 0: ex_dets = dets[_idx, :] overlap = iou(t_bbox.reshape(1, -1), ex_dets[:, :4]) weight[_idx - (i + 1)] = overlap weight = np.exp(-(weight * weight) / sigma) dets[i + 1:, 4] *= weight score_mask = dets[i + 1:, 4] >= threshold indexes_mask = indexes[i + 1:][score_mask] dets[i + 1:, :] = dets[indexes_mask, :] indexes[i + 1:] = indexes[indexes_mask] if len(indexes_mask) == 0: break keep = np.zeros(N, dtype=np.intp) keep[indexes] = 1 return np.where(keep == 1)[0] ``` 这里,soft_nms函数接受一个包含检测结果的二维数组dets,其中每个元素包含四个坐标[x1, y1, x2, y2]和一个分数score。函数的其他参数包括sigma、Nt、threshold和method等。 对于每个检测框,软NMS算法首先检查其余检测框的IoU。如果另一个检测框的IoU高于给定的Nt阈值,则它的分数将被调整到小于原始分数因子*(1-IoU)。然后,score_mask掩码被用来删掉具有较低分数的检测框,直到达到需要的最终物体数或所有框都不再死亡。最后,保留下来的检测框的索引被返回给用户。 这个PyTorch代码是一个很好的例子,可以帮助我们了解如何使用Python来自定义各种优化算法,包括Soft-NMS。 ### 回答3: Soft-NMS是一种基于非极大抑制(NMS)的目标检测算法。与传统的NMS算法不同,Soft-NMS通过对重叠框之间的分数进行加权来降低框的置信度,而不是直接将框丢弃。这样可以在一定程度上保留重叠框中的信息。 PyTorch是一种流行的深度学习框架,提供了供用户使用的许多预定义算法和函数。PyTorch实现Soft-NMS的源代码,这使得用户可以直接使用这种功能而不必手动实现。 在PyTorch中,Soft-NMS实现是在torchvision.ops.nms.soft_nms中完成的。它有以下四个参数: - boxes: 包含所有检测框的tensor。 - scores: 所有检测框的分数。 - iou_threshold: 重叠IOU阈值。如果两个框之间的IOU低于此阈值,则不会执行Soft-NMS。 - score_threshold: 分数阈值。如果框的分数低于此阈值,则不会执行Soft-NMS。 整个NMS算法的流程如下: 1. 将分数从高到低排序,并将相应的边界框记录下来。 2. 选择分数最高的框,并记录下该框的索引号。 3. 计算该框与所有其他框之间的IOU值,并使用Soft-NMS算法对其进行加权。 4. 删除其余重叠框,并在剩下的框中重复步骤2-4,直到所有框都被处理。 调用soft_nms函数后,将返回保留下来框的索引列表。然后可以使用这个索引列表来关联boxes和scores,并对检测到的目标进行可视化或其他后处理操作。 在使用Soft-NMS时,需要根据具体情况调整两个阈值。如果两个阈值设置得太低,可能会误保留一些低质量的检测框。如果设置得过高,可能会丢失一些真实的目标。因此,建议在数据集上进行测量和调整。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值