1 // This program converts a set of images to a lmdb/leveldb by storing them
2 // as Datum proto buffers.
4 // convert_imageset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME
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
12 #include <fstream> // NOLINT(readability/streams)
17 #include "boost/scoped_ptr.hpp"
18 #include "gflags/gflags.h"
19 #include "glog/logging.h"
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"
27 using namespace caffe; // NOLINT(build/namespaces)
29 using boost::scoped_ptr;
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',...).");
46 int main(int argc, char** argv) {
48 ::google::InitGoogleLogging(argv[0]);
49 // Print output to stderr (while still logging)
50 FLAGS_alsologtostderr = 1;
52 #ifndef GFLAGS_GFLAGS_H_
53 namespace gflags = google;
56 gflags::SetUsageMessage("Convert a set of images to the leveldb/lmdb\n"
57 "format used as input for Caffe.\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);
65 gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/convert_imageset");
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;
74 std::ifstream infile(argv[2]);
75 std::vector<std::pair<std::string, int> > lines;
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));
85 // randomly shuffle data
86 LOG(INFO) << "Shuffling data";
87 shuffle(lines.begin(), lines.end());
89 LOG(INFO) << "A total of " << lines.size() << " images.";
91 if (encode_type.size() && !encoded)
92 LOG(INFO) << "encode_type specified, assuming encoded=true.";
94 int resize_height = std::max<int>(0, FLAGS_resize_height);
95 int resize_width = std::max<int>(0, FLAGS_resize_width);
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());
103 std::string root_folder(argv[1]);
107 bool data_size_initialized = false;
109 for (int line_id = 0; line_id < lines.size(); ++line_id) {
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('.');
117 LOG(WARNING) << "Failed to guess the encoding of '" << fn << "'";
119 std::transform(enc.begin(), enc.end(), enc.begin(), ::tolower);
121 status = ReadImageToDatum(root_folder + lines[line_id].first,
122 lines[line_id].second, resize_height, resize_width, is_color,
124 if (status == false) continue;
126 if (!data_size_initialized) {
127 data_size = datum.channels() * datum.height() * datum.width();
128 data_size_initialized = true;
130 const std::string& data = datum.data();
131 CHECK_EQ(data.size(), data_size) << "Incorrect data field size "
136 string key_str = caffe::format_int(line_id, 8) + "_" + lines[line_id].first;
140 CHECK(datum.SerializeToString(&out));
141 txn->Put(key_str, out);
143 if (++count % 1000 == 0) {
146 txn.reset(db->NewTransaction());
147 LOG(INFO) << "Processed " << count << " files.";
150 // write the last batch
151 if (count % 1000 != 0) {
153 LOG(INFO) << "Processed " << count << " files.";
156 LOG(FATAL) << "This tool requires OpenCV; compile with USE_OPENCV.";