Added an iscolor flag to io.cpp method ReadImageToDatum to handle grayscale images and a corresponding commandline flag [-g] to convert_imageset.cpp.

1 files changed, 21 insertions, 14 deletions

diff --git a/tools/convert_imageset.cpp b/tools/convert_imageset.cpp
index 2420f953..bbf848ad 100644
--- a/tools/convert_imageset.cpp
+++ b/tools/convert_imageset.cpp
@@ -2,7 +2,7 @@
 // This program converts a set of images to a leveldb by storing them as Datum
 // proto buffers.
 // Usage:
-//    convert_imageset ROOTFOLDER/ LISTFILE DB_NAME RANDOM_SHUFFLE[0 or 1] \
+//   convert_imageset [-g] ROOTFOLDER/ LISTFILE DB_NAME RANDOM_SHUFFLE[0 or 1] \
 //                     [resize_height] [resize_width]
 // where ROOTFOLDER is the root folder that holds all the images, and LISTFILE
 // should be a list of files as well as their labels, in the format as
@@ -10,6 +10,9 @@
 //   ....
 // if RANDOM_SHUFFLE is 1, a random shuffle will be carried out before we
 // process the file lines.
+// Optional flag -g indicates the images should be read as
+// single-channel grayscale. If omitted, grayscale images will be
+// converted to color.
 
 #include <glog/logging.h>
 #include <leveldb/db.h>
@@ -30,24 +33,28 @@ using std::string;
 
 int main(int argc, char** argv) {
   ::google::InitGoogleLogging(argv[0]);
-  if (argc < 4 || argc > 7) {
+  if (argc < 4 || argc > 8) {
     printf("Convert a set of images to the leveldb format used\n"
         "as input for Caffe.\n"
         "Usage:\n"
-        "    convert_imageset ROOTFOLDER/ LISTFILE DB_NAME"
+        "    convert_imageset [-g] ROOTFOLDER/ LISTFILE DB_NAME"
         " RANDOM_SHUFFLE_DATA[0 or 1] [resize_height] [resize_width]\n"
         "The ImageNet dataset for the training demo is at\n"
         "    http://www.image-net.org/download-images\n");
     return 1;
   }
-  std::ifstream infile(argv[2]);
+
+  // Test whether argv[1] == "-g"
+  bool iscolor= !(string("-g") == string(argv[1]));
+  int  arg_offset = (iscolor ? 0 : 1);
+  std::ifstream infile(argv[arg_offset+2]);
   std::vector<std::pair<string, int> > lines;
   string filename;
   int label;
   while (infile >> filename >> label) {
     lines.push_back(std::make_pair(filename, label));
   }
-  if (argc >= 5 && argv[4][0] == '1') {
+  if (argc >= (arg_offset+5) && argv[arg_offset+4][0] == '1') {
     // randomly shuffle data
     LOG(INFO) << "Shuffling data";
     std::random_shuffle(lines.begin(), lines.end());
@@ -55,11 +62,11 @@ int main(int argc, char** argv) {
   LOG(INFO) << "A total of " << lines.size() << " images.";
   int resize_height = 0;
   int resize_width = 0;
-  if (argc >= 6) {
-    resize_height = atoi(argv[5]);
+  if (argc >= (arg_offset+6)) {
+    resize_height = atoi(argv[arg_offset+5]);
   }
-  if (argc >= 7) {
-    resize_width = atoi(argv[6]);
+  if (argc >= (arg_offset+7)) {
+    resize_width = atoi(argv[arg_offset+6]);
   }
 
   leveldb::DB* db;
@@ -67,12 +74,12 @@ int main(int argc, char** argv) {
   options.error_if_exists = true;
   options.create_if_missing = true;
   options.write_buffer_size = 268435456;
-  LOG(INFO) << "Opening leveldb " << argv[3];
+  LOG(INFO) << "Opening leveldb " << argv[arg_offset+3];
   leveldb::Status status = leveldb::DB::Open(
-      options, argv[3], &db);
-  CHECK(status.ok()) << "Failed to open leveldb " << argv[3];
+      options, argv[arg_offset+3], &db);
+  CHECK(status.ok()) << "Failed to open leveldb " << argv[arg_offset+3];
 
-  string root_folder(argv[1]);
+  string root_folder(argv[arg_offset+1]);
   Datum datum;
   int count = 0;
   const int kMaxKeyLength = 256;
@@ -82,7 +89,7 @@ int main(int argc, char** argv) {
   bool data_size_initialized = false;
   for (int line_id = 0; line_id < lines.size(); ++line_id) {
     if (!ReadImageToDatum(root_folder + lines[line_id].first,
-        lines[line_id].second, resize_height, resize_width, &datum)) {
+         lines[line_id].second, resize_height, resize_width, iscolor, &datum)) {
       continue;
     }
     if (!data_size_initialized) {