tools/compute_image_mean.cpp

   1 #include <stdint.h>
   2 #include <algorithm>
   3 #include <string>
   4 #include <utility>
   5 #include <vector>
   6
   7 #include "boost/scoped_ptr.hpp"
   8 #include "gflags/gflags.h"
   9 #include "glog/logging.h"
  10
  11 #include "caffe/proto/caffe.pb.h"
  12 #include "caffe/util/db.hpp"
  13 #include "caffe/util/io.hpp"
  14
  15 using namespace caffe;  // NOLINT(build/namespaces)
  16
  17 using std::max;
  18 using std::pair;
  19 using boost::scoped_ptr;
  20
  21 DEFINE_string(backend, "lmdb",
  22         "The backend {leveldb, lmdb} containing the images");
  23
  24 int main(int argc, char** argv) {
  25   ::google::InitGoogleLogging(argv[0]);
  26
  27 #ifdef USE_OPENCV
  28 #ifndef GFLAGS_GFLAGS_H_
  29   namespace gflags = google;
  30 #endif
  31
  32   gflags::SetUsageMessage("Compute the mean_image of a set of images given by"
  33         " a leveldb/lmdb\n"
  34         "Usage:\n"
  35         "    compute_image_mean [FLAGS] INPUT_DB [OUTPUT_FILE]\n");
  36
  37   gflags::ParseCommandLineFlags(&argc, &argv, true);
  38
  39   if (argc < 2 || argc > 3) {
  40     gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/compute_image_mean");
  41     return 1;
  42   }
  43
  44   scoped_ptr<db::DB> db(db::GetDB(FLAGS_backend));
  45   db->Open(argv[1], db::READ);
  46   scoped_ptr<db::Cursor> cursor(db->NewCursor());
  47
  48   BlobProto sum_blob;
  49   int count = 0;
  50   // load first datum
  51   Datum datum;
  52   datum.ParseFromString(cursor->value());
  53
  54   if (DecodeDatumNative(&datum)) {
  55     LOG(INFO) << "Decoding Datum";
  56   }
  57
  58   sum_blob.set_num(1);
  59   sum_blob.set_channels(datum.channels());
  60   sum_blob.set_height(datum.height());
  61   sum_blob.set_width(datum.width());
  62   const int data_size = datum.channels() * datum.height() * datum.width();
  63   int size_in_datum = std::max<int>(datum.data().size(),
  64                                     datum.float_data_size());
  65   for (int i = 0; i < size_in_datum; ++i) {
  66     sum_blob.add_data(0.);
  67   }
  68   LOG(INFO) << "Starting Iteration";
  69   while (cursor->valid()) {
  70     Datum datum;
  71     datum.ParseFromString(cursor->value());
  72     DecodeDatumNative(&datum);
  73
  74     const std::string& data = datum.data();
  75     size_in_datum = std::max<int>(datum.data().size(),
  76         datum.float_data_size());
  77     CHECK_EQ(size_in_datum, data_size) << "Incorrect data field size " <<
  78         size_in_datum;
  79     if (data.size() != 0) {
  80       CHECK_EQ(data.size(), size_in_datum);
  81       for (int i = 0; i < size_in_datum; ++i) {
  82         sum_blob.set_data(i, sum_blob.data(i) + (uint8_t)data[i]);
  83       }
  84     } else {
  85       CHECK_EQ(datum.float_data_size(), size_in_datum);
  86       for (int i = 0; i < size_in_datum; ++i) {
  87         sum_blob.set_data(i, sum_blob.data(i) +
  88             static_cast<float>(datum.float_data(i)));
  89       }
  90     }
  91     ++count;
  92     if (count % 10000 == 0) {
  93       LOG(INFO) << "Processed " << count << " files.";
  94     }
  95     cursor->Next();
  96   }
  97
  98   if (count % 10000 != 0) {
  99     LOG(INFO) << "Processed " << count << " files.";
 100   }
 101   for (int i = 0; i < sum_blob.data_size(); ++i) {
 102     sum_blob.set_data(i, sum_blob.data(i) / count);
 103   }
 104   // Write to disk
 105   if (argc == 3) {
 106     LOG(INFO) << "Write to " << argv[2];
 107     WriteProtoToBinaryFile(sum_blob, argv[2]);
 108   }
 109   const int channels = sum_blob.channels();
 110   const int dim = sum_blob.height() * sum_blob.width();
 111   std::vector<float> mean_values(channels, 0.0);
 112   LOG(INFO) << "Number of channels: " << channels;
 113   for (int c = 0; c < channels; ++c) {
 114     for (int i = 0; i < dim; ++i) {
 115       mean_values[c] += sum_blob.data(dim * c + i);
 116     }
 117     LOG(INFO) << "mean_value channel [" << c << "]:" << mean_values[c] / dim;
 118   }
 119 #else
 120   LOG(FATAL) << "This tool requires OpenCV; compile with USE_OPENCV.";
 121 #endif  // USE_OPENCV
 122   return 0;
 123 }