IVGCVSW-2662 Create ISerializer class and refactor Serializer

Change-Id: I8a500182193cd218d88c2bb406bf22f496fe43d1
Signed-off-by: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>

1 files changed, 47 insertions, 12 deletions

diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp
index 8aec3ca42a..376f6a87d3 100644
--- a/src/armnnSerializer/Serializer.hpp
+++ b/src/armnnSerializer/Serializer.hpp
@@ -6,17 +6,40 @@
 
 #include <armnn/ILayerVisitor.hpp>
 #include <armnn/LayerVisitorBase.hpp>
+
+#include <armnnSerializer/ISerializer.hpp>
+
 #include <iostream>
 #include <Schema_generated.h>
 
 namespace armnnSerializer
 {
 
-class Serializer : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy>
+class SerializerVisitor : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy>
 {
 public:
-    Serializer() {};
-    ~Serializer() {};
+    SerializerVisitor() {};
+    ~SerializerVisitor() {};
+
+    flatbuffers::FlatBufferBuilder& GetFlatBufferBuilder()
+    {
+        return m_flatBufferBuilder;
+    }
+
+    std::vector<unsigned int>& GetInputIds()
+    {
+        return m_inputIds;
+    }
+
+    std::vector<unsigned int>& GetOutputIds()
+    {
+        return m_outputIds;
+    }
+
+    std::vector<flatbuffers::Offset<armnn::armnnSerializer::AnyLayer>>& GetSerializedLayers()
+    {
+        return m_serializedLayers;
+    }
 
     void VisitAdditionLayer(const armnn::IConnectableLayer* layer,
                             const char* name = nullptr) override;
@@ -32,15 +55,6 @@ public:
     void VisitMultiplicationLayer(const armnn::IConnectableLayer* layer,
                                   const char* name = nullptr) override;
 
-    /// Serializes the network to ArmNN SerializedGraph.
-    /// @param [in] inNetwork The network to be serialized.
-    void Serialize(const armnn::INetwork& inNetwork);
-
-    /// Serializes the SerializedGraph to the stream.
-    /// @param [stream] the stream to save to
-    /// @return true if graph is Serialized to the Stream, false otherwise
-    bool SaveSerializedToStream(std::ostream& stream);
-
 private:
 
     /// Creates the Input Slots and Output Slots and LayerBase for the layer.
@@ -72,4 +86,25 @@ private:
     std::vector<unsigned int> m_outputIds;
 };
 
+class Serializer : public ISerializer
+{
+public:
+    Serializer() {};
+    ~Serializer() {};
+
+    /// Serializes the network to ArmNN SerializedGraph.
+    /// @param [in] inNetwork The network to be serialized.
+    void Serialize(const armnn::INetwork& inNetwork) override;
+
+    /// Serializes the SerializedGraph to the stream.
+    /// @param [stream] the stream to save to
+    /// @return true if graph is Serialized to the Stream, false otherwise
+    bool SaveSerializedToStream(std::ostream& stream) override;
+
+private:
+
+    /// Visitor to contruct serialized network
+    SerializerVisitor m_SerializerVisitor;
+};
+
 } //namespace armnnSerializer