From 8218d82b7ed838d490b6fb39ec7988da5d46829e Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Mon, 12 Mar 2018 15:33:06 +0800
Subject: [PATCH] Add model header template and move all generated file to same
 place.

---
 python/tools/model_header.template   | 22 ++++++++++++++++++++++
 python/tools/source_converter_lib.py |  8 ++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 python/tools/model_header.template

diff --git a/python/tools/model_header.template b/python/tools/model_header.template
new file mode 100644
index 00000000..9f5c776d
--- /dev/null
+++ b/python/tools/model_header.template
@@ -0,0 +1,22 @@
+//
+// Copyright (c) 2017 XiaoMi All rights reserved.
+// Generated by the mace converter.  DO NOT EDIT!
+//
+
+#include <string>
+
+#include "mace/public/mace.h"
+
+namespace mace {
+namespace {{tag}} {
+
+extern const unsigned char *LoadModelData(const char *model_data_file);
+
+extern void UnloadModelData(const unsigned char *model_data);
+
+extern NetDef CreateNet(const unsigned char *model_data);
+
+extern const std::string ModelChecksum();
+
+}  // namespace {{ tag }}
+}  // namespace mace
diff --git a/python/tools/source_converter_lib.py b/python/tools/source_converter_lib.py
index 48620344..8dd1cd7c 100644
--- a/python/tools/source_converter_lib.py
+++ b/python/tools/source_converter_lib.py
@@ -177,3 +177,11 @@ def convert_to_source(net_def, mode_pb_checksum, template_dir, obfuscate, model_
   )
   with open(output, "wb") as f:
     f.write(source)
+
+  # generate model header file
+  template_name = 'model_header.template'
+  source = j2_env.get_template(template_name).render(
+    tag = model_tag,
+  )
+  with open(output_dir + model_tag + '.h', "wb") as f:
+    f.write(source)
-- 
GitLab