From 2cd37cb745be35b86a9b146773c52e17989ab578 Mon Sep 17 00:00:00 2001 From: =?utf8?q?=EB=B0=95=EC=84=B8=ED=9D=AC/On-Device=20Lab=28SR=29/Princip?= =?utf8?q?al=20Engineer/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= Date: Fri, 5 Apr 2019 07:02:19 +0900 Subject: [PATCH] [tfkit] use arguments for unpack and pack (#3169) This will enable using file arguments for unpack and pack command Signed-off-by: SaeHie Park --- contrib/tfkit/src/PackCommand.cpp | 13 +++++++------ contrib/tfkit/src/UnpackCommand.cpp | 13 +++++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/contrib/tfkit/src/PackCommand.cpp b/contrib/tfkit/src/PackCommand.cpp index 37c5b61..4230aad 100644 --- a/contrib/tfkit/src/PackCommand.cpp +++ b/contrib/tfkit/src/PackCommand.cpp @@ -24,7 +24,6 @@ #include #include -#include #include #include @@ -84,12 +83,15 @@ void pack(tensorflow::GraphDef &graph_def) namespace tfkit { -int PackCommand::run(int, const char *const *) const +int PackCommand::run(int argc, const char *const *argv) const { tensorflow::GraphDef graph_def; - // Load text from standard input - google::protobuf::io::IstreamInputStream is{&std::cin}; + CmdArguments cmdargs(argc, argv); + + auto ioconfig = make_ioconfig(cmdargs); + + google::protobuf::io::IstreamInputStream is{ioconfig->in()}; if (!google::protobuf::TextFormat::Parse(&is, &graph_def)) { @@ -100,8 +102,7 @@ int PackCommand::run(int, const char *const *) const // convert float_val to tensor_content pack(graph_def); - // Write text into standard output - google::protobuf::io::OstreamOutputStream os{&std::cout}; + google::protobuf::io::OstreamOutputStream os{ioconfig->out()}; google::protobuf::TextFormat::Print(graph_def, &os); return 0; diff --git a/contrib/tfkit/src/UnpackCommand.cpp b/contrib/tfkit/src/UnpackCommand.cpp index 13d6e1e..8e51172 100644 --- a/contrib/tfkit/src/UnpackCommand.cpp +++ b/contrib/tfkit/src/UnpackCommand.cpp @@ -24,7 +24,6 @@ #include #include -#include #include namespace @@ -80,12 +79,15 @@ void unpack(tensorflow::GraphDef &graph_def) namespace tfkit { -int UnpackCommand::run(int, const char *const *) const +int UnpackCommand::run(int argc, const char *const *argv) const { tensorflow::GraphDef graph_def; - // Load text from standard input - google::protobuf::io::IstreamInputStream is{&std::cin}; + CmdArguments cmdargs(argc, argv); + + auto ioconfig = make_ioconfig(cmdargs); + + google::protobuf::io::IstreamInputStream is{ioconfig->in()}; if (!google::protobuf::TextFormat::Parse(&is, &graph_def)) { @@ -96,8 +98,7 @@ int UnpackCommand::run(int, const char *const *) const // convert tensor_content to float_val unpack(graph_def); - // Write text into standard output - google::protobuf::io::OstreamOutputStream os{&std::cout}; + google::protobuf::io::OstreamOutputStream os{ioconfig->out()}; google::protobuf::TextFormat::Print(graph_def, &os); return 0; -- 2.7.4