fix decode_png function.

Problem: In the case of palette png & channels=0, force into grayscale.

1 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..2f6f94e24c 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