tf.image.non_max_suppression 例子

import tensorflow as tf
import numpy as np
from keras import backend as K

rects=np.asarray([[1,2,3,4],[1,3,3,4],[1,3,4,4],[1,1,4,4],[1,1,3,4]],dtype=np.float32)

scores=np.asarray([0.4,0.5,0.72,0.9,0.45],dtype=np.float32)

import datetime
with tf.Session() as sess: for i in range(50): old=datetime.datetime.now() nms = tf.image.non_max_suppression(rects,scores, max_output_size=5,iou_threshold=0.5) print("cost time",(datetime.datetime.now()-old).microseconds) # print('face detectd', len(nms.eval()))
 for index, value in enumerate(nms.eval()): rect = rects[value] print(rect) # face = img[rect[1]:rect[1] + rect[3], rect[0]:rect[0] + rect[2], :]
 # face = imresize(face, shape_raw[0:2])
 # if preview: cv2.rectangle(img, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (0, 255, 0), 2)

实现极大值抑制non max suppression，其中boxes是不同boxes的坐标，scores是不同boxes预测的分数，max_boxes是保留的最大box的个数。

iou_threshold是一个阈值，去掉大于这个阈值的所有boxes。

下面是和python对比的效果：

#!/usr/bin/env python
# _*_ coding: utf-8 _*_

import numpy as np
import tensorflow as tf
def nms(bounding_boxes, score, threshold): # If no bounding boxes, return empty list
 if len(bounding_boxes) == 0: return [], [] # coordinates of bounding boxes
 start_x = bounding_boxes[:, 0] start_y = bounding_boxes[:, 1] end_x = bounding_boxes[:, 2] end_y = bounding_boxes[:, 3] # Picked bounding boxes
 picked_boxes = [] picked_score = [] # Compute areas of bounding boxes
 areas = (end_x - start_x + 1) * (end_y - start_y + 1) # Sort by confidence score of bounding boxes
 order = np.argsort(score) # Iterate bounding boxes
 while order.size > 0: # The index of largest confidence score
 index = order[-1] # Pick the bounding box with largest confidence score
 picked_boxes.append(bounding_boxes[index]) picked_score.append(confidence_score[index]) # Compute ordinates of intersection-over-union(IOU)
 x1 = np.maximum(start_x[index], start_x[order[:-1]]) x2 = np.minimum(end_x[index], end_x[order[:-1]]) y1 = np.maximum(start_y[index], start_y[order[:-1]]) y2 = np.minimum(end_y[index], end_y[order[:-1]]) # Compute areas of intersection-over-union
 w = np.maximum(0.0, x2 - x1 + 1) h = np.maximum(0.0, y2 - y1 + 1) intersection = w * h # Compute the ratio between intersection and union
 ratio = intersection / (areas[index] + areas[order[:-1]] - intersection) left = np.where(ratio < threshold) order = order[left] return picked_boxes, picked_score

bounding_boxes =np.asarray([[187, 82, 337, 317],[150, 67, 305, 282],[246, 121, 368, 304]],dtype=np.float32)
confidence_score = np.asarray([0.9, 0.75, 0.8],dtype=np.float32)

# Draw parameters
font_scale = 1
thickness = 2

# IoU threshold
threshold = 0.5

picked_boxes, picked_score = nms(bounding_boxes, confidence_score, threshold)

print(picked_boxes,picked_score)

import datetime
with tf.Session() as sess: for i in range(3): old=datetime.datetime.now() nms = tf.image.non_max_suppression(bounding_boxes,confidence_score, max_output_size=5,iou_threshold=threshold) print("cost time",(datetime.datetime.now()-old).microseconds) # print('face detectd', len(nms.eval()))
 for index, value in enumerate(nms.eval()): rect = bounding_boxes[value] print(rect)