From 09f190f5423bdb1ecfa6330abb3de5b937009bb0 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=A2=85=ED=98=84/=EB=8F=99=EC=9E=91=EC=A0=9C?=
 =?utf8?q?=EC=96=B4Lab=28SR=29/Senior=20Engineer/=EC=82=BC=EC=84=B1?=
 =?utf8?q?=EC=A0=84=EC=9E=90?= <jh1302.park@samsung.com>
Date: Thu, 19 Apr 2018 13:27:15 +0900
Subject: [PATCH] [caffegen] Support 'encode' (#95)

This commit implements 'encode' command which allows users to convert
prototxt (text format) to caffemodel (binary format).

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/caffegen/CMakeLists.txt        |  8 ++++++--
 contrib/caffegen/src/Driver.cpp        | 13 +++++++++++--
 contrib/caffegen/src/EncodeCommand.cpp | 35 ++++++++++++++++++++++++++++++++++
 contrib/caffegen/src/EncodeCommand.h   | 11 +++++++++++
 4 files changed, 63 insertions(+), 4 deletions(-)
 create mode 100644 contrib/caffegen/src/EncodeCommand.cpp
 create mode 100644 contrib/caffegen/src/EncodeCommand.h

diff --git a/contrib/caffegen/CMakeLists.txt b/contrib/caffegen/CMakeLists.txt
index c32ba24..9527a66 100644
--- a/contrib/caffegen/CMakeLists.txt
+++ b/contrib/caffegen/CMakeLists.txt
@@ -1,3 +1,7 @@
-file(GLOB_RECURSE SOURCES "src/*.cpp")
+if(TARGET nncc_frontend_caffe_core)
+  file(GLOB_RECURSE SOURCES "src/*.cpp")
 
-add_executable(caffegen ${SOURCES})
+  add_executable(caffegen ${SOURCES})
+  target_include_directories(caffegen PRIVATE include)
+  target_link_libraries(caffegen nncc_foundation nncc_frontend_caffe_core)
+endif(TARGET nncc_frontend_caffe_core)
diff --git a/contrib/caffegen/src/Driver.cpp b/contrib/caffegen/src/Driver.cpp
index b20690a..2d1b3b1 100644
--- a/contrib/caffegen/src/Driver.cpp
+++ b/contrib/caffegen/src/Driver.cpp
@@ -1,6 +1,15 @@
-#include <stdexcept>
+#include "EncodeCommand.h"
+
+#include <nncc/foundation/Memory.h>
+
+#include <map>
+#include <string>
 
 int main(int argc, char **argv)
 {
-  throw std::runtime_error{"Not implemented, yet"};
+  std::map<std::string, std::unique_ptr<Command>> commands;
+
+  commands["encode"] = nncc::foundation::make_unique<EncodeCommand>();
+
+  return commands.at(argv[1])->run();
 }
diff --git a/contrib/caffegen/src/EncodeCommand.cpp b/contrib/caffegen/src/EncodeCommand.cpp
new file mode 100644
index 0000000..382a270
--- /dev/null
+++ b/contrib/caffegen/src/EncodeCommand.cpp
@@ -0,0 +1,35 @@
+#include "EncodeCommand.h"
+
+#include <caffe.pb.h>
+
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/io/zero_copy_stream_impl.h>
+#include <google/protobuf/text_format.h>
+
+#include <iostream>
+
+int EncodeCommand::run(void) const
+{
+  caffe::NetParameter param;
+
+  // Load text from standard input
+  google::protobuf::io::FileInputStream is{0};
+
+  if (!google::protobuf::TextFormat::Parse(&is, &param))
+  {
+    std::cerr << "ERROR: Failed to parse prototxt" << std::endl;
+    return 255;
+  }
+
+  // Write binary into standard input
+  google::protobuf::io::FileOutputStream os{1};
+  google::protobuf::io::CodedOutputStream coded_os{&os};
+
+  if (!param.SerializeToCodedStream(&coded_os))
+  {
+    std::cerr << "ERROR: Failed to serialize" << std::endl;
+    return 255;
+  }
+
+  return 0;
+}
diff --git a/contrib/caffegen/src/EncodeCommand.h b/contrib/caffegen/src/EncodeCommand.h
new file mode 100644
index 0000000..4371fd1
--- /dev/null
+++ b/contrib/caffegen/src/EncodeCommand.h
@@ -0,0 +1,11 @@
+#ifndef __ENCODE_COMMAND_H__
+#define __ENCODE_COMMAND_H__
+
+#include "Command.h"
+
+struct EncodeCommand final : public Command
+{
+  int run(void) const override;
+};
+
+#endif // __ENCODE_COMMAND_H__
-- 
2.7.4