compute data mean for float_data

1 files changed, 11 insertions, 4 deletions

diff --git a/tools/compute_image_mean.cpp b/tools/compute_image_mean.cpp
index 65ff538f..771f5dc7 100644
--- a/tools/compute_image_mean.cpp
+++ b/tools/compute_image_mean.cpp
@@ -10,6 +10,7 @@
 
 using caffe::Datum;
 using caffe::BlobProto;
+using std::max;
 
 int main(int argc, char** argv) {
   ::google::InitGoogleLogging(argv[0]);
@@ -40,7 +41,8 @@ int main(int argc, char** argv) {
   sum_blob.set_height(datum.height());
   sum_blob.set_width(datum.width());
   const int data_size = datum.channels() * datum.height() * datum.width();
-  for (int i = 0; i < datum.data().size(); ++i) {
+  int size_in_datum = std::max<int>(datum.data().size(), datum.float_data_size());
+  for (int i = 0; i < size_in_datum; ++i) {
     sum_blob.add_data(0.);
   }
   LOG(INFO) << "Starting Iteration";
@@ -48,10 +50,15 @@ int main(int argc, char** argv) {
     // just a dummy operation
     datum.ParseFromString(it->value().ToString());
     const string& data = datum.data();
-    CHECK_EQ(data.size(), data_size) << "Incorrect data field size "
+    size_in_datum = std::max<int>(datum.data().size(), datum.float_data_size());
+    CHECK_EQ(size_in_datum, data_size) << "Incorrect data field size " << size_in_datum;
         << data.size();
-    for (int i = 0; i < data.size(); ++i) {
-      sum_blob.set_data(i, sum_blob.data(i) + (uint8_t)data[i]);
+    for (int i = 0; i < size_in_datum; ++i) {
+      if (data.size() != 0) {
+	sum_blob.set_data(i, sum_blob.data(i) + (uint8_t)data[i]);
+      } else {
+	sum_blob.set_data(i, sum_blob.data(i) + (float)datum.float_data(i));
+      }
     }
     ++count;
     if (count % 10000 == 0) {