Add separated Convert, Save and GetBuf

JDAI-CV · May 13, 2019 · 60ad90b · 60ad90b
1 parent 3ec2638
commit 60ad90b
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/tools/onnx2daq/OnnxConverter.cpp b/tools/onnx2daq/OnnxConverter.cpp
@@ -738,7 +738,6 @@ OnnxConverter::ConvertQuantInfosToFbs() {
 }
 
 void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto &model_proto,
-                            const std::string &filepath,
                             const css &table_file) {
     GOOGLE_PROTOBUF_VERIFY_VERSION;
 
@@ -1016,11 +1015,6 @@ void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto &model_proto,
     LOG(INFO) << "Shapes: ";
     LOG(INFO) << shaper_;
 
-    std::ofstream ofs(filepath);
-    ofs.write(reinterpret_cast<char *>(builder_.GetBufferPointer()),
-              builder_.GetSize());
-    ofs.close();
-
     skipped_act_.clear();
     layers_.clear();
     operands_.clear();
@@ -1030,3 +1024,16 @@ void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto &model_proto,
     onnx_tensors_.clear();
     shaper_.Clear();
 }
+
+void OnnxConverter::Save(const std::string &filename) {
+    std::ofstream ofs(filename);
+    ofs.write(reinterpret_cast<char *>(builder_.GetBufferPointer()),
+              builder_.GetSize());
+    ofs.close();
+}
+
+std::unique_ptr<uint8_t []> OnnxConverter::GetBuf() {
+    std::unique_ptr<uint8_t []> ptr(new uint8_t[builder_.GetSize()]);
+    memcpy(ptr.get(), builder_.GetBufferPointer(), builder_.GetSize());
+    return std::move(ptr);
+}
diff --git a/tools/onnx2daq/OnnxConverter.h b/tools/onnx2daq/OnnxConverter.h
@@ -171,5 +171,7 @@ class OnnxConverter {
 
    public:
     void Convert(const ONNX_NAMESPACE::ModelProto &model,
-                 const std::string &filepath, const css &table_file = "");
+                       const css &table_file = "");
+    void Save(const std::string &filename);
+    std::unique_ptr<uint8_t []> GetBuf();
 };
diff --git a/tools/onnx2daq/onnx2daq.cpp b/tools/onnx2daq/onnx2daq.cpp
@@ -32,7 +32,8 @@ int main(int argc, char **argv) {
     }
 
     OnnxConverter converter;
-    converter.Convert(model_proto, argv[2], table_file);
+    converter.Convert(model_proto, table_file);
+    converter.Save(argv[2]);
 
     google::protobuf::ShutdownProtobufLibrary();
     return 0;