Initial commit

app.py

@@ -0,0 +1,564 @@
+#!/usr/bin/env python
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""This tool creates an html visualization of a TensorFlow Lite graph.
+
+Example usage:
+
+python visualize.py foo.tflite foo.html
+"""
+
+import json
+import os
+import re
+import sys
+import numpy as np
+
+# pylint: disable=g-import-not-at-top
+if not os.path.splitext(__file__)[0].endswith(
+    os.path.join("tflite_runtime", "visualize")):
+  # This file is part of tensorflow package.
+  from tensorflow.lite.python import schema_py_generated as schema_fb
+else:
+  # This file is part of tflite_runtime package.
+  from tflite_runtime import schema_py_generated as schema_fb
+import gradio as gr
+from html import escape
+
+# A CSS description for making the visualizer
+# body {font-family: sans-serif; background-color: #fa0;}
+# # font-family: sans-serif;
+"""<style>
+table {background-color: #eca;}
+th {background-color: black; color: white;}
+h1 {
+  background-color: ffaa00;
+  padding:5px;
+  color: black;
+}
+
+svg {
+  margin: 10px;
+  border: 2px;
+  border-style: solid;
+  border-color: black;
+  background: white;
+}
+
+div {
+  border-radius: 5px;
+  background-color: #fec;
+  padding:5px;
+  margin:5px;
+}
+
+.tooltip {color: blue;}
+.tooltip .tooltipcontent  {
+    visibility: hidden;
+    color: black;
+    background-color: yellow;
+    padding: 5px;
+    border-radius: 4px;
+    position: absolute;
+    z-index: 1;
+}
+.tooltip:hover .tooltipcontent {
+    visibility: visible;
+}
+
+.edges line {
+  stroke: #333;
+}
+
+text {
+  font-weight: bold;
+}
+
+.nodes text {
+  color: black;
+  pointer-events: none;
+  font-size: 11px;
+}
+</style>"""
+
+_CSS = """
+<script src="https://d3js.org/d3.v4.min.js"></script>
+"""
+
+_D3_HTML_TEMPLATE = """
+  <script>
+    function buildGraph() {
+      // Build graph data
+      var graph = %s;
+
+      var svg = d3.select("#subgraph%d")
+      var width = svg.attr("width");
+      var height = svg.attr("height");
+      // Make the graph scrollable.
+      svg = svg.call(d3.zoom().on("zoom", function() {
+        svg.attr("transform", d3.event.transform);
+      })).append("g");
+
+
+      var color = d3.scaleOrdinal(d3.schemeDark2);
+
+      var simulation = d3.forceSimulation()
+          .force("link", d3.forceLink().id(function(d) {return d.id;}))
+          .force("charge", d3.forceManyBody())
+          .force("center", d3.forceCenter(0.5 * width, 0.5 * height));
+
+      var edge = svg.append("g").attr("class", "edges").selectAll("line")
+        .data(graph.edges).enter().append("path").attr("stroke","black").attr("fill","none")
+
+      // Make the node group
+      var node = svg.selectAll(".nodes")
+        .data(graph.nodes)
+        .enter().append("g")
+        .attr("x", function(d){return d.x})
+        .attr("y", function(d){return d.y})
+        .attr("transform", function(d) {
+          return "translate( " + d.x + ", " + d.y + ")"
+        })
+        .attr("class", "nodes")
+          .call(d3.drag()
+              .on("start", function(d) {
+                if(!d3.event.active) simulation.alphaTarget(1.0).restart();
+                d.fx = d.x;d.fy = d.y;
+              })
+              .on("drag", function(d) {
+                d.fx = d3.event.x; d.fy = d3.event.y;
+              })
+              .on("end", function(d) {
+                if (!d3.event.active) simulation.alphaTarget(0);
+                d.fx = d.fy = null;
+              }));
+      // Within the group, draw a box for the node position and text
+      // on the side.
+
+      var node_width = 150;
+      var node_height = 30;
+
+      node.append("rect")
+          .attr("r", "5px")
+          .attr("width", node_width)
+          .attr("height", node_height)
+          .attr("rx", function(d) { return d.group == 1 ? 1 : 10; })
+          .attr("stroke", "#000000")
+          .attr("fill", function(d) { return d.group == 1 ? "#dddddd" : "#000000"; })
+      node.append("text")
+          .text(function(d) { return d.name; })
+          .attr("x", 5)
+          .attr("y", 20)
+          .attr("fill", function(d) { return d.group == 1 ? "#000000" : "#eeeeee"; })
+      // Setup force parameters and update position callback
+
+
+      var node = svg.selectAll(".nodes")
+        .data(graph.nodes);
+
+      // Bind the links
+      var name_to_g = {}
+      node.each(function(data, index, nodes) {
+        console.log(data.id)
+        name_to_g[data.id] = this;
+      });
+
+      function proc(w, t) {
+        return parseInt(w.getAttribute(t));
+      }
+      edge.attr("d", function(d) {
+        function lerp(t, a, b) {
+          return (1.0-t) * a + t * b;
+        }
+        var x1 = proc(name_to_g[d.source],"x") + node_width /2;
+        var y1 = proc(name_to_g[d.source],"y") + node_height;
+        var x2 = proc(name_to_g[d.target],"x") + node_width /2;
+        var y2 = proc(name_to_g[d.target],"y");
+        var s = "M " + x1 + " " + y1
+            + " C " + x1 + " " + lerp(.5, y1, y2)
+            + " " + x2 + " " + lerp(.5, y1, y2)
+            + " " + x2  + " " + y2
+      return s;
+    });
+  }
+  console.log("Helllo!");
+  buildGraph();
+</script>
+"""
+
+
+def TensorTypeToName(tensor_type):
+  """Converts a numerical enum to a readable tensor type."""
+  for name, value in schema_fb.TensorType.__dict__.items():
+    if value == tensor_type:
+      return name
+  return None
+
+
+def BuiltinCodeToName(code):
+  """Converts a builtin op code enum to a readable name."""
+  for name, value in schema_fb.BuiltinOperator.__dict__.items():
+    if value == code:
+      return name
+  return None
+
+
+def NameListToString(name_list):
+  """Converts a list of integers to the equivalent ASCII string."""
+  if isinstance(name_list, str):
+    return name_list
+  else:
+    result = ""
+    if name_list is not None:
+      for val in name_list:
+        result = result + chr(int(val))
+    return result
+
+
+class OpCodeMapper:
+  """Maps an opcode index to an op name."""
+
+  def __init__(self, data):
+    self.code_to_name = {}
+    for idx, d in enumerate(data["operator_codes"]):
+      self.code_to_name[idx] = BuiltinCodeToName(d["builtin_code"])
+      if self.code_to_name[idx] == "CUSTOM":
+        self.code_to_name[idx] = NameListToString(d["custom_code"])
+
+  def __call__(self, x):
+    if x not in self.code_to_name:
+      s = "<UNKNOWN>"
+    else:
+      s = self.code_to_name[x]
+    return "%s (%d)" % (s, x)
+
+
+class DataSizeMapper:
+  """For buffers, report the number of bytes."""
+
+  def __call__(self, x):
+    if x is not None:
+      return "%d bytes" % len(x)
+    else:
+      return "--"
+
+
+class TensorMapper:
+  """Maps a list of tensor indices to a tooltip hoverable indicator of more."""
+
+  def __init__(self, subgraph_data):
+    self.data = subgraph_data
+
+  def __call__(self, x):
+    html = ""
+    if x is None:
+      return html
+
+    html += "<span class='tooltip'><span class='tooltipcontent'>"
+    for i in x:
+      tensor = self.data["tensors"][i]
+      html += str(i) + " "
+      html += NameListToString(tensor["name"]) + " "
+      html += TensorTypeToName(tensor["type"]) + " "
+      html += (repr(tensor["shape"]) if "shape" in tensor else "[]")
+      html += (repr(tensor["shape_signature"])
+               if "shape_signature" in tensor else "[]") + "<br>"
+    html += "</span>"
+    html += repr(x)
+    html += "</span>"
+    return html
+
+
+def GenerateGraph(subgraph_idx, g, opcode_mapper):
+  """Produces the HTML required to have a d3 visualization of the dag."""
+
+  def TensorName(idx):
+    return "t%d" % idx
+
+  def OpName(idx):
+    return "o%d" % idx
+
+  edges = []
+  nodes = []
+  first = {}
+  second = {}
+  pixel_mult = 200  # TODO(aselle): multiplier for initial placement
+  width_mult = 170  # TODO(aselle): multiplier for initial placement
+  for op_index, op in enumerate(g["operators"] or []):
+    if op["inputs"] is not None:
+      for tensor_input_position, tensor_index in enumerate(op["inputs"]):
+        if tensor_index not in first:
+          first[tensor_index] = ((op_index - 0.5 + 1) * pixel_mult,
+                                 (tensor_input_position + 1) * width_mult)
+        edges.append({
+            "source": TensorName(tensor_index),
+            "target": OpName(op_index)
+        })
+    if op["outputs"] is not None:
+      for tensor_output_position, tensor_index in enumerate(op["outputs"]):
+        if tensor_index not in second:
+          second[tensor_index] = ((op_index + 0.5 + 1) * pixel_mult,
+                                  (tensor_output_position + 1) * width_mult)
+        edges.append({
+            "target": TensorName(tensor_index),
+            "source": OpName(op_index)
+        })
+
+    nodes.append({
+        "id": OpName(op_index),
+        "name": opcode_mapper(op["opcode_index"]),
+        "group": 2,
+        "x": pixel_mult,
+        "y": (op_index + 1) * pixel_mult
+    })
+  for tensor_index, tensor in enumerate(g["tensors"]):
+    initial_y = (
+        first[tensor_index] if tensor_index in first else
+        second[tensor_index] if tensor_index in second else (0, 0))
+
+    nodes.append({
+        "id": TensorName(tensor_index),
+        "name": "%r (%d)" % (getattr(tensor, "shape", []), tensor_index),
+        "group": 1,
+        "x": initial_y[1],
+        "y": initial_y[0]
+    })
+  graph_str = json.dumps({"nodes": nodes, "edges": edges})
+
+  html = _D3_HTML_TEMPLATE % (graph_str, subgraph_idx)
+  return html
+
+
+def GenerateTableHtml(items, keys_to_print, display_index=True):
+  """Given a list of object values and keys to print, make an HTML table.
+
+  Args:
+    items: Items to print an array of dicts.
+    keys_to_print: (key, display_fn). `key` is a key in the object. i.e.
+      items[0][key] should exist. display_fn is the mapping function on display.
+      i.e. the displayed html cell will have the string returned by
+      `mapping_fn(items[0][key])`.
+    display_index: add a column which is the index of each row in `items`.
+
+  Returns:
+    An html table.
+  """
+  html = ""
+  # Print the list of  items
+  html += "<table><tr>\n"
+  html += "<tr>\n"
+  if display_index:
+    html += "<th>index</th>"
+  for h, mapper in keys_to_print:
+    html += "<th>%s</th>" % h
+  html += "</tr>\n"
+  for idx, tensor in enumerate(items):
+    html += "<tr>\n"
+    if display_index:
+      html += "<td>%d</td>" % idx
+    # print tensor.keys()
+    for h, mapper in keys_to_print:
+      val = tensor[h] if h in tensor else None
+      val = val if mapper is None else mapper(val)
+      html += "<td>%s</td>\n" % val
+
+    html += "</tr>\n"
+  html += "</table>\n"
+  return html
+
+
+def CamelCaseToSnakeCase(camel_case_input):
+  """Converts an identifier in CamelCase to snake_case."""
+  s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", camel_case_input)
+  return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
+
+
+def FlatbufferToDict(fb, preserve_as_numpy):
+  """Converts a hierarchy of FB objects into a nested dict.
+
+  We avoid transforming big parts of the flat buffer into python arrays. This
+  speeds conversion from ten minutes to a few seconds on big graphs.
+
+  Args:
+    fb: a flat buffer structure. (i.e. ModelT)
+    preserve_as_numpy: true if all downstream np.arrays should be preserved.
+      false if all downstream np.array should become python arrays
+  Returns:
+    A dictionary representing the flatbuffer rather than a flatbuffer object.
+  """
+  if isinstance(fb, int) or isinstance(fb, float) or isinstance(fb, str):
+    return fb
+  elif hasattr(fb, "__dict__"):
+    result = {}
+    for attribute_name in dir(fb):
+      attribute = fb.__getattribute__(attribute_name)
+      if not callable(attribute) and attribute_name[0] != "_":
+        snake_name = CamelCaseToSnakeCase(attribute_name)
+        preserve = True if attribute_name == "buffers" else preserve_as_numpy
+        result[snake_name] = FlatbufferToDict(attribute, preserve)
+    return result
+  elif isinstance(fb, np.ndarray):
+    return fb if preserve_as_numpy else fb.tolist()
+  elif hasattr(fb, "__len__"):
+    return [FlatbufferToDict(entry, preserve_as_numpy) for entry in fb]
+  else:
+    return fb
+
+
+def CreateDictFromFlatbuffer(buffer_data):
+  model_obj = schema_fb.Model.GetRootAsModel(buffer_data, 0)
+  model = schema_fb.ModelT.InitFromObj(model_obj)
+  return FlatbufferToDict(model, preserve_as_numpy=False)
+
+
+def create_html(tflite_input, input_is_filepath=True):  # pylint: disable=invalid-name
+  """Returns html description with the given tflite model.
+
+  Args:
+    tflite_input: TFLite flatbuffer model path or model object.
+    input_is_filepath: Tells if tflite_input is a model path or a model object.
+
+  Returns:
+    Dump of the given tflite model in HTML format.
+
+  Raises:
+    RuntimeError: If the input is not valid.
+  """
+
+  # Convert the model into a JSON flatbuffer using flatc (build if doesn't
+  # exist.
+  if input_is_filepath:
+    if not os.path.exists(tflite_input):
+      raise RuntimeError("Invalid filename %r" % tflite_input)
+    if tflite_input.endswith(".tflite") or tflite_input.endswith(".bin") or tflite_input.endswith(".tf_lite"):
+      with open(tflite_input, "rb") as file_handle:
+        file_data = bytearray(file_handle.read())
+      data = CreateDictFromFlatbuffer(file_data)
+    elif tflite_input.endswith(".json"):
+      data = json.load(open(tflite_input))
+    else:
+      raise RuntimeError("Input file was not .tflite or .json")
+  else:
+    data = CreateDictFromFlatbuffer(tflite_input)
+  html = ""
+  # html += _CSS
+  html += "<h1>TensorFlow Lite Model</h2>"
+
+  data["filename"] = tflite_input if input_is_filepath else (
+      "Null (used model object)")  # Avoid special case
+
+  toplevel_stuff = [("filename", None), ("version", None),
+                    ("description", None)]
+
+  html += "<table>\n"
+  for key, mapping in toplevel_stuff:
+    if not mapping:
+      mapping = lambda x: x
+    html += "<tr><th>%s</th><td>%s</td></tr>\n" % (key, mapping(data.get(key)))
+  html += "</table>\n"
+
+  # Spec on what keys to display
+  buffer_keys_to_display = [("data", DataSizeMapper())]
+  operator_keys_to_display = [("builtin_code", BuiltinCodeToName),
+                              ("custom_code", NameListToString),
+                              ("version", None)]
+
+  # Update builtin code fields.
+  for d in data["operator_codes"]:
+    d["builtin_code"] = max(d["builtin_code"], d["deprecated_builtin_code"])
+
+  for subgraph_idx, g in enumerate(data["subgraphs"]):
+    # Subgraph local specs on what to display
+    html += "<div class='subgraph'>"
+    tensor_mapper = TensorMapper(g)
+    opcode_mapper = OpCodeMapper(data)
+    op_keys_to_display = [("inputs", tensor_mapper), ("outputs", tensor_mapper),
+                          ("builtin_options", None),
+                          ("opcode_index", opcode_mapper)]
+    tensor_keys_to_display = [("name", NameListToString),
+                              ("type", TensorTypeToName), ("shape", None),
+                              ("shape_signature", None), ("buffer", None),
+                              ("quantization", None)]
+
+    html += "<h2>Subgraph %d</h2>\n" % subgraph_idx
+
+    # Inputs and outputs.
+    html += "<h3>Inputs/Outputs</h3>\n"
+    html += GenerateTableHtml([{
+        "inputs": g["inputs"],
+        "outputs": g["outputs"]
+    }], [("inputs", tensor_mapper), ("outputs", tensor_mapper)],
+                              display_index=False)
+
+    # Print the tensors.
+    html += "<h3>Tensors</h3>\n"
+    html += GenerateTableHtml(g["tensors"], tensor_keys_to_display)
+
+    # Print the ops.
+    if g["operators"]:
+      html += "<h3>Ops</h3>\n"
+      html += GenerateTableHtml(g["operators"], op_keys_to_display)
+
+    # Visual graph.
+    html += "<svg id='subgraph%d' width='1600' height='900'></svg>\n" % (
+        subgraph_idx,)
+    html += GenerateGraph(subgraph_idx, g, opcode_mapper)
+    html += "</div>"
+
+  # Buffers have no data, but maybe in the future they will
+  html += "<h2>Buffers</h2>\n"
+  html += GenerateTableHtml(data["buffers"], buffer_keys_to_display)
+
+  # Operator codes
+  html += "<h2>Operator Codes</h2>\n"
+  html += GenerateTableHtml(data["operator_codes"], operator_keys_to_display)
+
+  # html += "</body></html>\n"
+
+  # return f"<iframe src={escape(html)} ></iframe>"
+
+  html += """ <script src="https://d3js.org/d3.v4.min.js"></script> """
+  return html
+
+
+def main(argv):
+  try:
+    tflite_input = argv[1]
+    html_output = argv[2]
+  except IndexError:
+    print("Usage: %s <input tflite> <output html>" % (argv[0]))
+  else:
+    html = create_html(tflite_input)
+    with open(html_output, "w") as output_file:
+      output_file.write(html)
+
+def process_file(file):
+  try:
+    html = create_html(file.name)
+    return html
+  except Exception as e:
+    return f"Error: {str(e)}"
+
+with gr.Blocks(head=_CSS, ) as demo:
+  gr.Markdown("## TensorFlow Lite Model Visualizer")
+  file_input = gr.File(label="Upload TFLite File")
+  html_output = gr.HTML(label="Generated HTML", container=True)
+  file_input.change(process_file, inputs=file_input, outputs=html_output)
+
+demo.launch()
+
+
+# if __name__ == "__main__":
+#   main(sys.argv)