darknet框架中两个版本nms函数(do_nms_obj和do_nms_sort)的比较

前言

在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方式，大家可以根据自己的场景来选择。