1 files changed, 304 insertions, 0 deletions
diff --git a/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc b/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc
new file mode 100644
index 0000000000..23e72f9d57
--- /dev/null
+++ b/tensorflow/core/lib/jpeg/jpeg_mem_unittest.cc
@@ -0,0 +1,304 @@
+#include "tensorflow/core/lib/jpeg/jpeg_mem.h"
+
+#include <setjmp.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <memory>
+
+#include "tensorflow/core/lib/jpeg/jpeg_handle.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/port.h"
+#include "tensorflow/core/public/env.h"
+#include <gtest/gtest.h>
+
+#include "tensorflow/core/lib/core/casts.h"
+
+namespace tensorflow {
+namespace jpeg {
+namespace {
+
+const char kTestData[] = "tensorflow/core/lib/jpeg/testdata/";
+
+int ComputeSumAbsoluteDifference(const uint8* a, const uint8* b, int width,
+                                 int height, int a_stride, int b_stride) {
+  int totalerr = 0;
+  for (int i = 0; i < height; i++) {
+    const uint8* const pa = a + i * a_stride;
+    const uint8* const pb = b + i * b_stride;
+    for (int j = 0; j < 3 * width; j++) {
+      totalerr += abs(static_cast<int>(pa[j]) - static_cast<int>(pb[j]));
+    }
+  }
+  return totalerr;
+}
+
+// Reads the contents of the file into output
+void ReadFileToStringOrDie(Env* env, const string& filename, string* output) {
+  TF_CHECK_OK(ReadFileToString(env, filename, output));
+}
+
+void TestJPEG(Env* env, const string& jpegfile) {
+  // Read the data from the jpeg file into memory
+  string jpeg;
+  ReadFileToStringOrDie(Env::Default(), jpegfile, &jpeg);
+  const int fsize = jpeg.size();
+  const uint8* const temp = bit_cast<const uint8*>(jpeg.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
+  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
+  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
+  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) {
+  Env* env = Env::Default();
+  const string data_path = kTestData;
+
+  // Name of a valid jpeg file on the disk
+  TestJPEG(env, data_path + "jpeg_merge_test1.jpg");
+
+  // Exercise CMYK machinery as well
+  TestJPEG(env, data_path + "jpeg_merge_test1_cmyk.jpg");
+}
+
+TEST(JpegMemTest, Jpeg2) {
+  // create known data, for size in_w x in_h
+  const int in_w = 256;
+  const int in_h = 256;
+  const int stride1 = 3 * in_w;
+  const std::unique_ptr<uint8[]> refdata1(new uint8[stride1 * in_h]);
+  for (int i = 0; i < in_h; i++) {
+    for (int j = 0; j < in_w; j++) {
+      const int offset = i * stride1 + 3 * j;
+      refdata1[offset + 0] = i;
+      refdata1[offset + 1] = j;
+      refdata1[offset + 2] = static_cast<uint8>((i + j) >> 1);
+    }
+  }
+
+  // duplicate with weird input stride
+  const int stride2 = 3 * 357;
+  const std::unique_ptr<uint8[]> refdata2(new uint8[stride2 * in_h]);
+  for (int i = 0; i < in_h; i++) {
+    memcpy(&refdata2[i * stride2], &refdata1[i * stride1], 3 * in_w);
+  }
+
+  // Test compression
+  string cpdata1, cpdata2;
+  {
+    const string kXMP = "XMP_TEST_123";
+
+    // Compress it to JPEG
+    CompressFlags flags;
+    flags.format = FORMAT_RGB;
+    flags.quality = 97;
+    flags.xmp_metadata = kXMP;
+    cpdata1 = Compress(refdata1.get(), in_w, in_h, flags);
+    flags.stride = stride2;
+    cpdata2 = Compress(refdata2.get(), in_w, in_h, flags);
+    // Different input stride shouldn't change the output
+    CHECK_EQ(cpdata1, cpdata2);
+
+    // Verify valid XMP.
+    CHECK_NE(string::npos, cpdata1.find(kXMP));
+
+    // Test the other API, where a storage string is supplied
+    string cptest;
+    flags.stride = 0;
+    Compress(refdata1.get(), in_w, in_h, flags, &cptest);
+    CHECK_EQ(cptest, cpdata1);
+    flags.stride = stride2;
+    Compress(refdata2.get(), in_w, in_h, flags, &cptest);
+    CHECK_EQ(cptest, cpdata2);
+  }
+
+  // Uncompress twice: once with 3 components and once with autodetect
+  std::unique_ptr<uint8[]> imgdata1;
+  for (const int components : {0, 3}) {
+    // Uncompress it
+    UncompressFlags flags;
+    flags.components = components;
+    int w, h, c;
+    imgdata1.reset(
+        Uncompress(cpdata1.c_str(), cpdata1.length(), flags, &w, &h, &c, NULL));
+
+    // Check obvious formatting stuff
+    CHECK_EQ(w, in_w);
+    CHECK_EQ(h, in_h);
+    CHECK_EQ(c, 3);
+    CHECK(imgdata1.get());
+
+    // Compare the two images
+    const int totalerr = ComputeSumAbsoluteDifference(
+        imgdata1.get(), refdata1.get(), in_w, in_h, stride1, stride1);
+    CHECK_LE(totalerr, 85000);
+  }
+
+  // check the second image too. Should be bitwise identical to the first.
+  // uncompress using a weird stride
+  {
+    UncompressFlags flags;
+    flags.stride = 3 * 411;
+    const std::unique_ptr<uint8[]> imgdata2(new uint8[flags.stride * in_h]);
+    CHECK(imgdata2.get() == Uncompress(cpdata2.c_str(), cpdata2.length(), flags,
+                                       NULL, [&imgdata2](int w, int h, int c) {
+                                         CHECK_EQ(w, in_w);
+                                         CHECK_EQ(h, in_h);
+                                         CHECK_EQ(c, 3);
+                                         return imgdata2.get();
+                                       }));
+    const int totalerr = ComputeSumAbsoluteDifference(
+        imgdata1.get(), imgdata2.get(), in_w, in_h, stride1, flags.stride);
+    CHECK_EQ(totalerr, 0);
+  }
+}
+
+// Takes JPEG data and reads its headers to determine whether or not the JPEG
+// was chroma downsampled.
+bool IsChromaDownsampled(const string& jpegdata) {
+  // Initialize libjpeg structures to have a memory source
+  // Modify the usual jpeg error manager to catch fatal errors.
+  struct jpeg_decompress_struct cinfo;
+  struct jpeg_error_mgr jerr;
+  jmp_buf jpeg_jmpbuf;
+  cinfo.err = jpeg_std_error(&jerr);
+  cinfo.client_data = &jpeg_jmpbuf;
+  jerr.error_exit = CatchError;
+  if (setjmp(jpeg_jmpbuf)) return false;
+
+  // set up, read header, set image parameters, save size
+  jpeg_create_decompress(&cinfo);
+  SetSrc(&cinfo, jpegdata.c_str(), jpegdata.size(), false);
+
+  jpeg_read_header(&cinfo, TRUE);
+  jpeg_start_decompress(&cinfo);  // required to transfer image size to cinfo
+  const int components = cinfo.output_components;
+  if (components == 1) return false;
+
+  // Check validity
+  CHECK_EQ(3, components);
+  CHECK_EQ(cinfo.comp_info[1].h_samp_factor, cinfo.comp_info[2].h_samp_factor)
+      << "The h sampling factors should be the same.";
+  CHECK_EQ(cinfo.comp_info[1].v_samp_factor, cinfo.comp_info[2].v_samp_factor)
+      << "The v sampling factors should be the same.";
+  for (int i = 0; i < components; ++i) {
+    CHECK_GT(cinfo.comp_info[i].h_samp_factor, 0) << "Invalid sampling factor.";
+    CHECK_EQ(cinfo.comp_info[i].h_samp_factor, cinfo.comp_info[i].v_samp_factor)
+        << "The sampling factor should be the same in both directions.";
+  }
+
+  // We're downsampled if we use fewer samples for color than for brightness.
+  // Do this before deallocating cinfo.
+  const bool downsampled =
+      cinfo.comp_info[1].h_samp_factor < cinfo.comp_info[0].h_samp_factor;
+
+  jpeg_destroy_decompress(&cinfo);
+  return downsampled;
+}
+
+TEST(JpegMemTest, ChromaDownsampling) {
+  // Read the data from a test jpeg file into memory
+  const string jpegfile = string(kTestData) + "jpeg_merge_test1.jpg";
+  string jpeg;
+  ReadFileToStringOrDie(Env::Default(), jpegfile, &jpeg);
+
+  // Verify that compressing the JPEG with chroma downsampling works.
+  //
+  // First, uncompress the JPEG.
+  UncompressFlags unflags;
+  unflags.components = 3;
+  int w, h, c, num_warnings;
+  std::unique_ptr<uint8[]> uncompressed(Uncompress(
+      jpeg.c_str(), jpeg.size(), unflags, &w, &h, &c, &num_warnings));
+  CHECK(uncompressed.get() != NULL);
+  CHECK_EQ(num_warnings, 0);
+
+  // Recompress the JPEG with and without chroma downsampling
+  for (const bool downsample : {false, true}) {
+    CompressFlags flags;
+    flags.format = FORMAT_RGB;
+    flags.quality = 85;
+    flags.chroma_downsampling = downsample;
+    string recompressed;
+    Compress(uncompressed.get(), w, h, flags, &recompressed);
+    CHECK(!recompressed.empty());
+    CHECK_EQ(IsChromaDownsampled(recompressed), downsample);
+  }
+}
+
+void TestBadJPEG(Env* env, const string& bad_jpeg_file, int expected_width,
+                 int expected_height, const string& reference_RGB_file,
+                 const bool try_recover_truncated_jpeg) {
+  string jpeg;
+  ReadFileToStringOrDie(env, bad_jpeg_file, &jpeg);
+
+  UncompressFlags flags;
+  flags.components = 3;
+  flags.try_recover_truncated_jpeg = try_recover_truncated_jpeg;
+
+  int width, height, components;
+  std::unique_ptr<uint8[]> imgdata;
+  imgdata.reset(Uncompress(jpeg.c_str(), jpeg.size(), flags, &width, &height,
+                           &components, NULL));
+  if (expected_width > 0) {  // we expect the file to decode into 'something'
+    CHECK_EQ(width, expected_width);
+    CHECK_EQ(height, expected_height);
+    CHECK_EQ(components, 3);
+    CHECK(imgdata.get());
+    if (!reference_RGB_file.empty()) {
+      string ref;
+      ReadFileToStringOrDie(env, reference_RGB_file, &ref);
+      CHECK(!memcmp(ref.data(), imgdata.get(), ref.size()));
+    }
+  } else {  // no decodable
+    CHECK(!imgdata.get()) << "file:" << bad_jpeg_file;
+  }
+}
+
+TEST(JpegMemTest, BadJpeg) {
+  Env* env = Env::Default();
+  const string data_path = kTestData;
+
+  // Test corrupt file
+  TestBadJPEG(env, data_path + "bad_huffman.jpg", 1024, 768, "", false);
+  TestBadJPEG(env, data_path + "corrupt.jpg", 0 /*120*/, 90, "", false);
+
+  // Truncated files, undecodable because of missing lines:
+  TestBadJPEG(env, data_path + "corrupt34_2.jpg", 0, 3300, "", false);
+  TestBadJPEG(env, data_path + "corrupt34_3.jpg", 0, 3300, "", false);
+  TestBadJPEG(env, data_path + "corrupt34_4.jpg", 0, 3300, "", false);
+
+  // Try in 'recover' mode now:
+  TestBadJPEG(env, data_path + "corrupt34_2.jpg", 2544, 3300, "", true);
+  TestBadJPEG(env, data_path + "corrupt34_3.jpg", 2544, 3300, "", true);
+  TestBadJPEG(env, data_path + "corrupt34_4.jpg", 2544, 3300, "", true);
+}
+
+}  // namespace
+}  // namespace jpeg
+}  // namespace tensorflow