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