本博客记录使用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)结果：
