Pytorch实现nms (torchvision.ops.nms)和numpy实现nms
nms理解起来很简单:
将所有的boxes按照置信度从小到大排序,然后从boxes中删除置信度最大的box 将剩下的boxes与置信度最大的box,分别计算iou,去掉iou大于阈值(iou_threshold)的boxes 重复1,2直到索引为空
pytorch实现
from torch import Tensor
import torch
def box_area ( boxes: Tensor) - > Tensor:
"""
Computes the area of a set of bounding boxes, which are specified by its
(x1, y1, x2, y2) coordinates.
Arguments:
boxes (Tensor[N, 4]): boxes for which the area will be computed. They
are expected to be in (x1, y1, x2, y2) format
Returns:
area (Tensor[N]): area for each box
"""
return ( boxes[ : , 2 ] - boxes[ : , 0 ] ) * ( boxes[ : , 3 ] - boxes[ : , 1 ] )
def box_iou ( boxes1: Tensor, boxes2: Tensor) - > Tensor:
"""
Return intersection-over-union (Jaccard index) of boxes.
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
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
"""
area1 = box_area( boxes1)
area2 = box_area( boxes2)
lt = torch. max ( boxes1[ : , None , : 2 ] , boxes2[ : , : 2 ] )
rb = torch. min ( boxes1[ : , None , 2 : ] , boxes2[ : , 2 : ] )
wh = ( rb - lt) . clamp( min = 0 )
inter = wh[ : , : , 0 ] * wh[ : , : , 1 ]
iou = inter / ( area1[ : , None ] + area2 - inter)
return iou
def nms ( boxes: Tensor, scores: Tensor, iou_threshold: float ) :
"""
:param boxes: [N, 4], 此处传进来的框,是经过筛选(NMS之前选取过得分TopK)之后, 在传入之前处理好的;
:param scores: [N]
:param iou_threshold: 0.7
:return:
"""
keep = [ ]
idxs = scores. argsort( )
while idxs. numel( ) > 0 :
max_score_index = idxs[ - 1 ]
max_score_box = boxes[ max_score_index] [ None , : ]
keep. append( max_score_index)
if idxs. size( 0 ) == 1 :
break
idxs = idxs[ : - 1 ]
other_boxes = boxes[ idxs]
ious = box_iou( max_score_box, other_boxes)
idxs = idxs[ ious[ 0 ] <= iou_threshold]
keep = idxs. new( keep)
return keep
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 ] )
output = nms( boxes= box, scores= score, iou_threshold= 0.3 )
print ( output)
numpy 实现
import numpy as np
from numpy import array
def box_area ( boxes : array) :
"""
:param boxes: [N, 4]
:return: [N]
"""
return ( boxes[ : , 2 ] - boxes[ : , 0 ] ) * ( boxes[ : , 3 ] - boxes[ : , 1 ] )
def box_iou ( box1 : array, box2: array) :
"""
:param box1: [N, 4]
:param box2: [M, 4]
:return: [N, M]
"""
area1 = box_area( box1)
area2 = box_area( box2)
lt = np. maximum( box1[ : , np. newaxis, : 2 ] , box2[ : , : 2 ] )
rb = np. minimum( box1[ : , np. newaxis, 2 : ] , box2[ : , 2 : ] )
wh = rb - lt
wh = np. maximum( 0 , wh)
inter = wh[ : , : , 0 ] * wh[ : , : , 1 ]
iou = inter / ( area1[ : , np. newaxis] + area2 - inter)
return iou
def numpy_nms ( boxes : array, scores : array, iou_threshold : float ) :
idxs = scores. argsort( )
keep = [ ]
while idxs. size > 0 :
max_score_index = idxs[ - 1 ]
max_score_box = boxes[ max_score_index] [ None , : ]
keep. append( max_score_index)
if idxs. size == 1 :
break
idxs = idxs[ : - 1 ]
other_boxes = boxes[ idxs]
ious = box_iou( max_score_box, other_boxes)
idxs = idxs[ ious[ 0 ] <= iou_threshold]
keep = np. array( keep)
return keep
box = np. array( [ [ 2 , 3.1 , 7 , 5 ] , [ 3 , 4 , 8 , 4.8 ] , [ 4 , 4 , 5.6 , 7 ] , [ 0.1 , 0 , 8 , 1 ] ] )
score = np. array( [ 0.5 , 0.3 , 0.2 , 0.4 ] )
output = numpy_nms( boxes= box, scores= score, iou_threshold= 0.3 )
print ( output)