30 #include <opencv2/opencv.hpp> 39 #include "caffe2/core/common.h" 40 #include "caffe2/core/db.h" 41 #include "caffe2/core/init.h" 42 #include "caffe2/proto/caffe2.pb.h" 43 #include "caffe2/core/logging.h" 45 CAFFE2_DEFINE_bool(shuffle,
false,
46 "Randomly shuffle the order of images and their labels");
47 CAFFE2_DEFINE_string(input_folder,
"",
"The input image file name.");
51 "The text file containing the list of images.");
52 CAFFE2_DEFINE_string(output_db_name,
"",
"The output training leveldb name.");
53 CAFFE2_DEFINE_string(db,
"leveldb",
"The db type.");
54 CAFFE2_DEFINE_bool(raw,
false,
55 "If set, we pre-read the images and store the raw buffer.");
56 CAFFE2_DEFINE_bool(color,
true,
"If set, load images in color.");
60 "If caffe2::FLAGS_raw is set, scale the shorter edge to the given value.");
61 CAFFE2_DEFINE_bool(warp,
false,
"If warp is set, warp the images to square.");
65 "Number of image parsing and conversion threads.");
72 data_ = protos_.add_protos();
73 label_ = protos_.add_protos();
74 if (caffe2::FLAGS_raw) {
75 data_->set_data_type(TensorProto::BYTE);
78 if (caffe2::FLAGS_color) {
82 data_->set_data_type(TensorProto::STRING);
84 data_->add_string_data(
"");
86 label_->set_data_type(TensorProto::INT32);
88 label_->add_int32_data(0);
92 if (thread_.joinable()) {
97 void queue(
const std::pair<std::string, int>& pair) {
102 thread_ = std::thread(&Converter::run,
this);
106 std::unique_lock<std::mutex> lock(mutex_);
107 while (out_.empty()) {
111 auto value = out_.front();
118 const auto& input_folder = caffe2::FLAGS_input_folder;
119 std::unique_lock<std::mutex> lock(mutex_);
121 while (!in_.empty()) {
122 auto pair = in_.front();
126 label_->set_int32_data(0, pair.second);
129 if (!caffe2::FLAGS_raw) {
130 std::ifstream image_file_stream(input_folder + pair.first);
131 if (!image_file_stream) {
132 LOG(ERROR) <<
"Cannot open " << input_folder << pair.first
135 data_->mutable_string_data(0)->assign(
136 std::istreambuf_iterator<char>(image_file_stream),
137 std::istreambuf_iterator<char>());
141 cv::Mat img = cv::imread(
142 input_folder + pair.first,
143 caffe2::FLAGS_color ? CV_LOAD_IMAGE_COLOR
144 : CV_LOAD_IMAGE_GRAYSCALE);
148 int scaled_width, scaled_height;
149 if (caffe2::FLAGS_warp) {
150 scaled_width = caffe2::FLAGS_scale;
151 scaled_height = caffe2::FLAGS_scale;
152 }
else if (img.rows > img.cols) {
153 scaled_width = caffe2::FLAGS_scale;
155 static_cast<float>(img.rows) * caffe2::FLAGS_scale / img.cols;
157 scaled_height = caffe2::FLAGS_scale;
159 static_cast<float>(img.cols) * caffe2::FLAGS_scale / img.rows;
164 cv::Size(scaled_width, scaled_height),
168 data_->set_dims(0, scaled_height);
169 data_->set_dims(1, scaled_width);
172 DCHECK(resized_img.isContinuous());
173 auto nbytes = resized_img.total() * resized_img.elemSize();
174 data_->set_byte_data(resized_img.ptr(), nbytes);
177 protos_.SerializeToString(&value);
181 while (!out_.empty()) {
190 TensorProtos protos_;
193 std::queue<std::pair<std::string, int>> in_;
194 std::queue<std::string> out_;
197 std::condition_variable cv_;
201 void ConvertImageDataset(
202 const string& input_folder,
203 const string& list_filename,
204 const string& output_db_name,
206 std::ifstream list_file(list_filename);
207 std::vector<std::pair<std::string, int> > lines;
208 std::string filename;
210 while (list_file >> filename >> file_label) {
211 lines.push_back(std::make_pair(filename, file_label));
214 if (caffe2::FLAGS_shuffle) {
215 LOG(INFO) <<
"Shuffling data";
216 std::shuffle(lines.begin(), lines.end(), std::default_random_engine(1701));
219 auto num_threads = caffe2::FLAGS_num_threads;
220 if (num_threads < 1) {
221 num_threads = std::thread::hardware_concurrency();
224 LOG(INFO) <<
"Processing " << lines.size() <<
" images...";
225 LOG(INFO) <<
"Opening DB " << output_db_name;
227 auto db = db::CreateDB(caffe2::FLAGS_db, output_db_name, db::NEW);
228 auto transaction = db->NewTransaction();
230 LOG(INFO) <<
"Using " << num_threads <<
" processing threads...";
231 std::vector<Converter> converters(num_threads);
234 for (
auto i = 0; i < lines.size(); i++) {
235 converters[i % converters.size()].queue(lines[i]);
239 for (
auto& converter : converters) {
243 constexpr
auto key_max_length = 256;
244 char key_cstr[key_max_length];
247 for (
auto i = 0; i < lines.size(); i++) {
249 auto value = converters[i % converters.size()].get();
252 auto key_len = snprintf(
253 key_cstr,
sizeof(key_cstr),
"%08d_%s", i, lines[i].first.c_str());
254 DCHECK_LE(key_len,
sizeof(key_cstr));
257 transaction->Put(
string(key_cstr), value);
259 if (++count % 1000 == 0) {
261 transaction->Commit();
262 LOG(INFO) <<
"Processed " << count <<
" files.";
267 transaction->Commit();
268 LOG(INFO) <<
"Processed " << count <<
" files.";
274 int main(
int argc,
char** argv) {
276 caffe2::ConvertImageDataset(
277 caffe2::FLAGS_input_folder, caffe2::FLAGS_list_file,
278 caffe2::FLAGS_output_db_name, caffe2::FLAGS_shuffle);
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...