转载自:darknet框架中两个版本nms函数(do_nms_obj和do_nms_sort)的比较
作者:ltshan139
前言
在darknet框架中,有两个nms函数:do_nms_obj和do_nms_sort()。 在test_detector()函数中用的是do_nms_sort(), 而在darknet.py中用的是 do_nms_obj()。 值得注意的是,它们的运行结果会导致最终检测框不太一样。 有必要在这里结合代码解释一下,希望引起大家注意。
do_nms_obj
该函数的实现代码如下所示:
void do_nms_obj(detection *dets, int total, int classes, float thresh)
{
int i, j, k;
k = total-1;
for(i = 0; i <= k; ++i){
if(dets[i].objectness == 0){
detection swap = dets[i];
dets[i] = dets[k];
dets[k] = swap;
--k;
--i;
}
}
total = k+1;
for(i = 0; i < total; ++i){
dets[i].sort_class = -1;
}
qsort(dets, total, sizeof(detection), nms_comparator);
for(i = 0; i < total; ++i){
if(dets[i].objectness == 0) continue;
box a = dets[i].bbox;
for(j = i+1; j < total; ++j){
if(dets[j].objectness == 0) continue;
box b = dets[j].bbox;
if (box_iou(a, b) > thresh){
dets[j].objectness = 0;
for(k = 0; k < classes; ++k){
dets[j].prob[k] = 0;
}
}
}
}
}
上面 代码的意思是先按每个检测框(大于thresh。小于thresh的框在做nms之前其实已经过滤掉了)的objectness值由大到小顺序排列框;然后从第一个框开始轮询它和其它框的IOU,如果iou值大于nms thresh,那么其他框的的objectness和class probs都设为0。
这个函数代码简单明了,比较容易理解。但是有两个问题 1)对于密集目标群,不同目标会重叠,使用该nms函数很有可能把另外一个不同种类的检测框给过滤掉。 2)即使对没有重叠的目标,如果只考虑objectness值是不够客观的。相反,考虑class_probs更加客观,毕竟该值综合考虑了原始class_prob和objectness,如下所示
int get_yolo_detections(layer l, int w, int h, int netw, int neth, float thresh, int map, int relative, detection dets)
{
... ...
for(n = 0; n < l.n; ++n){
... ....
float objectness = predictions[obj_index];
if(objectness <= thresh) continue;
... ...
for(j = 0; j < l.classes; ++j){
... ...
float prob = objectness*predictions[class_index];
dets[count].prob[j] = (prob > thresh) ? prob : 0;
}
++count;
}
}
... ....
return count;
}
do_nms_sort
该函数代码如下
void do_nms_sort(detection *dets, int total, int classes, float thresh)
{
int i, j, k;
k = total-1;
for(i = 0; i <= k; ++i){
if(dets[i].objectness == 0){
detection swap = dets[i];
dets[i] = dets[k];
dets[k] = swap;
--k;
--i;
}
}
total = k+1;
for(k = 0; k < classes; ++k){
for(i = 0; i < total; ++i){
dets[i].sort_class = k;
}
qsort(dets, total, sizeof(detection), nms_comparator);
for(i = 0; i < total; ++i){
if(dets[i].prob[k] == 0) continue;
box a = dets[i].bbox;
for(j = i+1; j < total; ++j){
box b = dets[j].bbox;
if (box_iou(a, b) > thresh){
dets[j].prob[k] = 0;
}
}
}
}
}
该函数代码实现要复杂些。 它是对每个class都遍历一次所有框(上面函数只遍历一次所有框)。 在遍历框时,先按每个框在该类的probs来排序(这点也是和前面函数不一样),然后再使用iou来淘汰掉重复框,最后保留的是该类得分较高且没有太多重叠的框。 这样做显然能够解决上面函数所遇到的问题。 在实际检测中效果也相对好些。
结论
所以在用darknet框架推理时,要用do_nms_sort()。 当然近来也有一些新的nms方式,大家可以根据自己的场景来选择。
推荐阅读
更多海思AI芯片方案学习笔记欢迎关注海思AI芯片方案学习。