NMS的numpy写法
简单易懂的NMS的numpy写法
目标检测中的NMS,输入:boxes形如N*5,N个(x1,y1,x2,y2,score), thresh阈值为float型。
计算时,首先获取基于分数的降序排序order,计算全部box的面积area,对每个得分最高的boxes[i],计算与其余低分boxes[order[1:]]的交并比ious,去除ious高于阈值的box。
1 | def nms_plain(boxes, threshold): |
简洁的NMS的numpy写法
下面是一个简洁的NMS写法,主要是prod操作计算面积免去了许多多余行。
1 | def nms_neat(boxes, threshold): |
Soft-NMS的简洁的numpy写法
和NMS相比,Soft-NMS把高于thresh的box置信度置为score*weight,而非直接置为0。weight有多种计算方式,linear和Gaussian都有。给出一个Soft-NMS的简洁写法:
1 | def soft_nms(boxes, method = 'linear', thresh = 0.001, Nt = 0.3, sigma = 0.5): |
最后是画图和主函数验证:
1 | def plot_bbox(boxes, c = 'r'): |
NMS的pytorch写法
NMS的pytorch写法和Numpy类似,稍有不同,代码如下
1 | def nms_pytorch(bboxes, scores, threshold=0.5): |
NMS的cpp写法
本文的cpp写法是基于libtorch(pytorch c++ api),仿照pytorch,输入boxes和scores以及thresh变量。
1 | //输入boxes:Nx4; socres: N; thresh:介于(0,1)的float变量 |
事实上,除了使用libtorch的写法,将libtorch张量存储到c++的其他数据结构中,比如vector或者数组中同样可以实现NMS功能,本文不再赘述。