Merge pull request #20029 from yoya:master

PiperOrigin-RevId: 208076674

3 files changed, 11 insertions, 3 deletions

diff --git a/tensorflow/core/lib/png/png_io.cc b/tensorflow/core/lib/png/png_io.cc
index 62c803afb2..e226a15ccc 100644
--- a/tensorflow/core/lib/png/png_io.cc
+++ b/tensorflow/core/lib/png/png_io.cc
@@ -232,11 +232,19 @@ bool CommonInitDecode(StringPiece png_string, int desired_channels,
     CommonFreeDecode(context);
     return false;
   }
-  if (context->channels == 0) {  // Autodetect number of channels
-    context->channels = png_get_channels(context->png_ptr, context->info_ptr);
-  }
   const bool has_tRNS =
       (png_get_valid(context->png_ptr, context->info_ptr, PNG_INFO_tRNS)) != 0;
+  if (context->channels == 0) {  // Autodetect number of channels
+    if (context->color_type == PNG_COLOR_TYPE_PALETTE) {
+      if (has_tRNS) {
+        context->channels = 4;  // RGB + A(tRNS)
+      } else {
+        context->channels = 3;  // RGB
+      }
+    } else {
+      context->channels = png_get_channels(context->png_ptr, context->info_ptr);
+    }
+  }
   const bool has_alpha = (context->color_type & PNG_COLOR_MASK_ALPHA) != 0;
   if ((context->channels & 1) == 0) {  // We desire alpha
     if (has_alpha) {                   // There is alpha
diff --git a/tensorflow/core/lib/png/testdata/lena_palette.png b/tensorflow/core/lib/png/testdata/lena_palette.png
new file mode 100644
index 0000000000..d19ec04895
--- /dev/null
+++ b/tensorflow/core/lib/png/testdata/lena_palette.png
diff --git a/tensorflow/core/lib/png/testdata/lena_palette_trns.png b/tensorflow/core/lib/png/testdata/lena_palette_trns.png
new file mode 100644
index 0000000000..c298fee9ff
--- /dev/null
+++ b/tensorflow/core/lib/png/testdata/lena_palette_trns.png