Keep the colors functionality for both v1 and v2

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

tensorflow/python/kernel_tests/draw_bounding_box_op_test.py

@@ -61,15 +61,14 @@ class DrawBoundingBoxOpTest(test.TestCase):
     Args:
       img: 3-D numpy image on which to draw.
     """
+    color_table = colors
+    if colors is None:
       # THIS TABLE MUST MATCH draw_bounding_box_op.cc
-    default_color_table = np.asarray([[1, 1, 0, 1], [0, 0, 1, 1],
+      color_table = np.asarray([[1, 1, 0, 1], [0, 0, 1, 1],
                                 [1, 0, 0, 1], [0, 1, 0, 1],
                                 [0.5, 0, 0.5, 1], [0.5, 0.5, 0, 1],
                                 [0.5, 0, 0, 1], [0, 0, 0.5, 1],
                                 [0, 1, 1, 1], [1, 0, 1, 1]])
-    color_table = default_color_table
-    if colors is not None:
-      color_table = colors
     assert len(img.shape) == 3
     depth = img.shape[2]
     assert depth <= color_table.shape[1]
@@ -90,10 +89,7 @@ class DrawBoundingBoxOpTest(test.TestCase):
       image = ops.convert_to_tensor(image)
       image = image_ops_impl.convert_image_dtype(image, dtypes.float32)
       image = array_ops.expand_dims(image, 0)
-      if colors is None:
+      image = image_ops.draw_bounding_boxes(image, bboxes, colors=colors)
-        image = image_ops.draw_bounding_boxes(image, bboxes)
-      else:
-        image = image_ops.draw_bounding_boxes_v2(image, bboxes, colors)
       with self.cached_session(use_gpu=False) as sess:
         op_drawn_image = np.squeeze(sess.run(image), 0)
         self.assertAllEqual(test_drawn_image, op_drawn_image)

tensorflow/python/ops/image_ops_impl.py

@@ -3533,51 +3533,15 @@ def combined_non_max_suppression(boxes,
         boxes, scores, max_output_size_per_class, max_total_size, iou_threshold,
         score_threshold, pad_per_class)
 
-@tf_export(v1=['image.draw_bounding_boxes'])
-def draw_bounding_boxes(
-    images,
-    boxes,
-    name=None):
-  """Draw bounding boxes on a batch of images.
-
-  Outputs a copy of `images` but draws on top of the pixels zero or more bounding
-  boxes specified by the locations in `boxes`. The coordinates of the each
-  bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`. The
-  bounding box coordinates are floats in `[0.0, 1.0]` relative to the width and
-  height of the underlying image.
-
-  For example, if an image is 100 x 200 pixels (height x width) and the bounding
-  box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
-  the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
-
-  Parts of the bounding box may fall outside the image.
-
-  Args:
-    images: A `Tensor`. Must be one of the following types: `float32`, `half`.
-      4-D with shape `[batch, height, width, depth]`. A batch of images.
-    boxes: A `Tensor` of type `float32`.
-      3-D with shape `[batch, num_bounding_boxes, 4]` containing bounding
-      boxes.
-    name: A name for the operation (optional).
-
-  Returns:
-    A `Tensor`. Has the same type as `images`.
-  """
-  return gen_image_ops.draw_bounding_boxes(images, boxes, name)
-
 @tf_export('image.draw_bounding_boxes', v1=[])
-def draw_bounding_boxes_v2(
+def draw_bounding_boxes_v2(images, boxes, colors, name=None):
-    images,
-    boxes,
-    colors,
-    name=None):
   """Draw bounding boxes on a batch of images.
 
-  Outputs a copy of `images` but draws on top of the pixels zero or more bounding
+  Outputs a copy of `images` but draws on top of the pixels zero or more
-  boxes specified by the locations in `boxes`. The coordinates of the each
+  bounding boxes specified by the locations in `boxes`. The coordinates of the
-  bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`. The
+  each bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`.
-  bounding box coordinates are floats in `[0.0, 1.0]` relative to the width and
+  The bounding box coordinates are floats in `[0.0, 1.0]` relative to the width
-  height of the underlying image.
+  and height of the underlying image.
 
   For example, if an image is 100 x 200 pixels (height x width) and the bounding
   box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
@@ -3598,4 +3562,35 @@ def draw_bounding_boxes_v2(
   Returns:
     A `Tensor`. Has the same type as `images`.
   """
+  if colors is None and not compat.forward_compatible(2019, 5, 1):
+    return gen_image_ops.draw_bounding_boxes(images, boxes, name)
   return gen_image_ops.draw_bounding_boxes_v2(images, boxes, colors, name)
+
+@tf_export(v1=['image.draw_bounding_boxes'])
+def draw_bounding_boxes(images, boxes, name=None, colors=None):
+  """Draw bounding boxes on a batch of images.
+
+  Outputs a copy of `images` but draws on top of the pixels zero or more
+  bounding boxes specified by the locations in `boxes`. The coordinates of the
+  each bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`.
+  The bounding box coordinates are floats in `[0.0, 1.0]` relative to the width
+  and height of the underlying image.
+
+  For example, if an image is 100 x 200 pixels (height x width) and the bounding
+  box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
+  the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
+
+  Parts of the bounding box may fall outside the image.
+
+  Args:
+    images: A `Tensor`. Must be one of the following types: `float32`, `half`.
+      4-D with shape `[batch, height, width, depth]`. A batch of images.
+    boxes: A `Tensor` of type `float32`.
+      3-D with shape `[batch, num_bounding_boxes, 4]` containing bounding
+      boxes.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `Tensor`. Has the same type as `images`.
+  """
+  return draw_bounding_boxes_v2(images, boxes, colors, name)