diff options
Diffstat (limited to 'tensorflow/python/ops/image_ops_impl.py')
| -rw-r--r-- | tensorflow/python/ops/image_ops_impl.py | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py index 9440bab9ee..855a4d0c33 100644 --- a/tensorflow/python/ops/image_ops_impl.py +++ b/tensorflow/python/ops/image_ops_impl.py @@ -20,6 +20,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.compat import compat from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -2110,6 +2111,64 @@ def non_max_suppression(boxes, iou_threshold, score_threshold) +@tf_export('image.non_max_suppression_padded') +def non_max_suppression_padded(boxes, + scores, + max_output_size, + iou_threshold=0.5, + score_threshold=float('-inf'), + pad_to_max_output_size=False, + name=None): + """Greedily selects a subset of bounding boxes in descending order of score. + + Performs algorithmically equivalent operation to tf.image.non_max_suppression, + with the addition of an optional parameter which zero-pads the output to + be of size `max_output_size`. + The output of this operation is a tuple containing the set of integers + indexing into the input collection of bounding boxes representing the selected + boxes and the number of valid indices in the index set. The bounding box + coordinates corresponding to the selected indices can then be obtained using + the `tf.slice` and `tf.gather` operations. For example: + selected_indices_padded, num_valid = tf.image.non_max_suppression_padded( + boxes, scores, max_output_size, iou_threshold, + score_threshold, pad_to_max_output_size=True) + selected_indices = tf.slice( + selected_indices_padded, tf.constant([0]), num_valid) + selected_boxes = tf.gather(boxes, selected_indices) + + Args: + boxes: A 2-D float `Tensor` of shape `[num_boxes, 4]`. + scores: A 1-D float `Tensor` of shape `[num_boxes]` representing a single + score corresponding to each box (each row of boxes). + max_output_size: A scalar integer `Tensor` representing the maximum number + of boxes to be selected by non max suppression. + iou_threshold: A float representing the threshold for deciding whether boxes + overlap too much with respect to IOU. + score_threshold: A float representing the threshold for deciding when to + remove boxes based on score. + pad_to_max_output_size: bool. If True, size of `selected_indices` output + is padded to `max_output_size`. + name: A name for the operation (optional). + + Returns: + selected_indices: A 1-D integer `Tensor` of shape `[M]` representing the + selected indices from the boxes tensor, where `M <= max_output_size`. + valid_outputs: A scalar integer `Tensor` denoting how many elements in + `selected_indices` are valid. Valid elements occur first, then padding. + """ + with ops.name_scope(name, 'non_max_suppression_padded'): + iou_threshold = ops.convert_to_tensor(iou_threshold, name='iou_threshold') + score_threshold = ops.convert_to_tensor( + score_threshold, name='score_threshold') + if compat.forward_compatible(2018, 8, 7) or pad_to_max_output_size: + return gen_image_ops.non_max_suppression_v4( + boxes, scores, max_output_size, iou_threshold, score_threshold, + pad_to_max_output_size) + else: + return gen_image_ops.non_max_suppression_v3( + boxes, scores, max_output_size, iou_threshold, score_threshold) + + @tf_export('image.non_max_suppression_overlaps') def non_max_suppression_with_overlaps(overlaps, scores, |