From 2e05c95ecd8077e348078235c072b159a27ba95e Mon Sep 17 00:00:00 2001
From: Pavel P <pavlov.pavel@gmail.com>
Date: Sun, 15 Jul 2018 20:33:46 -0700
Subject: [PATCH] Add multi-file input mode

---
 butteraugli/butteraugli_main.cc | 60 ++++++++++++++++++++++++++++-----
 1 file changed, 51 insertions(+), 9 deletions(-)

diff --git a/butteraugli/butteraugli_main.cc b/butteraugli/butteraugli_main.cc
index bd8ae09..072477e 100755
--- a/butteraugli/butteraugli_main.cc
+++ b/butteraugli/butteraugli_main.cc
@@ -358,18 +358,56 @@ void CreateHeatMapImage(const ImageF& distmap, double good_threshold,
   }
 }
 
+static bool ends_with(const std::string& str,
+                     const std::string& suffix) {
+  return str.size() >= suffix.size() && 0 ==
+    str.compare(str.size() - suffix.size(), suffix.size(), suffix);
+}
+
+// check for multi-file mode (if 3rd argument is png|jpg|jpeg
+// input file
+static bool isMultiFileMode(int argc, char* argv[]) {
+  if (argc >= 4 && (ends_with(argv[3], ".png") ||
+    ends_with(argv[3], ".jpg") || ends_with(argv[3], ".jpeg"))) {
+    FILE* const fl = fopen(argv[3], "rb");
+    if (fl != NULL) {
+      fclose(fl);
+      return true;
+    }
+  }
+  return false;
+}
+
+int ProcessFile(bool multiFile, std::vector<Image8>& rgb1,
+    const char* fileName, const char* heatMapFile);
+
 // main() function, within butteraugli namespace for convenience.
 int Run(int argc, char* argv[]) {
   if (argc != 3 && argc != 4) {
     fprintf(stderr,
             "Usage: %s {image1.(png|jpg|jpeg)} {image2.(png|jpg|jpeg)} "
-            "[heatmap.ppm]\n",
-            argv[0]);
+            "[heatmap.ppm]\n"
+            "or: %s {image1.(png|jpg|jpeg)} {image2.(png|jpg|jpeg)} "
+            "{image3.(png|jpg|jpeg)} ... {imageN.(png|jpg|jpeg)}\n",
+            argv[0], argv[0]);
     return 1;
   }
 
+  int imageN = 2;
+  bool multiFile = isMultiFileMode(argc, argv);
   std::vector<Image8> rgb1 = ReadImageOrDie(argv[1]);
-  std::vector<Image8> rgb2 = ReadImageOrDie(argv[2]);
+  do {
+    if (int ret = ProcessFile(multiFile, rgb1, argv[imageN],
+      multiFile || argc!=4 ? NULL : argv[imageN+1]))
+      return ret;
+  } while (multiFile && ++imageN < argc);
+
+  return 0;
+}
+
+int ProcessFile(bool multiFile, std::vector<Image8>& rgb1,
+    const char* fileName, const char* heatMapFile) {
+  std::vector<Image8> rgb2 = ReadImageOrDie(fileName);
 
   if (rgb1.size() == 3 && rgb2.size() == 4) {
     // Adding a missing alpha channel to one of the images.
@@ -378,7 +416,7 @@ int Run(int argc, char* argv[]) {
     // Adding a missing alpha channel to one of the images.
     rgb2.push_back(Image8(rgb2[0].xsize(), rgb2[0].ysize(), 255));
   } else if (rgb1.size() != rgb2.size()) {
-    fprintf(stderr, "Different number of channels: %lu vs %lu\n", rgb1.size(),
+    fprintf(stderr, "Different number of channels: %zu vs %zu\n", rgb1.size(),
             rgb2.size());
     exit(1);
   }
@@ -387,7 +425,7 @@ int Run(int argc, char* argv[]) {
     if (rgb1[c].xsize() != rgb2[c].xsize() ||
         rgb1[c].ysize() != rgb2[c].ysize()) {
       fprintf(
-          stderr, "The images are not equal in size: (%lu,%lu) vs (%lu,%lu)\n",
+          stderr, "The images are not equal in size: (%zu,%zu) vs (%zu,%zu)\n",
           rgb1[c].xsize(), rgb2[c].xsize(), rgb1[c].ysize(), rgb2[c].ysize());
       return 1;
     }
@@ -424,22 +462,26 @@ int Run(int argc, char* argv[]) {
       diff_map_ptr = &diff_map_on_white;
     }
   }
+  if (multiFile) {
+    printf("%-16lf%s\n", diff_value, fileName);
+    return 0;
+  }
   printf("%lf\n", diff_value);
 
-  if (argc == 4) {
+  if (heatMapFile) {
     const double good_quality = ::butteraugli::ButteraugliFuzzyInverse(1.5);
     const double bad_quality = ::butteraugli::ButteraugliFuzzyInverse(0.5);
     std::vector<uint8_t> rgb;
     CreateHeatMapImage(*diff_map_ptr, good_quality, bad_quality,
                        rgb1[0].xsize(), rgb2[0].ysize(), &rgb);
-    FILE* const fmap = fopen(argv[3], "wb");
+    FILE* const fmap = fopen(heatMapFile, "wb");
     if (fmap == NULL) {
-      fprintf(stderr, "Cannot open %s\n", argv[3]);
+      fprintf(stderr, "Cannot open %s\n", heatMapFile);
       perror("fopen");
       return 1;
     }
     bool ok = true;
-    if (fprintf(fmap, "P6\n%lu %lu\n255\n",
+    if (fprintf(fmap, "P6\n%zu %zu\n255\n",
                       rgb1[0].xsize(), rgb1[0].ysize()) < 0){
       perror("fprintf");
       ok = false;