First tests done

Author: MKaptein

sclblonnx/__init__.py

@@ -27,7 +27,8 @@
 from .node import \
     node, \
     add_node, \
-    add_nodes
+    add_nodes, \
+    delete_node
 
 from .constant import \
     constant, \

sclblonnx/input.py

@@ -1,18 +1,17 @@
 from onnx import helper as xhelp
 from onnx import onnx_ml_pb2 as xpb2
-
-from sclblonnx.utils import _value, _data_type, _parse_element
+from sclblonnx.utils import _value, _data_type, _parse_element, _print
 
 
 # list_inputs lists all inputs of a graph
 def list_inputs(graph: xpb2.GraphProto):
-    """ Tries to list the outputs of a given graph.
+    """ Tries to list the inputs of a given graph.
 
     Args:
         graph the ONNX graph
     """
     if type(graph) is not xpb2.GraphProto:
-        print("graph is not a valid ONNX graph.")
+        _print("graph is not a valid ONNX graph.")
         return False
 
     i = 1
@@ -21,6 +20,9 @@ def list_inputs(graph: xpb2.GraphProto):
         print("Input {}: Name: '{}', Type: {}, Dimension: {}".format(i, name, dtype, shape))
         i += 1
 
+    if i == 1:
+        print("No inputs found.")
+
     return True
 
 
@@ -38,13 +40,14 @@ def add_input(
         name: String, the name of the input as used to determine the graph topology.
         data_type: String, the data type of the input. Run list_data_types() for an overview.
         dimensions: List[] specifying the dimensions of the input.
+        **kwargs
 
     Returns:
         The extended graph.
 
     """
     if type(graph) is not xpb2.GraphProto:
-        print("graph is not a valid ONNX graph.")
+        _print("graph is not a valid ONNX graph.")
         return False
 
     dtype = _data_type(data_type)
