YOLOV4學習系列（三） 結合代碼講解DIOU-NMS

前言

NMS（非最大抑制）是目標檢測算法後處理中常用的技術，用來將redundant檢測框給過濾掉。YOLOV4沒有用經典的NMS，取而代之的是DIOU-NMS。本博客接下來會講解其原理和代碼實現。

原理

在經典的NMS中，得分最高的檢測框和其它檢測框逐一算出一個對應的IOU值，並將該值超過NMS threshold的框全部過濾掉。可以看出，在經典NMS算法中，IOU是唯一考量的因素。

但是在實際應用場景中，當兩個不同物體捱得很近時，由於IOU值比較大，往往經過NMS處理後，只剩下一個檢測框，這樣導致漏檢的錯誤情況發生。

基於此，DIOU-NMS就不僅僅考慮IOU，還考慮兩個框中心點之間的距離。如果兩個框之間IOU比較大，但是兩個框的距離比較大時，可能會認爲這是兩個物體的框而不會被過濾掉。 其公式如下：

 這裏其實沒有列出DIOU是具體如何計算的（後面會結合代碼講解）。公式大體上講了得分最高的預測框M和其它框Bi的（IOU-DIOU）值比較小時，Bi的得分值Si仍然保持，否則，當（IOU-DIOU）大於NMS threshold值時，Si值就設成0了，即被過濾掉。

代碼

我們來結合代碼來看下YOLOV4的DIOU-NMS的具體實現。

// https://github.com/Zzh-tju/DIoU-darknet
// https://arxiv.org/abs/1911.08287
void diounms_sort(detection *dets, int total, int classes, float thresh, NMS_KIND nms_kind, float beta1)
{
    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_v3);
        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 && nms_kind == CORNERS_NMS)
                {
                    float sum_prob = pow(dets[i].prob[k], 2) + pow(dets[j].prob[k], 2);
                    float alpha_prob = pow(dets[i].prob[k], 2) / sum_prob;
                    float beta_prob = pow(dets[j].prob[k], 2) / sum_prob;
                    dets[j].prob[k] = 0;
                }
                else if (box_iou(a, b) > thresh && nms_kind == GREEDY_NMS) {
                    dets[j].prob[k] = 0;
                }
                else {
                    if (box_diounms(a, b, beta1) > thresh && nms_kind == DIOU_NMS) {
                        dets[j].prob[k] = 0;
                    }
                }
            }

        }
    }
}

上面代碼重點在下面這段：

                    if (box_diounms(a, b, beta1) > thresh && nms_kind == DIOU_NMS) {
                        dets[j].prob[k] = 0;

box_diounms（）就是用來計算diou的值，如下所示。

float box_diounms(box a, box b, float beta1)
{
    boxabs ba = box_c(a, b);
    float w = ba.right - ba.left;
    float h = ba.bot - ba.top;
    float c = w * w + h * h;
    float iou = box_iou(a, b);
    if (c == 0) {
        return iou;
    }
    float d = (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
    float u = pow(d / c, beta1);
    float diou_term = u;
    return iou - diou_term;
}

box_c是a和b框的最小外接框 ，

c是該最小外接框的對角線長的平方；

iou是經典的iou求法，即a、b框的交併比；

d爲兩框中心點距離的平方；

u值則是d/c的beta1次方。注意算diou值時需要引入一個新的參數beta1。 這個需要在cfg裏面yolo層參數中指定，缺省爲0.6。 

最終得diou值爲 iou-u， 跟上面原理中得公式完全一致。