[NFC] Move loadImagesAndPreprocess() to Image.h

rdzhabarov · Mar 28, 2019 · 73bec4b · 73bec4b
1 parent 0f61083
commit 73bec4b
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 41 deletions.
diff --git a/include/glow/Base/Image.h b/include/glow/Base/Image.h
@@ -22,6 +22,9 @@
 
 #include "glow/Base/Tensor.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+
 #include <tuple>
 
 namespace glow {
@@ -91,6 +94,21 @@ Tensor readPngImageAndPreprocess(llvm::StringRef filename,
                                  ImageChannelOrder imageChannelOrder,
                                  ImageLayout imageLayout,
                                  bool useImagenetNormalization);
+
+/// Loads and normalizes all PNGs into a tensor in the NHWC format with the
+/// requested channel ordering.
+/// \param filenames list of filenames to read.
+/// \param inputImageData Tensor to save the resulting output.
+/// \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 useImagenetNormalization use special normalization for Imagenet.
+void loadImagesAndPreprocess(const llvm::ArrayRef<std::string> &filenames,
+                             Tensor *inputImageData,
+                             ImageNormalizationMode imageNormMode,
+                             ImageChannelOrder imageChannelOrder,
+                             ImageLayout imageLayout,
+                             bool useImagenetNormalization);
 } // namespace glow
 
 #endif // GLOW_BASE_IMAGE_H
diff --git a/lib/Base/Image.cpp b/lib/Base/Image.cpp
@@ -391,3 +391,45 @@ Tensor glow::readPngImageAndPreprocess(llvm::StringRef filename,
   }
   return imageData;
 }
+
+void glow::loadImagesAndPreprocess(const llvm::ArrayRef<std::string> &filenames,
+                                   Tensor *inputImageData,
+                                   ImageNormalizationMode imageNormMode,
+                                   ImageChannelOrder imageChannelOrder,
+                                   ImageLayout imageLayout,
+                                   bool useImagenetNormalization) {
+  assert(!filenames.empty() &&
+         "There must be at least one filename in filenames.");
+  unsigned numImages = filenames.size();
+
+  // Get image dimensions and check if grayscale or color.
+  size_t imgHeight;
+  size_t imgWidth;
+  bool isGray;
+  std::tie(imgHeight, imgWidth, isGray) = getPngInfo(filenames[0].c_str());
+  const size_t numChannels = isGray ? 1 : 3;
+
+  // Allocate a tensor for the batch.
+  ShapeVector batchDims;
+  switch (imageLayout) {
+  case ImageLayout::NCHW:
+    batchDims = {numImages, numChannels, imgHeight, imgWidth};
+    break;
+  case ImageLayout::NHWC:
+    batchDims = {numImages, imgHeight, imgWidth, numChannels};
+    break;
+  }
+  inputImageData->reset(ElemKind::FloatTy, batchDims);
+  auto IIDH = inputImageData->getHandle<>();
+
+  // Read images into local tensors and add to batch.
+  for (size_t n = 0; n < filenames.size(); n++) {
+    Tensor localCopy = readPngImageAndPreprocess(filenames[n], imageNormMode,
+                                                 imageChannelOrder, imageLayout,
+                                                 useImagenetNormalization);
+    assert(std::equal(localCopy.dims().begin(), localCopy.dims().end(),
+                      inputImageData->dims().begin() + 1) &&
+           "All images must have the same dimensions");
+    IIDH.insertSlice(localCopy, n);
+  }
+}
diff --git a/tools/loader/ImageClassifier.cpp b/tools/loader/ImageClassifier.cpp
@@ -80,46 +80,6 @@ llvm::cl::opt<bool> convertInAndOutToFp16(
 
 } // namespace
 
-/// Loads and normalizes all PNGs into a tensor in the NHWC format with the
-/// requested channel ordering.
-void loadImagesAndPreprocess(const llvm::cl::list<std::string> &filenames,
-                             Tensor *inputImageData) {
-  assert(!filenames.empty() &&
-         "There must be at least one filename in filenames.");
-  unsigned numImages = filenames.size();
-
-  // Get image dimensions and check if grayscale or color.
-  size_t imgHeight;
-  size_t imgWidth;
-  bool isGray;
-  std::tie(imgHeight, imgWidth, isGray) = getPngInfo(filenames[0].c_str());
-  const size_t numChannels = isGray ? 1 : 3;
-
-  // Allocate a tensor for the batch.
-  ShapeVector batchDims;
-  switch (imageLayout) {
-  case ImageLayout::NCHW:
-    batchDims = {numImages, numChannels, imgHeight, imgWidth};
-    break;
-  case ImageLayout::NHWC:
-    batchDims = {numImages, imgHeight, imgWidth, numChannels};
-    break;
-  }
-  inputImageData->reset(ElemKind::FloatTy, batchDims);
-  auto IIDH = inputImageData->getHandle<>();
-
-  // Read images into local tensors and add to batch.
-  for (size_t n = 0; n < filenames.size(); n++) {
-    Tensor localCopy = readPngImageAndPreprocess(filenames[n], imageNormMode,
-                                                 imageChannelOrder, imageLayout,
-                                                 useImagenetNormalization);
-    assert(std::equal(localCopy.dims().begin(), localCopy.dims().end(),
-                      inputImageData->dims().begin() + 1) &&
-           "All images must have the same dimensions");
-    IIDH.insertSlice(localCopy, n);
-  }
-}
-
 /// Write a prompt to stdout asking for filenames for classification. Read in
 /// those filenames and add them to \p filenames. \p filenames is cleared before
 /// adding the new set of filenames from stdin. \returns false if the passed in
@@ -380,7 +340,9 @@ int main(int argc, char **argv) {
           getNextImageFilenames(&inputImageFilenames)) ||
          isFirstRun) {
     // Load and process the image data into the inputImageData Tensor.
-    loadImagesAndPreprocess(inputImageFilenames, &inputImageData);
+    loadImagesAndPreprocess(inputImageFilenames, &inputImageData, imageNormMode,
+                            imageChannelOrder, imageLayout,
+                            useImagenetNormalization);
 
     // If this is the first run, then we need to build and compile the model.
     if (isFirstRun) {