tf.image.non_max_suppression 例子

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)