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