Merge pull request #5 from PINTO0309/restore_metadata

Fixed to restore metadata

Author: PINTO0309

sne4onnx/__init__.py

@@ -1,3 +1,3 @@
 from sne4onnx.onnx_network_extraction import extraction, main
 
-__version__ = '1.0.12'
+__version__ = '1.0.13'

sne4onnx/onnx_network_extraction.py

@@ -122,6 +122,10 @@ def extraction(
     # domain, ir_version
     domain: str = onnx_graph.domain
     ir_version: int = onnx_graph.ir_version
+    meta_data = {'domain': domain, 'ir_version': ir_version}
+    metadata_props = None
+    if hasattr(onnx_graph, 'metadata_props'):
+        metadata_props = onnx_graph.metadata_props
 
     graph = gs.import_onnx(onnx_graph)
     graph.cleanup().toposort()
@@ -176,9 +180,14 @@ def extraction(
     # Shape Estimation
     extracted_graph = None
     try:
-        extracted_graph = onnx.shape_inference.infer_shapes(gs.export_onnx(graph, do_type_check=False, **{'domain': domain, 'ir_version': ir_version}))
+        exported_onnx_graph = gs.export_onnx(graph, do_type_check=False, **meta_data)
+        if metadata_props is not None:
+            exported_onnx_graph.metadata_props.extend(metadata_props)
+        extracted_graph = onnx.shape_inference.infer_shapes(exported_onnx_graph)
     except Exception as e:
-        extracted_graph = gs.export_onnx(graph, do_type_check=False, **{'domain': domain, 'ir_version': ir_version})
+        extracted_graph = gs.export_onnx(graph, do_type_check=False, **meta_data)
+        if metadata_props is not None:
+            exported_onnx_graph.metadata_props.extend(metadata_props)
         if not non_verbose:
             print(
                 f'{Color.YELLOW}WARNING:{Color.RESET} '+