Enable partial JPEG decompression.

PiperOrigin-RevId: 167329034

1 files changed, 190 insertions, 39 deletions

diff --git a/tensorflow/core/lib/jpeg/jpeg_mem.cc b/tensorflow/core/lib/jpeg/jpeg_mem.cc
index 258793aa1e..3c7e5ca696 100644
--- a/tensorflow/core/lib/jpeg/jpeg_mem.cc
+++ b/tensorflow/core/lib/jpeg/jpeg_mem.cc
@@ -70,13 +70,24 @@ class FewerArgsForCompiler {
   int stride_;
 };
 
+// Check whether the crop window is valid, assuming crop is true.
+bool IsCropWindowValid(const UncompressFlags& flags, int input_image_width,
+                       int input_image_height) {
+  // Crop window is valid only if it is non zero and all the window region is
+  // within the original image.
+  return flags.crop_width > 0 && flags.crop_height > 0 && flags.crop_x >= 0 &&
+         flags.crop_y >= 0 &&
+         flags.crop_y + flags.crop_height <= input_image_height &&
+         flags.crop_x + flags.crop_width <= input_image_width;
+}
+
 uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   // unpack the argball
   const int datasize = argball->datasize_;
   const auto& flags = argball->flags_;
   const int ratio = flags.ratio;
   int components = flags.components;
-  int stride = flags.stride;            // may be 0
+  int stride = flags.stride;              // may be 0
   int64* const nwarn = argball->pnwarn_;  // may be NULL
 
   // Can't decode if the ratio is not recognized by libjpeg
@@ -159,8 +170,43 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
     return nullptr;
   }
 
+  JDIMENSION target_output_width = cinfo.output_width;
+  JDIMENSION target_output_height = cinfo.output_height;
+  JDIMENSION skipped_scanlines = 0;
+#if !defined(WIN32)
+  if (flags.crop) {
+    // Update target output height and width based on crop window.
+    target_output_height = flags.crop_height;
+    target_output_width = flags.crop_width;
+
+    // So far, cinfo holds the original input image information.
+    if (!IsCropWindowValid(flags, cinfo.output_width, cinfo.output_height)) {
+      LOG(ERROR) << "Invalid crop window: x=" << flags.crop_x
+                 << ", y=" << flags.crop_y << ", w=" << target_output_width
+                 << ", h=" << target_output_height
+                 << " for image_width: " << cinfo.output_width
+                 << " and image_height: " << cinfo.output_height;
+      jpeg_destroy_decompress(&cinfo);
+      return nullptr;
+    }
+
+    // Update cinfo.output_width. It is tricky that cinfo.output_width must
+    // fall on an Minimum Coded Unit (MCU) boundary; if it doesn't, then it will
+    // be moved left to the nearest MCU boundary, and width will be increased
+    // accordingly. Therefore, the final cinfo.crop_width might differ from the
+    // given flags.crop_width. Please see libjpeg library for details.
+    JDIMENSION crop_width = flags.crop_width;
+    JDIMENSION crop_x = flags.crop_x;
+    jpeg_crop_scanline(&cinfo, &crop_x, &crop_width);
+
+    // Update cinfo.output_scanline.
+    skipped_scanlines = jpeg_skip_scanlines(&cinfo, flags.crop_y);
+    CHECK_EQ(skipped_scanlines, flags.crop_y);
+  }
+#endif
+
   // check for compatible stride
-  const int min_stride = cinfo.output_width * components * sizeof(JSAMPLE);
+  const int min_stride = target_output_width * components * sizeof(JSAMPLE);
   if (stride == 0) {
     stride = min_stride;
   } else if (stride < min_stride) {
@@ -170,47 +216,88 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   }
 
   // Remember stride and height for use in Uncompress
-  argball->height_ = cinfo.output_height;
+  argball->height_ = target_output_height;
   argball->stride_ = stride;
 
-  uint8* const dstdata = argball->allocate_output_(
-      cinfo.output_width, cinfo.output_height, components);
+#if defined(WIN32)
+  uint8* dstdata = nullptr;
+  if (flags.crop) {
+    dstdata = new JSAMPLE[stride * target_output_height];
+  } else {
+    dstdata = argball->allocate_output_(target_output_width,
+                                        target_output_height, components);
+  }
+#else
+  uint8* dstdata = argball->allocate_output_(target_output_width,
+                                             target_output_height, components);
+#endif
   if (dstdata == nullptr) {
     jpeg_destroy_decompress(&cinfo);
     return nullptr;
   }
   JSAMPLE* output_line = static_cast<JSAMPLE*>(dstdata);
 
