目标检测之单类别NMS

小女孩真可爱2024-07-19 15:41

long time no see!

在目标检测中，常见的是多类别NMS，也就是只对相同类别的boxes来计算IOU；但现实场景中经常遇到同一个物体被识别成2个类别，也就是模型认为它既是类别1也是类别2.这时候通过多类别nms就过滤不掉这种重叠的框。所以就需要进行单类别NMS：即把所有的boxes都认为是一个类别，然后再计算IOU来过滤。

这个函数的三个输入参数分别是：模型检测得到的框（x,y,w,h）、每个框的得分、nms阈值

python 复制代码

def oneclass_nms(boxes, class_probs, nms_threshold):


    def get_iou(box1, box2):
        """
        计算两个边界框的IOU
        :param box1: 第一个边界框，格式为 [x1, y1, x2, y2]
        :param box2: 第二个边界框，格式为 [x1, y1, x2, y2]
        :return: IOU的值
        """
        x11, y11, x12, y12 = box1
        x21, y21, x22, y22 = box2

        # 计算边界框的交集
        inter_x1 = max(x11, x21)
        inter_y1 = max(y11, y21)
        inter_x2 = min(x12, x22)
        inter_y2 = min(y12, y22)

        # 计算交集面积
        inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

        # 计算边界框的总面积
        box1_area = (x12 - x11) * (y12 - y11)
        box2_area = (x22 - x21) * (y22 - y21)

        # 计算并集面积
        union_area = box1_area + box2_area - inter_area

        # 计算IOU
        iou = inter_area / union_area
        return iou

    # 初始化一个空列表来存储保留的边界
    boxes_list = copy.deepcopy(boxes.tolist())
    boxes_list_copy = copy.deepcopy(boxes.tolist())
    box_save = set()

    while boxes_list:
        box_a = boxes_list.pop(0)
        for box_b in boxes_list:
            if get_iou(box_a, box_b) > 0.1:
                box_save.add(boxes_list_copy.index(box_a))

    all_index = set(list(range(len(boxes_list_copy))))
    # 获取all_index中不在keep中的索引
    diff = all_index - box_save
    diff = list(diff)
    diff = sorted(diff, key=lambda x: x)

    return diff

在官方的代码中已经有boxes, class_probs, nms_threshold这三个参数的输出，我们只需把它传入上面的函数就可以了。在官方yolo的基础上修改代码如下（注释掉的是官方原始的代码）

在non_max_suppression这个函数里插入我们的单类别nms函数即可。把官方的nms注释掉换成自定义的nms就OK了