Make decode_jpeg automatically turn CMYK to RGB

Our core/lib/jpeg now refuses to output 4 channel jpegs.  It still accepts them
as input, but unconditionally converts them to grayscale or RGB.
Change: 114239759

2 files changed, 14 insertions, 16 deletions

diff --git a/tensorflow/core/lib/jpeg/jpeg_mem.cc b/tensorflow/core/lib/jpeg/jpeg_mem.cc
index 0dddae2459..8727f1367b 100644
--- a/tensorflow/core/lib/jpeg/jpeg_mem.cc
+++ b/tensorflow/core/lib/jpeg/jpeg_mem.cc
@@ -77,11 +77,16 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   int stride = flags.stride;            // may be 0
   int* const nwarn = argball->pnwarn_;  // may be NULL
 
-  // can't decode if the ratio is not recognized by libjpeg
+  // Can't decode if the ratio is not recognized by libjpeg
   if ((ratio != 1) && (ratio != 2) && (ratio != 4) && (ratio != 8)) {
     return nullptr;
   }
 
+  // Channels must be autodetect, grayscale, or rgb.
+  if (!(components == 0 || components == 1 || components == 3)) {
+    return nullptr;
+  }
+
   // if empty image, return
   if (datasize == 0 || srcdata == NULL) return nullptr;
 
@@ -105,8 +110,8 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   SetSrc(&cinfo, srcdata, datasize, flags.try_recover_truncated_jpeg);
   jpeg_read_header(&cinfo, TRUE);
 
-  // Set components automatically if desired
-  if (components == 0) components = cinfo.num_components;
+  // Set components automatically if desired, autoconverting cmyk to rgb.
+  if (components == 0) components = std::min(cinfo.num_components, 3);
 
   // set grayscale and ratio parameters
   switch (components) {
@@ -114,11 +119,10 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
       cinfo.out_color_space = JCS_GRAYSCALE;
       break;
     case 3:
-    case 4:
       if (cinfo.jpeg_color_space == JCS_CMYK ||
           cinfo.jpeg_color_space == JCS_YCCK) {
-        // always use cmyk for output in a 4 channel jpeg. libjpeg has a builtin
-        // decoder.
+        // Always use cmyk for output in a 4 channel jpeg. libjpeg has a builtin
+        // decoder.  We will further convert to rgb below.
         cinfo.out_color_space = JCS_CMYK;
       } else {
         cinfo.out_color_space = JCS_RGB;
diff --git a/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc b/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc
index 7022496e98..9cdfa3c749 100644
--- a/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc
+++ b/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc
@@ -61,33 +61,27 @@ void TestJPEG(Env* env, const string& jpegfile) {
   const int fsize = jpeg.size();
   const uint8* const temp = bit_cast<const uint8*>(jpeg.data());
 
-  // try partial decoding (half of the data)
+  // Try partial decoding (half of the data)
   int w, h, c;
   std::unique_ptr<uint8[]> imgdata;
 
   UncompressFlags flags;
   flags.components = 3;
 
-  // set min_acceptable_fraction to something insufficient
+  // Set min_acceptable_fraction to something insufficient
   flags.min_acceptable_fraction = 0.8;
   imgdata.reset(Uncompress(temp, fsize / 2, flags, &w, &h, &c, NULL));
   CHECK(imgdata.get() == NULL);
 
-  // now, use a value that makes fsize/2 be enough for a black-filling
+  // Now, use a value that makes fsize/2 be enough for a black-filling
   flags.min_acceptable_fraction = 0.01;
   imgdata.reset(Uncompress(temp, fsize / 2, flags, &w, &h, &c, NULL));
   CHECK(imgdata.get() != NULL);
 
-  // finally, uncompress the whole data
+  // Finally, uncompress the whole data
   flags.min_acceptable_fraction = 1.0;
   imgdata.reset(Uncompress(temp, fsize, flags, &w, &h, &c, NULL));
   CHECK(imgdata.get() != NULL);
-
-  // Uncompress the data to RGBA, too
-  flags.min_acceptable_fraction = 1.0;
-  flags.components = 4;
-  imgdata.reset(Uncompress(temp, fsize, flags, &w, &h, &c, NULL));
-  CHECK(imgdata.get() != NULL);
 }
 
 TEST(JpegMemTest, Jpeg) {