Repeal revert of #1878

2 files changed, 20 insertions, 14 deletions

diff --git a/tools/compute_image_mean.cpp b/tools/compute_image_mean.cpp
index dff63a09..b1fc7cae 100644
--- a/tools/compute_image_mean.cpp
+++ b/tools/compute_image_mean.cpp
@@ -50,7 +50,7 @@ int main(int argc, char** argv) {
   Datum datum;
   datum.ParseFromString(cursor->value());
 
-  if (DecodeDatum(&datum)) {
+  if (DecodeDatumNative(&datum)) {
     LOG(INFO) << "Decoding Datum";
   }
 
@@ -68,7 +68,7 @@ int main(int argc, char** argv) {
   while (cursor->valid()) {
     Datum datum;
     datum.ParseFromString(cursor->value());
-    DecodeDatum(&datum);
+    DecodeDatumNative(&datum);
 
     const std::string& data = datum.data();
     size_in_datum = std::max<int>(datum.data().size(),
diff --git a/tools/convert_imageset.cpp b/tools/convert_imageset.cpp
index 7fbf5b05..816a91f9 100644
--- a/tools/convert_imageset.cpp
+++ b/tools/convert_imageset.cpp
@@ -39,6 +39,8 @@ 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) {
   ::google::InitGoogleLogging(argv[0]);
@@ -63,6 +65,7 @@ int main(int argc, char** argv) {
   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;
 
   std::ifstream infile(argv[2]);
   std::vector<std::pair<std::string, int> > lines;
@@ -78,11 +81,8 @@ int main(int argc, char** argv) {
   }
   LOG(INFO) << "A total of " << lines.size() << " images.";
 
-  if (encoded) {
-    CHECK_EQ(FLAGS_resize_height, 0) << "With encoded don't resize images";
-    CHECK_EQ(FLAGS_resize_width, 0) << "With encoded don't resize images";
-    CHECK(!check_size) << "With encoded cannot check_size";
-  }
+  if (encode_type.size() && !encoded)
+    LOG(INFO) << "encode_type specified, assuming encoded=true.";
 
   int resize_height = std::max<int>(0, FLAGS_resize_height);
   int resize_width = std::max<int>(0, FLAGS_resize_width);
@@ -98,18 +98,24 @@ int main(int argc, char** argv) {
   int count = 0;
   const int kMaxKeyLength = 256;
   char key_cstr[kMaxKeyLength];
-  int data_size;
+  int data_size = 0;
   bool data_size_initialized = false;
 
   for (int line_id = 0; line_id < lines.size(); ++line_id) {
     bool status;
-    if (encoded) {
-      status = ReadFileToDatum(root_folder + lines[line_id].first,
-        lines[line_id].second, &datum);
-    } else {
-      status = ReadImageToDatum(root_folder + lines[line_id].first,
-          lines[line_id].second, resize_height, resize_width, is_color, &datum);
+    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);
     }
+    status = ReadImageToDatum(root_folder + lines[line_id].first,
+        lines[line_id].second, resize_height, resize_width, is_color,
+        enc, &datum);
     if (status == false) continue;
     if (check_size) {
       if (!data_size_initialized) {