NMS（Non-Maximum Suppression）非极大抑制的几种写法

NMS（Non-Maximum Suppression）非极大抑制的几种写法

• dets矩阵,  thresh阈值定义：

import numpy as np
import tensorflow as tf
import torch

dets = np.array([[30, 20, 230, 200, 1],
                 [50, 50, 260, 220, 0.9],
                 [210, 30, 420, 5, 0.8],
                 [430, 280, 460, 360, 0.7]])
thresh = 0.35

 

• numpy 版本NMS：

​
def nms_numpy(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    scores = dets[:, -1]

    ids_sorted = np.argsort(-1 * scores)
    ids = []

    while len(ids_sorted) > 0:
        i = ids_sorted[0]
        ids.append(ids_sorted[0])

        inter_x1 = np.maximum(x1[i], x1[ids_sorted[1:]])
        inter_y1 = np.maximum(y1[i], y1[ids_sorted[1:]])
        inter_x2 = np.minimum(x2[i], x2[ids_sorted[1:]])
        inter_y2 = np.minimum(y2[i], y2[ids_sorted[1:]])
        inter_w = np.maximum((inter_x2 - inter_x1), 0)
        inter_h = np.maximum((inter_y2 - inter_y1), 0)

        inters = inter_w * inter_h
        unions = areas[i] + areas[ids_sorted[1:]] - inters
        ious = inters / unions

        ids_keep = np.where(ious <= thresh)[0]
        if len(ids_keep) == 0:
            break
        ids_sorted = ids_sorted[ids_keep + 1]
    return ids
​

• PyTorch版NMS： 

def nms_torch(dets, thresh):
    dets = torch.tensor(dets)
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    areas = (x2 - x1) * (y2 - y1)

    scores = dets[:, -1]
    _, ids_sorted = scores.sort(0, descending=True)

    ids = []
    while ids_sorted.numel() > 0:
        i = ids_sorted.item() if (ids_sorted.numel() == 1) else ids_sorted[0]
        ids.append(i)
        
        if ids_sorted.numel() == 1:
            break
        
        inter_x1 = x1[ids_sorted[1:]].clamp(min=x1[i])
        inter_y1 = y1[ids_sorted[1:]].clamp(min=y1[i])
        inter_x2 = x2[ids_sorted[1:]].clamp(max=x2[i])
        inter_y2 = y2[ids_sorted[1:]].clamp(max=y2[i])
        inter_w = (inter_x2 - inter_x1).clamp(min=0)
        inter_h = (inter_y2 - inter_y1).clamp(min=0)
        
        inters = inter_w * inter_h
        unions = areas[i] + areas[ids_sorted[1:]] - inters
        ious = inters / unions
        
        ids_keep = (ious <= thresh).nonzero().squeeze()
        if ids_keep.numel() == 0:
            break
        ids_sorted = ids_sorted[ids_keep + 1]
    return torch.LongTensor(ids)

•  Tensorflow版NMS 版本1： 

def nms_tensorflow_1(dets, thresh):
    dets = tf.convert_to_tensor(dets, tf.float32)
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    areas = (x2 - x1) * (y2 - y1)

    scores = dets[:, 4]
    k = tf.shape(scores)[0]
    _, ids_sorted = tf.nn.top_k(scores, k, sorted=True)

    def my_cond(loop, ids_sorted, ids):
        cur_len = tf.shape(ids_sorted)[0]
        flag = tf.cond(tf.equal(cur_len, 0), lambda: False, lambda: True)
        return flag

    def my_body(loop, ids_sorted, ids):
        i = ids_sorted[0]
        inter_x1 = tf.maximum(x1[i], tf.gather(x1, ids_sorted[1:]))
        inter_y1 = tf.maximum(y1[i], tf.gather(y1, ids_sorted[1:]))
        inter_x2 = tf.minimum(x2[i], tf.gather(x2, ids_sorted[1:]))
        inter_y2 = tf.minimum(y2[i], tf.gather(y2, ids_sorted[1:]))
        inter_w = tf.maximum(0.0, inter_x2 - inter_x1)
        inter_h = tf.maximum(0.0, inter_y2 - inter_y1)

        inters = inter_w * inter_h
        unions = areas[i] + tf.gather(areas, ids_sorted[1:]) - inters
        ious = inters / unions

        ids_keep = tf.where(ious <= thresh)
        ids_keep = tf.reshape(ids_keep, [-1])
        ids_sorted = tf.gather(ids_sorted, ids_keep + 1)
        ids = tf.concat([ids, tf.reshape(i, [1])], 0)
        return loop+1, ids_sorted, ids

    loop = tf.constant(0)
    ids = tf.constant(0, shape=[1])
    _, _, ids_res = tf.while_loop(cond=my_cond, body=my_body, loop_vas=[loop, ids_sorted, ids],
                                  shape_invariants=[tf.TensorShape(None), tf.TensorShape(None), tf.TensorShape(None)])
    ids_res = ids_res[1:]
    
    with tf.Session() as sess:
        ids = sess.run(ids_res)
    return ids

•  Tensorflow版NMS 版本2：  

def nms_tensorflow_2(dets, thresh):
    boxes = np.array(dets[:, :4], np.float32)
    scores = np.array(dets[:, -1], np.float32)

    max_output_size = 10
    selected_indices = tf.image.non_max_suppression(boxes, scores, max_output_size, thresh)
    with tf.Session() as sess:
        selected_indices = sess.run(selected_indices)
    return selected_indices