From 07332c8f1f8120392df9a2a9143e9a934c4b4bd4 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EC=9C=A4=ED=98=84=EC=8B=9D/On-Device=20Lab=28SR=29/Princip?=
 =?utf8?q?al=20Engineer/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?=
 <hyunsik.yoon@samsung.com>
Date: Mon, 24 Jun 2019 19:40:45 +0900
Subject: [PATCH] [tfldump] Dumping attrs of custom op (#3949)

* [tfldump] Dumping attrs of custom op

This commit enables tfldump to print custom op's attribute information.

Signed-off-by: Hyun Sik Yoon <hyunsik.yoon@samsung.com>

* fix #include
---
 contrib/tfldump/CMakeLists.txt    |  1 +
 contrib/tfldump/src/OpPrinter.cpp | 43 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/contrib/tfldump/CMakeLists.txt b/contrib/tfldump/CMakeLists.txt
index 7e3f1b6..a526ae8 100644
--- a/contrib/tfldump/CMakeLists.txt
+++ b/contrib/tfldump/CMakeLists.txt
@@ -18,3 +18,4 @@ target_include_directories(tfldump PRIVATE include)
 target_link_libraries(tfldump tfldump_flatbuffer)
 target_link_libraries(tfldump safemain)
 target_link_libraries(tfldump stdex)
+target_link_libraries(tfldump flatbuffers)
diff --git a/contrib/tfldump/src/OpPrinter.cpp b/contrib/tfldump/src/OpPrinter.cpp
index cb67b35..9ecc4ff 100644
--- a/contrib/tfldump/src/OpPrinter.cpp
+++ b/contrib/tfldump/src/OpPrinter.cpp
@@ -19,6 +19,8 @@
 
 #include <stdex/Memory.h>
 
+#include <flatbuffers/flexbuffers.h>
+
 using stdex::make_unique;
 
 namespace tfldump
@@ -127,6 +129,46 @@ public:
   }
 };
 
+class CustomOpPrinter : public OpPrinter
+{
+public:
+  void options(const tflite::Operator *op, std::ostream &os) const override
+  {
+    if (op->custom_options_format() != tflite::CustomOptionsFormat::CustomOptionsFormat_FLEXBUFFERS)
+    {
+      os << "    ";
+      os << "Unknown custom option format";
+      return;
+    }
+
+    const flatbuffers::Vector<uint8_t> *option_buf = op->custom_options();
+
+    if (option_buf == nullptr || option_buf->size() == 0)
+    {
+      os << "No attrs found." << std::endl;
+      return;
+    }
+
+    // printing attrs
+    // attrs of custom ops are encoded in flexbuffer format
+    auto attr_map = flexbuffers::GetRoot(option_buf->data(), option_buf->size()).AsMap();
+
+    os << "    ";
+    auto keys = attr_map.Keys();
+    for (int i = 0; i < keys.size(); i++)
+    {
+      auto key = keys[i].ToString();
+      os << key << "(" << attr_map[key].ToString() << ") ";
+    }
+
+    // Note: attr in "Shape" type does not seem to be converted by tflite_convert.
+    // When the converted tflite file (with custom op) is opened with hexa editory,
+    // attrs names can be found but attr name in "Shape" type is not found.
+
+    os << std::endl;
+  }
+};
+
 OpPrinterRegistry::OpPrinterRegistry()
 {
   _op_map[tflite::BuiltinOperator_AVERAGE_POOL_2D] = make_unique<Pool2DPrinter>();
@@ -137,6 +179,7 @@ OpPrinterRegistry::OpPrinterRegistry()
   // There is no Option for ReLU and ReLU6
   _op_map[tflite::BuiltinOperator_RESHAPE] = make_unique<ReshapePrinter>();
   _op_map[tflite::BuiltinOperator_SOFTMAX] = make_unique<SoftmaxPrinter>();
+  _op_map[tflite::BuiltinOperator_CUSTOM] = make_unique<CustomOpPrinter>();
 }
 
 } // namespace tfldump
-- 
2.7.4