YOLOv5 | NMS | 手动实现NMS

这个NMS可以按照相同的类别进行nms，不同类别是不会进行nms的，可以加一个“开关”（参数）来控制是否使用进行同类别iou，但是懒得加了哈哈哈哈 ，就是一个if语句，offset这个简直绝了，学的源码的哈哈哈哈哈哈，不信你去看

YOLOv5 | NMS | 源码解析-CSDN博客

import numpy as np
import torch
 
def nms(boxes, score, conf_thres, iou_thres, classes):
 
    tf = score[:] > conf_thres
    score = score[tf]
 
    keep = []
    idx = np.argsort(score)[::-1]
 
    pred = torch.hstack((boxes,classes))
    offset = pred[:,4:5] * 4096
    boxes = pred[:,0:4] + offset
 
    # print(boxes)
 
    while len(idx) > 0:
 
        keep.append(idx[0])
        overlap = np.zeros_like(idx[1:], dtype= np.float32)
        
        for i, j in enumerate(idx[1:]):
            bbox1 = boxes[idx[0]]
            bbox2 = boxes[j]
            out = iou(bbox1, bbox2)
            overlap[i] = out
        idx = idx[1:][overlap < iou_thres]
 
    return keep
    # box = []
    # for i in keep:
    #     box.append(boxes[i])
    # return box
 
def iou(bbox1,bbox2):
    x1, y1, w1, h1 = bbox1 # left up right bottom
    x2, y2, w2, h2 = bbox2
    if (x1 <= ((x2+ w2)/2) <= w1) or ((y1 <= ((y2+ h2)/2) <= h1)): # 必须加上等号
        left_top_x = max(x1, x2)
        left_top_y = max(y1, y2) 
        right_bottom_x = min(w1, w2)
        right_bottom_y = min(h1, h2)
 
        I = (right_bottom_x - left_top_x) * (right_bottom_y - left_top_y)
        o = (w1 - x1) * (h1- y1) + (w2 - x2) * (h2 - y2) - I
 
        IoU = I / o
 
        return IoU
    else: # 不相交
        return 0
 
if __name__ == '__main__':
    boxes = torch.Tensor([
        [100, 100, 200, 200],
        [120, 110, 220, 210],
        [300, 320, 400, 400],
        [180, 100, 300, 180]
    ])
    scores = np.array([0.9, 0.8, 0.7, 0.6])
    classes = torch.Tensor([[0],[1],[1],[2]])
    # np.concatenate 先转置 可以拼接 
    out = nms(boxes, scores, 0.5, 0.25,classes)
    print(out)

写的比较简单，不会高级的写法