misc update
[jacinto-ai/caffe-jacinto.git] / src / programs / convert_dataset.cpp
1 // Copyright 2013 Yangqing Jia
2 // This program converts a set of images to a leveldb by storing them as Datum
3 // proto buffers.
4 // Usage:
5 //    convert_dataset 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 0
9 // ....
11 #include <glog/logging.h>
12 #include <leveldb/db.h>
14 #include <string>
15 #include <iostream>
16 #include <fstream>
18 #include "caffe/proto/caffe.pb.h"
19 #include "caffe/util/io.hpp"
21 using namespace caffe;
22 using std::string;
24 // A utility function to generate random strings
25 void GenerateRandomPrefix(const int n, string* key) {
26   const char* kCHARS = "abcdefghijklmnopqrstuvwxyz";
27   key->clear();
28   for (int i = 0; i < n; ++i) {
29     key->push_back(kCHARS[rand() % 26]);
30   }
31   key->push_back('_');
32 }
34 int main(int argc, char** argv) {
35   ::google::InitGoogleLogging(argv[0]);
36   std::ifstream infile(argv[2]);
37   leveldb::DB* db;
38   leveldb::Options options;
39   options.error_if_exists = true;
40   options.create_if_missing = true;
41   LOG(INFO) << "Opening leveldb " << argv[3];
42   leveldb::Status status = leveldb::DB::Open(
43       options, argv[3], &db);
44   CHECK(status.ok()) << "Failed to open leveldb " << argv[3];
46   string root_folder(argv[1]);
47   string filename;
48   int label;
49   Datum datum;
50   string key;
51   string value;
52   while (infile >> filename >> label) {
53     ReadImageToDatum(root_folder + filename, label, &datum);
54     // get the key, and add a random string so the leveldb will have permuted
55     // data
56     GenerateRandomPrefix(8, &key);
57     key += filename;
58     // get the value
59     datum.SerializeToString(&value);
60     db->Put(leveldb::WriteOptions(), key, value);
61     LOG(ERROR) << "Writing " << key;
62   }
64   delete db;
65   return 0;
66 }