4 #include "boost/algorithm/string.hpp"
5 #include "google/protobuf/text_format.h"
7 #include "caffe/blob.hpp"
8 #include "caffe/common.hpp"
9 #include "caffe/net.hpp"
10 #include "caffe/proto/caffe.pb.h"
11 #include "caffe/util/db.hpp"
12 #include "caffe/util/format.hpp"
13 #include "caffe/util/io.hpp"
20 namespace db = caffe::db;
22 template<typename Dtype>
23 int feature_extraction_pipeline(int argc, char** argv);
25 int main(int argc, char** argv) {
26 return feature_extraction_pipeline<float>(argc, argv);
27 // return feature_extraction_pipeline<double>(argc, argv);
30 template<typename Dtype>
31 int feature_extraction_pipeline(int argc, char** argv) {
32 ::google::InitGoogleLogging(argv[0]);
33 const int num_required_args = 7;
34 if (argc < num_required_args) {
36 "This program takes in a trained network and an input data layer, and then"
37 " extract features of the input data produced by the net.\n"
38 "Usage: extract_features pretrained_net_param"
39 " feature_extraction_proto_file extract_feature_blob_name1[,name2,...]"
40 " save_feature_dataset_name1[,name2,...] num_mini_batches db_type"
41 " [CPU/GPU] [DEVICE_ID=0]\n"
42 "Note: you can extract multiple features in one pass by specifying"
43 " multiple feature blob names and dataset names separated by ','."
44 " The names cannot contain white space characters and the number of blobs"
45 " and datasets must be equal.";
48 int arg_pos = num_required_args;
50 arg_pos = num_required_args;
51 if (argc > arg_pos && strcmp(argv[arg_pos], "GPU") == 0) {
52 LOG(ERROR)<< "Using GPU";
54 if (argc > arg_pos + 1) {
55 device_id = atoi(argv[arg_pos + 1]);
56 CHECK_GE(device_id, 0);
58 LOG(ERROR) << "Using Device_id=" << device_id;
59 Caffe::SetDevice(device_id);
60 Caffe::set_mode(Caffe::GPU);
62 LOG(ERROR) << "Using CPU";
63 Caffe::set_mode(Caffe::CPU);
66 arg_pos = 0; // the name of the executable
67 std::string pretrained_binary_proto(argv[++arg_pos]);
69 // Expected prototxt contains at least one data layer such as
70 // the layer data_layer_name and one feature blob such as the
71 // fc7 top blob to extract features.
74 name: "data_layer_name"
77 source: "/path/to/your/images/to/extract/feature/images_leveldb"
78 mean_file: "/path/to/your/image_mean.binaryproto"
84 top: "label_blob_name"
96 std::string feature_extraction_proto(argv[++arg_pos]);
97 boost::shared_ptr<Net<Dtype> > feature_extraction_net(
98 new Net<Dtype>(feature_extraction_proto, caffe::TEST));
99 feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto);
101 std::string extract_feature_blob_names(argv[++arg_pos]);
102 std::vector<std::string> blob_names;
103 boost::split(blob_names, extract_feature_blob_names, boost::is_any_of(","));
105 std::string save_feature_dataset_names(argv[++arg_pos]);
106 std::vector<std::string> dataset_names;
107 boost::split(dataset_names, save_feature_dataset_names,
108 boost::is_any_of(","));
109 CHECK_EQ(blob_names.size(), dataset_names.size()) <<
110 " the number of blob names and dataset names must be equal";
111 size_t num_features = blob_names.size();
113 for (size_t i = 0; i < num_features; i++) {
114 CHECK(feature_extraction_net->has_blob(blob_names[i]))
115 << "Unknown feature blob name " << blob_names[i]
116 << " in the network " << feature_extraction_proto;
119 int num_mini_batches = atoi(argv[++arg_pos]);
121 std::vector<boost::shared_ptr<db::DB> > feature_dbs;
122 std::vector<boost::shared_ptr<db::Transaction> > txns;
123 const char* db_type = argv[++arg_pos];
124 for (size_t i = 0; i < num_features; ++i) {
125 LOG(INFO)<< "Opening dataset " << dataset_names[i];
126 boost::shared_ptr<db::DB> db(db::GetDB(db_type));
127 db->Open(dataset_names.at(i), db::NEW);
128 feature_dbs.push_back(db);
129 boost::shared_ptr<db::Transaction> txn(db->NewTransaction());
133 LOG(ERROR)<< "Extracting Features";
136 std::vector<int> image_indices(num_features, 0);
137 for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) {
138 feature_extraction_net->Forward();
139 for (int i = 0; i < num_features; ++i) {
140 const boost::shared_ptr<Blob<Dtype> > feature_blob =
141 feature_extraction_net->blob_by_name(blob_names[i]);
142 int batch_size = feature_blob->num();
143 int dim_features = feature_blob->count() / batch_size;
144 const Dtype* feature_blob_data;
145 for (int n = 0; n < batch_size; ++n) {
146 datum.set_height(feature_blob->height());
147 datum.set_width(feature_blob->width());
148 datum.set_channels(feature_blob->channels());
150 datum.clear_float_data();
151 feature_blob_data = feature_blob->cpu_data() +
152 feature_blob->offset(n);
153 for (int d = 0; d < dim_features; ++d) {
154 datum.add_float_data(feature_blob_data[d]);
156 string key_str = caffe::format_int(image_indices[i], 10);
159 CHECK(datum.SerializeToString(&out));
160 txns.at(i)->Put(key_str, out);
162 if (image_indices[i] % 1000 == 0) {
163 txns.at(i)->Commit();
164 txns.at(i).reset(feature_dbs.at(i)->NewTransaction());
165 LOG(ERROR)<< "Extracted features of " << image_indices[i] <<
166 " query images for feature blob " << blob_names[i];
168 } // for (int n = 0; n < batch_size; ++n)
169 } // for (int i = 0; i < num_features; ++i)
170 } // for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index)
171 // write the last batch
172 for (int i = 0; i < num_features; ++i) {
173 if (image_indices[i] % 1000 != 0) {
174 txns.at(i)->Commit();
176 LOG(ERROR)<< "Extracted features of " << image_indices[i] <<
177 " query images for feature blob " << blob_names[i];
178 feature_dbs.at(i)->Close();
181 LOG(ERROR)<< "Successfully extracted the features!";