diff --git a/data/coco/README.md b/data/coco/README.md new file mode 100644 index 0000000..7589a7c --- /dev/null +++ b/data/coco/README.md @@ -0,0 +1,38 @@ +### Preparation +1. Download Images and Annotations from [MSCOCO](http://mscoco.org/dataset/#download). By default, we assume the data is stored in `$HOME/data/coco` + +2. Get the coco code. We will call the directory that you cloned coco into `$COCO_ROOT` + ```Shell + git clone https://github.com/weiliu89/coco.git + cd coco + git checkout dev + ``` + +3. Build the coco code. + ```Shell + cd PythonAPI + python setup.py build_ext --inplace + ``` + +4. Split the annotation to many files per image and get the image size info. + ```Shell + # Check scripts/batch_split_annotation.py and change settings accordingly. + python scripts/batch_split_annotation.py + # Create the minival2014_name_size.txt and test-dev2015_name_size.txt in $CAFFE_ROOT/data/coco + python scripts/batch_get_image_size.py + ``` + +5. Create the LMDB file. + ```Shell + cd $CAFFE_ROOT + # Create the minival.txt, testdev.txt, test.txt, train.txt in data/coco/ + python data/coco/create_list.py + # You can modify the parameters in create_data.sh if needed. + # It will create lmdb files for minival, testdev, test, and train with encoded original image: + # - $HOME/data/coco/lmdb/coco_minival_lmdb + # - $HOME/data/coco/lmdb/coco_testdev_lmdb + # - $HOME/data/coco/lmdb/coco_test_lmdb + # - $HOME/data/coco/lmdb/coco_train_lmdb + # and make soft links at examples/coco/ + ./data/coco/create_data.sh + ``` diff --git a/data/coco/create_data.sh b/data/coco/create_data.sh new file mode 100644 index 0000000..3ff38de --- /dev/null +++ b/data/coco/create_data.sh @@ -0,0 +1,29 @@ +cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd ) +root_dir=$cur_dir/../.. + +cd $root_dir + +redo=false +data_root_dir="$HOME/data/coco" +dataset_name="coco" +mapfile="$root_dir/data/$dataset_name/labelmap_coco.prototxt" +anno_type="detection" +label_type="json" +db="lmdb" +min_dim=0 +max_dim=0 +width=0 +height=0 + +extra_cmd="--encode-type=jpg --encoded" +if $redo +then + extra_cmd="$extra_cmd --redo" +fi +#for subset in minival testdev train test +for subset in minival testdev train test +do + python $root_dir/scripts/create_annoset.py --anno-type=$anno_type --label-type=$label_type --label-map-file=$mapfile --min-dim=$min_dim --max-dim=$max_dim \ + --resize-width=$width --resize-height=$height --check-label $extra_cmd $data_root_dir $root_dir/data/$dataset_name/$subset.txt \ + $data_root_dir/$db/$dataset_name"_"$subset"_"$db examples/$dataset_name 2>&1 | tee $root_dir/data/$dataset_name/$subset.log +done diff --git a/data/coco/create_list.py b/data/coco/create_list.py new file mode 100644 index 0000000..27d1968 --- /dev/null +++ b/data/coco/create_list.py @@ -0,0 +1,121 @@ +import argparse +import os +from random import shuffle +import shutil +import subprocess +import sys + +HOMEDIR = os.path.expanduser("~") +CURDIR = os.path.dirname(os.path.realpath(__file__)) + +# If true, re-create all list files. +redo = True +# The root directory which holds all information of the dataset. +data_dir = "{}/data/coco".format(HOMEDIR) +# The directory name which holds the image sets. +imgset_dir = "ImageSets" +# The direcotry which contains the images. +img_dir = "images" +img_ext = "jpg" +# The directory which contains the annotations. +anno_dir = "Annotations" +anno_ext = "json" + +train_list_file = "{}/train.txt".format(CURDIR) +minival_list_file = "{}/minival.txt".format(CURDIR) +testdev_list_file = "{}/testdev.txt".format(CURDIR) +test_list_file = "{}/test.txt".format(CURDIR) + +# Create training set. +# We follow Ross Girschick's split. +if redo or not os.path.exists(train_list_file): + datasets = ["train2014", "valminusminival2014"] + img_files = [] + anno_files = [] + for dataset in datasets: + imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset) + with open(imgset_file, "r") as f: + for line in f.readlines(): + name = line.strip("\n") + subset = name.split("_")[1] + img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext) + assert os.path.exists("{}/{}".format(data_dir, img_file)), \ + "{}/{} does not exist".format(data_dir, img_file) + anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext) + assert os.path.exists("{}/{}".format(data_dir, anno_file)), \ + "{}/{} does not exist".format(data_dir, anno_file) + img_files.append(img_file) + anno_files.append(anno_file) + # Shuffle the images. + idx = [i for i in xrange(len(img_files))] + shuffle(idx) + with open(train_list_file, "w") as f: + for i in idx: + f.write("{} {}\n".format(img_files[i], anno_files[i])) + +if redo or not os.path.exists(minival_list_file): + datasets = ["minival2014"] + subset = "val2014" + img_files = [] + anno_files = [] + for dataset in datasets: + imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset) + with open(imgset_file, "r") as f: + for line in f.readlines(): + name = line.strip("\n") + img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext) + assert os.path.exists("{}/{}".format(data_dir, img_file)), \ + "{}/{} does not exist".format(data_dir, img_file) + anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext) + assert os.path.exists("{}/{}".format(data_dir, anno_file)), \ + "{}/{} does not exist".format(data_dir, anno_file) + img_files.append(img_file) + anno_files.append(anno_file) + with open(minival_list_file, "w") as f: + for i in xrange(len(img_files)): + f.write("{} {}\n".format(img_files[i], anno_files[i])) + +if redo or not os.path.exists(testdev_list_file): + datasets = ["test-dev2015"] + subset = "test2015" + img_files = [] + anno_files = [] + for dataset in datasets: + imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset) + with open(imgset_file, "r") as f: + for line in f.readlines(): + name = line.strip("\n") + img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext) + assert os.path.exists("{}/{}".format(data_dir, img_file)), \ + "{}/{} does not exist".format(data_dir, img_file) + anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext) + assert os.path.exists("{}/{}".format(data_dir, anno_file)), \ + "{}/{} does not exist".format(data_dir, anno_file) + img_files.append(img_file) + anno_files.append(anno_file) + with open(testdev_list_file, "w") as f: + for i in xrange(len(img_files)): + f.write("{} {}\n".format(img_files[i], anno_files[i])) + +if redo or not os.path.exists(test_list_file): + datasets = ["test2015"] + subset = "test2015" + img_files = [] + anno_files = [] + for dataset in datasets: + imgset_file = "{}/{}/{}.txt".format(data_dir, imgset_dir, dataset) + with open(imgset_file, "r") as f: + for line in f.readlines(): + name = line.strip("\n") + img_file = "{}/{}/{}.{}".format(img_dir, subset, name, img_ext) + assert os.path.exists("{}/{}".format(data_dir, img_file)), \ + "{}/{} does not exist".format(data_dir, img_file) + anno_file = "{}/{}/{}.{}".format(anno_dir, subset, name, anno_ext) + assert os.path.exists("{}/{}".format(data_dir, anno_file)), \ + "{}/{} does not exist".format(data_dir, anno_file) + img_files.append(img_file) + anno_files.append(anno_file) + with open(test_list_file, "w") as f: + for i in xrange(len(img_files)): + f.write("{} {}\n".format(img_files[i], anno_files[i])) + diff --git a/data/coco/labelmap_coco.prototxt b/data/coco/labelmap_coco.prototxt new file mode 100644 index 0000000..82252d2 --- /dev/null +++ b/data/coco/labelmap_coco.prototxt @@ -0,0 +1,405 @@ +item { + name: "none_of_the_above" + label: 0 + display_name: "background" +} +item { + name: "1" + label: 1 + display_name: "person" +} +item { + name: "2" + label: 2 + display_name: "bicycle" +} +item { + name: "3" + label: 3 + display_name: "car" +} +item { + name: "4" + label: 4 + display_name: "motorcycle" +} +item { + name: "5" + label: 5 + display_name: "airplane" +} +item { + name: "6" + label: 6 + display_name: "bus" +} +item { + name: "7" + label: 7 + display_name: "train" +} +item { + name: "8" + label: 8 + display_name: "truck" +} +item { + name: "9" + label: 9 + display_name: "boat" +} +item { + name: "10" + label: 10 + display_name: "traffic light" +} +item { + name: "11" + label: 11 + display_name: "fire hydrant" +} +item { + name: "13" + label: 12 + display_name: "stop sign" +} +item { + name: "14" + label: 13 + display_name: "parking meter" +} +item { + name: "15" + label: 14 + display_name: "bench" +} +item { + name: "16" + label: 15 + display_name: "bird" +} +item { + name: "17" + label: 16 + display_name: "cat" +} +item { + name: "18" + label: 17 + display_name: "dog" +} +item { + name: "19" + label: 18 + display_name: "horse" +} +item { + name: "20" + label: 19 + display_name: "sheep" +} +item { + name: "21" + label: 20 + display_name: "cow" +} +item { + name: "22" + label: 21 + display_name: "elephant" +} +item { + name: "23" + label: 22 + display_name: "bear" +} +item { + name: "24" + label: 23 + display_name: "zebra" +} +item { + name: "25" + label: 24 + display_name: "giraffe" +} +item { + name: "27" + label: 25 + display_name: "backpack" +} +item { + name: "28" + label: 26 + display_name: "umbrella" +} +item { + name: "31" + label: 27 + display_name: "handbag" +} +item { + name: "32" + label: 28 + display_name: "tie" +} +item { + name: "33" + label: 29 + display_name: "suitcase" +} +item { + name: "34" + label: 30 + display_name: "frisbee" +} +item { + name: "35" + label: 31 + display_name: "skis" +} +item { + name: "36" + label: 32 + display_name: "snowboard" +} +item { + name: "37" + label: 33 + display_name: "sports ball" +} +item { + name: "38" + label: 34 + display_name: "kite" +} +item { + name: "39" + label: 35 + display_name: "baseball bat" +} +item { + name: "40" + label: 36 + display_name: "baseball glove" +} +item { + name: "41" + label: 37 + display_name: "skateboard" +} +item { + name: "42" + label: 38 + display_name: "surfboard" +} +item { + name: "43" + label: 39 + display_name: "tennis racket" +} +item { + name: "44" + label: 40 + display_name: "bottle" +} +item { + name: "46" + label: 41 + display_name: "wine glass" +} +item { + name: "47" + label: 42 + display_name: "cup" +} +item { + name: "48" + label: 43 + display_name: "fork" +} +item { + name: "49" + label: 44 + display_name: "knife" +} +item { + name: "50" + label: 45 + display_name: "spoon" +} +item { + name: "51" + label: 46 + display_name: "bowl" +} +item { + name: "52" + label: 47 + display_name: "banana" +} +item { + name: "53" + label: 48 + display_name: "apple" +} +item { + name: "54" + label: 49 + display_name: "sandwich" +} +item { + name: "55" + label: 50 + display_name: "orange" +} +item { + name: "56" + label: 51 + display_name: "broccoli" +} +item { + name: "57" + label: 52 + display_name: "carrot" +} +item { + name: "58" + label: 53 + display_name: "hot dog" +} +item { + name: "59" + label: 54 + display_name: "pizza" +} +item { + name: "60" + label: 55 + display_name: "donut" +} +item { + name: "61" + label: 56 + display_name: "cake" +} +item { + name: "62" + label: 57 + display_name: "chair" +} +item { + name: "63" + label: 58 + display_name: "couch" +} +item { + name: "64" + label: 59 + display_name: "potted plant" +} +item { + name: "65" + label: 60 + display_name: "bed" +} +item { + name: "67" + label: 61 + display_name: "dining table" +} +item { + name: "70" + label: 62 + display_name: "toilet" +} +item { + name: "72" + label: 63 + display_name: "tv" +} +item { + name: "73" + label: 64 + display_name: "laptop" +} +item { + name: "74" + label: 65 + display_name: "mouse" +} +item { + name: "75" + label: 66 + display_name: "remote" +} +item { + name: "76" + label: 67 + display_name: "keyboard" +} +item { + name: "77" + label: 68 + display_name: "cell phone" +} +item { + name: "78" + label: 69 + display_name: "microwave" +} +item { + name: "79" + label: 70 + display_name: "oven" +} +item { + name: "80" + label: 71 + display_name: "toaster" +} +item { + name: "81" + label: 72 + display_name: "sink" +} +item { + name: "82" + label: 73 + display_name: "refrigerator" +} +item { + name: "84" + label: 74 + display_name: "book" +} +item { + name: "85" + label: 75 + display_name: "clock" +} +item { + name: "86" + label: 76 + display_name: "vase" +} +item { + name: "87" + label: 77 + display_name: "scissors" +} +item { + name: "88" + label: 78 + display_name: "teddy bear" +} +item { + name: "89" + label: 79 + display_name: "hair drier" +} +item { + name: "90" + label: 80 + display_name: "toothbrush" +} diff --git a/tools/convert_annoset.cpp b/tools/convert_annoset.cpp new file mode 100644 index 0000000..32ca258 --- /dev/null +++ b/tools/convert_annoset.cpp @@ -0,0 +1,207 @@ +// This program converts a set of images and annotations to a lmdb/leveldb by +// storing them as AnnotatedDatum proto buffers. +// Usage: +// convert_annoset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME +// +// where ROOTFOLDER is the root folder that holds all the images and +// annotations, and LISTFILE should be a list of files as well as their labels +// or label files. +// For classification task, the file should be in the format as +// imgfolder1/img1.JPEG 7 +// .... +// For detection task, the file should be in the format as +// imgfolder1/img1.JPEG annofolder1/anno1.xml +// .... + +#include +#include // NOLINT(readability/streams) +#include +#include +#include +#include + +#include "boost/scoped_ptr.hpp" +#include "boost/variant.hpp" +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "caffe/proto/caffe.pb.h" +#include "caffe/util/db.hpp" +#include "caffe/util/format.hpp" +#include "caffe/util/io.hpp" +#include "caffe/util/rng.hpp" + +using namespace caffe; // NOLINT(build/namespaces) +using std::pair; +using boost::scoped_ptr; + +DEFINE_bool(gray, false, + "When this option is on, treat images as grayscale ones"); +DEFINE_bool(shuffle, false, + "Randomly shuffle the order of images and their labels"); +DEFINE_string(backend, "lmdb", + "The backend {lmdb, leveldb} for storing the result"); +DEFINE_string(anno_type, "classification", + "The type of annotation {classification, detection}."); +DEFINE_string(label_type, "xml", + "The type of annotation file format."); +DEFINE_string(label_map_file, "", + "A file with LabelMap protobuf message."); +DEFINE_bool(check_label, false, + "When this option is on, check that there is no duplicated name/label."); +DEFINE_int32(min_dim, 0, + "Minimum dimension images are resized to (keep same aspect ratio)"); +DEFINE_int32(max_dim, 0, + "Maximum dimension images are resized to (keep same aspect ratio)"); +DEFINE_int32(resize_width, 0, "Width images are resized to"); +DEFINE_int32(resize_height, 0, "Height images are resized to"); +DEFINE_bool(check_size, false, + "When this option is on, check that all the datum have the same size"); +DEFINE_bool(encoded, false, + "When this option is on, the encoded image will be save in datum"); +DEFINE_string(encode_type, "", + "Optional: What type should we encode the image as ('png','jpg',...)."); + +int main(int argc, char** argv) { +#ifdef USE_OPENCV + ::google::InitGoogleLogging(argv[0]); + // Print output to stderr (while still logging) + FLAGS_alsologtostderr = 1; + +#ifndef GFLAGS_GFLAGS_H_ + namespace gflags = google; +#endif + + gflags::SetUsageMessage("Convert a set of images and annotations to the " + "leveldb/lmdb format used as input for Caffe.\n" + "Usage:\n" + " convert_annoset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME\n"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (argc < 4) { + gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/convert_annoset"); + return 1; + } + + const bool is_color = !FLAGS_gray; + const bool check_size = FLAGS_check_size; + const bool encoded = FLAGS_encoded; + const string encode_type = FLAGS_encode_type; + const string anno_type = FLAGS_anno_type; + AnnotatedDatum_AnnotationType type; + const string label_type = FLAGS_label_type; + const string label_map_file = FLAGS_label_map_file; + const bool check_label = FLAGS_check_label; + std::map name_to_label; + + std::ifstream infile(argv[2]); + std::vector > > lines; + std::string filename; + int label; + std::string labelname; + if (anno_type == "classification") { + while (infile >> filename >> label) { + lines.push_back(std::make_pair(filename, label)); + } + } else if (anno_type == "detection") { + type = AnnotatedDatum_AnnotationType_BBOX; + LabelMap label_map; + CHECK(ReadProtoFromTextFile(label_map_file, &label_map)) + << "Failed to read label map file."; + CHECK(MapNameToLabel(label_map, check_label, &name_to_label)) + << "Failed to convert name to label."; + while (infile >> filename >> labelname) { + lines.push_back(std::make_pair(filename, labelname)); + } + } + if (FLAGS_shuffle) { + // randomly shuffle data + LOG(INFO) << "Shuffling data"; + shuffle(lines.begin(), lines.end()); + } + LOG(INFO) << "A total of " << lines.size() << " images."; + + if (encode_type.size() && !encoded) + LOG(INFO) << "encode_type specified, assuming encoded=true."; + + int min_dim = std::max(0, FLAGS_min_dim); + int max_dim = std::max(0, FLAGS_max_dim); + int resize_height = std::max(0, FLAGS_resize_height); + int resize_width = std::max(0, FLAGS_resize_width); + + // Create new DB + scoped_ptr db(db::GetDB(FLAGS_backend)); + db->Open(argv[3], db::NEW); + scoped_ptr txn(db->NewTransaction()); + + // Storing to db + std::string root_folder(argv[1]); + AnnotatedDatum anno_datum; + Datum* datum = anno_datum.mutable_datum(); + int count = 0; + int data_size = 0; + bool data_size_initialized = false; + + for (int line_id = 0; line_id < lines.size(); ++line_id) { + bool status = true; + std::string enc = encode_type; + if (encoded && !enc.size()) { + // Guess the encoding type from the file name + string fn = lines[line_id].first; + size_t p = fn.rfind('.'); + if ( p == fn.npos ) + LOG(WARNING) << "Failed to guess the encoding of '" << fn << "'"; + enc = fn.substr(p); + std::transform(enc.begin(), enc.end(), enc.begin(), ::tolower); + } + filename = root_folder + lines[line_id].first; + if (anno_type == "classification") { + label = boost::get(lines[line_id].second); + status = ReadImageToDatum(filename, label, resize_height, resize_width, + min_dim, max_dim, is_color, enc, datum); + } else if (anno_type == "detection") { + labelname = root_folder + boost::get(lines[line_id].second); + status = ReadRichImageToAnnotatedDatum(filename, labelname, resize_height, + resize_width, min_dim, max_dim, is_color, enc, type, label_type, + name_to_label, &anno_datum); + anno_datum.set_type(AnnotatedDatum_AnnotationType_BBOX); + } + if (status == false) { + LOG(WARNING) << "Failed to read " << lines[line_id].first; + continue; + } + if (check_size) { + if (!data_size_initialized) { + data_size = datum->channels() * datum->height() * datum->width(); + data_size_initialized = true; + } else { + const std::string& data = datum->data(); + CHECK_EQ(data.size(), data_size) << "Incorrect data field size " + << data.size(); + } + } + // sequential + string key_str = caffe::format_int(line_id, 8) + "_" + lines[line_id].first; + + // Put in db + string out; + CHECK(anno_datum.SerializeToString(&out)); + txn->Put(key_str, out); + + if (++count % 1000 == 0) { + // Commit db + txn->Commit(); + txn.reset(db->NewTransaction()); + LOG(INFO) << "Processed " << count << " files."; + } + } + // write the last batch + if (count % 1000 != 0) { + txn->Commit(); + LOG(INFO) << "Processed " << count << " files."; + } +#else + LOG(FATAL) << "This tool requires OpenCV; compile with USE_OPENCV."; +#endif // USE_OPENCV + return 0; +} diff --git a/tools/create_label_map.cpp b/tools/create_label_map.cpp new file mode 100644 index 0000000..4e16e6f --- /dev/null +++ b/tools/create_label_map.cpp @@ -0,0 +1,65 @@ +// This program reads in pairs label names and optionally ids and display names +// and store them in LabelMap proto buffer. +// Usage: +// create_label_map [FLAGS] MAPFILE OUTFILE +// where MAPFILE is a list of label names and optionally label ids and +// displaynames, and OUTFILE stores the information in LabelMap proto. +// Example: +// ./build/tools/create_label_map --delimiter=" " --include_background=true +// data/VOC2007/map.txt data/VOC2007/labelmap_voc.prototxt +// The format of MAPFILE is like following: +// class1 [1] [someclass1] +// ... +// The format of OUTFILE is like following: +// item { +// name: "class1" +// label: 1 +// display_name: "someclass1" +// } +// ... + +#include // NOLINT(readability/streams) +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "caffe/proto/caffe.pb.h" +#include "caffe/util/io.hpp" + +using namespace caffe; // NOLINT(build/namespaces) + +DEFINE_bool(include_background, false, + "When this option is on, include none_of_the_above as class 0."); +DEFINE_string(delimiter, " ", + "The delimiter used to separate fields in label_map_file."); + +int main(int argc, char** argv) { + ::google::InitGoogleLogging(argv[0]); + // Print output to stderr (while still logging) + FLAGS_alsologtostderr = 1; + +#ifndef GFLAGS_GFLAGS_H_ + namespace gflags = google; +#endif + + gflags::SetUsageMessage("Read in pairs label names and optionally ids and " + "display names and store them in LabelMap proto buffer.\n" + "Usage:\n" + " create_label_map [FLAGS] MAPFILE OUTFILE\n"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (argc < 3) { + gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/create_label_map"); + return 1; + } + + const bool include_background = FLAGS_include_background; + const string delimiter = FLAGS_delimiter; + + const string& map_file = argv[1]; + LabelMap label_map; + ReadLabelFileToLabelMap(map_file, include_background, delimiter, &label_map); + + WriteProtoToTextFile(label_map, argv[2]); +} diff --git a/tools/device_query.cpp b/tools/device_query.cpp new file mode 100644 index 0000000..03799e5 --- /dev/null +++ b/tools/device_query.cpp @@ -0,0 +1,7 @@ +#include "caffe/common.hpp" + +int main(int argc, char** argv) { + LOG(FATAL) << "Deprecated. Use caffe device_query " + "[--device_id=0] instead."; + return 0; +} diff --git a/tools/get_image_size.cpp b/tools/get_image_size.cpp new file mode 100644 index 0000000..94a0ac3 --- /dev/null +++ b/tools/get_image_size.cpp @@ -0,0 +1,113 @@ +// This program retrieves the sizes of a set of images. +// Usage: +// get_image_size [FLAGS] ROOTFOLDER/ LISTFILE OUTFILE +// +// where ROOTFOLDER is the root folder that holds all the images and +// annotations, and LISTFILE should be a list of files as well as their labels +// or label files. +// For classification task, the file should be in the format as +// imgfolder1/img1.JPEG 7 +// .... +// For detection task, the file should be in the format as +// imgfolder1/img1.JPEG annofolder1/anno1.xml +// .... + +#include // NOLINT(readability/streams) +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "caffe/util/io.hpp" + +using namespace caffe; // NOLINT(build/namespaces) + +DEFINE_string(name_id_file, "", + "A file which maps image_name to image_id."); + +int main(int argc, char** argv) { +#ifdef USE_OPENCV + ::google::InitGoogleLogging(argv[0]); + // Print output to stderr (while still logging) + FLAGS_alsologtostderr = 1; + +#ifndef GFLAGS_GFLAGS_H_ + namespace gflags = google; +#endif + + gflags::SetUsageMessage("Get sizes of a set of images.\n" + "Usage:\n" + " get_image_size ROOTFOLDER/ LISTFILE OUTFILE\n"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (argc < 4) { + gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/get_image_size"); + return 1; + } + + std::ifstream infile(argv[2]); + if (!infile.good()) { + LOG(FATAL) << "Failed to open file: " << argv[2]; + } + std::vector > lines; + std::string filename, label; + while (infile >> filename >> label) { + lines.push_back(std::make_pair(filename, label)); + } + infile.close(); + LOG(INFO) << "A total of " << lines.size() << " images."; + + const string name_id_file = FLAGS_name_id_file; + std::map map_name_id; + if (!name_id_file.empty()) { + std::ifstream nameidfile(name_id_file.c_str()); + if (!nameidfile.good()) { + LOG(FATAL) << "Failed to open name_id_file: " << name_id_file; + } + std::string name; + int id; + while (nameidfile >> name >> id) { + CHECK(map_name_id.find(name) == map_name_id.end()); + map_name_id[name] = id; + } + CHECK_EQ(map_name_id.size(), lines.size()); + } + + // Storing to outfile + boost::filesystem::path root_folder(argv[1]); + std::ofstream outfile(argv[3]); + if (!outfile.good()) { + LOG(FATAL) << "Failed to open file: " << argv[3]; + } + int height, width; + int count = 0; + for (int line_id = 0; line_id < lines.size(); ++line_id) { + boost::filesystem::path img_file = root_folder / lines[line_id].first; + GetImageSize(img_file.string(), &height, &width); + std::string img_name = img_file.stem().string(); + if (map_name_id.size() == 0) { + outfile << img_name << " " << height << " " << width << std::endl; + } else { + CHECK(map_name_id.find(img_name) != map_name_id.end()); + int img_id = map_name_id.find(img_name)->second; + outfile << img_id << " " << height << " " << width << std::endl; + } + + if (++count % 1000 == 0) { + LOG(INFO) << "Processed " << count << " files."; + } + } + // write the last batch + if (count % 1000 != 0) { + LOG(INFO) << "Processed " << count << " files."; + } + outfile.flush(); + outfile.close(); +#else + LOG(FATAL) << "This tool requires OpenCV; compile with USE_OPENCV."; +#endif // USE_OPENCV + return 0; +}