【目標檢測】 IoU 計算及 NMS 計算

IoU 計算

計算 IoU 的代碼因爲不是很複雜，所以在一些偏深度學習的崗位面試時比較容易遇到。一般都是給定一個點的形式爲，[x1,y1,x2,y2] 或者是 [x1,y1,w1,h1]。

'''
Input: 
p_x=[x1,y1,w1,h1]
p_y=[x2,y2,w2,h2]

'''

def IoU(p_x,p_y):
    area_x = p_x[2]*p_x[3]
    area_y = p_y[2]*p_y[3]
    
    x_tl = max(p_x[0],p_y[0])
    y_tl = max(p_x[1],p_y[1])
    x_br = min(p_x[0]+p_x[2],p_y[0]+p_y[2])
    y_br = min(p_x[1]+p_x[3],p_y[1]+p_y[3])

    inter_w = max(x_br-x_tl,0)   # 與 0 比較是爲了防止兩個不相交的情況
    inter_h = max(y_br-y_tl,0)   # 與 0 比較是爲了防止兩個不相交的情況

    inter_area = inter_h*inter_w

    return inter_area*1.0 / (area_x+area_y - inter_area)

NMS 計算

雖然 NMS 需要基於上面的計算 IoU 公式，但是上面是簡單情況。只有兩個框，這裏需要藉助 Numpy，直接多維運算，否則程序寫起來會很麻煩，運行起來也會因爲多個 for 循環慢很多，這裏直接參考：

import numpy as np

'''
input:
dets=[[x1,y1,x2,y2],[x1,y1,x2,y2],...] 
thresh: float
'''

def nms(dets, thresh):
   # 以下 5 個都是多維的數組
    x1 = dets[:, 0] 
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]
   
   # 用數組存儲所有 box 的面積
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    
    # 將數組的類別分數從大到小排序，並得到它們排序後的索引順序 
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]  # np.where() 直接將小於閾值的去除了，返回 tuple（[index1,index2,...]）
        order = order[inds + 1]

    return keep