-  // Temporary buffer used for CMYK -> RGB conversion.
+  // jpeg_read_scanlines requires the buffers to be allocated based on
+  // cinfo.output_width, but the target image width might be different if crop
+  // is enabled and crop_width is not MCU aligned. In this case, we need to
+  // realign the scanline output to achieve the exact cropping.  Notably, only
+  // cinfo.output_width needs to fall on MCU boundary, while cinfo.output_height
+  // has no such constraint.
+  const bool need_realign_cropped_scanline =
+      (target_output_width != cinfo.output_width);
   const bool use_cmyk = (cinfo.out_color_space == JCS_CMYK);
-  tempdata = use_cmyk ? new JSAMPLE[cinfo.output_width * 4] : nullptr;
+
+  if (use_cmyk) {
+    // Temporary buffer used for CMYK -> RGB conversion.
+    tempdata = new JSAMPLE[cinfo.output_width * 4];
+  } else if (need_realign_cropped_scanline) {
+    // Temporary buffer used for MCU-aligned scanline data.
+    tempdata = new JSAMPLE[cinfo.output_width * components];
+  }
 
   // If there is an error reading a line, this aborts the reading.
   // Save the fraction of the image that has been read.
-  argball->height_read_ = cinfo.output_height;
-  while (cinfo.output_scanline < cinfo.output_height) {
+  argball->height_read_ = target_output_height;
+
+  // These variables are just to avoid repeated computation in the loop.
+  const int max_scanlines_to_read = skipped_scanlines + target_output_height;
+  const int mcu_align_offset =
+      (cinfo.output_width - target_output_width) * (use_cmyk ? 4 : components);
+  while (cinfo.output_scanline < max_scanlines_to_read) {
     int num_lines_read = 0;
-    if (cinfo.out_color_space == JCS_CMYK) {
+    if (use_cmyk) {
       num_lines_read = jpeg_read_scanlines(&cinfo, &tempdata, 1);
-      // Convert CMYK to RGB
-      for (size_t i = 0; i < cinfo.output_width; ++i) {
-        int c = tempdata[4 * i + 0];
-        int m = tempdata[4 * i + 1];
-        int y = tempdata[4 * i + 2];
-        int k = tempdata[4 * i + 3];
-        int r, g, b;
-        if (cinfo.saw_Adobe_marker) {
-          r = (k * c) / 255;
-          g = (k * m) / 255;
-          b = (k * y) / 255;
-        } else {
-          r = (255 - k) * (255 - c) / 255;
-          g = (255 - k) * (255 - m) / 255;
-          b = (255 - k) * (255 - y) / 255;
+      if (num_lines_read > 0) {
+        // Convert CMYK to RGB if scanline read succeeded.
+        for (size_t i = 0; i < target_output_width; ++i) {
+          int offset = 4 * i;
+          if (need_realign_cropped_scanline) {
+            // Align the offset for MCU boundary.
+            offset += mcu_align_offset;
+          }
+          const int c = tempdata[offset + 0];
+          const int m = tempdata[offset + 1];
+          const int y = tempdata[offset + 2];
+          const int k = tempdata[offset + 3];
+          int r, g, b;
+          if (cinfo.saw_Adobe_marker) {
+            r = (k * c) / 255;
+            g = (k * m) / 255;
+            b = (k * y) / 255;
+          } else {
+            r = (255 - k) * (255 - c) / 255;
+            g = (255 - k) * (255 - m) / 255;
+            b = (255 - k) * (255 - y) / 255;
+          }
+          output_line[3 * i + 0] = r;
+          output_line[3 * i + 1] = g;
+          output_line[3 * i + 2] = b;
         }
-        output_line[3 * i + 0] = r;
-        output_line[3 * i + 1] = g;
-        output_line[3 * i + 2] = b;
+      }
+    } else if (need_realign_cropped_scanline) {
+      num_lines_read = jpeg_read_scanlines(&cinfo, &tempdata, 1);
+      if (num_lines_read > 0) {
+        memcpy(output_line, tempdata + mcu_align_offset, min_stride);
       }
     } else {
       num_lines_read = jpeg_read_scanlines(&cinfo, &output_line, 1);
@@ -218,12 +305,13 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
     // Handle error cases
     if (num_lines_read == 0) {
       LOG(ERROR) << "Premature end of JPEG data. Stopped at line "
-                 << cinfo.output_scanline << "/" << cinfo.output_height;
+                 << cinfo.output_scanline - skipped_scanlines << "/"
+                 << target_output_height;
       if (!flags.try_recover_truncated_jpeg) {
-        argball->height_read_ = cinfo.output_scanline;
+        argball->height_read_ = cinfo.output_scanline - skipped_scanlines;
         error = JPEGERRORS_UNEXPECTED_END_OF_DATA;
       } else {
-        for (size_t line = cinfo.output_scanline; line < cinfo.output_height;
+        for (size_t line = cinfo.output_scanline; line < max_scanlines_to_read;
              ++line) {
           if (line == 0) {
             // If even the first line is missing, fill with black color
@@ -235,9 +323,9 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
           output_line += stride;
         }
         argball->height_read_ =
-            cinfo.output_height;  // consider all lines as read
+            target_output_height;  // consider all lines as read
         // prevent error-on-exit in libjpeg:
-        cinfo.output_scanline = cinfo.output_height;
+        cinfo.output_scanline = max_scanlines_to_read;
       }
       break;
     }
@@ -248,23 +336,33 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
   delete[] tempdata;
   tempdata = nullptr;
 
+#if !defined(WIN32)
+  if (flags.crop && cinfo.output_scanline < cinfo.output_height) {
+    // Skip the rest of scanlines, required by jpeg_destroy_decompress.
+    jpeg_skip_scanlines(&cinfo,
+                        cinfo.output_height - flags.crop_y - flags.crop_height);
+    // After this, cinfo.output_height must be equal to cinfo.output_height;
+    // otherwise, jpeg_destroy_decompress would fail.
+  }
+#endif
+
   // Convert the RGB data to RGBA, with alpha set to 0xFF to indicate
   // opacity.
   // RGBRGBRGB... --> RGBARGBARGBA...
   if (components == 4) {
     // Start on the last line.
     JSAMPLE* scanlineptr = static_cast<JSAMPLE*>(
-        dstdata + static_cast<int64>(cinfo.output_height - 1) * stride);
+        dstdata + static_cast<int64>(target_output_height - 1) * stride);
     const JSAMPLE kOpaque = -1;  // All ones appropriate for JSAMPLE.
-    const int right_rgb = (cinfo.output_width - 1) * 3;
-    const int right_rgba = (cinfo.output_width - 1) * 4;
+    const int right_rgb = (target_output_width - 1) * 3;
+    const int right_rgba = (target_output_width - 1) * 4;
 
-    for (int y = cinfo.output_height; y-- > 0;) {
+    for (int y = target_output_height; y-- > 0;) {
       // We do all the transformations in place, going backwards for each row.
       const JSAMPLE* rgb_pixel = scanlineptr + right_rgb;
       JSAMPLE* rgba_pixel = scanlineptr + right_rgba;
       scanlineptr -= stride;
-      for (int x = cinfo.output_width; x-- > 0;
+      for (int x = target_output_width; x-- > 0;
            rgba_pixel -= 4, rgb_pixel -= 3) {
         // We copy the 3 bytes at rgb_pixel into the 4 bytes at rgba_pixel
         // The "a" channel is set to be opaque.
@@ -319,8 +417,61 @@ uint8* UncompressLow(const void* srcdata, FewerArgsForCompiler* argball) {
       LOG(ERROR) << "Unhandled case " << error;
       break;
   }
-  jpeg_destroy_decompress(&cinfo);
 
+#if defined(WIN32)
+  // TODO(tanmingxing): delete all these code after migrating to libjpeg_turbo
+  // for Windows.
+  if (flags.crop) {
+    // Update target output height and width based on crop window.
+    target_output_height = flags.crop_height;
+    target_output_width = flags.crop_width;
+
+    // cinfo holds the original input image information.
+    if (!IsCropWindowValid(flags, cinfo.output_width, cinfo.output_height)) {
+      LOG(ERROR) << "Invalid crop window: x=" << flags.crop_x
+                 << ", y=" << flags.crop_y << ", w=" << target_output_width
+                 << ", h=" << target_output_height
+                 << " for image_width: " << cinfo.output_width
+                 << " and image_height: " << cinfo.output_height;
+      delete[] dstdata;
+      jpeg_destroy_decompress(&cinfo);
+      return nullptr;
+    }
+
+    const uint8* full_image = dstdata;
+    dstdata = argball->allocate_output_(target_output_width,
+                                        target_output_height, components);
+    if (dstdata == nullptr) {
+      delete[] full_image;
+      jpeg_destroy_decompress(&cinfo);
+      return nullptr;
+    }
+
+    const int full_image_stride = stride;
+    // Update stride and hight for crop window.
+    const int min_stride = target_output_width * components * sizeof(JSAMPLE);
+    if (flags.stride == 0) {
+      stride = min_stride;
+    }
+    argball->height_ = target_output_height;
+    argball->stride_ = stride;
+
+    if (argball->height_read_ > target_output_height) {
+      argball->height_read_ = target_output_height;
+    }
+    const int crop_offset = flags.crop_x * components * sizeof(JSAMPLE);
+    const uint8* full_image_ptr = full_image + flags.crop_y * full_image_stride;
+    uint8* crop_image_ptr = dstdata;
+    for (int i = 0; i < argball->height_read_; i++) {
+      memcpy(crop_image_ptr, full_image_ptr + crop_offset, min_stride);
+      crop_image_ptr += stride;
+      full_image_ptr += full_image_stride;
+    }
+    delete[] full_image;
+  }
+#endif
+
+  jpeg_destroy_decompress(&cinfo);
   return dstdata;
 }