[Model Parser] Prepare argument (#2068)

Prepare argument input

- verbose level (-v, -vv)
- tensor list (-t ...)
- operator list (-o ...)

Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>

diff --git a/tools/tflitefile_tool/model_parser.py b/tools/tflitefile_tool/model_parser.py
index 5e9c0d3..35b8f87 100755 (executable)
--- a/tools/tflitefile_tool/model_parser.py
+++ b/tools/tflitefile_tool/model_parser.py
@@ -21,6 +21,42 @@ class TFLiteModelFileParser(object):
         # Read flatbuffer file descriptor using argument
         self.tflite_file = args.input_file
 
+        # Set print level (0 ~ 2)
+        # TODO: print information based on level
+        self.print_level = args.verbose
+        if (args.verbose > 2):
+            self.print_level = 2
+        if (args.verbose < 0):
+            self.print_level = 0
+
+        # Set tensor index list to print information
+        # TODO:
+        #   Print tensors in list only
+        #   Print all tensors if argument used and not specified index number
+        if (args.tensor != None):
+            if (len(args.tensor) == 0):
+                self.print_all_tensor = True
+            else:
+                self.print_all_tensor = False
+                self.print_tensor_index = []
+
+                for tensor_index in args.tensor:
+                    self.print_tensor_index.append(int(tensor_index))
+
+        # Set operator index list to print information
+        # TODO:
+        #   Print operators in list only
+        #   Print all operators if argument used and not specified index number
+        if (args.operator != None):
+            if (len(args.operator) == 0):
+                self.print_all_oeprator = True
+            else:
+                self.print_all_oeprator = False
+                self.print_operator_index = []
+
+                for operator_index in args.operator:
+                    self.print_operator_index.append(int(operator_index))
+
     def main(self):
         # Generate Model: top structure of tflite model file
         buf = self.tflite_file.read()
@@ -55,6 +91,19 @@ if __name__ == '__main__':
     arg_parser = argparse.ArgumentParser()
     arg_parser.add_argument(
         "input_file", type=argparse.FileType('rb'), help="tflite file to read")
+    arg_parser.add_argument(
+        '-v',
+        '--verbose',
+        action='count',
+        default=0,
+        help="set print level (0~2, default: 0)")
+    arg_parser.add_argument(
+        '-t', '--tensor', nargs='*', help="tensor ID to print information (default: all)")
+    arg_parser.add_argument(
+        '-o',
+        '--operator',
+        nargs='*',
+        help="operator ID to print information (default: all)")
     args = arg_parser.parse_args()
 
     # Call main function