Merge pull request #57 from JDAI-CV/support_onnx_model_output

support onnx model output
JDAI-CV · Jun 24, 2019 · 90cb0f8 · 90cb0f8
2 parents 99baa7b + a75548d
commit 90cb0f8
Show file tree

Hide file tree

Showing 6 changed files with 231 additions and 106 deletions.
diff --git a/.clang-format b/.clang-format
@@ -3,3 +3,4 @@ IndentWidth: 4
 AllowShortBlocksOnASingleLine: false
 AllowShortLoopsOnASingleLine: false
 AllowShortFunctionsOnASingleLine: false
+AlignEscapedNewlines: Left
diff --git a/common/daq.fbs b/common/daq.fbs
@@ -190,6 +190,7 @@ table Model {
     initializers:[Tensor];
     inputs:[Input];
     quant_infos:[QuantInfo];
+    outputs:[string];
 }
 
 root_type Model;
diff --git a/dnnlibrary/src/DaqReader.cpp b/dnnlibrary/src/DaqReader.cpp
@@ -177,6 +177,18 @@ void AddInputs(const DNN::Model &model, ModelBuilder &builder) {
     }
 }
 
+void AddOutputs(const DNN::Model &model, ModelBuilder &builder) {
+    using namespace android::nn::wrapper;
+    if (model.outputs() == nullptr) {
+        return;
+    }
+    for (const auto &output : *model.outputs()) {
+        PNT(output->str());
+        css output_name = output->str();
+        builder.AddOutput(output_name);
+    }
+}
+
 void AddLayers(const DNN::Model &model, ModelBuilder &builder) {
     for (const auto layer : *model.layers()) {
         switch (layer->type()) {
@@ -359,5 +371,6 @@ void ReadDaqImpl(const uint8_t *buf, ModelBuilder &builder) {
     AddInitializersFromBuffer(*model, builder);
     AddInputs(*model, builder);
     AddLayers(*model, builder);
+    AddOutputs(*model, builder);
 }
 }  // namespace dnn