Add a version for readPngImageAndPreprocess that can also take a tens…

…or as input. (#2896)

Summary:
*Description*:
Implement a version for readPngImageAndPreprocess that can also take a tensor as input.
*Testing*:
Add a test case in ImageTest: readPngImageAndPreprocessWithAndWithoutInputTensor() that verifies the new version works correctly.
[Optional Fixes #issue]
#2358
Pull Request resolved: #2896

Differential Revision: D15343817

Pulled By: ZchiPitt

fbshipit-source-id: c010e08190fe4625bc9d4484dd1b89c1437cf5c3

include/glow/Base/Image.h

@@ -93,7 +93,8 @@ bool writePngImage(Tensor *T, const char *filename,
                    llvm::ArrayRef<float> mean = zeroMean,
                    llvm::ArrayRef<float> stddev = oneStd);
 
-/// Read a png image and preprocess it according to several parameters.
+/// Read a png image and preprocess it according to several parameters. Create a
+/// tensor and store the preprocessed image data into this tensor.
 /// \param filename the png file to read.
 /// \param imageNormMode normalize values to this range.
 /// \param imageChannelOrder the order of color channels.
@@ -107,6 +108,24 @@ Tensor readPngImageAndPreprocess(llvm::StringRef filename,
                                  llvm::ArrayRef<float> mean = zeroMean,
                                  llvm::ArrayRef<float> stddev = oneStd);
 
+/// Read a png image and preprocess it according to several parameters. Take a
+/// tensor as a parameter and store the preprocessed image data into this
+/// tensor.
+/// \param imageData the tensor into which the preprocessed image data
+///  will be stored.
+/// \param filename the png file to read.
+/// \param imageNormMode normalize values to this range.
+/// \param imageChannelOrder the order of color channels.
+/// \param imageLayout the order of dimensions (channel, height, and width).
+/// \param mean use special mean to normalize.
+/// \param stdev use special stddev to normalize.
+void readPngImageAndPreprocess(Tensor &imageData, llvm::StringRef filename,
+                               ImageNormalizationMode imageNormMode,
+                               ImageChannelOrder imageChannelOrder,
+                               ImageLayout imageLayout,
+                               llvm::ArrayRef<float> mean = zeroMean,
+                               llvm::ArrayRef<float> stddev = oneStd);
+
 /// Loads and normalizes all PNGs into a tensor in the NHWC format with the
 /// requested channel ordering.
 /// \param filenames list of filenames to read.

lib/Base/Image.cpp

@@ -362,6 +362,19 @@ Tensor glow::readPngImageAndPreprocess(llvm::StringRef filename,
                                        llvm::ArrayRef<float> mean,
                                        llvm::ArrayRef<float> stddev) {
   Tensor imageData;
+  readPngImageAndPreprocess(imageData, filename, imageNormMode,
+                            imageChannelOrder, imageLayout, mean, stddev);
+  return imageData;
+}
+
+void glow::readPngImageAndPreprocess(Tensor &imageData,
+                                     llvm::StringRef filename,
+                                     ImageNormalizationMode imageNormMode,
+                                     ImageChannelOrder imageChannelOrder,
+                                     ImageLayout imageLayout,
+                                     llvm::ArrayRef<float> mean,
+                                     llvm::ArrayRef<float> stddev) {
+
   auto range = normModeToRange(imageNormMode);
   bool loadSuccess =
       !readPngImage(&imageData, filename.data(), range, mean, stddev);
@@ -391,7 +404,6 @@ Tensor glow::readPngImageAndPreprocess(llvm::StringRef filename,
     imageData.transpose(&transposed, {2u, 0u, 1u});
     imageData = std::move(transposed);
   }
-  return imageData;
 }
 
 void glow::loadImagesAndPreprocess(const llvm::ArrayRef<std::string> &filenames,

tests/unittests/ImageTest.cpp

@@ -55,6 +55,40 @@ TEST(Image, readBadImages) {
   ASSERT_FALSE(loadSuccess);
 }
 
+TEST(Image, readPngImageAndPreprocessWithAndWithoutInputTensor) {
+  auto image1 = readPngImageAndPreprocess(
+      "tests/images/imagenet/cat_285.png", ImageNormalizationMode::k0to1,
+      ImageChannelOrder::RGB, ImageLayout::NHWC, imagenetNormMean,
+      imagenetNormStd);
+  Tensor image2;
+  readPngImageAndPreprocess(image2, "tests/images/imagenet/cat_285.png",
+                            ImageNormalizationMode::k0to1,
+                            ImageChannelOrder::BGR, ImageLayout::NCHW,
+                            imagenetNormMean, imagenetNormStd);
+
+  // Test if the preprocess actually happened.
+  size_t imgHeight = image1.dims()[0];
+  size_t imgWidth = image1.dims()[1];
+  size_t numChannels = image1.dims()[2];
+
+  Tensor transposed;
+  image2.transpose(&transposed, {1u, 2u, 0u});
+  image2 = std::move(transposed);
+
+  Tensor swizzled(image1.getType());
+  auto IH = image1.getHandle();
+  auto SH = swizzled.getHandle();
+  for (size_t z = 0; z < numChannels; z++) {
+    for (size_t y = 0; y < imgHeight; y++) {
+      for (size_t x = 0; x < imgWidth; x++) {
+        SH.at({x, y, numChannels - 1 - z}) = IH.at({x, y, z});
+      }
+    }
+  }
+  image1 = std::move(swizzled);
+  EXPECT_TRUE(image1.isEqual(image2, 0.01));
+}
+
 TEST(Image, writePngImage) {
   auto range = std::make_pair(0.f, 1.f);
   Tensor localCopy;