Upperlan
0%

NMS和IOU代码实现

发表于 分类于

NMS

非极大值抑制的具体实现代码如下面的nms函数的定义，需要说明的是数据集中含有多个类别的物体，所以这里需要做多分类非极大值抑制，其实现原理与非极大值抑制相同，区别在于需要对每个类别都做非极大值抑制，实现代码如下面的multiclass_nms所示。

nms逻辑

  1. 将scores排序，从大到小
    while inds.size > 0:
    选一个最大score_m
    score_m < 阈值： break
    否则
    for(bbbox ： 所有确认保留的bbox索引 keen_inds)：
     计算所有bbox和已保留的bbox的IoU
 如果很大，就认为重叠，不保留
 如果保留，push到确认保留的bbox，
 pop一个最大score的bbox

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
# 非极大值抑制
# 假设score：[0.2, 0.3, 0.1, 0.8, 0.7]
# bboxes[nums_bbox, 4]
#i：batch size
def nms(bboxes, scores, score_thresh, nms_thresh, pre_nms_topk, i=0, c=0):
    """
    nms
    """
    inds = np.argsort(scores) 
    # 将scores中的元素从小到大排列，提取其对应的index(索引)，然后输出到y。例如：score[2]=0.1最小，所以y[0]=2,
    # score[4]=0.8最大，所以inds[5]=0.8。
    inds = inds[::-1] # 倒序，变成从大到小
    keep_inds = []
    while(len(inds) > 0):
        cur_ind = inds[0] # 选出当前最大score的索引
        cur_score = scores[cur_ind] # 从索引取出值来
        # if score of the box is less than score_thresh, just drop it
        if cur_score < score_thresh:
            break

        keep = True
        for ind in keep_inds:
            current_box = bboxes[cur_ind] # 选出当前最大的score的bbox
            remain_box = bboxes[ind] # 剩下的所有bbox一个一个来
            iou = box_iou_xyxy(current_box, remain_box) # 剩下所有bbox和最大scorebbox的IoU
            if iou > nms_thresh: # 大于一定IoU的直接去掉,不保存索引，也就是不加入keep_ind
                keep = False
                break
#         if i == 0 and c == 4 and cur_ind == 951:
#             print('suppressed, ', keep, i, c, cur_ind, ind, iou)
        if keep:
            keep_inds.append(cur_ind)
        inds = inds[1:] # 每个循环结束，就pop掉一个最大score

    return np.array(keep_inds)

# 多分类非极大值抑制
def multiclass_nms(bboxes, scores, score_thresh=0.01, nms_thresh=0.45, pre_nms_topk=1000, pos_nms_topk=100):
    """
    This is for multiclass_nms
    """
    batch_size = bboxes.shape[0]
    class_num = scores.shape[1]
    rets = []
    for i in range(batch_size):
        bboxes_i = bboxes[i]
        scores_i = scores[i]
        ret = []
        for c in range(class_num):
            scores_i_c = scores_i[c]
            keep_inds = nms(bboxes_i, scores_i_c, score_thresh, nms_thresh, pre_nms_topk, i=i, c=c)
            if len(keep_inds) < 1:
                continue
            keep_bboxes = bboxes_i[keep_inds]
            keep_scores = scores_i_c[keep_inds]
            keep_results = np.zeros([keep_scores.shape[0], 6])
            keep_results[:, 0] = c
            keep_results[:, 1] = keep_scores[:]
            keep_results[:, 2:6] = keep_bboxes[:, :]
            ret.append(keep_results)
        if len(ret) < 1:
            rets.append(ret)
            continue
        ret_i = np.concatenate(ret, axis=0)
        scores_i = ret_i[:, 1]
        if len(scores_i) > pos_nms_topk:
            inds = np.argsort(scores_i)[::-1]
            inds = inds[:pos_nms_topk]
            ret_i = ret_i[inds]

        rets.append(ret_i)

    return rets

IoU

找到相交区域的斜对角两个边界点的坐标，计算面积。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import nump as np

def get_iou(pred_bboxes, gt_bbox):
    """
    (x1,y1,x2,y2)
    """
    ximin = np.maximum(pred[:, 0], gt_bbox[0])
    yimin = np.maximum(pred[:, 1], gt_bbox[1])
    ximax = np.minimum(pred[:,2], gt_bbox[2])
    yimax = np.minimum(pred[:,3], gt_bbox[3])
    wi = np.maximum(ximax - ximin, 0.)
    hi = np.maximum(yimax - yimin, 0.)
    
    inter = wi * yi
    
    uni = (pred[:,2] - pred[:, 0] + 1) * (pred[:,3] - pred[:, 1] + 1) + \
        (gt_bbox[2] - gt_bbox[0] + 1) * (gt_bbox[3] - gt_bbox[1] + 1) - inter
    
    IoU = inter / uni
    
    return IoU