From fcde19e77424c8c716a7bfc6e1e1cb363b22f20e Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=A2=85=ED=98=84/On-Device=20Lab=28SR=29/Staff?=
 =?utf8?q?=20Engineer/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?=
 <jh1302.park@samsung.com>
Date: Mon, 15 Jul 2019 15:32:02 +0900
Subject: [PATCH] [moco/tf] Importer with User-Defined GraphBuilderSource
 (#4259)

This commit introduces a new Importer constructor which allows users to
inject their-own GraphBuilderSource to Importer.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/moco-tf/src/Importer.cpp |  6 ++++++
 contrib/moco-tf/src/Importer.h   | 11 +++++++++++
 2 files changed, 17 insertions(+)

diff --git a/contrib/moco-tf/src/Importer.cpp b/contrib/moco-tf/src/Importer.cpp
index 3a4d148..70087aa 100644
--- a/contrib/moco-tf/src/Importer.cpp
+++ b/contrib/moco-tf/src/Importer.cpp
@@ -270,6 +270,12 @@ std::unique_ptr<loco::Graph> Importer::import(const ModelSignature &signature,
 
   const GraphBuilderSource *source_ptr = &moco::tf::GraphBuilderRegistry::get();
 
+  if (_source != nullptr)
+  {
+    // Use user-defined GraphBuilderSource
+    source_ptr = _source;
+  }
+
   convert_graph(*source_ptr, signature, tf_graph_def, graph.get());
 
   transform_graph(graph.get());
diff --git a/contrib/moco-tf/src/Importer.h b/contrib/moco-tf/src/Importer.h
index 6875e50..e5faafd 100644
--- a/contrib/moco-tf/src/Importer.h
+++ b/contrib/moco-tf/src/Importer.h
@@ -20,6 +20,8 @@
 #include <moco/tf/Frontend.h>
 #include <moco/tf/Names.h>
 
+#include "GraphBuilderRegistry.h"
+
 #include <loco.h>
 
 #include <tensorflow/core/framework/graph.pb.h>
@@ -37,7 +39,16 @@ public:
   Importer();
 
 public:
+  explicit Importer(const GraphBuilderSource *source) : _source{source}
+  {
+    // DO NOTHING
+  }
+
+public:
   std::unique_ptr<loco::Graph> import(const ModelSignature &, tensorflow::GraphDef &) const;
+
+private:
+  const GraphBuilderSource *_source = nullptr;
 };
 
 } // namespace tf
-- 
2.7.4