1 files changed, 144 insertions, 0 deletions
diff --git a/tensorflow/examples/android/jni/jni_utils.cc b/tensorflow/examples/android/jni/jni_utils.cc
new file mode 100644
index 0000000000..3fffc19cb6
--- /dev/null
+++ b/tensorflow/examples/android/jni/jni_utils.cc
@@ -0,0 +1,144 @@
+#include "tensorflow/examples/android/jni/jni_utils.h"
+
+#include <android/asset_manager.h>
+#include <android/asset_manager_jni.h>
+#include <jni.h>
+#include <stdlib.h>
+
+#include <string>
+#include <vector>
+#include <fstream>
+#include <sstream>
+
+#include "tensorflow/core/platform/logging.h"
+#include "google/protobuf/src/google/protobuf/io/zero_copy_stream_impl.h"
+#include "google/protobuf/src/google/protobuf/io/zero_copy_stream_impl_lite.h"
+#include "google/protobuf/src/google/protobuf/io/coded_stream.h"
+#include "google/protobuf/src/google/protobuf/message_lite.h"
+
+static const char* const ASSET_PREFIX = "file:///android_asset/";
+
+namespace {
+class IfstreamInputStream : public ::google::protobuf::io::CopyingInputStream {
+ public:
+  explicit IfstreamInputStream(const std::string& file_name)
+      : ifs_(file_name.c_str(), std::ios::in | std::ios::binary) {}
+  ~IfstreamInputStream() { ifs_.close(); }
+
+  int Read(void* buffer, int size) {
+    if (!ifs_) {
+      return -1;
+    }
+    ifs_.read(static_cast<char*>(buffer), size);
+    return ifs_.gcount();
+  }
+
+ private:
+  std::ifstream ifs_;
+};
+}  // namespace
+
+bool PortableReadFileToProto(const std::string& file_name,
+                             ::google::protobuf::MessageLite* proto) {
+  ::google::protobuf::io::CopyingInputStreamAdaptor stream(
+      new IfstreamInputStream(file_name));
+  stream.SetOwnsCopyingStream(true);
+  // TODO(jiayq): the following coded stream is for debugging purposes to allow
+  // one to parse arbitrarily large messages for MessageLite. One most likely
+  // doesn't want to put protobufs larger than 64MB on Android, so we should
+  // eventually remove this and quit loud when a large protobuf is passed in.
+  ::google::protobuf::io::CodedInputStream coded_stream(&stream);
+  // Total bytes hard limit / warning limit are set to 1GB and 512MB
+  // respectively.
+  coded_stream.SetTotalBytesLimit(1024LL << 20, 512LL << 20);
+  return proto->ParseFromCodedStream(&coded_stream);
+}
+
+bool IsAsset(const char* const filename) {
+  return strstr(filename, ASSET_PREFIX) == filename;
+}
+
+void ReadFileToProto(AAssetManager* const asset_manager,
+                     const char* const filename,
+                     google::protobuf::MessageLite* message) {
+  if (!IsAsset(filename)) {
+    VLOG(0) << "Opening file: " << filename;
+    CHECK(PortableReadFileToProto(filename, message));
+    return;
+  }
+
+  CHECK_NOTNULL(asset_manager);
+
+  const char* const asset_filename = filename + strlen(ASSET_PREFIX);
+  AAsset* asset = AAssetManager_open(asset_manager,
+                                     asset_filename,
+                                     AASSET_MODE_STREAMING);
+  CHECK_NOTNULL(asset);
+
+  off_t start;
+  off_t length;
+  const int fd = AAsset_openFileDescriptor(asset, &start, &length);
+
+  if (fd >= 0) {
+    // If it has a file descriptor that means it can be memmapped directly
+    // from the APK.
+    VLOG(0) << "Opening asset " << asset_filename
+            << " from disk with zero-copy.";
+    google::protobuf::io::FileInputStream is(fd);
+    google::protobuf::io::LimitingInputStream lis(&is, start + length);
+    lis.Skip(start);
+    CHECK(message->ParseFromZeroCopyStream(&lis));
+    is.Close();
+  } else {
+    // It may be compressed, in which case we have to uncompress
+    // it to memory first.
+    VLOG(0) << "Opening asset " << asset_filename
+            << " from disk with copy.";
+    const off_t data_size = AAsset_getLength(asset);
+    const void* const memory = AAsset_getBuffer(asset);
+    CHECK(message->ParseFromArray(memory, data_size));
+  }
+  AAsset_close(asset);
+}
+
+void ReadFileToString(AAssetManager* const asset_manager,
+                      const char* const filename, std::string* str) {
+  if (!IsAsset(filename)) {
+    VLOG(0) << "Opening file: " << filename;
+    std::ifstream t(filename);
+    std::string tmp((std::istreambuf_iterator<char>(t)),
+                     std::istreambuf_iterator<char>());
+    tmp.swap(*str);
+    t.close();
+    return;
+  }
+
+  CHECK_NOTNULL(asset_manager);
+  const char* const asset_filename = filename + strlen(ASSET_PREFIX);
+  AAsset* asset = AAssetManager_open(asset_manager,
+                                     asset_filename,
+                                     AASSET_MODE_STREAMING);
+  CHECK_NOTNULL(asset);
+  VLOG(0) << "Opening asset " << asset_filename << " from disk with copy.";
+  const off_t data_size = AAsset_getLength(asset);
+  const char* memory = reinterpret_cast<const char*>(AAsset_getBuffer(asset));
+
+  std::string tmp(memory, memory + data_size);
+  tmp.swap(*str);
+  AAsset_close(asset);
+}
+
+void ReadFileToVector(AAssetManager* const asset_manager,
+                      const char* const filename,
+                      std::vector<std::string>* str_vector) {
+  std::string labels_string;
+  ReadFileToString(asset_manager, filename, &labels_string);
+  std::istringstream ifs(labels_string);
+  str_vector->clear();
+  std::string label;
+  while (std::getline(ifs, label)) {
+    str_vector->push_back(label);
+  }
+  VLOG(0) << "Read " << str_vector->size() << " values from " << filename;
+}
+