Merge branch 'feature/onnx_upload' into develop

Author: Alex

backend/requirements.txt

@@ -65,3 +65,4 @@ websocket-client==0.58.0
 Werkzeug==1.0.1
 wrapt==1.12.1
 zipp==3.4.1
+onnx==1.10.1

backend/server.py

@@ -7,10 +7,11 @@
 import cairosvg
 from flask import Flask, send_file, request, jsonify
 from translate.translate_keras import translate_keras
+from translate.translate_onnx import translate_onnx
 from translate.graph import Graph
 
 app = Flask(__name__)
-app.config['UPLOAD_EXTENSIONS'] = ['.h5']
+app.config['UPLOAD_EXTENSIONS'] = ['.h5', '.onnx']
 app.config['UPLOAD_PATH'] = 'models'
 ok_status = 200
 error_status = 400
@@ -168,6 +169,9 @@ def get_network(identifier):
     if os.path.exists(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.h5')):
         graph = translate_keras(os.path.join(
             app.config['UPLOAD_PATH'], identifier, 'model.h5'))
+    elif os.path.exists(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.onnx')):
+        graph = translate_onnx(os.path.join(
+            app.config['UPLOAD_PATH'], identifier, 'model.onnx'))
     else:
         graph = translate_keras(os.path.join(
             app.config['UPLOAD_PATH'], identifier, 'model_current.py'))
@@ -190,7 +194,7 @@ def get_code(identifier):
         object -- a http response containing the code for the network
     """
     check_upload_path_exists(identifier)
-    if not (os.path.exists(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.h5'))):
+    if not (os.path.exists(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.h5')) or os.path.exists(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.onnx'))):
         with open(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model_current.py'), 'r') as myfile:
             keras_code = myfile.read()
             return keras_code, ok_status, text_type
@@ -231,7 +235,8 @@ def upload_model(identifier):
     file_ext = os.path.splitext(filename)[1]
     if file_ext not in app.config['UPLOAD_EXTENSIONS']:
         return "", error_status, text_type
-    file_path = os.path.join(app.config['UPLOAD_PATH'], identifier, 'model.h5')
+    file_path = os.path.join(
+        app.config['UPLOAD_PATH'], identifier, f'model{file_ext}')
     uploaded_file.save(file_path)
     return "No code loaded since model file is present.", ok_status, text_type
 
@@ -247,12 +252,14 @@ def delete_model(identifier):
         object -- a http response signaling the change worked
     """
     check_upload_path_exists(identifier)
-    model_path = os.path.join(
+    keras_model_path = os.path.join(
         app.config['UPLOAD_PATH'], identifier, 'model.h5')
-    print(model_path)
-    if (os.path.exists(model_path)):
-        print('test2')
-        os.remove(model_path)
+    onnx_model_path = os.path.join(
+        app.config['UPLOAD_PATH'], identifier, 'model.onnx')
+    if (os.path.exists(keras_model_path)):
+        os.remove(keras_model_path)
+    if (os.path.exists(onnx_model_path)):
+        os.remove(onnx_model_path)
     with open(os.path.join(app.config['UPLOAD_PATH'], identifier, 'model_current.py'), 'r') as myfile:
         keras_code = myfile.read()
         return keras_code, ok_status, text_type

backend/translate/layer.py

@@ -1,52 +1,82 @@
 """Used to handle layers as objects."""
+
+
 class Layer:
-  """Representation of Layers."""
-  # Initialize the Properties.
-  def __init__(self, class_name, name, layer):
-    self.properties = {}
-    self.input = []
-    self.input_names = []
-    self.output = []
-    self.name = name
-    self.type = class_name
-    self.dimensions = {
-        'in': 0,
-        'out': layer.get_output_at(0).get_shape().as_list()[1:]
-    }
-    if isinstance(layer.get_input_at(0), list):
-      self.dimensions['in'] = layer.get_input_at(0)[0].get_shape().as_list()[1:]
-    else:
-      self.dimensions['in'] = layer.get_input_at(0).get_shape().as_list()[1:]
-    if isinstance(layer.get_output_at(0), list):
-      self.dimensions['out'] = layer.get_output_at(0)[0].get_shape().as_list()[1:]
-    else:
-      self.dimensions['out'] = layer.get_output_at(0).get_shape().as_list()[1:]
-
-
-  def add_specs(self, specs):
-    """Add Specifications.
-
-    Arguments:
-        specs {list} -- the specifications to be added to the layer
-    """
-    # Update Properties based on Specs. Try to Parse the Specs.
-    for prop in specs:
-      if isinstance(specs[prop], dict):
-        self.properties[prop] = specs[prop]['class_name']
-      else:
-        self.properties[prop] = specs[prop]
-
-
-  def add_input_names(self, nodes):
-    """Adds all the Names of the Input Nodes to the Layer.
-
-    Arguments:
-        nodes {list} -- all nodes that are inputs to the layer
-    """
-    if (len(nodes) > 0):
-      for node in nodes[0]:
-        self.input_names.append(node[0])
-
-  # String Representation of the Layer.
-  def __repr__(self):
-    return "%s(properties: %r)" % (self.__class__, self.properties)
+    """Representation of Layers."""
+    # Initialize the Properties.
+
+    def __init__(self, class_name, name, dimensions):
+        self.properties = {}
+        self.input = []
+        self.input_names = []
+        self.output = []
+        self.name = name
+        self.type = class_name
+        self.dimensions = dimensions
+
+    @classmethod
+    def from_keras(cls, class_name, name, layer):
+        dimensions = {
+            'in': 0,
+            'out': layer.get_output_at(0).get_shape().as_list()[1:]
+        }
+        if isinstance(layer.get_input_at(0), list):
+            dimensions['in'] = layer.get_input_at(
+                0)[0].get_shape().as_list()[1:]
+        else:
+            dimensions['in'] = layer.get_input_at(0).get_shape().as_list()[1:]
+        if isinstance(layer.get_output_at(0), list):
+            dimensions['out'] = layer.get_output_at(
+                0)[0].get_shape().as_list()[1:]
+        else:
+            dimensions['out'] = layer.get_output_at(
+                0).get_shape().as_list()[1:]
+        return cls(class_name, name, dimensions)
+
+    @classmethod
+    def from_onnx(cls, layer, onnx_graph):
+        dimensions = {
+            'in': 0,
+            'out': 0
+        }
+        if len(layer.input) > 0:
+            for value_info in onnx_graph.value_info:
+                if (layer.input[0] == value_info.name):
+                    dimensions['in'] = list(
+                        map(lambda x: x.dim_value, value_info.type.tensor_type.shape.dim[1:]))
+                if (layer.output[0] == value_info.name):
+                    dimensions['out'] = list(
+                        map(lambda x: x.dim_value, value_info.type.tensor_type.shape.dim[1:]))
+            for value_info in onnx_graph.input:
+                if (layer.input[0] == value_info.name):
+                    dimensions['in'] = list(
+                        map(lambda x: x.dim_value, value_info.type.tensor_type.shape.dim[1:]))
+            for value_info in onnx_graph.output:
+                if (layer.output[0] == value_info.name):
+                    dimensions['out'] = list(
+                        map(lambda x: x.dim_value, value_info.type.tensor_type.shape.dim[1:]))
+        return cls(layer.op_type, layer.name, dimensions)
+
+    def add_specs(self, specs):
+        """Add Specifications.
+        Arguments:
+            specs {list} -- the specifications to be added to the layer
+        """
+        # Update Properties based on Specs. Try to Parse the Specs.
+        for prop in specs:
+            if isinstance(specs[prop], dict):
+                self.properties[prop] = specs[prop]['class_name']
+            else:
+                self.properties[prop] = specs[prop]
+
+    def add_input_names_from_node(self, nodes):
+        """Adds all the Names of the Input Nodes to the Layer.
+        Arguments:
+            nodes {list} -- all nodes that are inputs to the layer
+        """
+        if (len(nodes) > 0):
+            for node in nodes[0]:
+                self.input_names.append(node[0])
+
+    def __repr__(self):
+        return "%s(properties: %r)" % (self.__class__, self.properties)

backend/translate/translate.py

@@ -17,17 +17,17 @@ def translate_keras(filename):
     Returns:
         object -- the result of this translation, can be an error
     """
-    epicbox.configure(profiles=[
-        epicbox.Profile('python', 'tf_plus_keras:latest')])
-    general_reader = open('translate/keras_loader.txt', 'rb')
-    general_code = general_reader.read()
     if keras_ext in filename:
         try:
             return graph_from_model_file(filename)
         except Exception as err:
             return {'error_class': '', 'line_number': 1,
                     'detail': "Model could not be loaded correctly. Error: " + str(err)}
     else:
+        epicbox.configure(profiles=[
+            epicbox.Profile('python', 'tf_plus_keras:latest')])
+        general_reader = open('translate/keras_loader.txt', 'rb')
+        general_code = general_reader.read()
         with open(filename, 'rb') as myfile:
             keras_code = myfile.read()
             try:
@@ -99,8 +99,8 @@ def add_layer_type(layer_json, model_layer, graph, previous_node):
     Returns:
         String -- name of the new layer
     """
-    new_layer = layer.Layer(layer_json['class_name'],
-                            layer_json['config']['name'], model_layer)
+    new_layer = layer.Layer.from_keras(layer_json['class_name'],
+                                       layer_json['config']['name'], model_layer)
     new_layer.add_specs(layer_json['config'])
     return add_to_graph(new_layer, layer_json, graph, previous_node)
 
@@ -118,9 +118,9 @@ def add_to_graph(new_layer, model_layer, graph, previous_node):
         String -- name of the new layer
     """
     try:
-        new_layer.add_input_names(model_layer['inbound_nodes'])
+        new_layer.add_input_names_from_node(model_layer['inbound_nodes'])
     except Exception:
         if previous_node != '':
-            new_layer.add_input_names([[[previous_node, 0, 0, {}]]])
+            new_layer.add_input_names_from_node([[[previous_node, 0, 0, {}]]])
     graph.add_layer(new_layer)
     return new_layer.name

backend/translate/translate_keras.py

@@ -0,0 +1,70 @@
+import onnx
+from onnx import shape_inference
+from translate.graph import Graph
+import translate.layer as layer
+
+
+def translate_onnx(file):
+    """Get a graph from an external file defining the neural network.
+    Arguments:
+        file {Path} -- the path to the onnx file to be translated
+    Returns:
+        object -- the result of this translation, can be an error
+    """
+    onnx_model = onnx.load(file)
+    onnx_model = shape_inference.infer_shapes(onnx_model)
+    graph = Graph()
+    previous_node = ''
+    for index, node in enumerate(onnx_model.graph.node):
+        if index >= 0:
+            add_layer_type(node, graph, previous_node, onnx_model.graph)
+    graph.resolve_input_names()
+    return graph
+
+
+def add_layer_type(node, graph, previous_node, onnx_graph):
+    """Add a Layer. Layers are identified by name and equipped using the spec.
+    Arguments:
+        node {onnx.Node} -- the node that represents the layer
+        graph {object} -- neural network graph
+        previous_node {String} -- name of the node before this one
+        onnx_graph {onnx.Graph} -- the graph that represents this network
+    Returns:
+        String -- name of the new layer
+    """
+    new_layer = layer.Layer.from_onnx(node, onnx_graph)
+    add_specs_onnx(new_layer, node.attribute)
+    return add_to_graph(new_layer, node, onnx_graph, graph, previous_node)
+
+
+def add_to_graph(new_layer, node, onnx_graph, graph, previous_node):
+    """Takes new layer, adds the Properties and then adds the Layer to the Graph.
+    Arguments:
+        new_layer {object} -- the layer to be added to the graph
+        model_layer {dict} -- json dict of the layer props
+        graph {object} -- the neural network graph
+        previous_node {String} -- name of the previous layer
+    Returns:
+        String -- name of the new layer
+    """
+    for input in node.input:
+        for graph_node in onnx_graph.node:
+            for output in graph_node.output:
+                if (input == output):
+                    new_layer.input_names.append(graph_node.name)
+    graph.add_layer(new_layer)
+    return new_layer.name
+
+
+def add_specs_onnx(new_layer, attributes):
+    for attribute in attributes:
+        if (attribute.type == onnx.AttributeProto.AttributeType.FLOAT):
+            new_layer.properties[attribute.name] = attribute.f
+        elif (attribute.type == onnx.AttributeProto.AttributeType.FLOATS):
+            new_layer.properties[attribute.name] = str(attribute.floats)
+        elif (attribute.type == onnx.AttributeProto.AttributeType.INT):
+            new_layer.properties[attribute.name] = attribute.i
+        elif (attribute.type == onnx.AttributeProto.AttributeType.INTS):
+            new_layer.properties[attribute.name] = str(attribute.ints)
+        else:
+            print(attribute.type, 'not supported yet.')