torchvision.opt.nms/torchvision.ops.box_iou使用(计算nms与计算iou)

 本博客记录使用torchvision.opt.nms()函数求解nms用法及torchvision.ops.box_iou()函数求解iou用法。

 torchvision.opt.nms()参数如下(来源源码)：

Parameters
    ----------
    boxes : Tensor[N, 4])
        boxes to perform NMS on. They
        are expected to be in (x1, y1, x2, y2) format
    scores : Tensor[N]
        scores for each one of the boxes
    iou_threshold : float
        discards all overlapping
        boxes with IoU > iou_threshold

    Returns
    -------
    keep : Tensor
        int64 tensor with the indices
        of the elements that have been kept
        by NMS, sorted in decreasing order of scores
    """

torchvision.ops.box_iou()参数如下(来源源码)：

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
    """

 

用法代码如下：

import torch
import torchvision

if __name__ == '__main__':
    box = torch.tensor([[2, 3.1, 7, 5], [3, 4, 8, 4.8], [4, 4, 5.6, 7], [0.1, 0, 8, 1]])
    score = torch.tensor([0.5, 0.3, 0.2, 0.4])
    # iou值超过iou_threshold阈值则排除，否则判定不相交
    output_index = torchvision.ops.nms(boxes=box, scores=score, iou_threshold=0.1)
    print(output_index)

    box1 = torch.tensor([[2, 3.1, 7, 5], [3, 4, 8, 4.8], [4, 4, 5.6, 7]])
    box2 = torch.tensor([[2, 4, 7, 5], [3, 4, 8, 4.8], [4, 4, 5.6, 7], [0.1, 0, 8, 1]])
    '''
    boxes1 (Tensor[N, 4])
    boxes2 (Tensor[M, 4])
    return : iou (Tensor[N, M])
    '''
    iou = torchvision.ops.box_iou(box1, box2)  # 计算交叉值 
    print('IOU of bboxes:')
    print(iou)

结果：