tools/convert_imageset.cpp

   1 // This program converts a set of images to a lmdb/leveldb by storing them
   2 // as Datum proto buffers.
   3 // Usage:
   4 //   convert_imageset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME
   5 //
   6 // where ROOTFOLDER is the root folder that holds all the images, and LISTFILE
   7 // should be a list of files as well as their labels, in the format as
   8 //   subfolder1/file1.JPEG 7
   9 //   ....
  10
  11 #include <algorithm>
  12 #include <fstream>  // NOLINT(readability/streams)
  13 #include <string>
  14 #include <utility>
  15 #include <vector>
  16
  17 #include "boost/scoped_ptr.hpp"
  18 #include "gflags/gflags.h"
  19 #include "glog/logging.h"
  20
  21 #include "caffe/proto/caffe.pb.h"
  22 #include "caffe/util/db.hpp"
  23 #include "caffe/util/format.hpp"
  24 #include "caffe/util/io.hpp"
  25 #include "caffe/util/rng.hpp"
  26
  27 using namespace caffe;  // NOLINT(build/namespaces)
  28 using std::pair;
  29 using boost::scoped_ptr;
  30
  31 DEFINE_bool(gray, false,
  32     "When this option is on, treat images as grayscale ones");
  33 DEFINE_bool(shuffle, false,
  34     "Randomly shuffle the order of images and their labels");
  35 DEFINE_string(backend, "lmdb",
  36         "The backend {lmdb, leveldb} for storing the result");
  37 DEFINE_int32(resize_width, 0, "Width images are resized to");
  38 DEFINE_int32(resize_height, 0, "Height images are resized to");
  39 DEFINE_bool(check_size, false,
  40     "When this option is on, check that all the datum have the same size");
  41 DEFINE_bool(encoded, false,
  42     "When this option is on, the encoded image will be save in datum");
  43 DEFINE_string(encode_type, "",
  44     "Optional: What type should we encode the image as ('png','jpg',...).");
  45
  46 int main(int argc, char** argv) {
  47 #ifdef USE_OPENCV
  48   ::google::InitGoogleLogging(argv[0]);
  49   // Print output to stderr (while still logging)
  50   FLAGS_alsologtostderr = 1;
  51
  52 #ifndef GFLAGS_GFLAGS_H_
  53   namespace gflags = google;
  54 #endif
  55
  56   gflags::SetUsageMessage("Convert a set of images to the leveldb/lmdb\n"
  57         "format used as input for Caffe.\n"
  58         "Usage:\n"
  59         "    convert_imageset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME\n"
  60         "The ImageNet dataset for the training demo is at\n"
  61         "    http://www.image-net.org/download-images\n");
  62   gflags::ParseCommandLineFlags(&argc, &argv, true);
  63
  64   if (argc < 4) {
  65     gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/convert_imageset");
  66     return 1;
  67   }
  68
  69   const bool is_color = !FLAGS_gray;
  70   const bool check_size = FLAGS_check_size;
  71   const bool encoded = FLAGS_encoded;
  72   const string encode_type = FLAGS_encode_type;
  73
  74   std::ifstream infile(argv[2]);
  75   std::vector<std::pair<std::string, int> > lines;
  76   std::string line;
  77   size_t pos;
  78   int label;
  79   while (std::getline(infile, line)) {
  80     pos = line.find_last_of(' ');
  81     label = atoi(line.substr(pos + 1).c_str());
  82     lines.push_back(std::make_pair(line.substr(0, pos), label));
  83   }
  84   if (FLAGS_shuffle) {
  85     // randomly shuffle data
  86     LOG(INFO) << "Shuffling data";
  87     shuffle(lines.begin(), lines.end());
  88   }
  89   LOG(INFO) << "A total of " << lines.size() << " images.";
  90
  91   if (encode_type.size() && !encoded)
  92     LOG(INFO) << "encode_type specified, assuming encoded=true.";
  93
  94   int resize_height = std::max<int>(0, FLAGS_resize_height);
  95   int resize_width = std::max<int>(0, FLAGS_resize_width);
  96
  97   // Create new DB
  98   scoped_ptr<db::DB> db(db::GetDB(FLAGS_backend));
  99   db->Open(argv[3], db::NEW);
 100   scoped_ptr<db::Transaction> txn(db->NewTransaction());
 101
 102   // Storing to db
 103   std::string root_folder(argv[1]);
 104   Datum datum;
 105   int count = 0;
 106   int data_size = 0;
 107   bool data_size_initialized = false;
 108
 109   for (int line_id = 0; line_id < lines.size(); ++line_id) {
 110     bool status;
 111     std::string enc = encode_type;
 112     if (encoded && !enc.size()) {
 113       // Guess the encoding type from the file name
 114       string fn = lines[line_id].first;
 115       size_t p = fn.rfind('.');
 116       if ( p == fn.npos )
 117         LOG(WARNING) << "Failed to guess the encoding of '" << fn << "'";
 118       enc = fn.substr(p+1);
 119       std::transform(enc.begin(), enc.end(), enc.begin(), ::tolower);
 120     }
 121     status = ReadImageToDatum(root_folder + lines[line_id].first,
 122         lines[line_id].second, resize_height, resize_width, is_color,
 123         enc, &datum);
 124     if (status == false) continue;
 125     if (check_size) {
 126       if (!data_size_initialized) {
 127         data_size = datum.channels() * datum.height() * datum.width();
 128         data_size_initialized = true;
 129       } else {
 130         const std::string& data = datum.data();
 131         CHECK_EQ(data.size(), data_size) << "Incorrect data field size "
 132             << data.size();
 133       }
 134     }
 135     // sequential
 136     string key_str = caffe::format_int(line_id, 8) + "_" + lines[line_id].first;
 137
 138     // Put in db
 139     string out;
 140     CHECK(datum.SerializeToString(&out));
 141     txn->Put(key_str, out);
 142
 143     if (++count % 1000 == 0) {
 144       // Commit db
 145       txn->Commit();
 146       txn.reset(db->NewTransaction());
 147       LOG(INFO) << "Processed " << count << " files.";
 148     }
 149   }
 150   // write the last batch
 151   if (count % 1000 != 0) {
 152     txn->Commit();
 153     LOG(INFO) << "Processed " << count << " files.";
 154   }
 155 #else
 156   LOG(FATAL) << "This tool requires OpenCV; compile with USE_OPENCV.";
 157 #endif  // USE_OPENCV
 158   return 0;
 159 }