From 9e8a730f08d5d9275dfed767d4ae15a44639b950 Mon Sep 17 00:00:00 2001
From: Akshay Modi <nareshmodi@google.com>
Date: Tue, 26 Nov 2019 13:52:58 -0800
Subject: [PATCH] Add a wrapper around encode_png which converts to tensor
 before calling into the gen ops code.

This way if there is uint16 np array, we'll use uint16 as the dtype correctly
instead of using uint8 (which would happen before due to EncodePng having a
default type of uint8).

PiperOrigin-RevId: 282633227
Change-Id: I46d5bda1f94fda3e0a4b15c57591ef86dd1649a4
---
 .../python_api/api_def_EncodePng.pbtxt        |  1 +
 tensorflow/python/ops/image_ops_impl.py       | 29 +++++++++++++++++++
 tensorflow/python/ops/image_ops_test.py       |  8 +++++
 3 files changed, 38 insertions(+)

diff --git a/tensorflow/core/api_def/python_api/api_def_EncodePng.pbtxt b/tensorflow/core/api_def/python_api/api_def_EncodePng.pbtxt
index 42717ba7d56..97f0017d7a7 100644
--- a/tensorflow/core/api_def/python_api/api_def_EncodePng.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_EncodePng.pbtxt
@@ -3,4 +3,5 @@ op {
   endpoint {
     name: "image.encode_png"
   }
+  visibility: HIDDEN
 }
diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index 84c3e38bd0d..f9252fdd547 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -2232,6 +2232,35 @@ tf_export(
         gen_image_ops.extract_jpeg_shape)
 
 
+@tf_export('image.encode_png')
+def encode_png(image, compression=-1, name=None):
+  r"""PNG-encode an image.
+
+  `image` is a 3-D uint8 or uint16 Tensor of shape `[height, width, channels]`
+  where `channels` is:
+
+  *   1: for grayscale.
+  *   2: for grayscale + alpha.
+  *   3: for RGB.
+  *   4: for RGBA.
+
+  The ZLIB compression level, `compression`, can be -1 for the PNG-encoder
+  default or a value from 0 to 9.  9 is the highest compression level,
+  generating the smallest output, but is slower.
+
+  Args:
+    image: A `Tensor`. Must be one of the following types: `uint8`, `uint16`.
+      3-D with shape `[height, width, channels]`.
+    compression: An optional `int`. Defaults to `-1`. Compression level.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `Tensor` of type `string`.
+  """
+  return gen_image_ops.encode_png(
+      ops.convert_to_tensor(image), compression, name)
+
+
 @tf_export(
     'io.decode_image',
     'image.decode_image',
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 7ee86bd096e..f4c0f6926b1 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -5119,6 +5119,7 @@ class SobelEdgesTest(test_util.TensorFlowTestCase):
       self.assertAllClose(expected_batch, actual_sobel)
 
 
+@test_util.run_all_in_graph_and_eager_modes
 class DecodeImageTest(test_util.TensorFlowTestCase):
 
   def testJpegUint16(self):
@@ -5141,6 +5142,13 @@ class DecodeImageTest(test_util.TensorFlowTestCase):
       image0, image1 = self.evaluate([image0, image1])
       self.assertAllEqual(image0, image1)
 
+      # NumPy conversions should happen before
+      x = np.random.randint(256, size=(4, 4, 3), dtype=np.uint16)
+      x_str = image_ops_impl.encode_png(x)
+      x_dec = image_ops_impl.decode_image(
+          x_str, channels=3, dtype=dtypes.uint16)
+      self.assertAllEqual(x, x_dec)
+
   def testGifUint16(self):
     with self.cached_session(use_gpu=True) as sess:
       base = "tensorflow/core/lib/gif/testdata"