Hide NonMaxSuppression and NonMaxSuppressionV2 ops and add a python wrapper that sets a backwards compatible default value for iou_threshold.

PiperOrigin-RevId: 163844703

tensorflow/python/ops/hidden_ops.txt

@@ -158,13 +158,15 @@ TensorArrayWriteV3
 SymbolicGradient
 
 # image_ops
+AdjustContrastv2
+NonMaxSuppression
+NonMaxSuppressionV2
 RandomCrop
 ResizeBilinearGrad
 ResizeNearestNeighborGrad
-AdjustContrastv2
-ScaleImageGrad
 SampleDistortedBoundingBox
 SampleDistortedBoundingBoxV2
+ScaleImageGrad
 
 # io_ops
 FixedLengthRecordReader

tensorflow/python/ops/image_ops_impl.py

@@ -1573,3 +1573,49 @@ def sample_distorted_bounding_box(image_size, bounding_boxes, seed=None,
                 max_attempts=max_attempts,
                 use_image_if_no_bounding_boxes=use_image_if_no_bounding_boxes,
                 name=name)
+
+
+def non_max_suppression(boxes,
+                        scores,
+                        max_output_size,
+                        iou_threshold=0.5,
+                        name=None):
+  """Greedily selects a subset of bounding boxes in descending order of score.
+
+  Prunes away boxes that have high intersection-over-union (IOU) overlap
+  with previously selected boxes.  Bounding boxes are supplied as
+  [y1, x1, y2, x2], where (y1, x1) and (y2, x2) are the coordinates of any
+  diagonal pair of box corners and the coordinates can be provided as normalized
+  (i.e., lying in the interval [0, 1]) or absolute.  Note that this algorithm
+  is agnostic to where the origin is in the coordinate system.  Note that this
+  algorithm is invariant to orthogonal transformations and translations
+  of the coordinate system; thus translating or reflections of the coordinate
+  system result in the same boxes being selected by the algorithm.
+  The output of this operation is a set of integers indexing into the input
+  collection of bounding boxes representing the selected boxes.  The bounding
+  box coordinates corresponding to the selected indices can then be obtained
+  using the `tf.gather operation`.  For example:
+    selected_indices = tf.image.non_max_suppression(
+        boxes, scores, max_output_size, iou_threshold)
+    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.
+    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`.
+  """
+  with ops.name_scope(name, 'non_max_suppression'):
+    iou_threshold = ops.convert_to_tensor(iou_threshold, name='iou_threshold')
+    # pylint: disable=protected-access
+    return gen_image_ops._non_max_suppression_v2(boxes, scores, max_output_size,
+                                                 iou_threshold)
+    # pylint: enable=protected-access

tensorflow/python/ops/image_ops_test.py

@@ -2859,5 +2859,23 @@ class FormatTest(test_util.TensorFlowTestCase):
           decode(io_ops.read_file(path)).eval()
 
 
+class NonMaxSuppressionTest(test_util.TensorFlowTestCase):
+
+  def testSelectFromThreeClusters(self):
+    boxes_np = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
+                [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
+    scores_np = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
+    max_output_size_np = 3
+    iou_threshold_np = 0.5
+    with self.test_session():
+      boxes = constant_op.constant(boxes_np)
+      scores = constant_op.constant(scores_np)
+      max_output_size = constant_op.constant(max_output_size_np)
+      iou_threshold = constant_op.constant(iou_threshold_np)
+      selected_indices = image_ops.non_max_suppression(
+          boxes, scores, max_output_size, iou_threshold).eval()
+      self.assertAllClose(selected_indices, [3, 0, 5])
+
+
 if __name__ == "__main__":
   googletest.main()

tensorflow/tools/api/golden/tensorflow.image.pbtxt

@@ -94,7 +94,7 @@ tf_module {
   }
   member_method {
     name: "non_max_suppression"
-    argspec: "args=[\'boxes\', \'scores\', \'max_output_size\', \'iou_threshold\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+    argspec: "args=[\'boxes\', \'scores\', \'max_output_size\', \'iou_threshold\', \'name\'], varargs=None, keywords=None, defaults=[\'0.5\', \'None\'], "
   }
   member_method {
     name: "pad_to_bounding_box"