leap-model-parser 0.1.224__py3-none-any.whl → 0.1.225.dev1__py3-none-any.whl

leap_model_parser/contract/graph.py

@@ -47,6 +47,7 @@ class Node:
     outputs: Dict[str, ConnectionOutput] = field(default_factory=dict)
     pruning_plan_id: Optional[str] = None
     wrapper: Optional[WrapperData] = None
+    shape: Optional[List[Optional[int]]] = None
 
     def __key(self):
         return (self.id, self.name)

leap_model_parser/keras_json_model_import.py

@@ -5,7 +5,6 @@ from typing import Set, Dict, Any, List, Type, Optional, Tuple
 
 import tensorflow as tf  # type: ignore
 from keras.engine.keras_tensor import KerasTensor  # type: ignore
-from keras.engine.node import Node  # type: ignore
 from keras.engine.node import Node as keras_node # type: ignore
 from keras.layers import Layer # type: ignore
 from keras.layers.convolutional.base_conv import Conv # type: ignore
@@ -221,7 +220,7 @@ class KerasJsonModelImport:
 
         self.layer_data_adjustments(data, layer_metadata)
         node = Node(id=str(self.id), name=layer_metadata.get(
-            "class_name", layer["class_name"]), data=data, position=[0, 0])
+            "class_name", layer["class_name"]), data=data, position=[0, 0], shape=layer['output_shape'])
         if 'wrapper' in layer:
             node.wrapper = layer['wrapper']
         self.nodes_cache[node_key] = node

leap_model_parser/model_parser.py

@@ -54,6 +54,10 @@ class ModelParser:
             ImportModelTypeEnum.PB_TF2.value: self.convert_pb_model,
         }
 
+    def _add_output_node_shape_to_model_schema(self, model_schema: Dict, keras_model: Model):
+        for i, layer in enumerate(keras_model.layers):
+            model_schema['config']['layers'][i]['output_shape'] = list(layer.output_shape)
+
     def get_keras_model_and_model_graph(
             self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[Dict[str, Node], List[InputInfo], Optional[Model], Optional[str]]:
         model_to_keras_converter: Optional[Callable[[str], Tuple[Dict[str, Node], Model, Optional[str]]]] = \
@@ -64,6 +68,9 @@ class ModelParser:
 
         file_path = str(model_path)
         model_schema, keras_model_with_weights, error_info = model_to_keras_converter(file_path)
+
+        self._add_output_node_shape_to_model_schema(model_schema, keras_model_with_weights)
+
         model_generator = KerasJsonModelImport(self.custom_layers)
 
         keras_model = keras_model_with_weights

{leap_model_parser-0.1.224.dist-info → leap_model_parser-0.1.225.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.224
+Version: 0.1.225.dev1
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.224.dist-info → leap_model_parser-0.1.225.dev1.dist-info}/RECORD

@@ -1,13 +1,13 @@
 LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
 leap_model_parser/__init__.py,sha256=OAU7rFHAVVWUM-cDtQ4Ohum567KN8M-YTkHZp5KiYbo,132
 leap_model_parser/contract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-leap_model_parser/contract/graph.py,sha256=c5rSlJJQ3JzLhyeniglzrP5GH0fp8UZDQKZw1fcrNe8,1132
+leap_model_parser/contract/graph.py,sha256=uE0WecEB8WjDFNMZbRLEzmYq7w8nnyjiLU5iA9O9eqw,1180
 leap_model_parser/contract/importmodelresponse.py,sha256=GlvnKS8xrebU2Sj0dxqtEhAOOo3DlOtT7AMJ2BlzH9E,145
 leap_model_parser/contract/nodedata.py,sha256=1_ML0nzp3QUZ0_9mGSLhfO4_hqjYMwi0DWLwymUnWEs,43326
 leap_model_parser/contract/ui_components.json,sha256=0lsxwOLElW1E-imCcdh3zKPWgzFuQ_bApG6aHvYfTvo,410591
-leap_model_parser/keras_json_model_import.py,sha256=MD9ksO54VywikEAb2lxaNHq8Nbp8oiJab2Yn8UB-DRI,16841
+leap_model_parser/keras_json_model_import.py,sha256=Bt1wZX5PK6YFUKGk2wNq03M5sckeBb3HdDpzMpMGWrA,16819
 leap_model_parser/leap_graph_editor.py,sha256=0SDHNVkMq94PXWaTqZzMlzJMRLWebOj4uyz64c_KCTQ,13559
-leap_model_parser/model_parser.py,sha256=Je37-FjrNgi8w17rBIt_b71xm-BoiopBAvgW30qfb7o,6614
+leap_model_parser/model_parser.py,sha256=ovXrVQqdbWJPMikJp96Llim-8fmcVM_Udoif5Tcf_kg,6949
 leap_model_parser/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 leap_model_parser/utils/layerpedia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 leap_model_parser/utils/layerpedia/layerpedia.py,sha256=1syubfXBTB630TVkgcQ-Ge7Qe9Zbr6EtZRreuqCJnQ8,9292
@@ -18,8 +18,8 @@ leap_model_parser/utils/uicomponents/generatenodedata.py,sha256=LRaPlO5jJ9pUtkvL
 leap_model_parser/utils/uicomponents/tensorflowinscpection.py,sha256=ym613z9iQKPDBpr0RYD35bTABdm1L-Ez86G47BYT7qw,6775
 leap_model_parser/utils/uicomponents/ui_components.json,sha256=0lsxwOLElW1E-imCcdh3zKPWgzFuQ_bApG6aHvYfTvo,410591
 leap_model_parser/utils/uicomponents/ui_components_config.yaml,sha256=cRH8T-c3TAL0nfefRvt9pFsjbTWNEg38NRyHR7RpJsk,19534
-leap_model_parser-0.1.224.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
-leap_model_parser-0.1.224.dist-info/METADATA,sha256=XzCmWdsrO-N07khELKcWtfS5G7e-vVj5rF4hIe9qJ-0,1104
-leap_model_parser-0.1.224.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-leap_model_parser-0.1.224.dist-info/entry_points.txt,sha256=ZvV6EuQt1uAqwapNg5Lo2qjJM9ZG5g2wfzZoLh_Ztyk,77
-leap_model_parser-0.1.224.dist-info/RECORD,,
+leap_model_parser-0.1.225.dev1.dist-info/LICENSE,sha256=qIwWjdspQeSMTtnFZBC8MuT-95L02FPvzRUdWFxrwJY,1067
+leap_model_parser-0.1.225.dev1.dist-info/METADATA,sha256=RK88uZWeYk9yba7RknolL0wNTsI72PK0pIP8_Ur9e5g,1109
+leap_model_parser-0.1.225.dev1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+leap_model_parser-0.1.225.dev1.dist-info/entry_points.txt,sha256=ZvV6EuQt1uAqwapNg5Lo2qjJM9ZG5g2wfzZoLh_Ztyk,77
+leap_model_parser-0.1.225.dev1.dist-info/RECORD,,