@@ -54,21 +57,34 @@ def add_input(
     try:
         graph.input.append(xhelp.make_tensor_value_info(name, dtype, dimensions, **kwargs), *kwargs)
     except Exception as e:
-        print("Unable to add the input: " + str(e))
+        _print("Unable to add the input: " + str(e))
         return False
     return graph
 
 
 # rename_input renames an input
 def rename_input(graph, current_name, new_name):
-    # We have to rename the input itself:
+    """ Rename an input to a graph
+
+    Args:
+        graph: A graph, onnx.onnx_ml_pb2.GraphProto.
+        current_name: String, the current input name.
+        new_name: String, the name desired input name.
+
+    Returns:
+        The changed graph.
+    """
+    if type(graph) is not xpb2.GraphProto:
+        _print("graph is not a valid ONNX graph.")
+        return False
+
     found = False
     for input in graph.input:
         if input.name == current_name:
             input.name = new_name
             found = True
     if not found:
-        print("err")
+        _print("Unable to find the input to rename.")
         return False
 
     # And rename it in every nodes that takes this as input:
@@ -93,14 +109,15 @@ def replace_input(
         graph: A graph, onnx.onnx_ml_pb2.GraphProto.
         name: String, the name of the input as used to determine the graph topology.
         data_type: String, the data type of the input. Run list_data_types() for an overview.
-        dimensions: List[] specifying the dimensions of the input.
+        dimensions: List[] specifying the dimensions of the input.,
+        **kwargs
 
     Returns:
         The extended graph.
 
     """
     if type(graph) is not xpb2.GraphProto:
-        print("graph is not a valid ONNX graph.")
+        _print("graph is not a valid ONNX graph.")
         return graph
 
     # Remove the named input
@@ -111,33 +128,32 @@ def replace_input(
                 graph.input.remove(elem)
                 found = True
     except Exception as e:
-        print("Unable to iterate the inputs. " + str(e))
+        _print("Unable to iterate the inputs. " + str(e))
         return False
     if not found:
-        print("Unable to find the input by name.")
+        _print("Unable to find the input by name.")
 
     # Create the new value
     try:
         val = _value(name, data_type, dimensions, **kwargs)
     except Exception as e:
-        print("Unable to create value. " + str(e))
+        _print("Unable to create value. " + str(e))
         return False
 
     # Add the value to the input
     try:
         graph.input.append(val, *kwargs)
     except Exception as e:
-        print("Unable to add the input: " + str(e))
+        _print("Unable to add the input: " + str(e))
         return False
 
     return graph
 
 
 # delete_input deletes an existing input
 def delete_input(
-            graph: xpb2.GraphProto,
-            name: str,
-            **kwargs):
+        graph: xpb2.GraphProto,
+        name: str):
     """ Removes an existing input of a graph by name
 
     Args:
@@ -159,10 +175,10 @@ def delete_input(
                 graph.input.remove(elem)
                 found = True
     except Exception as e:
-        print("Unable to iterate the inputs. " + str(e))
+        _print("Unable to iterate the inputs. " + str(e))
         return False
     if not found:
-        print("Unable to find the input by name.")
+        _print("Unable to find the input by name.")
         return False
 
-    return graph
+    return graph

sclblonnx/node.py

@@ -1,7 +1,7 @@
 from onnx import helper as xhelp
 from onnx import onnx_ml_pb2 as xpb2
-
 import sclblonnx._globals as glob
+from sclblonnx.utils import _print
 
 
 # Node creates a new node
@@ -18,6 +18,7 @@ def node(
         inputs: [] list of inputs (names to determine the graph topology)
         outputs: [] list of outputs (names to determine the graph topology)
         name: The name of this node (Optional)
+        **kwargs
     """
     if not name:
         name = "sclbl-onnx-node" + str(glob.NODE_COUNT)
@@ -26,7 +27,7 @@ def node(
     try:
         node = xhelp.make_node(op_type, inputs, outputs, name, **kwargs)
     except Exception as e:
-        print("Unable to create node: " + str(e))
+        _print("Unable to create node: " + str(e))
         return False
     return node
 
@@ -43,20 +44,23 @@ def add_node(
     Args:
         graph: A graph, onnx.onnx_ml_pb2.GraphProto.
         node: A node, onnx.onnx_ml_pb2.NodeProto.
+        **kwargs
 
     Returns:
         The extended graph.
     """
     if type(graph) is not xpb2.GraphProto:
-        print("graph is not a valid ONNX graph.")
+        _print("The graph is not a valid ONNX graph.")
         return False
+
     if type(node) is not xpb2.NodeProto:
-        print("node is not a valid ONNX node.")
+        _print("The node is not a valid ONNX node.")
         return False
+
     try:
         graph.node.append(node, **kwargs)
     except Exception as e:
-        print("Unable to extend graph: " + str(e))
+        _print("Unable to extend graph: " + str(e))
         return False
     return graph
 
@@ -72,6 +76,7 @@ def add_nodes(
     Args:
         graph: A graph, onnx.onnx_ml_pb2.GraphProto.
         nodes: A list of nodes, [onnx.onnx_ml_pb2.NodeProto].
+        **kwargs
 
     Returns:
         The extended graph.
@@ -80,8 +85,50 @@ def add_nodes(
         print("graph is not a valid ONNX graph.")
         return False
 
-    for node in nodes:
+    for node in nodes:  # error handling in add_node
         graph = add_node(graph, node, **kwargs)
         if not graph:
             return False
-    return graph
+
+    return graph
+
+
+# delete_node deletes a node from a graph
+def delete_node(
+        graph: xpb2.GraphProto,
+        node_name: str = "",
+        **kwargs):
+    """ Add node appends a node to graph g and returns the extended graph
+
+    Prints a message and returns False if fails.
+
+    Args:
+        graph: A graph, onnx.onnx_ml_pb2.GraphProto.
+        node_name: Name of the node to remove.
+        **kwargs
+
+    Returns:
+        The extended graph.
+    """
+    if type(graph) is not xpb2.GraphProto:
+        _print("The graph is not a valid ONNX graph.")
+        return False
+
+    if not node_name:
+        _print("Please specify a node name.")
+        return False
+
+    found = False
+    try:
+        for elem in graph.node:
+            if elem.name == node_name:
+                graph.node.remove(elem)
+                found = True
+    except Exception as e:
+        _print("Unable to iterate the nodes. " + str(e))
+        return False
+    if not found:
+        _print("Unable to find the node by name.")
+        return False
+
+    return graph