NMS代码基础

参考链接

参考链接

import torch
def iou(box,boxes):
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    bosex_area = (boxes[:,2] - boxes[:,0]) * (boxes[:,3] - boxes[:,1])

    xx1 = torch.maximum(box[0],boxes[:,0])
    xx2 = torch.minimum(box[2],boxes[:,2])
    yy1 = torch.maximum(box[1],boxes[:,1])
    yy2 = torch.minimum(box[3],boxes[:,3])

    w,h = torch.maximum(torch.tensor(0),(xx2 - xx1)),torch.maximum(torch.tensor(0),(yy2 - yy1))

    area = w * h

    return area / (box_area + bosex_area - area)


def nms(boxes,thres):
    new_boxes = boxes[boxes[:,5].argsort(descending=True)]
    keep_box = []
    while len(new_boxes) > 0:
        _box = new_boxes[0]
        keep_box.append(_box)
        if len(new_boxes) > 1:
            _boxes = new_boxes[1:]
            new_boxes = _boxes[torch.where(iou(_box,_boxes) < thres)]
    
    return torch.stack(keep_box)



if __name__ == "__main__":
    box = torch.tensor([0,0,4,4])
    boxes = torch.tensor([[4,4,7,7],[2,2,6,6]])
    iou_ = iou(box,boxes)